Anormalib - new

https://github.com/openvinotoolkit/anomalib?tab=readme-ov-file

GitHub - openvinotoolkit/anomalib: An anomaly detection library comprising state-of-the-art algorithms and features such as expe

결함이 있는 이미지(이상 이미지)를 탐지하고, 모델을 학습하고 평가한 후, 새로운 이미지에서 결함을 찾아내는 추론까지 수행

 

• 데이터셋 로드:
  • MVTec AD 데이터셋을 사용하여 "hazelnut" (헤이즐넛) 카테고리의 이미지를 불러옵니다. 이 데이터셋은 산업용 결함 탐지에서 사용되며, 정상 및 비정상(결함이 있는) 이미지를 포함합니다.
  • 이미지를 (256, 256) 크기로 조정하고, 배치 크기 32로 학습 및 테스트 데이터를 준비합니다.
• PatchCore 모델 설정:
  • PatchCore라는 이상 탐지 모델을 사용합니다. 이 모델은 ResNet18 백본을 사용하며, 사전 학습된 가중치를 불러옵니다. 이 모델은 작은 이상(예: 스크래치, 균열 등)을 탐지하는 데 유리한 모델입니다.
• 모델 학습:
  • PyTorch Lightning의 Trainer를 사용하여 PatchCore 모델을 학습합니다. 이 단계에서 주어진 데이터셋으로 모델이 학습됩니다. 학습된 모델은 이미지에서 정상과 비정상을 구분하는 능력을 가지게 됩니다.
• 모델 평가:
  • 학습된 모델이 테스트 데이터셋에서 얼마나 잘 동작하는지 평가합니다. 즉, 모델이 새로운 이미지에서 결함을 얼마나 정확하게 찾아내는지 확인합니다.
• 모델 저장:
  • 학습된 모델을 저장하여, 이후에 추론 단계에서 사용할 수 있도록 합니다.
• 추론:
  • 새로운 데이터가 주어졌을 때, 모델이 그 데이터에서 결함을 탐지하는 작업을 수행합니다. 저장된 모델을 불러와 결함이 있는지 없는지를 예측합니다.
• 시각화:
  • 첫 번째 이미지의 이상 탐지 결과를 시각화합니다. 결과는 anomaly_map이라는 형태로 나와서, 이미지의 각 부분에서 결함이 있는지 없는지를 색상으로 표시합니다. 'hot' 컬러맵을 사용해 결함이 있는 부분을 강조하여 보여줍니다.

설치 모듈

nvidia-smi

nvcc --version

가상환경 및 설치

# 1. Conda 가상 환경 생성 및 활성화
conda create -n anomaly_detection python=3.9
conda activate anomaly_detection

# 2. PyTorch 설치 (CUDA 11.8 기준)
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

# 3. Anomalib 및 기타 패키지 설치
pip install anomalib matplotlib

또는

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

pip install anomalib pytorch-lightning matplotlib

코드

# 필요한 모듈을 가져옵니다.
from anomalib.data import MVTec
from anomalib.models import Patchcore
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint  # 체크포인트 콜백 추가
import matplotlib.pyplot as plt

# 데이터셋 로드: MVTec AD 데이터셋을 로드합니다.
datamodule = MVTec(
    root="path_to_actual_dataset",  # 실제 데이터 경로로 변경
    category="hazelnut",  # 분석하고자 하는 카테고리
    image_size=(256, 256),
    train_batch_size=32,
    test_batch_size=32,
)

# 모델 초기화: PatchCore 모델을 초기화합니다.
model = Patchcore(
    input_size=(256, 256),
    backbone="resnet18",  # 사용할 백본 네트워크
    pre_trained=True,     # 사전 학습된 가중치 사용 여부
)

# Checkpoint 콜백 생성: 모델을 저장할 경로 설정
checkpoint_callback = ModelCheckpoint(
    dirpath="path_to_save_model",  # 저장 경로
    filename="best_model",  # 저장 파일명
    save_top_k=1,  # 최고 성능 모델 1개만 저장
    monitor="val_loss",  # 검증 손실을 기준으로 최적 모델 선택
    mode="min",  # 손실이 최소일 때 저장
)

