Create app.py

app.py

@@ -0,0 +1,65 @@
+import torch
+from torchvision import models, transforms
+from PIL import Image
+import gradio as gr
+import pandas as pd
+
+# Load pre-trained models
+resnet18 = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
+resnet18.eval()
+
+convnext_tiny = models.convnext_tiny(weights=models.ConvNeXt_Tiny_Weights.DEFAULT)
+convnext_tiny.eval()
+
+# Image preprocessing
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225]),
+])
+
+# Load class labels
+labels = models.ResNet18_Weights.DEFAULT.meta['categories']
+
+# Function to get top-k predictions
+def get_topk(output, k=3):
+    probs = torch.nn.functional.softmax(output, dim=1)
+    topk_probs, topk_indices = torch.topk(probs, k)
+    return [(labels[idx], float(prob) * 100) for idx, prob in zip(topk_indices[0], topk_probs[0])]
+
+# Inference function
+def classify_image(image):
+    image = Image.fromarray(image)
+    input_tensor = preprocess(image).unsqueeze(0)  # Add batch dimension
+
+    # ResNet18 top-3 predictions
+    with torch.no_grad():
+        resnet_output = resnet18(input_tensor)
+    resnet_top3 = get_topk(resnet_output)
+
+    # ConvNeXt-Tiny top-3 predictions
+    with torch.no_grad():
+        convnext_output = convnext_tiny(input_tensor)
+    convnext_top3 = get_topk(convnext_output)
+
+    # Create DataFrame for table display
+    df = pd.DataFrame({
+        "Rank": [1, 2, 3],
+        "ResNet-18": [f"{label} ({prob:.2f}%)" for label, prob in resnet_top3],
+        "ConvNeXt-Tiny": [f"{label} ({prob:.2f}%)" for label, prob in convnext_top3]
+    })
+    return df
+
+# Gradio interface
+iface = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Dataframe(headers=["Rank", "ResNet-18", "ConvNeXt-Tiny"], type="pandas"),
+    title="Image Classification Validator",
+    description="Upload an AI-generated image to see top-3 predictions from ResNet-18 and ConvNeXt-Tiny with probabilities."
+)
+
+if __name__ == "__main__":
+    iface.launch()