pytorch lightningでのpredictエラー

anomaly_maps = self.model(batch["image"])

IndexError: too many indices for tensor of dimension 4

import os
from argparse import ArgumentParser, Namespace

from pytorch_lightning import Trainer
import torch
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
import pytorch_lightning as pl

from util.src.anomalib.config import get_configurable_parameters
from util.src.anomalib.data import get_datamodule
from util.src.anomalib.models import get_model
from util.src.anomalib.utils.callbacks import get_callbacks
from util.auto_edit import auto_edit_other_app

import json
from PIL import Image
import numpy as np
import yaml
import datetime
import shutil
import glob


# 画像ファイルのパスを指定
image_path = "1.jpg"

# PIL Imageとして画像を読み込む
your_image = Image.open(image_path)
# 1. 画像をテンソルに変換
transform = transforms.Compose([transforms.ToTensor()])
image = transform(your_image)  # your_imageは推論したい画像

# 2. それを含むデータセットを作成
class SingleImageDataset(Dataset):
    def __init__(self, image):
        self.data = [image]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def main():
    os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
    test(1)
    # while True:
    #     predict_type = input("onnx、ckptどちらで推論しますか\n 'onnx' or 'ckpt': ").lower()
    #     if predict_type in ['o', 'on', 'onn', 'onnx']:
    #         test(0)
    #         break
    #     elif predict_type in ['c', 'ck', 'ckp', 'ckpt']:
    #         test(1)
    #         break


def test(file_type):
    """Test an anomaly classification and segmentation model that is initially trained via `tools/train.py`.

    The script is able to write the results into both filesystem and a logger such as Tensorboard.
    """
    # args = get_args()

    with open("train_result_file/train_settings.json") as settings_file:
        file_contents = json.load(settings_file)
        model_type = file_contents['ai_model'] # train_type
        task_type = file_contents['task']
        image_size = file_contents['image_size']
        print("model_type", model_type, task_type, image_size)

    CONFIG_PATH = f"util/src/anomalib/models/{model_type.lower()}/config.yaml"

    # necessary for image_size to change
    with open(CONFIG_PATH, 'r') as ymlfile:
        data = yaml.safe_load(ymlfile)
        data['dataset']['image_size'] = int(image_size)
    with open(CONFIG_PATH, "w") as ymlfile:
        yaml.dump(data, ymlfile)


    if file_type == 1:
        files = glob.glob("train_result_file/model*.ckpt")
        weight_file = files[0]
    else:
        files = glob.glob("train_result_file/model*.onnx")
        weight_file = files[0]

    config = get_configurable_parameters(
        model_name=model_type.lower(),
        config_path=CONFIG_PATH,
        weight_file=weight_file,
    )

    image_folder = "images"
    edit_json_path = "train_result_file/image_edit.json"
    if os.path.isfile(edit_json_path):
        os.makedirs(f"{image_folder}/edited", exist_ok=True)
        if model_type.lower() == "padim":
            os.makedirs(f"{image_folder}/edited/OK", exist_ok=True)
            os.makedirs(f"{image_folder}/edited/NG", exist_ok=True)
            os.makedirs(f"{image_folder}/edited/testOK", exist_ok=True)

            edit_files(f"{image_folder}/OK", edit_json_path, f"{image_folder}/edited/OK")
            edit_files(f"{image_folder}/NG", edit_json_path, f"{image_folder}/edited/NG")
            edit_files(f"{image_folder}/testOK", edit_json_path, f"{image_folder}/edited/testOK")

        else:
            edit_files(image_folder, edit_json_path, f"{image_folder}/edited")

        image_folder = f"{image_folder}/edited"

    fol_time = "{0:%Y%m%d-%H%M%S}".format(datetime.datetime.now())
    output_folder = f"results/{fol_time}"
    os.makedirs(output_folder, exist_ok=True)

    config["dataset"]["path"] = f"{image_folder}/"
    config["dataset"]["task"] = task_type #"classification" "segmentation"
    config["visualization"]["image_save_path"] = f"{output_folder}/"

    datamodule = get_datamodule(config)
    model = get_model(config)

    callbacks = get_callbacks(config)

    
    # trainer.test(model=model, datamodule=datamodule)

    dataset = SingleImageDataset(image)

    # 3. データセットを用いてデータローダーを作成
    data_loader = DataLoader(dataset, batch_size=1)

    # 4. Trainer.predict()メソッドを呼び出して推論
    trainer = Trainer(callbacks=callbacks, **config.trainer)
    predictions = trainer.predict(model=model, dataloaders=data_loader)

    # 予測結果を取得
    for prediction in predictions:
        # 予測結果を処理する（具体的な処理はモデルに依存します）
        print(prediction)
    
    # trainer.predict(model=model, dataloaders=datamodule)

    if os.path.exists(f"results/{model_type.lower()}") and os.path.isdir(f"results/{model_type.lower()}"):
        shutil.rmtree(f"results/{model_type.lower()}")


def get_args() -> Namespace:
    """Get CLI arguments.

    Returns:
        Namespace: CLI arguments.
    """
    parser = ArgumentParser()
    parser.add_argument("--model", type=str, default="stfpm",
                        help="Name of the algorithm to train/test")
    parser.add_argument("--config", type=str, required=False,
                        help="Path to a model config file")
    parser.add_argument("--weight_file", type=str,
                        default="weights/model.ckpt")

    args = parser.parse_args()
    return args


