Update app.py

app.py

@@ -3,15 +3,15 @@ import numpy as np
 import torch
 import torch.nn as nn
 import gradio as gr
-from dataclasses import dataclass
+import time
 from torchvision.models import efficientnet_v2_m, EfficientNet_V2_M_Weights
 from torchvision.ops import nms, box_iou
 import torch.nn.functional as F
 from torchvision import transforms
 from PIL import Image, ImageDraw, ImageFont, ImageFilter
-from dog_database import get_dog_description
 from breed_health_info import breed_health_info
 from breed_noise_info import breed_noise_info
+from dog_database import get_dog_description
 from scoring_calculation_system import UserPreferences
 from recommendation_html_format import format_recommendation_html, get_breed_recommendations
 from history_manager import UserHistoryManager
@@ -42,19 +42,19 @@ model_yolo = YOLO('yolov8l.pt')
 history_manager = UserHistoryManager()
 
 dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staffordshire_Terrier",
-              "Appenzeller", "Australian_Terrier", "Bedlington_Terrier", "Bernese_Mountain_Dog",
+              "Appenzeller", "Australian_Terrier", "Bedlington_Terrier", "Bernese_Mountain_Dog", "Bichon_Frise",
               "Blenheim_Spaniel", "Border_Collie", "Border_Terrier", "Boston_Bull", "Bouvier_Des_Flandres",
               "Brabancon_Griffon", "Brittany_Spaniel", "Cardigan", "Chesapeake_Bay_Retriever",
-              "Chihuahua", "Dandie_Dinmont", "Doberman", "English_Foxhound", "English_Setter",
+              "Chihuahua", "Dachshund", "Dandie_Dinmont", "Doberman", "English_Foxhound", "English_Setter",
               "English_Springer", "EntleBucher", "Eskimo_Dog", "French_Bulldog", "German_Shepherd",
               "German_Short-Haired_Pointer", "Gordon_Setter", "Great_Dane", "Great_Pyrenees",
-              "Greater_Swiss_Mountain_Dog", "Ibizan_Hound", "Irish_Setter", "Irish_Terrier",
+              "Greater_Swiss_Mountain_Dog","Havanese", "Ibizan_Hound", "Irish_Setter", "Irish_Terrier",
               "Irish_Water_Spaniel", "Irish_Wolfhound", "Italian_Greyhound", "Japanese_Spaniel",
               "Kerry_Blue_Terrier", "Labrador_Retriever", "Lakeland_Terrier", "Leonberg", "Lhasa",
               "Maltese_Dog", "Mexican_Hairless", "Newfoundland", "Norfolk_Terrier", "Norwegian_Elkhound",
               "Norwich_Terrier", "Old_English_Sheepdog", "Pekinese", "Pembroke", "Pomeranian",
               "Rhodesian_Ridgeback", "Rottweiler", "Saint_Bernard", "Saluki", "Samoyed",
-              "Scotch_Terrier", "Scottish_Deerhound", "Sealyham_Terrier", "Shetland_Sheepdog",
+              "Scotch_Terrier", "Scottish_Deerhound", "Sealyham_Terrier", "Shetland_Sheepdog", "Shiba_Inu",
               "Shih-Tzu", "Siberian_Husky", "Staffordshire_Bullterrier", "Sussex_Spaniel",
               "Tibetan_Mastiff", "Tibetan_Terrier", "Walker_Hound", "Weimaraner",
               "Welsh_Springer_Spaniel", "West_Highland_White_Terrier", "Yorkshire_Terrier",
@@ -68,6 +68,7 @@ dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staff
               "Standard_Schnauzer", "Toy_Poodle", "Toy_Terrier", "Vizsla", "Whippet",
               "Wire-Haired_Fox_Terrier"]
 
+
 class MultiHeadAttention(nn.Module):
 
     def __init__(self, in_dim, num_heads=8):
@@ -122,15 +123,19 @@ class BaseModel(nn.Module):
         logits = self.classifier(attended_features)
         return logits, attended_features
 
-
-num_classes = 120
+# Initialize model
+num_classes = len(dog_breeds)
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = BaseModel(num_classes=num_classes, device=device)
 
-checkpoint = torch.load('best_model_81_dog.pth', map_location=torch.device('cpu'))
-model.load_state_dict(checkpoint['model_state_dict'])
+# Initialize base model
+model = BaseModel(num_classes=num_classes, device=device).to(device)
 
-# evaluation mode
+# Load model path
+model_path = "124_best_model_dog.pth"
+checkpoint = torch.load(model_path, map_location=device)
+
+# Load model state
+model.load_state_dict(checkpoint["base_model"], strict=False)
 model.eval()
 
 # Image preprocessing function
@@ -149,24 +154,38 @@ def preprocess_image(image):
     return transform(image).unsqueeze(0)
 
 async def predict_single_dog(image):
-    image_tensor = preprocess_image(image)
+    """
+    Predicts the dog breed using only the classifier.
+    Args:
+        image: PIL Image or numpy array
+    Returns:
+        tuple: (top1_prob, topk_breeds, relative_probs)
+    """
+    image_tensor = preprocess_image(image).to(device)
+    
     with torch.no_grad():
