Update app.py

app.py

@@ -6,7 +6,10 @@ import numpy as np
 
 model = YOLO("model.pt")
 
-def predict(img, dot_color, text_color_choice, dot_size, font_size):
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
+def predict(img, target_class, dot_color, text_color_choice, dot_size, font_size):
     path = img.split("\\")[-1].split(".")[0]
     print("path", path)
 
@@ -15,11 +18,7 @@ def predict(img, dot_color, text_color_choice, dot_size, font_size):
     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:
@@ -30,7 +29,7 @@ def predict(img, dot_color, text_color_choice, dot_size, font_size):
     else:
         text_rgb = (255, 255, 255)
 
-    # 檢查輸入的大小是否合理
+    # 大小設定
     try:
         r = int(dot_size)
     except:
@@ -43,34 +42,51 @@ def predict(img, dot_color, text_color_choice, dot_size, font_size):
         fsize = 36
     if fsize < 8: fsize = 8
 
+    count = 0
+    annot_img = None
+    annot_img_numbered = None
+
+    # 轉換成 YOLO 類別索引
+    if target_class == "膠囊":
+        class_idx = 0
+    else:
+        class_idx = 1
+
     for i in results:
-        count = len(i.boxes)
-        img_pil = Image.fromarray(i.orig_img[..., ::-1])  # BGR → RGB
+        # 過濾出該類別的框
+        selected_boxes = []
+        if hasattr(i.boxes, "cls"):
+            for b, c in zip(i.boxes.xyxy, i.boxes.cls):
+                if int(c) == class_idx:
+                    selected_boxes.append(b)
+
+        count = len(selected_boxes)
+        img_pil = Image.fromarray(i.orig_img[..., ::-1])
         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)
+            font = ImageFont.truetype("DejaVuSans.ttf", fsize)
         except:
-            font = ImageFont.load_default()
+            try:
+                font = ImageFont.truetype("arial.ttf", fsize)
+            except:
+                font = ImageFont.load_default()
 
-        for idx, box in enumerate(i.boxes.xyxy, start=1):
+        for idx, box in enumerate(selected_boxes, 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)
@@ -78,6 +94,7 @@ def predict(img, dot_color, text_color_choice, dot_size, font_size):
 
     return str(count), annot_img, annot_img_numbered
 
+
 # 自定義 CSS 樣式（背景改為白色、簡潔風格）
 custom_css = """
 .gradio-container {
@@ -212,6 +229,13 @@ with gr.Blocks(
                 label="請選擇或拖放圖片",
                 elem_classes=["image-upload"]
             )
+            target_class = gr.Radio(
+                choices=["膠囊", "錠劑"],
+                value="膠囊",
+                label="要計算的類別",
+                info="選擇只計算錠劑或膠囊的總數"
+            )
+
             gr.HTML('<div style="text-align: center; margin-top: 15px;">')
             button = gr.Button(
                 "開始計算",
@@ -281,7 +305,7 @@ with gr.Blocks(
     
     button.click(
     fn=predict,
-    inputs=[img, dot_color, text_color_choice, dot_size, font_size],
+    inputs=[img, target_class, dot_color, text_color_choice, dot_size, font_size],
     outputs=[data_output, img_output, img_output_numbered]
 )