Upload app.py

app.py

@@ -7,7 +7,7 @@ from torchvision import models
 import torch.nn as nn
 
 # Define the LightningModule class (should match the training code)
-class ResNet50Lightning(pl.LightningModule):
+class ResNet50Image2k(pl.LightningModule):
     def __init__(self, num_classes=1000):
         super().__init__()
         self.model = models.resnet50(pretrained=False)
@@ -17,8 +17,8 @@ class ResNet50Lightning(pl.LightningModule):
         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)
+checkpoint_path = "./resnet50_exp.ckpt"  # Replace with your checkpoint file path
+model = ResNet50Image2k.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 = model.to(device)
@@ -52,9 +52,11 @@ def predict_top5(image):
     return top5_results
 
 examples = [
-    ["examples/espresso.jpg.webp"],  # Example 1
-    ["examples/american_bullfrog.jpg"],  # Example 2
-    ["examples/tiger_shark.jpg"],  # Example 3
+    ["Images/Bird.JPEG"],  # Example 1
+    ["Images/Chamelion.JPEG"],  # Example 2
+    ["Images/Lizard.JPEG"],  # Example 3
+    ["Images/Shark.JPEG"],  # Example 4
+    ["Images/Turtle.JPEG"],  # Example 5
 ]
 
 # Create the Gradio interface