def edit_files(image_folder, edit_json_path, dst_folder):
    with open(edit_json_path, "r") as edit_file:
        json_load = json.load(edit_file)

        # auto_edit_other_appは、学習とテストで使用しているauto_editをベースとしているので変更したくない
        # auto_edit_other_appを使用するには、マスクと位置合わせの画像パスを前で指定する必要あり
        edit_param_data = {}
        for edit_count in range(1, len(json_load.keys())):
            if json_load[str(edit_count)]["event"] == "masked":
                mask_name = json_load[str(edit_count)]["event_data"]["mask_name"]
                edit_param_data["mask_image_path"] = f"train_result_file/{mask_name}"
            if json_load[str(edit_count)]["event"] == "position":
                master_image_file_name = json_load[str(edit_count)]["event_data"]["master_image_file_name"]
                edit_param_data["position_image_path"] = f"train_result_file/{master_image_file_name}"
        print("edit_param_data",edit_param_data)

        for image_name in os.listdir(image_folder):
            image_path = os.path.join(image_folder, image_name)

            if os.path.isfile(image_path):
                img = Image.open(image_path)
                img = img.convert('RGB')
                img = np.array(img, np.float32)
                edit_img = auto_edit_other_app(json_load, img, edit_param_data)# 出力はuint8
                
                db_image = Image.fromarray(edit_img)# uint8で保存
                result_img = f"{dst_folder}/edited_{image_name}.png"
                db_image.save(result_img, quality=95)

main()

"""FastFlow Lightning Model Implementation."""

# Copyright (C) 2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import torch
from omegaconf import DictConfig, ListConfig
from pytorch_lightning.callbacks import EarlyStopping
from pytorch_lightning.utilities.types import STEP_OUTPUT
from torch import Tensor, optim

from anomalib.models.components import AnomalyModule
from anomalib.models.fastflow.loss import FastflowLoss
from anomalib.models.fastflow.torch_model import FastflowModel


class Fastflow(AnomalyModule):
    """PL Lightning Module for the FastFlow algorithm.

    Args:
        input_size (tuple[int, int]): Model input size.
        backbone (str): Backbone CNN network
        pre_trained (bool, optional): Boolean to check whether to use a pre_trained backbone.
        flow_steps (int, optional): Flow steps.
        conv3x3_only (bool, optinoal): Use only conv3x3 in fast_flow model. Defaults to False.
        hidden_ratio (float, optional): Ratio to calculate hidden var channels. Defaults to 1.0.
    """

    def __init__(
        self,
        input_size: tuple[int, int],
        backbone: str,
        pre_trained: bool = True,
        flow_steps: int = 8,
        conv3x3_only: bool = False,
        hidden_ratio: float = 1.0,
    ) -> None:
        super().__init__()

        self.model = FastflowModel(
            input_size=input_size,
            backbone=backbone,
            pre_trained=pre_trained,
            flow_steps=flow_steps,
            conv3x3_only=conv3x3_only,
            hidden_ratio=hidden_ratio,
        )
        self.loss = FastflowLoss()

    def training_step(self, batch: dict[str, str | Tensor], *args, **kwargs) -> STEP_OUTPUT:
        """Forward-pass input and return the loss.

        Args:
            batch (batch: dict[str, str | Tensor]): Input batch
            _batch_idx: Index of the batch.

        Returns:
            STEP_OUTPUT: Dictionary containing the loss value.
        """
        del args, kwargs  # These variables are not used.

        hidden_variables, jacobians = self.model(batch["image"])
        loss = self.loss(hidden_variables, jacobians)
        self.log("train_loss", loss.item(), on_epoch=True, prog_bar=True, logger=True)
        return {"loss": loss}

    def validation_step(self, batch: dict[str, str | Tensor], *args, **kwargs) -> STEP_OUTPUT:
        """Forward-pass the input and return the anomaly map.

        Args:
            batch (dict[str, str | Tensor]): Input batch

        Returns:
            STEP_OUTPUT | None: batch dictionary containing anomaly-maps.
        """
        del args, kwargs  # These variables are not used.

        anomaly_maps = self.model(batch["image"])
        batch["anomaly_maps"] = anomaly_maps
        return batch


class FastflowLightning(Fastflow):
    """PL Lightning Module for the FastFlow algorithm.

    Args:
        hparams (DictConfig | ListConfig): Model params
    """

    def __init__(self, hparams: DictConfig | ListConfig) -> None:
        super().__init__(
            input_size=hparams.model.input_size,
            backbone=hparams.model.backbone,
            pre_trained=hparams.model.pre_trained,
            flow_steps=hparams.model.flow_steps,
            conv3x3_only=hparams.model.conv3x3_only,
            hidden_ratio=hparams.model.hidden_ratio,
        )
        self.hparams: DictConfig | ListConfig  # type: ignore
        self.save_hyperparameters(hparams)

    def configure_callbacks(self) -> list[EarlyStopping]:
        """Configure model-specific callbacks.

        Note:
            This method is used for the existing CLI.
            When PL CLI is introduced, configure callback method will be
                deprecated, and callbacks will be configured from either
                config.yaml file or from CLI.
        """
        early_stopping = EarlyStopping(
            monitor=self.hparams.model.early_stopping.metric,
            patience=self.hparams.model.early_stopping.patience,
            mode=self.hparams.model.early_stopping.mode,
        )
        return [early_stopping]

    def configure_optimizers(self) -> torch.optim.Optimizer:
        """Configures optimizers for each decoder.

        Note:
            This method is used for the existing CLI.
            When PL CLI is introduced, configure optimizers method will be
                deprecated, and optimizers will be configured from either
                config.yaml file or from CLI.

        Returns:
            Optimizer: Adam optimizer for each decoder
        """
        return optim.Adam(
            params=self.model.parameters(),
            lr=self.hparams.model.lr,
            weight_decay=self.hparams.model.weight_decay,
        )