added gradioApp.py to main

gradioApp.py

@@ -0,0 +1,88 @@
+import gradio as gr
+import torch
+from torchvision import transforms, models
+from PIL import Image
+import torch.nn as nn
+import os
+os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
+
+
+# Use the model architecture
+class ResNet18(nn.Module):
+    def __init__(self, num_classes):
+        super(ResNet18, self).__init__()
+        self.resnet18 = models.resnet18(weights='ResNet18_Weights.DEFAULT')
+        self.resnet18.fc = nn.Linear(self.resnet18.fc.in_features, num_classes)
+
+    def forward(self, x):
+        return self.resnet18(x)
+
+# Load the pretrained classifier
+num_classes = 2
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = ResNet18(num_classes=num_classes)
+model.load_state_dict(torch.load('resnet_state_dict.pth', map_location=device)) # Load trained state path from resnet_state_dict.pth
+model = model.to(device)
+model.eval()
+
+# Transform
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,), (0.5,))
+])
+
+# Define classes
+class_names = ["Yes, it is a hotdog :)", "No, it isn't a hotdog! :("]
+
+# Prediction function
+def predict(image):
+    try:
+        if isinstance(image, Image.Image):
+            image = image.convert("RGB")
+        else:
+            raise ValueError("Input is not a PIL Image")
+
+        image = transform(image).unsqueeze(0)
+        image = image.to(device)
+
+        # Perform inference
+        with torch.no_grad():
+            output = model(image)
+            _, predicted = torch.max(output, 1)
+
+        return class_names[predicted.item()]
+    except Exception as e:
+        return str(e)
+
+# Use one of the preset images if not for an uploaded hotdog image
+preset_images = [
+    'data/test/hot_dog/133012.jpg',
+    'data/test/hot_dog/133015.jpg',
+    'data/test/hot_dog/133245.jpg',
+    'data/test/hot_dog/135628.jpg',
+    'data/test/hot_dog/138933.jpg',
+    'data/test/not_hot_dog/6229.jpg',
+    'data/test/not_hot_dog/6261.jpg',
+    'data/test/not_hot_dog/6709.jpg',
+    'data/test/not_hot_dog/6926.jpg',
+    'data/test/not_hot_dog/7056.jpg']
+
+# Gradio interface
+iface = gr.Interface(
+    fn=predict, 
+    inputs=gr.Image(type="pil", label="Upload your image"),
+    theme='gstaff/xkcd',
+    outputs=gr.Textbox(label="Is it a hotodog?"),  # Show the predicted class name
+    live=True,
+    description="Your friendly hotdog/nothotdog classifier"
+)
+
+header = gr.Markdown("""
+    # Welcome to the Hotdog Classifier! 🍔
+    This app classifies whether an image shows a hotdog or not.
+    Upload an image or choose from the preset images below.
+""")
+
+# Launch the app, currently share set to True
+iface.launch()