Update app.py

app.py

@@ -43,7 +43,7 @@ dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staff
               "Miniature_Poodle", "Miniature_Schnauzer", "Otterhound", "Papillon", "Pug", "Redbone", 
               "Schipperke", "Silky_Terrier", "Soft-Coated_Wheaten_Terrier", "Standard_Poodle", 
               "Standard_Schnauzer", "Toy_Poodle", "Toy_Terrier", "Vizsla", "Whippet", 
-              "Wire-Haired_Fox_Terrier","Havanese", "Dachshund", "Shiba_Inu", "Bichon_Frise"]
+              "Wire-Haired_Fox_Terrier"]
 
 class MultiHeadAttention(nn.Module):
 
@@ -100,11 +100,11 @@ class BaseModel(nn.Module):
         return logits, attended_features
 
 
-num_classes = 124
+num_classes = 120
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = BaseModel(num_classes=num_classes, device=device)
 
-checkpoint = torch.load('[124_82.30]_best_model_dog.pth', map_location=torch.device('cpu'))
+checkpoint = torch.load('best_model_81_dog.pth', map_location=torch.device('cpu'))
 model.load_state_dict(checkpoint['model_state_dict'])
 
 # evaluation mode
@@ -207,40 +207,67 @@ async def process_single_dog(image):
     
     # Case 1: Low confidence - unclear image or breed not in dataset
     if top1_prob < 0.15:
+        error_message = '''
+            <div class="dog-info-card">
+                <div class="breed-info">
+                    <p class="warning-message">
+                        <span class="icon">⚠️</span>
+                        The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.
+                    </p>
+                </div>
+            </div>
+        '''
         initial_state = {
-            "explanation": "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.",
+            "explanation": error_message,
             "image": None,
             "is_multi_dog": False
         }
-        return initial_state["explanation"], None, initial_state
+        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)
-        formatted_description = format_description(description, breed)
+        formatted_description = format_description_html(description, breed)  # 使用 format_description_html
+        html_content = f'''
+            <div class="dog-info-card">
+                <div class="breed-info">
+                    {formatted_description}
+                </div>
+            </div>
+        '''
         initial_state = {
-            "explanation": formatted_description,
+            "explanation": html_content,
             "image": image,
             "is_multi_dog": False
         }
-        return formatted_description, image, initial_state
+        return html_content, image, initial_state
         
     # Case 3: Medium confidence - show top 3 breeds with relative probabilities
     else:
-        breeds_info = ""
+        breeds_html = ""
         for i, (breed, prob) in enumerate(zip(topk_breeds, relative_probs)):
             description = get_dog_description(breed)
-            formatted_description = format_description(description, breed)
-            breeds_info += f"\n\nBreed {i+1}: **{breed}** (Confidence: {prob})\n{formatted_description}"
+            formatted_description = format_description_html(description, breed)  # 使用 format_description_html
+            breeds_html += f'''
+                <div class="dog-info-card">
+                    <div class="breed-info">
+                        <div class="breed-header">
+                            <span class="breed-name">Breed {i+1}: {breed}</span>
+                            <span class="confidence-badge">Confidence: {prob}</span>
+                        </div>
+                        {formatted_description}
+                    </div>
+                </div>
+            '''
 
         initial_state = {
-            "explanation": breeds_info,
+            "explanation": breeds_html,
             "image": image,
             "is_multi_dog": False
         }
-        return breeds_info, image, initial_state
+        return breeds_html, image, initial_state
 
         
 async def predict(image):