app.py

import torch
from torchvision import transforms, models
from PIL import Image
import gradio as gr
import os

# Use CPU
device = torch.device('cpu')

# Define ResNet-50 Architecture
model = models.resnet50(weights=None)

# Chanege model ouputs to fit this data (num_classes=37)
model.fc = torch.nn.Linear(2048, 37)

# Load model's weight
model.load_state_dict(torch.load('./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 (埃及貓)', 
               '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 predict images function
def classify_image(image):
    image = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        outputs = model(image)
        probabilities, indices = torch.topk(outputs, k=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:.2f}" for class_name, prob in predictions}

examples_path = './examples'

if os.path.exists(examples_path):
    print(f"[INFO] Found examples folder at {examples_path}")
else:
    print(f"[ERROR] Examples folder not found at {examples_path}")

examples = [[examples_path + "/" + img] for img in os.listdir(examples_path)]

# Create the reference list
breed_list_text = """
### Recognizable Breeds:

- 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 (約克夏犬)
"""


# Gradio Interface
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=f'A ResNet50-based model for classifying 37 different pet breeds.\n\n{breed_list_text}', 
    article='[Oxford Project](https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/The%20Oxford-IIIT%20Pet%20Project)'
)

demo.launch()