Update app.py

app.py

@@ -4,26 +4,26 @@ from PIL import Image
 import gradio as gr
 import os
 
-# 使用 CPU
+# Use CPU
 device = torch.device('cpu')
 
-# 定義 ResNet-50 模型架構
-model = models.resnet50(weights=None)  # 不使用預訓練權重，僅構建模型架構
+# define ResNet-50 architecture
+model = models.resnet50(weights=None)
 
-# 加載模型權重到模型架構
-model.load_state_dict(torch.load('resnet50_model_weights.pth', map_location=device))
+# change the output to 37 (num_classes)
+model.fc = torch.nn.Linear(2048, 37)
+
+# Load model weights
+model.load_state_dict(torch.load('/content/Oxford_Pet_classifier/resnet50_model_weights.pth', map_location=device))
 
-# 設置模型為評估模式
 model.eval()
 
-# 定義影像預處理
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# 定義類別名稱
 class_names = ['Abyssinian (阿比西尼亞貓)', 'American Bulldog (美國鬥牛犬)', 'American Pit Bull Terrier (美國比特鬥牛梗)',
                'Basset Hound (巴吉度獵犬)', 'Beagle (米格魯)', 'Bengal (孟加拉貓)', 'Birman (緬甸貓)', 'Bombay (孟買貓)', 
                'Boxer (拳師犬)', 'British Shorthair (英國短毛貓)', 'Chihuahua (吉娃娃)', 'Egyptian Mau (埃及貓)', 
@@ -35,19 +35,27 @@ class_names = ['Abyssinian (阿比西尼亞貓)', 'American Bulldog (美國鬥
                'Siamese (暹羅貓)', 'Sphynx (無毛貓)', 'Staffordshire Bull Terrier (史塔福郡鬥牛犬)', 
                'Wheaten Terrier (小麥色梗)', 'Yorkshire Terrier (約克夏犬)']
 
-# 定義預測函數
+# predict function
 def classify_image(image):
-    image = transform(image).unsqueeze(0).to(device)  # 確保影像資料處理在 CPU 上
+    image = transform(image).unsqueeze(0).to(device)  # make sure on the cpu
     with torch.no_grad():
         outputs = model(image)
-        probabilities, indices = torch.topk(outputs, k=3)  # 取得前3個預測
-        probabilities = torch.nn.functional.softmax(probabilities, dim=1)  # 將結果轉換為機率
+        probabilities, indices = torch.topk(outputs, k=3)  # top 3 predictions
+        probabilities = torch.nn.functional.softmax(probabilities, dim=1)  # turn into probabilities
         predictions = [(class_names[idx], prob.item()) for idx, prob in zip(indices[0], probabilities[0])]
         return {class_name: f"{prob * 100:.2f}%" for class_name, prob in predictions}
 
-# Gradio 介面
-examples = [["examples/" + img] for img in os.listdir('examples')]
+# define examples_path
+examples_path = '/content/Oxford_Pet_classifier/examples'
+
+# make sure examples is exists
+if os.path.exists(examples_path):
+    examples = [[examples_path + "/" + img] for img in os.listdir(examples_path)]
+else:
+    print(f"[ERROR] Examples folder not found at {examples_path}")
+    examples = []  
 
+# Gradio Interface
 demo = gr.Interface(fn=classify_image,  
                     inputs=gr.Image(type="pil"),  
                     outputs=[gr.Label(num_top_classes=3, label="Top 3 Predictions")],  
@@ -56,5 +64,4 @@ demo = gr.Interface(fn=classify_image,
                     description='A ResNet50-based model for classifying 37 different pet breeds.',
                     article='https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/The%20Oxford-IIIT%20Pet%20Project')  
 
-# 啟動 Gradio demo
-demo.launch()
+demo.launch()