Update app.py

app.py

@@ -8,6 +8,8 @@ import torch.nn.functional as F
 from torchvision import transforms
 from PIL import Image
 from data_manager import get_dog_description
+import wikipedia
+from urllib.parse import quote
 
 dog_breeds = ["Afghan_Hound(阿富汗獵犬)", "African_Hunting_Dog(非洲野犬)", "Airedale(艾爾谷犬)",
             "American_Staffordshire_Terrier(美國斯塔福郡梗)", "Appenzeller(亞賓澤爾犬)",
@@ -139,18 +141,53 @@ def preprocess_image(image):
 
     return transform(image).unsqueeze(0)
 
+# def predict(image):
+#     try:
+#         image_tensor = preprocess_image(image)
+#         with torch.no_grad():
+#             logits, _ = model(image_tensor)
+#             _, predicted = torch.max(logits, 1)
+
+#         breed = dog_breeds[predicted.item()]  # Map label to breed name
+
+#         # Retrieve breed description
+#         description = get_dog_description(breed)
+
+#         # Formatting the description for better display
+#         if isinstance(description, dict):
+#             description_str = f"**Breed**: {description['Breed']}\n\n"
+#             description_str += f"**Size**: {description['Size']}\n\n"
+#             description_str += f"**Lifespan**: {description['Lifespan']}\n\n"
+#             description_str += f"**Temperament**: {description['Temperament']}\n\n"
+#             description_str += f"**Care Level**: {description['Care Level']}\n\n"
+#             description_str += f"**Good with Children**: {description['Good with Children']}\n\n"
+#             description_str += f"**Exercise Needs**: {description['Exercise Needs']}\n\n"
+#             description_str += f"**Grooming Needs**: {description['Grooming Needs']}\n\n"
+#             description_str += f"**Description**: {description['Description']}\n\n"
+#         else:
+#             description_str = description
+
+#         return description_str
+#     except Exception as e:
+#         return f"An error occurred: {e}"
+
 def predict(image):
     try:
         image_tensor = preprocess_image(image)
         with torch.no_grad():
-            logits, _ = model(image_tensor)
+            logits, * = model(image_tensor)
             _, predicted = torch.max(logits, 1)
-
         breed = dog_breeds[predicted.item()]  # Map label to breed name
-
+        
         # Retrieve breed description
         description = get_dog_description(breed)
-
+        
+        # Generate Wikipedia link
+        try:
+            wiki_link = wikipedia.page(f"{breed} dog").url
+        except:
+            wiki_link = f"https://en.wikipedia.org/wiki/Special:Search?search={quote(breed)}+dog"
+        
         # Formatting the description for better display
         if isinstance(description, dict):
             description_str = f"**Breed**: {description['Breed']}\n\n"
@@ -164,89 +201,70 @@ def predict(image):
             description_str += f"**Description**: {description['Description']}\n\n"
         else:
             description_str = description
-
+        
+        # Add Wikipedia link
+        description_str += f"\n\n[Click here to view the Wikipedia page for {breed}]({wiki_link})"
+        
         return description_str
     except Exception as e:
         return f"An error occurred: {e}"
 
+
 iface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(label="Upload an image", type="numpy"),  # 支持拖放和圖片編輯
-    outputs="markdown",
-    title="<span style='font-family:Roboto; font-weight:bold; color:#2C3E50;'>Dog Breed Classifier</span>",
-    description="<span style='font-family:Open Sans; color:#34495E;'>Upload an image, and the system will predict the breed and provide detailed information from the database.</span>",
+    inputs=gr.Image(label="Upload a dog image", type="numpy"),
+    outputs=gr.Markdown(label="Prediction Results"),
+    title="<h1 style='font-family:Roboto; font-weight:bold; color:#2C3E50; text-align:center;'>🐶 Dog Breed Classifier 🔍</h1>",
+    description="<p style='font-family:Open Sans; color:#34495E; text-align:center;'>Upload a picture of a dog, and AI will predict its breed, provide detailed information, and include a Wikipedia link!</p>",
     examples=['Border_Collie.jpg',
               'Golden_Retriever.jpeg',
               'Saint_Bernard.jpeg',
               'French_Bulldog.jpeg',
               'Samoyed.jpg'],
-    css = """
-/* 美化按鈕樣式 */
-.gr-button {
-    background-color: #1ABC9C;
-    color: white;
-    border-radius: 12px;
-    box-shadow: 0px 5px 10px rgba(0, 0, 0, 0.15);
-    padding: 12px 24px;
-    transition: background-color 0.3s ease, transform 0.2s ease;
-    font-weight: bold;
-}
-.gr-button:hover {
-    background-color: #16A085;
-    transform: scale(1.05);
-}
-
-/* 美化標題 */
-h1, h2, h3, h4, h5 {
-    font-family: 'Roboto', sans-serif;
-    text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.1);
-    color: #2C3E50;
-    font-weight: bold;
-}
-
-/* 設定背景漸層 */
-body {
-    background: linear-gradient(to bottom right, #ECF0F1, #FFFFFF);
-    overflow: hidden;  /* 隱藏滾動條 */
-}
-
-/* 美化圖片上傳區塊 */
-.gr-box {
-    background: #F9F9F9;
-    border: 2px solid #BDC3C7;
-    border-radius: 15px;
-    padding: 20px;
-    box-shadow: 0px 5px 15px rgba(0, 0, 0, 0.1);
-    transition: box-shadow 0.3s ease;
-}
-.gr-box:hover {
-    box-shadow: 0px 8px 20px rgba(0, 0, 0, 0.2);
-}
-
-/* 修改 markdown 區塊字體及行高 */
-.output-markdown {
-    font-family: 'Noto Sans', sans-serif;
-    line-height: 1.75;
-    color: #34495E;
-    overflow: hidden;
-    max-height: none;
-}
-
-/* 調整範例圖片區塊樣式 */
-.examples img {
-    border-radius: 10px;
-    transition: transform 0.3s ease, box-shadow 0.3s ease;
-}
-.examples img:hover {
-    transform: scale(1.05);
-    box-shadow: 0px 5px 10px rgba(0, 0, 0, 0.15);
-}
-
-/* 去掉範例圖片滾動條 */
-.examples {
-    overflow: hidden;
-}
-""", theme='default')
+     css = """
+    /* 新增樣式 */
+    .container {
+        max-width: 900px;
+        margin: 0 auto;
+        padding: 20px;
+        background-color: rgba(255, 255, 255, 0.9);
+        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;
+    }
+    
+    .output-markdown, .output-image {
+        margin-top: 20px;
+        padding: 15px;
+        background-color: #f5f5f5;
+        border-radius: 10px;
+    }
+    
+    .examples {
+        display: flex;
+        justify-content: center;
+        flex-wrap: wrap;
+        gap: 10px;
+        margin-top: 20px;
+    }
+    
+    .examples img {
+        width: 100px;
+        height: 100px;
+        object-fit: cover;
+    }
+    """,
+    theme='default')