Updated app.py

app.py

@@ -3,19 +3,10 @@ import traceback
 import gradio as gr
 import torch
 from torchvision.models import get_model
-from torchvision.transforms import v2
+from torchvision.transforms import transforms
 from torchvision.transforms.functional import InterpolationMode
 
 
-# Imagenet-1k classes
-if not os.path.exists("imagenet_classes.txt"):
-    os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")
-
-# Download an example image from the pytorch website
-if not os.path.exists("dog.jpg"):
-    torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
-
-
 # Function to load the model with custom weights
 def load_model(weights_path):
     model = get_model("resnet50", num_classes=1000)
@@ -48,13 +39,12 @@ def classify_image(image):
     return result
 
 # Define image transformation to match the model input
-transform = v2.Compose([
-    v2.Resize(256, interpolation=InterpolationMode.BILINEAR, antialias=True),
-    v2.CenterCrop(224),
-    v2.PILToTensor(),
-    v2.ToDtype(torch.float, scale=True),
-    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    v2.ToPureTensor(),
+transform = transforms.Compose([
+    transforms.Resize(256, interpolation=InterpolationMode.BILINEAR, antialias=True),
+    transforms.CenterCrop(224),
+    transforms.PILToTensor(),
+    transforms.ConvertImageDtype(torch.float),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
 # Path to the pre-trained model weights (should be set by the user)
@@ -67,13 +57,17 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil"),  # Image input (in PIL format)
     outputs=gr.Label(num_top_classes=5),  # Output will be the predicted top 5 classes with confidence scores
     title = "Image Recognition using ResNet-50 trained on Imagenet-1K",
+    live = True, 
     description = "<p style='text-align: center'> Gradio demo for ResNet, Deep residual networks pre-trained on ImageNet. To use it, simply upload your image, or click one of the examples to load them. </p>",
     article = "<p style='text-align: center'> \
         <a href='https://arxiv.org/abs/1512.03385' target='_blank'>Deep Residual Learning for Image Recognition</a> | \
             <a href='https://github.com/KD1994/session-9-imagenet-resnet50' target='_blank'>Github Repo</a> \
                 </p>",
     examples = [
-        ['dog.jpg']
+        ['examples/dog.jpg'], 
+        ['examples/great-white-shark.jpg'],
+        ['examples/american-goldfinch.jpg'],
+        ['examples/hognose-snake.jpg']
     ]
 )