DawnC commited on
Commit
4c3f6bd
·
verified ·
1 Parent(s): 1d9fa1f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -15
app.py CHANGED
@@ -10,10 +10,9 @@ device = torch.device('cpu')
10
  # Define ResNet-50 Architecture
11
  model = models.resnet50(weights=None)
12
 
13
- # revised full connected layer to 37 (num_classes)
14
  model.fc = torch.nn.Linear(2048, 37)
15
 
16
- # Load Model weights
17
  model.load_state_dict(torch.load('./resnet50_model_weights.pth', map_location=device))
18
 
19
  model.eval()
@@ -24,6 +23,7 @@ transform = transforms.Compose([
24
  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
25
  ])
26
 
 
27
  class_names = ['Abyssinian (阿比西尼亞貓)', 'American Bulldog (美國鬥牛犬)', 'American Pit Bull Terrier (美國比特鬥牛梗)',
28
  'Basset Hound (巴吉度獵犬)', 'Beagle (米格魯)', 'Bengal (孟加拉貓)', 'Birman (緬甸貓)', 'Bombay (孟買貓)',
29
  'Boxer (拳師犬)', 'British Shorthair (英國短毛貓)', 'Chihuahua (吉娃娃)', 'Egyptian Mau (埃及貓)',
@@ -35,38 +35,48 @@ class_names = ['Abyssinian (阿比西尼亞貓)', 'American Bulldog (美國鬥
35
  'Siamese (暹羅貓)', 'Sphynx (無毛貓)', 'Staffordshire Bull Terrier (史塔福郡鬥牛犬)',
36
  'Wheaten Terrier (小麥色梗)', 'Yorkshire Terrier (約克夏犬)']
37
 
38
- # predict function
39
  def classify_image(image):
40
- image = transform(image).unsqueeze(0).to(device) # make sure prediction on cpu
 
41
  with torch.no_grad():
 
42
  outputs = model(image)
43
- probabilities, indices = torch.topk(outputs, k=3) # top 3 predictions
44
- probabilities = torch.nn.functional.softmax(probabilities, dim=1)
 
 
 
45
  predictions = [(class_names[idx], prob.item()) for idx, prob in zip(indices[0], probabilities[0])]
46
- return {class_name: prob for class_name, prob in predictions}
47
-
48
 
 
49
  examples_path = './examples'
50
 
 
51
  if os.path.exists(examples_path):
52
  print(f"[INFO] Found examples folder at {examples_path}")
53
  else:
54
  print(f"[ERROR] Examples folder not found at {examples_path}")
55
 
56
- # Gradio Interface
 
57
  examples = [[examples_path + "/" + img] for img in os.listdir(examples_path)]
58
 
59
- # drop down list
60
- dropdown = gr.Dropdown(choices=class_names, label="Select a breed", type="value")
61
 
 
62
  demo = gr.Interface(
63
  fn=classify_image,
64
- inputs=[gr.Image(type="pil"), dropdown],
65
- outputs=[gr.Label(num_top_classes=3, label="Top 3 Predictions")],
66
  examples=examples,
67
  title='Oxford Pet 🐈🐕',
68
- description='A ResNet50-based model for classifying 37 different pet breeds.',
69
- article='[Oxford Project](https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/The%20Oxford-IIIT%20Pet%20Project)'
 
70
  )
71
 
 
72
  demo.launch()
 
10
  # Define ResNet-50 Architecture
11
  model = models.resnet50(weights=None)
12
 
13
+ # Revise fully connected layer to output 37 classes (num_classes = 37)
14
  model.fc = torch.nn.Linear(2048, 37)
15
 
 
16
  model.load_state_dict(torch.load('./resnet50_model_weights.pth', map_location=device))
17
 
18
  model.eval()
 
23
  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
24
  ])
25
 
26
+ # List of class names (37 dog and cat breeds)
27
  class_names = ['Abyssinian (阿比西尼亞貓)', 'American Bulldog (美國鬥牛犬)', 'American Pit Bull Terrier (美國比特鬥牛梗)',
28
  'Basset Hound (巴吉度獵犬)', 'Beagle (米格魯)', 'Bengal (孟加拉貓)', 'Birman (緬甸貓)', 'Bombay (孟買貓)',
29
  'Boxer (拳師犬)', 'British Shorthair (英國短毛貓)', 'Chihuahua (吉娃娃)', 'Egyptian Mau (埃及貓)',
 
35
  'Siamese (暹羅貓)', 'Sphynx (無毛貓)', 'Staffordshire Bull Terrier (史塔福郡鬥牛犬)',
36
  'Wheaten Terrier (小麥色梗)', 'Yorkshire Terrier (約克夏犬)']
37
 
38
+ # Prediction function
39
  def classify_image(image):
40
+ # Apply transformation and add batch dimension
41
+ image = transform(image).unsqueeze(0).to(device)
42
  with torch.no_grad():
43
+ # Make predictions using the model
44
  outputs = model(image)
45
+ # Apply softmax to get probabilities
46
+ probabilities = torch.nn.functional.softmax(outputs, dim=1)
47
+ # Get the top 3 predictions
48
+ probabilities, indices = torch.topk(probabilities, k=3)
49
+ # Return the class names with their corresponding probabilities
50
  predictions = [(class_names[idx], prob.item()) for idx, prob in zip(indices[0], probabilities[0])]
51
+ return {class_name: prob for class_name, prob in predictions} # Return raw float numbers
 
52
 
53
+ # Path to the folder containing example images
54
  examples_path = './examples'
55
 
56
+ # Check if the example images folder exists
57
  if os.path.exists(examples_path):
58
  print(f"[INFO] Found examples folder at {examples_path}")
59
  else:
60
  print(f"[ERROR] Examples folder not found at {examples_path}")
61
 
62
+ # Gradio interface
63
+ # Load example images from the folder
64
  examples = [[examples_path + "/" + img] for img in os.listdir(examples_path)]
65
 
66
+ # Create dropdown menu for users to see available classes (as reference, no direct connection to prediction)
67
+ dropdown = gr.Dropdown(choices=class_names, label="Recognizable Breeds", type="value")
68
 
69
+ # Define Gradio Interface
70
  demo = gr.Interface(
71
  fn=classify_image,
72
+ inputs=[gr.Image(type="pil")], # Only image input is used for prediction
73
+ outputs=[gr.Label(num_top_classes=3, label="Top 3 Predictions")], # Outputs top 3 predictions with probabilities
74
  examples=examples,
75
  title='Oxford Pet 🐈🐕',
76
+ description='A ResNet50-based model for classifying 37 different pet breeds.',
77
+ article='[Oxford Project](https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/The%20Oxford-IIIT%20Pet%20Project)',
78
+ inputs_dropdown=dropdown # Dropdown used as a reference list
79
  )
80
 
81
+ # Launch Gradio demo
82
  demo.launch()