Update app.py

app.py

@@ -271,52 +271,27 @@ def predict(image):
         image_tensor = preprocess_image(image)
         with torch.no_grad():
             output = model(image_tensor)
-            if isinstance(output, tuple):
-                logits = output[0]
-            else:
-                logits = output
+            logits = output[0] if isinstance(output, tuple) else output
 
-            # 取得預測的top k結果
             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]]
+            topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
 
             if top1_prob >= 0.5:
-                # 正確辨識時，返回該品種資訊
-                predicted = topk_indices[0][0]
-                breed = dog_breeds[predicted.item()]
+                # High confidence prediction
+                breed = topk_breeds[0]
                 description = get_dog_description(breed)
-                akc_link = get_akc_breeds_link()
-
-                if isinstance(description, dict):
-                    description_str = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
-                else:
-                    description_str = description
-
-                # 添加AKC連結
-                description_str += f"\n\n**Want to learn more about dog breeds?** [Visit the AKC dog breeds page]({akc_link}) and search for {breed} to find detailed information."
-
-                # 添加免責聲明
-                disclaimer = ("\n\n*Disclaimer: The external link provided leads to the American Kennel Club (AKC) dog breeds page. "
-                              "You may need to search for the specific breed on that page. "
-                              "I am not responsible for the content on external sites. "
-                              "Please refer to the AKC's terms of use and privacy policy.*")
-                description_str += disclaimer
-
-                return description_str, gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False)
+                return format_description(description, breed), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
             elif top1_prob < 0.1:
-                # 如果信心度低於 0.1，返回提示請上傳更清晰的圖片
-                return "The image is too unclear or the dog breed is not in the dataset. Please upload a clearer image of the dog.", gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False)
-
+                # Very low confidence prediction
+                return ("The image is too unclear or the dog breed is not in the dataset. Please upload a clearer image of the dog.",
+                        gr.update(visible=False), gr.update(visible=False), gr.update(visible=False))
             else:
-                # 不確定時，返回top 3的預測結果，並且允許點擊查看詳細資訊
-                topk_breeds = [dog_breeds[idx.item()] for idx in topk_indices[0]]
-                topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
-
-                # 提供說明
+                # Medium confidence prediction
                 explanation = (
                     f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n"
                     f"1. {topk_breeds[0]} ({topk_probs_percent[0]} confidence)\n"
@@ -324,24 +299,20 @@ def predict(image):
                     f"3. {topk_breeds[2]} ({topk_probs_percent[2]} confidence)\n\n"
                     "Click on a button to view more information about the breed."
                 )
+                return explanation, gr.update(visible=True, value=f"More about {topk_breeds[0]}"), gr.update(visible=True, value=f"More about {topk_breeds[1]}"), gr.update(visible=True, value=f"More about {topk_breeds[2]}")
 
-                return explanation, gr.Button.update(visible=True, value=f"More about {topk_breeds[0]}"), gr.Button.update(visible=True, value=f"More about {topk_breeds[1]}"), gr.Button.update(visible=True, value=f"More about {topk_breeds[2]}")
     except Exception as e:
-        return f"An error occurred: {e}", gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False)
+        return f"An error occurred: {e}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
-def show_details(breed):
-    description = get_dog_description(breed)
-    akc_link = get_akc_breeds_link()
-    
+def format_description(description, breed):
     if isinstance(description, dict):
         description_str = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
     else:
         description_str = description
 
-    # 添加AKC連結
+    akc_link = get_akc_breeds_link()
     description_str += f"\n\n**Want to learn more about dog breeds?** [Visit the AKC dog breeds page]({akc_link}) and search for {breed} to find detailed information."
 
