Create app.py

app.py

@@ -0,0 +1,145 @@
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import io
+
+from pathlib import Path
+import os, shutil
+from tqdm.auto import tqdm
+import torchvision
+from torch.utils.data import DataLoader
+from torchvision.datasets import ImageFolder
+from torchvision.transforms import transforms
+import torch.optim as optim
+from torchvision.models import resnet50, ResNet50_Weights
+import urllib.request
+import tarfile
+
+# Transform
+transform = transforms.Compose([
+    transforms.Resize((224,224)),
+    transforms.ToTensor()
+])
+
+# Dataset download
+urllib.request.urlretrieve(
+    "https://www.mydrive.ch/shares/38536/3830184030e49fe74747669442f0f282/download/420937484-1629951672/carpet.tar.xz",
+    "carpet.tar.xz"
+)
+
+with tarfile.open('carpet.tar.xz') as f:
+    f.extractall('.')
+
+# Feature extractor class
+class resnet_feature_extractor(torch.nn.Module):
+    def __init__(self):
+        super(resnet_feature_extractor, self).__init__()
+        self.model = resnet50(weights=ResNet50_Weights.DEFAULT)
+        self.model.eval()
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        def hook(module, input, output):
+            self.features.append(output)
+
+        self.model.layer2[-1].register_forward_hook(hook)            
+        self.model.layer3[-1].register_forward_hook(hook) 
+
+    def forward(self, input):
+        self.features = []
+        with torch.no_grad():
+            _ = self.model(input)
+
+        self.avg = torch.nn.AvgPool2d(3, stride=1)
+        fmap_size = self.features[0].shape[-2]
+        self.resize = torch.nn.AdaptiveAvgPool2d(fmap_size)
+
+        resized_maps = [self.resize(self.avg(fmap)) for fmap in self.features]
+        patch = torch.cat(resized_maps, 1)
+        patch = patch.reshape(patch.shape[1], -1).T
+        return patch
+
+# Initialize backbone
+backbone = resnet_feature_extractor()
+
+# Memory bank
+memory_bank = []
+folder_path = Path("carpet/train/good")
+for pth in tqdm(folder_path.iterdir(), leave=False):
+    with torch.no_grad():
+        data = transform(Image.open(pth)).unsqueeze(0)
+        features = backbone(data)
+        memory_bank.append(features.cpu().detach())
+memory_bank = torch.cat(memory_bank, dim=0)
+
+# Threshold
+y_score = []
+for pth in tqdm(folder_path.iterdir(), leave=False):
+    data = transform(Image.open(pth)).unsqueeze(0)
+    with torch.no_grad():
+        features = backbone(data)
+    distances = torch.cdist(features, memory_bank, p=2.0)
+    dist_score, _ = torch.min(distances, dim=1) 
+    s_star = torch.max(dist_score)
+    y_score.append(s_star.cpu().numpy())
+
+best_threshold = np.mean(y_score) + 2 * np.std(y_score)
+
+
+# Gradio Function
+
+def detect_fault(uploaded_image):
+    test_image = transform(uploaded_image).unsqueeze(0)
+
+    with torch.no_grad():
+        features = backbone(test_image)
+
+    distances = torch.cdist(features, memory_bank, p=2.0)
+    dist_score, _ = torch.min(distances, dim=1) 
+    s_star = torch.max(dist_score)
+    segm_map = dist_score.view(1, 1, 28, 28) 
+    segm_map = torch.nn.functional.interpolate(
+                segm_map,
+                size=(224, 224),
+                mode='bilinear'
+            ).cpu().squeeze().numpy()
+
+    y_score_image = s_star.cpu().numpy()
+    y_pred_image = 1*(y_score_image >= best_threshold)
+    class_label = ['Image Is OK','Image Is Not OK']
+
+    # Plot results
+    fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+    axs[0].imshow(test_image.squeeze().permute(1,2,0).cpu().numpy())
+    axs[0].set_title("Original Image")
+    axs[0].axis("off")
+
+    axs[1].imshow(segm_map, cmap='jet')
+    axs[1].set_title(f"Anomaly Score: {y_score_image / best_threshold:0.4f}\nPrediction: {class_label[y_pred_image]}")
+    axs[1].axis("off")
+
+    axs[2].imshow((segm_map > best_threshold*1.25), cmap='gray')
+    axs[2].set_title("Fault Segmentation Map")
+    axs[2].axis("off")
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png")
+    buf.seek(0)
+    result_image = Image.open(buf)
+    plt.close(fig)
+
+    return result_image
+
+# Launch Gradio App
+demo = gr.Interface(
+    fn=detect_fault,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Image(type="pil", label="Detection Result"),
+    title="Fault Detection in Images",
+    description="Upload an image and the model will detect if there are any faults and show the segmentation map."
+)
+
+if __name__ == "__main__":
+    demo.launch()