Update app.py

app.py

@@ -1,29 +1,93 @@
 from PIL import Image
 import numpy as np
+import cv2
 import gradio as gr
 from ultralytics import YOLO
 
 MODEL_PATH = "best.pt"
 model = YOLO(MODEL_PATH)
 
-def predict(image, conf):
+def draw_pose_clean(image, conf=0.25, show_labels=False, point_size=3):
     if image is None:
         return None
-    results = model.predict(source=np.array(image), conf=conf, imgsz=640, verbose=False)[0]
-    plotted = results.plot()  # BGR numpy array
-    plotted = plotted[:, :, ::-1]  # BGR -> RGB
-    return Image.fromarray(plotted)
+
+    img = np.array(image.convert("RGB"))
+    results = model.predict(source=img, conf=conf, imgsz=640, verbose=False)[0]
+
+    out = img.copy()
+    if results.keypoints is None or len(results.keypoints.xy) == 0:
+        return Image.fromarray(out)
+
+    kpts_xy = results.keypoints.xy.cpu().numpy()
+    kpts_conf = results.keypoints.conf.cpu().numpy() if results.keypoints.conf is not None else None
+
+    for det_i, det_kpts in enumerate(kpts_xy):
+        confs = kpts_conf[det_i] if kpts_conf is not None else np.ones(len(det_kpts))
+        for i, (x, y) in enumerate(det_kpts):
+            score = float(confs[i])
+            if score < conf:
+                continue
+
+            x_i, y_i = int(x), int(y)
+            cv2.circle(out, (x_i, y_i), int(point_size), (0, 255, 0), -1, lineType=cv2.LINE_AA)
+
+            # Optional: only label every 5th keypoint to reduce clutter
+            if show_labels and (i % 5 == 0):
+                cv2.putText(
+                    out,
+                    f"{i+1}",
+                    (x_i + 4, y_i - 4),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.35,
+                    (255, 255, 0),
+                    1,
+                    cv2.LINE_AA,
+                )
+
+    return Image.fromarray(out)
 
 demo = gr.Interface(
-    fn=predict,
+    fn=draw_pose_clean,
     inputs=[
-        gr.Image(type="pil", label="Input Image"),
-        gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence"),
+        gr.Image(type="pil", label="Input image"),
+        gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence threshold"),
+        gr.Checkbox(value=False, label="Show keypoint index labels (sparser)"),
+        gr.Slider(1, 8, value=3, step=1, label="Point size"),
     ],
-    outputs=gr.Image(type="pil", label="Pose Result"),
-    title="DogFLW YOLOv8 Pose",
-    description="Upload a dog image to detect 46 facial landmarks."
+    outputs=gr.Image(type="pil", label="Output"),
+    title="DogFLW Pose (Clean View)",
+    description="Less cluttered output: points only by default. Toggle sparse labels if needed.",
 )
 
 if __name__ == "__main__":
     demo.launch()
+
+# from PIL import Image
+# import numpy as np
+# import gradio as gr
+# from ultralytics import YOLO
+
+# MODEL_PATH = "best.pt"
+# model = YOLO(MODEL_PATH)
+
+# def predict(image, conf):
+#     if image is None:
+#         return None
+#     results = model.predict(source=np.array(image), conf=conf, imgsz=640, verbose=False)[0]
+#     plotted = results.plot()  # BGR numpy array
+#     plotted = plotted[:, :, ::-1]  # BGR -> RGB
+#     return Image.fromarray(plotted)
+
+# demo = gr.Interface(
+#     fn=predict,
+#     inputs=[
+#         gr.Image(type="pil", label="Input Image"),
+#         gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence"),
+#     ],
+#     outputs=gr.Image(type="pil", label="Pose Result"),
+#     title="DogFLW YOLOv8 Pose",
+#     description="Upload a dog image to detect 46 facial landmarks."
+# )
+
+# if __name__ == "__main__":
+#     demo.launch()