Add automatic download of imagenet_classes.txt

app.py

@@ -1,7 +1,9 @@
+import os
 import gradio as gr
 import glob
 import time
 import random
+import requests
 
 # Import necessary libraries
 from torchvision import models, transforms
@@ -20,8 +22,23 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
+# Function to download imagenet_classes.txt
+def download_imagenet_classes():
+    url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
+    response = requests.get(url)
+    if response.status_code == 200:
+        with open("imagenet_classes.txt", "wb") as f:
+            f.write(response.content)
+        print("imagenet_classes.txt downloaded successfully.")
+    else:
+        print("Failed to download imagenet_classes.txt")
+
+# Check if imagenet_classes.txt exists, if not, download it
+if not os.path.exists("imagenet_classes.txt"):
+    download_imagenet_classes()
+
 # Load class labels
-with open('imagenet_classes.txt') as f:
+with open('imagenet_classes.txt', 'r') as f:
     labels = [line.strip() for line in f.readlines()]
 
 def classify_image(image):

requirements.txt

@@ -2,4 +2,5 @@ gradio==3.23.0
 Pillow==9.5.0
 torch==1.13.1
 torchvision==0.14.1
-urllib3<1.27,>=1.25.4
+urllib3<1.27,>=1.25.4
+requests==2.31.0