-    # 添加免責聲明
     disclaimer = ("\n\n*Disclaimer: The external link provided leads to the American Kennel Club (AKC) dog breeds page. "
                   "You may need to search for the specific breed on that page. "
                   "I am not responsible for the content on external sites. "
@@ -350,20 +321,11 @@ def show_details(breed):
 
     return description_str
 
-# Gradio Interface
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(label="Upload a dog image", type="numpy"),
-    outputs=[gr.Markdown(label="Prediction Results"), gr.Button(value="View More 1", visible=False), gr.Button(value="View More 2", visible=False), gr.Button(value="View More 3", visible=False)],
-    title="<h1 style='font-family:Roboto; font-weight:bold; color:#2C3E50; text-align:center;'>🐶 Dog Breed Classifier 🔍</h1>",
-    article= 'For more details on this project and other work, feel free to visit my GitHub [Dog Breed Classifier](https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/Dog%20Breed%20Classifier)',
-    description="<p style='font-family:Open Sans; color:#34495E; text-align:center;'>Upload a picture of a dog, and model will predict its breed, provide detailed information, and include an extra information link!</p>",
-    examples=['Border_Collie.jpg',
-              'Golden_Retriever.jpeg',
-              'Saint_Bernard.jpeg',
-              'French_Bulldog.jpeg',
-              'Samoyed.jpg'],
-     css = """
+def show_details(breed):
+    description = get_dog_description(breed)
+    return format_description(description, breed)
+
+with gr.Blocks(css="""
     .container {
         max-width: 900px;
         margin: 0 auto;
@@ -372,25 +334,14 @@ iface = gr.Interface(
         border-radius: 15px;
         box-shadow: 0 0 20px rgba(0, 0, 0, 0.1);
     }
-    
-    .gr-form {
-        display: flex;
-        flex-direction: column;
-        align-items: center;
-    }
-    
-    .gr-box {
-        width: 100%;
-        max-width: 500px;
-    }
-    
+    .gr-form { display: flex; flex-direction: column; align-items: center; }
+    .gr-box { width: 100%; max-width: 500px; }
     .output-markdown, .output-image {
         margin-top: 20px;
         padding: 15px;
         background-color: #f5f5f5;
         border-radius: 10px;
     }
-    
     .examples {
         display: flex;
         justify-content: center;
@@ -398,22 +349,38 @@ iface = gr.Interface(
         gap: 10px;
         margin-top: 20px;
     }
-    
     .examples img {
         width: 100px;
         height: 100px;
         object-fit: cover;
     }
-    """,
-    theme='default'
-)
+    """) as iface:
+
+    gr.HTML("<h1 style='font-family:Roboto; font-weight:bold; color:#2C3E50; text-align:center;'>🐶 Dog Breed Classifier 🔍</h1>")
+    gr.HTML("<p style='font-family:Open Sans; color:#34495E; text-align:center;'>Upload a picture of a dog, and the model will predict its breed, provide detailed information, and include an extra information link!</p>")
+    
+    with gr.Row():
+        input_image = gr.Image(label="Upload a dog image", type="numpy")
+        output = gr.Markdown(label="Prediction Results")
+    
+    with gr.Row():
+        btn1 = gr.Button("View More 1", visible=False)
+        btn2 = gr.Button("View More 2", visible=False)
+        btn3 = gr.Button("View More 3", visible=False)
+
+    input_image.change(predict, inputs=input_image, outputs=[output, btn1, btn2, btn3])
+    
+    btn1.click(show_details, inputs=btn1, outputs=output)
+    btn2.click(show_details, inputs=btn2, outputs=output)
+    btn3.click(show_details, inputs=btn3, outputs=output)
+
+    gr.Examples(
+        examples=['Border_Collie.jpg', 'Golden_Retriever.jpeg', 'Saint_Bernard.jpeg', 'French_Bulldog.jpeg', 'Samoyed.jpg'],
+        inputs=input_image
+    )
 
-# 定義按鈕的點擊事件
-iface.add_event_listener('View More 1', fn=lambda: show_details(topk_breeds[0]))
-iface.add_event_listener('View More 2', fn=lambda: show_details(topk_breeds[1]))
-iface.add_event_listener('View More 3', fn=lambda: show_details(topk_breeds[2]))
+    gr.HTML('For more details on this project and other work, feel free to visit my GitHub <a href="https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/Dog%20Breed%20Classifier">Dog Breed Classifier</a>')
 
-# 啟動 Gradio 應用
 if __name__ == "__main__":
     iface.launch()