Question of custom dataset inference!

[cory619-seungwon]

Hello, all.

I trained my custom dataset with anomalib patchcore.
Training was good.

However, when inferencing my evaluation images(total 42 images), the number of predictions is only 4.
So, there's deviation between the number of evaluation images & predictions now...
Error occured as "Mismatch between number of predictions (4) and images (42).".

I think "datamodule = Folder()" make my custom dataset split so, these error occurred..
What should I modify in my code to run inference on my custom dataset?
In short, how to inference only 1 image path?

I don't wanna split my custom dataset when inferencing evaluation data.
With this idea, I will make inference sever file to inference only 1 image_path when client_file send 1 image_path.

Let me share my code as below with several information.

• train batch size : 2
• eval batch size : 1

Thanks in advance.

import argparse
import os
from pathlib import Path
from PIL import Image
import torch
import numpy as np
import cv2
import matplotlib.pyplot as plt
from anomalib.models import Patchcore
from anomalib.engine import Engine
from anomalib.utils.normalization import NormalizationMethod
from anomalib.data import Folder
from anomalib import TaskType


def load_image(image_path):
    image = Image.open(image_path).convert("RGB")
    image = torch.tensor(np.array(image)).permute(2, 0, 1).float()  # (H, W, C) -> (C, H, W)
    image = image.unsqueeze(0) 
    return image


def overlay_anomaly_map_on_image(image, anomaly_map):
    image_np = image.squeeze().permute(1, 2, 0).cpu().numpy().astype(np.uint8)  # (C, H, W) --> (H, W, C)
    anomaly_map_resized = cv2.resize(anomaly_map, (image_np.shape[1], image_np.shape[0]))
    min_val = np.min(anomaly_map_resized)
    max_val = np.max(anomaly_map_resized)
    anomaly_map_rescaled = ((anomaly_map_resized - min_val) / (max_val - min_val) * 255).astype(np.uint8)
    heatmap = cv2.applyColorMap(anomaly_map_rescaled, cv2.COLORMAP_JET)
    overlay_image = cv2.addWeighted(image_np, 0.6, heatmap, 0.3, 0)

    return overlay_image



def train_model(data_root, model_save_path):
    datamodule = Folder(
        name='dacon',
        root=data_root,
        normal_dir="good",
        abnormal_dir="crack",
        mask_dir="crack_mask",
        task=TaskType.CLASSIFICATION,
        train_batch_size=2,
        eval_batch_size=1,
        seed=0
    )

    model = Patchcore(
        backbone="wide_resnet101_2",
        pre_trained=True,
        coreset_sampling_ratio=0.1,
        layers=["layer2", "layer3"]
    )

    engine = Engine(
        normalization=NormalizationMethod.MIN_MAX,
        task=TaskType.CLASSIFICATION,
        accelerator="cuda"
    )

    engine.train(model=model, datamodule=datamodule)

    model_save_dir = Path(model_save_path)
    model_save_dir.mkdir(parents=True, exist_ok=True)
    model_save_file = model_save_dir / "trained_model.pth"
    torch.save(model.state_dict(), model_save_file)

    print(f"model saved in {model_save_file}.")


def test_model(data_root, model_load_path):
    
    evaluation_dir = Path("D:\\patchcore\\anomalib_2\\evaluation_images")
    inference_output_dir = Path("D:\\patchcore\\anomalib_2\\inference_results")
    inference_output_dir.mkdir(parents=True, exist_ok=True)

    if not evaluation_dir.exists() or len(list(evaluation_dir.glob('*.png'))) == 0:
        raise FileNotFoundError(f"No valid images found in {evaluation_dir}")

    model = Patchcore(
        backbone="wide_resnet101_2",
        pre_trained=False,  
        coreset_sampling_ratio=0.1,
        layers=["layer2", "layer3"]
    )

    # Load model checkpoint
    model_load_file = Path(model_load_path) / "trained_model.pth"
    if not model_load_file.exists():
        raise FileNotFoundError(f"Saved model not found in {model_load_file}")

    model.load_state_dict(torch.load(model_load_file))
    print(f"Model loaded from {model_load_file}")

    engine = Engine(
        normalization=NormalizationMethod.MIN_MAX,
        task=TaskType.CLASSIFICATION,
        accelerator="cuda"
    )

    datamodule = Folder(
        name='evaluation',
        root=str(evaluation_dir), 
        normal_dir=str(evaluation_dir), 
        task=TaskType.CLASSIFICATION,
        eval_batch_size=1,
    )


    predictions = engine.predict(datamodule=datamodule, model=model)

    image_paths = list(evaluation_dir.glob('*.png'))
    
    if len(predictions) != len(image_paths):
        raise ValueError(f"Mismatch between number of predictions ({len(predictions)}) and images ({len(image_paths)})")

    for idx, image_path in enumerate(image_paths):
        image_name = image_path.stem
        anomaly_map = predictions[idx]['anomaly_map'].squeeze().cpu().numpy()
        pred_score = predictions[idx]['pred_score'].item()


        prediction_map_path = inference_output_dir / f"{image_name}_prediction_map.png"

        plt.title(f"Prediction Map for {image_name}")
        plt.savefig(prediction_map_path)
        plt.close()


        anomaly_percentage = min(int(pred_score), 100)


        annotated_image = np.array(Image.open(image_path))
        cv2.putText(annotated_image, f"anomalous ({anomaly_percentage}%)", (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        annotated_image_path = inference_output_dir / f"{image_name}_anomalous.png"
        Image.fromarray(annotated_image).save(annotated_image_path)

    print(f"Inference results saved in {inference_output_dir}")




def main():
    parser = argparse.ArgumentParser(description="Patchcore anomaly detection training and inference")
    parser.add_argument('--train', action='store_true', help="Train the Patchcore model")
    parser.add_argument('--test', action='store_true', help="Test the Patchcore model")
    args = parser.parse_args()

    data_root = "D:\\patchcore\\anomalib_2\\dataset_patchcore"
    model_save_path = "D:\\patchcore\\anomalib_2\\models"
    image_dir = "D:\\patchcore\\anomalib_2\\dataset_patchcore\\crack"

    if args.train:
        train_model(data_root, model_save_path)
    elif args.test:
        test_model(data_root, model_save_path)
    else:
        print("Please specify --train or --test to proceed.")


if __name__ == '__main__':
    main()