Update app.py

app.py

@@ -165,165 +165,165 @@ def _predict_single_dog(image):
     return top1_prob, topk_breeds, topk_probs_percent
 
 
-# async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4):
-#     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
-#     dogs = []
-#     for box in results.boxes:
-#         if box.cls == 16:  # COCO 資料集中狗的類別是 16
-#             xyxy = box.xyxy[0].tolist()
-#             confidence = box.conf.item()
-#             cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
-#             dogs.append((cropped_image, confidence, xyxy))
-#     return dogs
+async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4):
+    results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+    dogs = []
+    for box in results.boxes:
+        if box.cls == 16:  # COCO 資料集中狗的類別是 16
+            xyxy = box.xyxy[0].tolist()
+            confidence = box.conf.item()
+            cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
+            dogs.append((cropped_image, confidence, xyxy))
+    return dogs
 
 
-# async def predict(image):
-#     if image is None:
-#         return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
+async def predict(image):
+    if image is None:
+        return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
 
-#     try:
-#         if isinstance(image, np.ndarray):
-#             image = Image.fromarray(image)
+    try:
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
 
-#         dogs = await detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4)
+        dogs = await detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4)
         
-#         if len(dogs) <= 1:
-#             return await process_single_dog(image)
+        if len(dogs) <= 1:
+            return await process_single_dog(image)
 
-#         # 多狗情境
-#         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
-#         explanations = []
-#         buttons = []
-#         annotated_image = image.copy()
-#         draw = ImageDraw.Draw(annotated_image)
-#         font = ImageFont.load_default()
+        # 多狗情境
+        color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
+        explanations = []
+        buttons = []
+        annotated_image = image.copy()
+        draw = ImageDraw.Draw(annotated_image)
+        font = ImageFont.load_default()
 
-#         for i, (cropped_image, _, box) in enumerate(dogs):
-#             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
-#             color = color_list[i % len(color_list)]
-#             draw.rectangle(box, outline=color, width=3)
-#             draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
+        for i, (cropped_image, _, box) in enumerate(dogs):
+            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
+            color = color_list[i % len(color_list)]
+            draw.rectangle(box, outline=color, width=3)
+            draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
 
-#             breed = topk_breeds[0]
-#             if top1_prob >= 0.5:
-#                 description = get_dog_description(breed)
-#                 formatted_description = format_description(description, breed)
-#                 explanations.append(f"Dog {i+1}: {formatted_description}")
-#             elif top1_prob >= 0.2:
-#                 dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
-#                 dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
-#                 explanations.append(dog_explanation)
-#                 buttons.extend([gr.update(visible=True, value=f"Dog {i+1}: More about {breed}") for breed in topk_breeds[:3]])
-#             else:
-#                 explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
+            breed = topk_breeds[0]
+            if top1_prob >= 0.5:
+                description = get_dog_description(breed)
+                formatted_description = format_description(description, breed)
+                explanations.append(f"Dog {i+1}: {formatted_description}")
+            elif top1_prob >= 0.2:
+                dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
+                dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
+                explanations.append(dog_explanation)
+                buttons.extend([gr.update(visible=True, value=f"Dog {i+1}: More about {breed}") for breed in topk_breeds[:3]])
+            else:
+                explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
 
-#         final_explanation = "\n\n".join(explanations)
-#         if buttons:
-#             final_explanation += "\n\nClick on a button to view more information about the breed."
-#             initial_state = {
-#                 "explanation": final_explanation,
-#                 "buttons": buttons,
-#                 "show_back": True
-#             }
-#             return (final_explanation, annotated_image, 
-#                     buttons[0] if len(buttons) > 0 else gr.update(visible=False),
-#                     buttons[1] if len(buttons) > 1 else gr.update(visible=False),
-#                     buttons[2] if len(buttons) > 2 else gr.update(visible=False),
-#                     gr.update(visible=True),
-#                     initial_state)
-#         else:
-#             initial_state = {
-#                 "explanation": final_explanation,
-#                 "buttons": [],
-#                 "show_back": False
-#             }
-#             return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
+        final_explanation = "\n\n".join(explanations)
+        if buttons:
+            final_explanation += "\n\nClick on a button to view more information about the breed."
+            initial_state = {
+                "explanation": final_explanation,
+                "buttons": buttons,
+                "show_back": True
+            }
+            return (final_explanation, annotated_image, 
+                    buttons[0] if len(buttons) > 0 else gr.update(visible=False),
+                    buttons[1] if len(buttons) > 1 else gr.update(visible=False),
+                    buttons[2] if len(buttons) > 2 else gr.update(visible=False),
+                    gr.update(visible=True),
+                    initial_state)
+        else:
+            initial_state = {
+                "explanation": final_explanation,
+                "buttons": [],
+                "show_back": False
+            }
+            return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
 
