Update app.py

app.py

@@ -1,15 +1,27 @@
 import gradio as gr
 import torch
-from torchvision import transforms, models
+import pytorch_lightning as pl
+from torchvision import transforms
 from PIL import Image
+from torchvision import models
+import torch.nn as nn
 
-# Load your trained ResNet50 model checkpoint
+# Define the LightningModule class (should match the training code)
+class ResNet50Lightning(pl.LightningModule):
+    def __init__(self, num_classes=1000):
+        super().__init__()
+        self.model = models.resnet50(pretrained=False)
+        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)
+    
+    def forward(self, x):
+        return self.model(x)
+
+# Load the model from PyTorch Lightning checkpoint
+checkpoint_path = "./resnet50_40epoch_imagenet1k.ckpt"  # Replace with your checkpoint file path
+model = ResNet50Lightning.load_from_checkpoint(checkpoint_path)
+model.eval()  # Set the model to evaluation mode
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = models.resnet50(num_classes=1000)  # Ensure the model matches your architecture
-checkpoint = torch.load("resnet50_40epoch_imagenet1k.ckpt", map_location=device)  # Replace with your checkpoint path
-model.load_state_dict(checkpoint['model_state_dict'])  # Load state_dict from your checkpoint
 model = model.to(device)
-model.eval()
 
 # Load ImageNet class labels
 with open("classes.txt") as f:
@@ -23,7 +35,7 @@ preprocess = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Function to predict top-5 classes
+# Define the prediction function
 def predict_top5(image):
     # Preprocess the image
     image = preprocess(image).unsqueeze(0).to(device)
@@ -35,7 +47,7 @@ def predict_top5(image):
 
     # Get top-5 predictions
     top5_prob, top5_catid = torch.topk(probabilities, 5)
-    top5_results = {class_labels[catid]: prob.item() for prob, catid in zip(top5_prob, top5_catid)}
+    top5_results = {class_labels[catid]: f"{prob.item():.4f}" for prob, catid in zip(top5_prob, top5_catid)}
 
     return top5_results