# Trainer 초기화: 학습 및 평가를 관리할 pytorch lightning의 Trainer 사용
trainer = Trainer(
    max_epochs=50,  # 원하는 에폭 수 설정
    callbacks=[checkpoint_callback]  # 모델 저장 콜백 추가
)

# 모델 학습: 데이터셋을 사용하여 모델을 학습시킵니다.
trainer.fit(
    model=model,
    datamodule=datamodule
)

# 모델 평가: 학습된 모델을 테스트 데이터로 평가합니다.
results = trainer.test(
    model=model,
    datamodule=datamodule
)

# 결과 출력: 결과를 확인합니다.
print(results)

# 모델 저장: 모델 저장 경로는 CheckpointCallback에서 처리됩니다.
# 추가 저장을 원한다면 아래 코드 사용 가능:
# trainer.save_checkpoint("path_to_save_model/best_model.ckpt")

# 새로운 데이터로 추론 수행: 새 데이터에 대한 이상 탐지를 수행합니다.
predictions = trainer.predict(
    model=model,
    datamodule=datamodule,
    ckpt_path="path_to_save_model/best_model.ckpt"  # 저장된 모델 경로
)

# 예측 결과를 출력합니다.
print(predictions)

# 예시에서 첫 번째 이미지의 anomaly_map을 추출하여 시각화
# predictions의 구조를 먼저 확인
if 'anomaly_map' in predictions[0]:  # anomaly_map 키 존재 여부 확인
    anomaly_map = predictions[0]['anomaly_map']  # 예측 결과 구조에 맞춰 변경

    # anomaly_map 시각화
    plt.imshow(anomaly_map, cmap='hot')  # 'hot' 컬러맵을 사용하여 이상점 강조
    plt.colorbar()  # 색상 막대 추가
    plt.title("Anomaly Map for sample_image_1.png")
    plt.show()
else:
    print("Anomaly map not found in predictions.")

 

dataset_root/

• train/
  • normal/
    • image1.png, image2.png, ...
• test/
  • normal/
    • image1.png, image2.png, ...
  • anomaly/
    • image1.png, image2.png, ...

train

# Step 1: Import necessary libraries
from anomalib.data import Folder
from anomalib.models import Patchcore
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
import matplotlib.pyplot as plt
import torch

# Step 2: DataModule Initialization
# Folder datamodule is used for loading data from custom directories.
datamodule = Folder(
    root="path_to_custom_dataset",  # Path to your dataset
    normal_dir="train/normal",      # Path to normal training images
    abnormal_dir="test/anomaly",    # Path to anomaly test images
    normal_test_dir="test/normal",  # Path to normal test images
    image_size=(256, 256),          # Resize the images
    train_batch_size=32,
    test_batch_size=32
)

# Step 3: Model Initialization
model = Patchcore(
    input_size=(256, 256),     # Input image size
    backbone="resnet18",       # Backbone architecture
    layers_to_extract_from=[2, 3],  # Feature extraction layers
    pre_trained=True           # Use pre-trained weights
)

# Step 4: Callbacks (For saving the best model based on validation performance)
checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',
    dirpath='path_to_save_model',
    filename='best_model',
    save_top_k=1,
    mode='min',
)

# Step 5: Trainer Initialization (for managing training process)
trainer = Trainer(
    max_epochs=50,  # Number of training epochs
    callbacks=[checkpoint_callback],  # Save model checkpoint
    gpus=1 if torch.cuda.is_available() else 0  # Use GPU if available
)

# Step 6: Train the model
trainer.fit(model, datamodule=datamodule)

# Step 7: Test the model (Evaluate on test dataset)
trainer.test(model, datamodule=datamodule)

# Step 8: Save the best model
trainer.save_checkpoint("path_to_save_model/best_model.ckpt")

inference

# Step 1: Load the trained model from checkpoint
model = Patchcore.load_from_checkpoint("path_to_save_model/best_model.ckpt")

# Step 2: Perform inference on new data
predictions = trainer.predict(
    model=model,
    datamodule=datamodule  # Reusing the datamodule with test data
)