-#     except Exception as e:
-#         error_msg = f"An error occurred: {str(e)}"
-#         print(error_msg)  # 添加日誌輸出
-#         return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
+    except Exception as e:
+        error_msg = f"An error occurred: {str(e)}"
+        print(error_msg)  # 添加日誌輸出
+        return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
 
 
-async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4, merge_threshold=0.5):
-    results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
-    dogs = []
+# async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4, merge_threshold=0.5):
+#     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+#     dogs = []
     
-    image_area = image.width * image.height
-    min_area_ratio = 0.005  # 最小檢測面積佔整個圖像的比例
+#     image_area = image.width * image.height
+#     min_area_ratio = 0.005  # 最小檢測面積佔整個圖像的比例
     
-    for box in results.boxes:
-        if box.cls == 16:  # COCO 數據集中狗的類別是 16
-            xyxy = box.xyxy[0].tolist()
-            area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
-            if area / image_area >= min_area_ratio:
-                confidence = box.conf.item()
-                dogs.append((xyxy, confidence))
+#     for box in results.boxes:
+#         if box.cls == 16:  # COCO 數據集中狗的類別是 16
+#             xyxy = box.xyxy[0].tolist()
+#             area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
+#             if area / image_area >= min_area_ratio:
+#                 confidence = box.conf.item()
+#                 dogs.append((xyxy, confidence))
     
-    if dogs:
-        boxes = torch.tensor([dog[0] for dog in dogs])
-        scores = torch.tensor([dog[1] for dog in dogs])
+#     if dogs:
+#         boxes = torch.tensor([dog[0] for dog in dogs])
+#         scores = torch.tensor([dog[1] for dog in dogs])
         
-        # 應用 NMS
-        keep = nms(boxes, scores, iou_threshold)
+#         # 應用 NMS
+#         keep = nms(boxes, scores, iou_threshold)
         
-        merged_dogs = []
-        for i in keep:
-            xyxy = boxes[i].tolist()
-            confidence = scores[i].item()
-            merged_dogs.append((xyxy, confidence))
+#         merged_dogs = []
+#         for i in keep:
+#             xyxy = boxes[i].tolist()
+#             confidence = scores[i].item()
+#             merged_dogs.append((xyxy, confidence))
         
-        # 後處��：分離過於接近的檢測框
-        final_dogs = []
-        while merged_dogs:
-            base_dog = merged_dogs.pop(0)
-            to_merge = [base_dog]
+#         # 後處理：分離過於接近的檢測框
+#         final_dogs = []
+#         while merged_dogs:
+#             base_dog = merged_dogs.pop(0)
+#             to_merge = [base_dog]
             
-            i = 0
-            while i < len(merged_dogs):
-                iou = box_iou(torch.tensor([base_dog[0]]), torch.tensor([merged_dogs[i][0]]))[0][0].item()
-                if iou > merge_threshold:
-                    to_merge.append(merged_dogs.pop(i))
-                else:
-                    i += 1
+#             i = 0
+#             while i < len(merged_dogs):
+#                 iou = box_iou(torch.tensor([base_dog[0]]), torch.tensor([merged_dogs[i][0]]))[0][0].item()
+#                 if iou > merge_threshold:
+#                     to_merge.append(merged_dogs.pop(i))
+#                 else:
+#                     i += 1
             