-        output = model(image_tensor)
-        logits = output[0] if isinstance(output, tuple) else output
-        probabilities = F.softmax(logits, dim=1)
-        topk_probs, topk_indices = torch.topk(probabilities, k=3)
-        top1_prob = topk_probs[0][0].item()
-        topk_breeds = [dog_breeds[idx.item()] for idx in topk_indices[0]]
-
-        # Calculate relative probabilities for display
-        raw_probs = [prob.item() for prob in topk_probs[0]]
-        sum_probs = sum(raw_probs)
-        relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in raw_probs]
-
-    return top1_prob, topk_breeds, relative_probs
-
-
-async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.45):
+        # Get model outputs (只使用logits，不需要features)
+        logits = model(image_tensor)[0]  # 如果model仍返回tuple，取第一個元素
+        probs = F.softmax(logits, dim=1)
+        
+        # Classifier prediction
+        top5_prob, top5_idx = torch.topk(probs, k=5)
+        breeds = [dog_breeds[idx.item()] for idx in top5_idx[0]]
+        probabilities = [prob.item() for prob in top5_prob[0]]
+        
+        # Calculate relative probabilities
+        sum_probs = sum(probabilities[:3])  # 只取前三個來計算相對概率
+        relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in probabilities[:3]]
+        
+        # Debug output
+        print("\nClassifier Predictions:")
+        for breed, prob in zip(breeds[:5], probabilities[:5]):
+            print(f"{breed}: {prob:.4f}")
+            
+        return probabilities[0], breeds[:3], relative_probs
+
+
+async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     boxes = []
@@ -193,7 +212,6 @@ async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.45):
 
     return dogs
 
-
 def non_max_suppression(boxes, iou_threshold):
     keep = []
     boxes = sorted(boxes, key=lambda x: x[1], reverse=True)
@@ -218,52 +236,6 @@ def calculate_iou(box1, box2):
     return iou
 
 
-async def process_single_dog(image):
-    """Process a single dog image and return breed predictions and HTML output."""
-    top1_prob, topk_breeds, relative_probs = await predict_single_dog(image)
-
-    # Case 1: Low confidence - unclear image or breed not in dataset
-    if top1_prob < 0.2:
-        error_message = format_warning_html(
-            'The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.'
-        )
-        initial_state = {
-            "explanation": error_message,
-            "image": None,
-            "is_multi_dog": False
-        }
-        return error_message, None, initial_state
-
-    breed = topk_breeds[0]
-
-    # Case 2: High confidence - single breed result
-    if top1_prob >= 0.45:
-        description = get_dog_description(breed)
-        html_content = format_single_dog_result(breed, description)
-        initial_state = {
-            "explanation": html_content,
-            "image": image,
-            "is_multi_dog": False
-        }
-        return html_content, image, initial_state
-
-    # Case 3: Medium confidence - show top 3 breeds with relative probabilities
-    description = get_dog_description(breed)
-    breeds_html = format_multiple_breeds_result(
-        topk_breeds=topk_breeds,
-        relative_probs=relative_probs,
-        color='#34C759',  # 使用單狗顏色
-        index=1,  # 因為是單狗處理，所以index為1
-        get_dog_description=get_dog_description
-    )
-
-    initial_state = {
-        "explanation": breeds_html,
-        "image": image,
-        "is_multi_dog": False
-    }
-    return breeds_html, image, initial_state
-
 
 def create_breed_comparison(breed1: str, breed2: str) -> dict:
     breed1_info = get_dog_description(breed1)
@@ -353,21 +325,46 @@ async def predict(image):
             top1_prob, topk_breeds, relative_probs = await predict_single_dog(cropped_image)
             combined_confidence = detection_confidence * top1_prob
 
-            # Format results based on confidence
-            if combined_confidence < 0.2:
+            # Format results based on confidence with error handling
+            try:
+                if combined_confidence < 0.2:
+                    dogs_info += format_error_message(color, i+1)
+                elif top1_prob >= 0.45:
+                    breed = topk_breeds[0]
+                    description = get_dog_description(breed)
+                    # Handle missing breed description
+                    if description is None:
+                        # 如果沒有描述，創建一個基本描述
+                        description = {
+                            "Name": breed,
+                            "Size": "Unknown",
+                            "Exercise Needs": "Unknown",
+                            "Grooming Needs": "Unknown",
+                            "Care Level": "Unknown",
+                            "Good with Children": "Unknown",
+                            "Description": f"Identified as {breed.replace('_', ' ')}"
+                        }
+                    dogs_info += format_single_dog_result(breed, description, color)
+                else:
+                    # 修改format_multiple_breeds_result的調用，包含錯誤處理
+                    dogs_info += format_multiple_breeds_result(
+                        topk_breeds,
+                        relative_probs,
+                        color,
+                        i+1,
+                        lambda breed: get_dog_description(breed) or {
+                            "Name": breed,
+                            "Size": "Unknown",
+                            "Exercise Needs": "Unknown",
+                            "Grooming Needs": "Unknown",
+                            "Care Level": "Unknown",
+                            "Good with Children": "Unknown",
+                            "Description": f"Identified as {breed.replace('_', ' ')}"
+                        }
+                    )
+            except Exception as e:
+                print(f"Error formatting results for dog {i+1}: {str(e)}")
                 dogs_info += format_error_message(color, i+1)
-            elif top1_prob >= 0.45:
-                breed = topk_breeds[0]
-                description = get_dog_description(breed)
-                dogs_info += format_single_dog_result(breed, description, color)
-            else:
-                dogs_info += format_multiple_breeds_result(
-                    topk_breeds,
-                    relative_probs,
-                    color,
-                    i+1,
-                    get_dog_description
-                )
 
         # Wrap final HTML output
         html_output = format_multi_dog_container(dogs_info)
@@ -422,6 +419,7 @@ def show_details_html(choice, previous_output, initial_state):
 def main():
     with gr.Blocks(css=get_css_styles()) as iface:
         # Header HTML
+
         gr.HTML("""
         <header style='text-align: center; padding: 20px; margin-bottom: 20px;'>
             <h1 style='font-size: 2.5em; margin-bottom: 10px; color: #2D3748;'>
@@ -467,6 +465,7 @@ def main():
                 history_component=history_component
             )
 
+
             # 4. 最後創建歷史記錄標籤頁
             create_history_tab(history_component)