# Step 3: Visualize the results
for idx, prediction in enumerate(predictions):
    # Extract anomaly map for each prediction
    anomaly_map = prediction['anomaly_map']

    # Visualize the anomaly map using matplotlib
    plt.imshow(anomaly_map, cmap='hot')  # Hot colormap highlights anomalies
    plt.colorbar()
    plt.title(f"Anomaly Map for test image {idx+1}")
    plt.show()

검증데이터셋 사용시 (테스트 데이터셋에서 일부 사용)

# Step 1: Import necessary libraries
from anomalib.data import Folder
from anomalib.models import Patchcore
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
import torch

# Step 2: DataModule Initialization with validation
datamodule = Folder(
    root="path_to_custom_dataset",   # Path to your dataset
    normal_dir="train/normal",       # Path to normal training images
    abnormal_dir="test/anomaly",     # Path to anomaly test images
    normal_test_dir="test/normal",   # Path to normal test images
    val_split=0.2,                   # Split 20% of training data for validation
    image_size=(256, 256),           # Resize the images
    train_batch_size=32,
    test_batch_size=32,
    val_batch_size=32                # Add validation batch size
)

# Step 3: Model Initialization
model = Patchcore(
    input_size=(256, 256),     # Input image size
    backbone="resnet18",       # Backbone architecture
    layers_to_extract_from=[2, 3],  # Feature extraction layers
    pre_trained=True           # Use pre-trained weights
)

# Step 4: Callbacks (For saving the best model based on validation performance)
checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',        # Monitor validation loss
    dirpath='path_to_save_model',
    filename='best_model',
    save_top_k=1,
    mode='min',
)

# Step 5: Trainer Initialization
trainer = Trainer(
    max_epochs=50,            # Number of training epochs
    callbacks=[checkpoint_callback],  # Save model checkpoint
    gpus=1 if torch.cuda.is_available() else 0,  # Use GPU if available
    check_val_every_n_epoch=1  # Perform validation every epoch
)

# Step 6: Train the model with validation
trainer.fit(model, datamodule=datamodule)

# Step 7: Test the model (Evaluate on test dataset)
trainer.test(model, datamodule=datamodule)

# Step 8: Save the best model
trainer.save_checkpoint("path_to_save_model/best_model.ckpt")

검증용 데이터를 사용할때 특정경로의 데이터 사용

path_to_validation_data/
   normal/
      image1.png, image2.png, ...
   anomaly/
      image1.png, image2.png, ...

# Step 1: Import necessary libraries
from anomalib.data import Folder
from anomalib.models import Patchcore
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
import torch

# Step 2: DataModule Initialization with explicit validation directory
datamodule = Folder(
    root="path_to_custom_dataset",      # Path to your dataset
    normal_dir="train/normal",          # Path to normal training images
    abnormal_dir="test/anomaly",        # Path to anomaly test images
    normal_test_dir="test/normal",      # Path to normal test images
    val_dir="path_to_validation_data",  # Path to the validation data
    image_size=(256, 256),              # Resize the images
    train_batch_size=32,
    val_batch_size=32,                  # Validation batch size
    test_batch_size=32                  # Test batch size
)

# Step 3: Model Initialization
model = Patchcore(
    input_size=(256, 256),     # Input image size
    backbone="resnet18",       # Backbone architecture
    layers_to_extract_from=[2, 3],  # Feature extraction layers
    pre_trained=True           # Use pre-trained weights
)

# Step 4: Callbacks (For saving the best model based on validation performance)
checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',        # Monitor validation loss
    dirpath='path_to_save_model',
    filename='best_model',
    save_top_k=1,
    mode='min',
)

# Step 5: Trainer Initialization
trainer = Trainer(
    max_epochs=50,            # Number of training epochs
    callbacks=[checkpoint_callback],  # Save model checkpoint
    gpus=1 if torch.cuda.is_available() else 0,  # Use GPU if available
    check_val_every_n_epoch=1  # Perform validation every epoch
)

# Step 6: Train the model with the custom validation set
trainer.fit(model, datamodule=datamodule)

# Step 7: Test the model (Evaluate on test dataset)
trainer.test(model, datamodule=datamodule)

# Step 8: Save the best model
trainer.save_checkpoint("path_to_save_model/best_model.ckpt")