-            if len(to_merge) == 1:
-                final_dogs.append(base_dog)
-            else:
-                # 如果檢測到多個重疊框，嘗試分離它們
-                centers = torch.tensor([[((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)] for box, _ in to_merge])
-                distances = torch.cdist(centers, centers)
+#             if len(to_merge) == 1:
+#                 final_dogs.append(base_dog)
+#             else:
+#                 # 如果檢測到多個重疊框，嘗試分離它們
+#                 centers = torch.tensor([[((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)] for box, _ in to_merge])
+#                 distances = torch.cdist(centers, centers)
                 
-                if torch.any(distances > 0):  # 確保不是完全重疊
-                    max_distance = distances.max()
-                    if max_distance > (base_dog[0][2] - base_dog[0][0]) * 0.5:  # 如果最大距離大於框寬度的一半
-                        final_dogs.extend(to_merge)
-                    else:
-                        # 合併為一個框
-                        merged_box = torch.tensor([box for box, _ in to_merge]).mean(dim=0)
-                        merged_confidence = max(conf for _, conf in to_merge)
-                        final_dogs.append((merged_box.tolist(), merged_confidence))
-                else:
-                    # 完全重疊的情況，保留置信度最高的
-                    best_dog = max(to_merge, key=lambda x: x[1])
-                    final_dogs.append(best_dog)
+#                 if torch.any(distances > 0):  # 確保不是完全重疊
+#                     max_distance = distances.max()
+#                     if max_distance > (base_dog[0][2] - base_dog[0][0]) * 0.5:  # 如果最大距離大於框寬度的一半
+#                         final_dogs.extend(to_merge)
+#                     else:
+#                         # 合併為一個框
+#                         merged_box = torch.tensor([box for box, _ in to_merge]).mean(dim=0)
+#                         merged_confidence = max(conf for _, conf in to_merge)
+#                         final_dogs.append((merged_box.tolist(), merged_confidence))
+#                 else:
+#                     # 完全重疊的情況，保留置信度最高的
+#                     best_dog = max(to_merge, key=lambda x: x[1])
+#                     final_dogs.append(best_dog)
         
-        # 擴展邊界框並創建剪裁的圖像
-        expanded_dogs = []
-        for xyxy, confidence in final_dogs:
-            expanded_xyxy = [
-                max(0, xyxy[0] - 20),
-                max(0, xyxy[1] - 20),
-                min(image.width, xyxy[2] + 20),
-                min(image.height, xyxy[3] + 20)
-            ]
-            cropped_image = image.crop(expanded_xyxy)
-            expanded_dogs.append((cropped_image, confidence, expanded_xyxy))
+#         # 擴展邊界框並創建剪裁的圖像
+#         expanded_dogs = []
+#         for xyxy, confidence in final_dogs:
+#             expanded_xyxy = [
+#                 max(0, xyxy[0] - 20),
+#                 max(0, xyxy[1] - 20),
+#                 min(image.width, xyxy[2] + 20),
+#                 min(image.height, xyxy[3] + 20)
+#             ]
+#             cropped_image = image.crop(expanded_xyxy)
+#             expanded_dogs.append((cropped_image, confidence, expanded_xyxy))
         
-        return expanded_dogs
+#         return expanded_dogs
     
-    # 如果沒有檢測到狗狗，返回整張圖片
-    return [(image, 1.0, [0, 0, image.width, image.height])]
+#     # 如果沒有檢測到狗狗，返回整張圖片
+#     return [(image, 1.0, [0, 0, image.width, image.height])]
 
 # async def predict(image):
 #     if image is None:
@@ -396,166 +396,6 @@ async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4, me
 #         return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
         
 
-# async def process_single_dog(image):
-#     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
-#     if top1_prob < 0.2:
-#         initial_state = {
-#             "explanation": "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.",
-#             "buttons": [],
-#             "show_back": False
-#         }
-#         return initial_state["explanation"], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
-
-#     breed = topk_breeds[0]
-#     description = get_dog_description(breed)
-
-#     if top1_prob >= 0.5:
-#         formatted_description = format_description(description, breed)
-#         initial_state = {
-#             "explanation": formatted_description,
-#             "buttons": [],
-#             "show_back": False
-#         }
-#         return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
-#     else:
-#         explanation = (
-#             f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
-#             f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
-#             f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]} confidence)\n"
-#             f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
-#             "Click on a button to view more information about the breed."
-#         )
-#         buttons = [
-#             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]}")
-#         ]
-#         initial_state = {
-#             "explanation": explanation,
-#             "buttons": buttons,
-#             "show_back": True
-#         }
-#         return explanation, image, buttons[0], buttons[1], buttons[2], gr.update(visible=True), initial_state
-
-# def show_details(choice, previous_output, initial_state):
-#     if not choice:
-#         return previous_output, gr.update(visible=True), initial_state
-
-#     try:
-#         breed = choice.split("More about ")[-1]
-#         description = get_dog_description(breed)
-#         formatted_description = format_description(description, breed)
-#         return formatted_description, gr.update(visible=True), initial_state
-#     except Exception as e:
-#         error_msg = f"An error occurred while showing details: {e}"
-#         print(error_msg)  # 添加日誌輸出
-#         return error_msg, gr.update(visible=True), initial_state
-
-# # 介面部分
-# with gr.Blocks() as iface:
-#     gr.HTML("<h1 style='text-align: center;'>🐶 Dog Breed Classifier 🔍</h1>")
-#     gr.HTML("<p style='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="pil")
-#         output_image = gr.Image(label="Annotated Image")
-    
-#     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)
-    
-#     back_button = gr.Button("Back", visible=False)
-    
-#     initial_state = gr.State()
-    
-#     input_image.change(
-#         predict,
-#         inputs=input_image,
-#         outputs=[output, output_image, btn1, btn2, btn3, back_button, initial_state]
-#     )
-
-#     for btn in [btn1, btn2, btn3]:
-#         btn.click(
-#             show_details,
-#             inputs=[btn, output, initial_state],
-#             outputs=[output, back_button, initial_state]
-#         )
-
-#     back_button.click(
-#         lambda state: (state["explanation"], 
-#                        state["buttons"][0] if len(state["buttons"]) > 0 else gr.update(visible=False),
-#                        state["buttons"][1] if len(state["buttons"]) > 1 else gr.update(visible=False),
-#                        gr.update(visible=state["show_back"])),
-#         inputs=[initial_state],
-#         outputs=[output, btn1, btn2, btn3, back_button]
-#     )
-
-#     gr.Examples(
-#         examples=['Border_Collie.jpg', 'Golden_Retriever.jpeg', 'Saint_Bernard.jpeg', 'French_Bulldog.jpeg', 'Samoyed.jpg'],
-#         inputs=input_image
-#     )
-
-#     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_Breed_Classifier">Dog Breed Classifier</a>')
-
-#     if __name__ == "__main__":
-#         iface.launch()
-
-
-async def predict(image):
-    if image is None:
-        return "Please upload an image to start.", None, [], gr.update(visible=False), None
-
-    try:
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-
-        dogs = await detect_multiple_dogs(image)
-        
-        if len(dogs) <= 1:
-            return await process_single_dog(image)
-
-        color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
-        explanations = []
-        buttons = []
-        annotated_image = image.copy()
-        draw = ImageDraw.Draw(annotated_image)
-        font = ImageFont.load_default()
-
-        for i, (cropped_image, _, box) in enumerate(dogs):
-            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
-            color = color_list[i % len(color_list)]
-            draw.rectangle(box, outline=color, width=3)
-            draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
-
-            breed = topk_breeds[0]
-            if top1_prob >= 0.5:
-                description = get_dog_description(breed)
-                formatted_description = format_description(description, breed)
-                explanations.append(f"Dog {i+1}: {formatted_description}")
-            else:
-                dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
-                dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
-                explanations.append(dog_explanation)
-                buttons.extend([f"Dog {i+1}: More about {breed}" for breed in topk_breeds[:3]])
-
-        final_explanation = "\n\n".join(explanations)
-        if buttons:
-            final_explanation += "\n\nClick on a button to view more information about the breed."
-        
-        initial_state = {
-            "explanation": final_explanation,
-            "buttons": buttons,
-            "show_back": bool(buttons)
-        }
-        return final_explanation, annotated_image, buttons, gr.update(visible=bool(buttons)), initial_state
-
-    except Exception as e:
-        error_msg = f"An error occurred: {str(e)}"
-        print(error_msg)  # 添加日誌輸出
-        return error_msg, None, [], gr.update(visible=False), None
 
 async def process_single_dog(image):
     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
@@ -565,7 +405,7 @@ async def process_single_dog(image):
             "buttons": [],
             "show_back": False
         }
