Update app.py

app.py

@@ -1,28 +1,82 @@
 import gradio as gr
 from ultralytics import YOLO
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
 
 model = YOLO("model.pt")
 
-def predict(img):
+def predict(img, dot_color, text_color_choice, dot_size, font_size):
     path = img.split("\\")[-1].split(".")[0]
     print("path", path)
 
     results = model.predict(source=img, save=False, show_labels=False, show_conf=False)
-    
-    # Handle both single result and list of results
+
     if not isinstance(results, (list, tuple)):
         results = [results]
-    
+
     count = 0
     annot_img = None
+    annot_img_numbered = None
+
+    # 轉換顏色文字為 RGB
+    if dot_color == "紅色":
+        dot_rgb = (255, 0, 0)
+    else:
+        dot_rgb = (57, 255, 20)  # 螢光綠
+
+    if text_color_choice == "黑色":
+        text_rgb = (0, 0, 0)
+    else:
+        text_rgb = (255, 255, 255)
+
+    # 檢查輸入的大小是否合理
+    try:
+        r = int(dot_size)
+    except:
+        r = 10
+    if r < 2: r = 2
+
+    try:
+        fsize = int(font_size)
+    except:
+        fsize = 36
+    if fsize < 8: fsize = 8
+
     for i in results:
         count = len(i.boxes)
-        annot_img = i.plot(labels=False, masks=False)
-        # Convert BGR to RGB if needed
-        if annot_img.shape[-1] == 3:  # Check if image has 3 channels
-            annot_img = annot_img[..., ::-1]  # Reverse the color channels
-    
-    return str(count), annot_img
+        img_pil = Image.fromarray(i.orig_img[..., ::-1])  # BGR → RGB
+        img_pil_numbered = img_pil.copy()
+
+        draw = ImageDraw.Draw(img_pil)
+        draw_num = ImageDraw.Draw(img_pil_numbered)
+
+        # 載入字體（若無 Arial 則 fallback）
+        try:
+            font = ImageFont.truetype("arial.ttf", fsize)
+        except:
+            font = ImageFont.load_default()
+
+        for idx, box in enumerate(i.boxes.xyxy, start=1):
+            x1, y1, x2, y2 = box[:4].tolist()
+            cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
+
+            # -------- 第一張（只有圓點）--------
+            draw.ellipse((cx - r, cy - r, cx + r, cy + r), fill=dot_rgb, outline=None)
+
+            # -------- 第二張（圓點 + 編號）--------
+            draw_num.ellipse((cx - r, cy - r, cx + r, cy + r), fill=dot_rgb, outline=None)
+
+            # 最後一顆用紅字，其餘依使用者設定
+            num_color = (255, 0, 0) if idx == count else text_rgb
+
+            # 在圓點右上方標註數字
+            draw_num.text((cx + r + 4, cy - r), str(idx), fill=num_color, font=font)
+
+        annot_img = np.array(img_pil)
+        annot_img_numbered = np.array(img_pil_numbered)
+
+    return str(count), annot_img, annot_img_numbered
 
 # 自定義 CSS 樣式（背景改為白色、簡潔風格）
 custom_css = """
@@ -166,6 +220,34 @@ with gr.Blocks(
             )
             gr.HTML('</div><div style="text-align:center; margin-top:8px; color:#888;">支援 JPG/PNG，建議 640×640 以上</div>')
             gr.HTML('</div>')
+            dot_color = gr.Radio(
+                choices=["紅色", "螢光綠"],
+                value="紅色",
+                label="圓點顏色",
+                info="選擇藥錠中心圓點顏色"
+            )
+
+            text_color_choice = gr.Radio(
+                choices=["黑色", "白色"],
+                value="黑色",
+                label="文字顏色",
+                info="選擇編號文字顏色"
+            )
+
+            dot_size = gr.Textbox(
+                label="圓點大小",
+                value="10",
+                info="設定圓點半徑大小（建議範圍 5～20）"
+            )
+
+            font_size = gr.Textbox(
+                label="字體大小",
+                value="36",
+                info="設定編號文字大小（建議範圍 20～60）"
+            )
+            
+            
+            
         
         with gr.Column(scale=1):
             
@@ -179,6 +261,11 @@ with gr.Blocks(
                 type="numpy",
                 elem_classes=["result-image"]
             )
+            gr.HTML('<div style="text-align:center; margin:20px 0;"><label style="font-size:14px;color:#555;">標註 + 編號結果</label></div>')
+            img_output_numbered = gr.Image(
+                type="numpy",
+                elem_classes=["result-image"]
+            )
             gr.HTML('</div>')
     
     gr.HTML("""
@@ -192,7 +279,11 @@ with gr.Blocks(
         </div>
     """)
     
-    button.click(fn=predict, inputs=img, outputs=[data_output, img_output])
+    button.click(
+    fn=predict,
+    inputs=[img, dot_color, text_color_choice, dot_size, font_size],
+    outputs=[data_output, img_output, img_output_numbered]
+)
 
 if __name__ == "__main__":
     demo.launch()