fixed bugs

app.py

@@ -1,57 +1,60 @@
 import torch
-from torchvision import transforms
-from torchvision import models
+from torchvision import transforms, models
 from PIL import Image
 import gradio as gr
 import os
 
-# Use CPU
+# 使用 CPU
 device = torch.device('cpu')
 
-# Load the model ResNet-50 model architecture
-model = models.resnet50(pretrained=False)
+# 定義 ResNet-50 模型架構
+model = models.resnet50(weights=None)  # 不使用預訓練權重，僅構建模型架構
 
-# Load model's weight to CPU
-model = torch.load('resnet50_model_weights.pth', map_location=device)
+# 加載模型權重到模型架構
+model.load_state_dict(torch.load('resnet50_model_weights.pth', map_location=device))
+
+# 設置模型為評估模式
 model.eval()
 
-# Define the image preprocessing
+# 定義影像預處理
 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])
 ])
 
-# Define the class names
-class_names = ['Abyssinian', 'American Bulldog', 'American Pit Bull Terrier', 'Basset Hound', 'Beagle', 'Bengal', 'Birman', 'Bombay', 
-               'Boxer', 'British Shorthair', 'Chihuahua', 'Egyptian Mau', 'English Cocker Spaniel', 'English Setter', 'German Shorthaired', 
-               'Great Pyrenees', 'Havanese', 'Japanese Chin', 'Keeshond', 'Leonberger', 'Maine Coon', 'Miniature Pinscher', 'Newfoundland', 
-               'Persian', 'Pomeranian', 'Pug', 'Ragdoll', 'Russian Blue', 'Saint Bernard', 'Samoyed', 'Scottish Terrier', 'Shiba Inu', 
-               'Siamese', 'Sphynx', 'Staffordshire Bull Terrier', 'Wheaten Terrier', 'Yorkshire Terrier']
-
-# Define the predict function
+# 定義類別名稱
+class_names = ['Abyssinian (阿比西尼亞貓)', 'American Bulldog (美國鬥牛犬)', 'American Pit Bull Terrier (美國比特鬥牛梗)',
+               'Basset Hound (巴吉度獵犬)', 'Beagle (米格魯)', 'Bengal (孟加拉貓)', 'Birman (緬甸貓)', 'Bombay (孟買貓)', 
+               'Boxer (拳師犬)', 'British Shorthair (英國短毛貓)', 'Chihuahua (吉娃娃)', 'Egyptian Mau (埃及貓)', 
+               'English Cocker Spaniel (英國可卡犬)', 'English Setter (英國設得蘭犬)', 'German Shorthaired (德國短毛犬)', 
+               'Great Pyrenees (大白熊犬)', 'Havanese (哈瓦那犬)', 'Japanese Chin (日本狆)', 'Keeshond (荷蘭毛獅犬)', 
+               'Leonberger (萊昂貝格犬)', 'Maine Coon (緬因貓)', 'Miniature Pinscher (迷你品犬)', 'Newfoundland (紐芬蘭犬)', 
+               'Persian (波斯貓)', 'Pomeranian (博美犬)', 'Pug (哈巴狗)', 'Ragdoll (布偶貓)', 'Russian Blue (俄羅斯藍貓)', 
+               'Saint Bernard (聖伯納犬)', 'Samoyed (薩摩耶)', 'Scottish Terrier (蘇格蘭梗)', 'Shiba Inu (柴犬)', 
+               'Siamese (暹羅貓)', 'Sphynx (無毛貓)', 'Staffordshire Bull Terrier (史塔福郡鬥牛犬)', 
+               'Wheaten Terrier (小麥色梗)', 'Yorkshire Terrier (約克夏犬)']
+
+# 定義預測函數
 def classify_image(image):
-    image = transform(image).unsqueeze(0).to(device)  # Ensure image data is processed on CPU
+    image = transform(image).unsqueeze(0).to(device)  # 確保影像資料處理在 CPU 上
     with torch.no_grad():
         outputs = model(image)
-        _, predicted = torch.max(outputs, 1)
-        return class_names[predicted.item()]
-
-# Custom Gradio interface title, description, and article
-title = 'Oxford Pet 🐈🐕'
-description = 'A ResNet50-based computer vision model for classifying images of pets from the Oxford-IIIT Pet Dataset. The model can recognize 37 different pet breeds, including cats and dogs.'
-article = 'https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/The%20Oxford-IIIT%20Pet%20Project'
+        probabilities, indices = torch.topk(outputs, k=3)  # 取得前3個預測
+        probabilities = torch.nn.functional.softmax(probabilities, dim=1)  # 將結果轉換為機率
+        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 interface
+# Gradio 介面
 examples = [["examples/" + img] for img in os.listdir('examples')]
 
-demo = gr.Interface(fn=classify_image,  # Map input to output function
-                    inputs=gr.Image(type="pil"),  # Image input
-                    outputs=[gr.Label(num_top_classes=1, label="Predictions")],  # Predicted label
-                    examples=examples,  # Example images
-                    title=title,  
-                    description=description,  
-                    article=article)  
+demo = gr.Interface(fn=classify_image,  
+                    inputs=gr.Image(type="pil"),  
+                    outputs=[gr.Label(num_top_classes=3, label="Top 3 Predictions")],  
+                    examples=examples,  
+                    title='Oxford Pet 🐈🐕',  
+                    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')  
 
-# Launch the demo
-demo.launch()
+# 啟動 Gradio demo
+demo.launch()