-        return initial_state["explanation"], None, [], gr.update(visible=False), initial_state
+        return initial_state["explanation"], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
 
     breed = topk_breeds[0]
     description = get_dog_description(breed)
@@ -577,7 +417,7 @@ async def process_single_dog(image):
             "buttons": [],
             "show_back": False
         }
-        return formatted_description, image, [], gr.update(visible=False), initial_state
+        return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
     else:
         explanation = (
             f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
@@ -586,13 +426,17 @@ async def process_single_dog(image):
             f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
             "Click on a button to view more information about the breed."
         )
-        buttons = [f"More about {breed}" for breed in topk_breeds[:3]]
+        buttons = [
+            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]}")
+        ]
         initial_state = {
             "explanation": explanation,
             "buttons": buttons,
             "show_back": True
         }
-        return explanation, image, buttons, gr.update(visible=True), initial_state
+        return explanation, image, buttons[0], buttons[1], buttons[2], gr.update(visible=True), initial_state
 
 def show_details(choice, previous_output, initial_state):
     if not choice:
@@ -608,6 +452,7 @@ def show_details(choice, previous_output, initial_state):
         print(error_msg)  # 添加日誌輸出
         return error_msg, gr.update(visible=True), initial_state
 
+# 介面部分
 with gr.Blocks() as iface:
     gr.HTML("<h1 style='text-align: center;'>🐶 Dog Breed Classifier 🔍</h1>")
     gr.HTML("<p style='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>")
