Update app.py

app.py

@@ -1,13 +1,15 @@
 import torch
 import torch.nn as nn
-from torchvision.models import resnet18, resnet34, resnet50
+from torchvision.models import resnet50
 from torchvision import transforms
 from PIL import Image
 import gradio as gr
 
-# -----------------------
-# CLASS LABELS
-# -----------------------
+'''
+This is only the first model, the one from CV, input an image, output will be top three results, 
+will later change to return best results only
+'''
+
 CLASS_LABELS = [
     'Corn_Common_Rust', 'Corn_Gray_Leaf_Spot', 'Corn_Healthy', 'Corn_Northern_Leaf_Blight',
     'Potato_Early_Blight', 'Potato_Healthy', 'Potato_Late_Blight',
@@ -18,21 +20,11 @@ CLASS_LABELS = [
 
 NUM_CLASSES = len(CLASS_LABELS)
 
-
-# -----------------------
-# MODEL ARCHITECTURE
-# -----------------------
 class ResNetPlantDisease(nn.Module):
     def __init__(self, num_classes=17, model_name='resnet50', pretrained=False):
         super().__init__()
 
-        if model_name == 'resnet18':
-            self.backbone = resnet18(weights=None)
-            num_features = 512
-        elif model_name == 'resnet34':
-            self.backbone = resnet34(weights=None)
-            num_features = 512
-        elif model_name == 'resnet50':
+        if model_name == 'resnet50':
             self.backbone = resnet50(weights=None)
             num_features = 2048
         else:
@@ -50,18 +42,11 @@ class ResNetPlantDisease(nn.Module):
         return self.backbone(x)
 
 
-# -----------------------
-# LOAD MODEL
-# -----------------------
 model = ResNetPlantDisease(num_classes=NUM_CLASSES, model_name='resnet50')
 state = torch.load("plant_disease_resnet_model.pth", map_location="cpu")
 model.load_state_dict(state)
 model.eval()
 
-
-# -----------------------
-# TRANSFORMS
-# -----------------------
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -71,10 +56,6 @@ transform = transforms.Compose([
     )
 ])
 
-
-# -----------------------
-# PREDICT FUNCTION
-# -----------------------
 def predict(image):
     img = Image.fromarray(image)
     img = transform(img).unsqueeze(0)
@@ -86,10 +67,6 @@ def predict(image):
     result = {CLASS_LABELS[i]: float(probs[i]) for i in range(NUM_CLASSES)}
     return result
 
-
-# -----------------------
-# GRADIO UI + API
-# -----------------------
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy"),