@@ -618,54 +463,43 @@ with gr.Blocks() as iface:
     
     output = gr.Markdown(label="Prediction Results")
     
-    button_container = gr.Column()
+    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)
     
     back_button = gr.Button("Back", visible=False)
     
     initial_state = gr.State()
     
-    def create_buttons(button_texts):
-        buttons = []
-        for text in button_texts:
-            button = gr.Button(text)
-            button.click(
-                show_details,
-                inputs=[button, output, initial_state],
-                outputs=[output, back_button, initial_state]
-            )
-            buttons.append(button)
-        return buttons
-
-    def update_ui(explanation, image, button_texts, show_back, state):
-        button_container.clear()
-        buttons = create_buttons(button_texts)
-        return explanation, image, *buttons, gr.update(visible=show_back), state
-
     input_image.change(
         predict,
         inputs=input_image,
-        outputs=[output, output_image, button_container, back_button, initial_state]
-    ).then(
-        update_ui,
-        inputs=[output, output_image, button_container, back_button, initial_state],
-        outputs=[output, output_image] + [button_container] * 9 + [back_button, initial_state]
+        outputs=[output, output_image, btn1, btn2, btn3, back_button, initial_state]
     )
-    
-    def go_back(state):
-        return update_ui(state["explanation"], None, state["buttons"], state["show_back"], state)
+
+    for btn in [btn1, btn2, btn3]:
+        btn.click(
+            show_details,
+            inputs=[btn, output, initial_state],
+            outputs=[output, back_button, initial_state]
+        )
 
     back_button.click(
-        go_back,
+        lambda state: (state["explanation"], 
+                       state["buttons"][0] if len(state["buttons"]) > 0 else gr.update(visible=False),
+                       state["buttons"][1] if len(state["buttons"]) > 1 else gr.update(visible=False),
+                       gr.update(visible=state["show_back"])),
         inputs=[initial_state],
-        outputs=[output, output_image] + [button_container] * 9 + [back_button, initial_state]
+        outputs=[output, btn1, btn2, btn3, back_button]
     )
-    
+
     gr.Examples(
         examples=['Border_Collie.jpg', 'Golden_Retriever.jpeg', 'Saint_Bernard.jpeg', 'French_Bulldog.jpeg', 'Samoyed.jpg'],
         inputs=input_image
     )
-    
+
     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_Breed_Classifier">Dog Breed Classifier</a>')
 
-if __name__ == "__main__":
-    iface.launch()
+    if __name__ == "__main__":
+        iface.launch()