scripts for heart and star vertex annotator

src/utils.py

@@ -0,0 +1,45 @@
+import os
+
+from moviepy import ImageSequenceClip
+from natsort import natsorted
+
+
+def create_video_from_images(folder_path, output_video_file, fps):
+    """
+    Creates a video file from a sequence of images in a folder.
+
+    Args:
+        folder_path (str): The path to the folder containing the images.
+        output_video_file (str): The name of the output video file (e.g., 'my_video.mp4').
+        fps (int): The frames per second for the output video.
+    """
+    if not os.path.isdir(folder_path):
+        print(f"Error: The folder '{folder_path}' does not exist.")
+        return
+
+    # List all image files in the folder.
+    # We use natsorted to ensure files with numerical names (e.g., image-1.png, image-10.png)
+    # are sorted in a human-friendly way.
+    supported_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif')
+    image_files = [
+        os.path.join(folder_path, f)
+        for f in natsorted(os.listdir(folder_path))
+        if f.lower().endswith(supported_extensions)
+    ]
+
+    if not image_files:
+        print(f"Error: No supported image files found in '{folder_path}'.")
+        return
+
+    print(f"Found {len(image_files)} images. Creating video...")
+
+    try:
+        # Create a video clip from the list of image files.
+        clip = ImageSequenceClip(image_files, fps=fps)
+
+        # Write the video file to the specified path.
+        clip.write_videofile(output_video_file, fps=fps)
+
+        print(f"Successfully created video: '{output_video_file}'")
+    except Exception as e:
+        print(f"An error occurred while creating the video: {e}")

src/vit_pose/main_vit_pose.py

@@ -0,0 +1,81 @@
+import torch
+import requests
+import numpy as np
+import supervision as sv
+import cv2
+import os
+import sys
+import imageio.v2 as imageio
+from PIL import Image
+from transformers import AutoProcessor, RTDetrForObjectDetection, VitPoseForPoseEstimation, infer_device
+from vit_pose.vertex_annotator_heart import VertexAnnotatorHeart
+from vit_pose.vertex_annotator_star import VertexAnnotatorStar
+
+
+device = infer_device()
+
+# # Detect humans in the image
+person_image_processor = AutoProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
+person_model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365", device_map=device)
+
+  # Detect keypoints for each person found
+image_processor = AutoProcessor.from_pretrained("usyd-community/vitpose-base-simple")
+model = VitPoseForPoseEstimation.from_pretrained("usyd-community/vitpose-base-simple", device_map=device)
+
+def vit_pose_estimation(image, frame_count: int):
+  inputs = person_image_processor(images=image, return_tensors="pt").to(person_model.device)
+
+  with torch.no_grad():
+      outputs = person_model(**inputs)
+
+  results = person_image_processor.post_process_object_detection(
+      outputs, target_sizes=torch.tensor([(image.height, image.width)]), threshold=0.3
+  )
+  result = results[0]
+
+  # Human label refers 0 index in COCO dataset
+  person_boxes = result["boxes"][result["labels"] == 0]
+  person_boxes = person_boxes.cpu().numpy()
+
+  # Convert boxes from VOC (x1, y1, x2, y2) to COCO (x1, y1, w, h) format
+  person_boxes[:, 2] = person_boxes[:, 2] - person_boxes[:, 0]
+  person_boxes[:, 3] = person_boxes[:, 3] - person_boxes[:, 1]
+
+  inputs = image_processor(image, boxes=[person_boxes], return_tensors="pt").to(model.device)
+
+  with torch.no_grad():
+      outputs = model(**inputs)
+
+  pose_results = image_processor.post_process_pose_estimation(outputs, boxes=[person_boxes])
+  image_pose_result = pose_results[0]
+
+  xy = torch.stack([pose_result['keypoints'] for pose_result in image_pose_result]).cpu().numpy()
+  scores = torch.stack([pose_result['scores'] for pose_result in image_pose_result]).cpu().numpy()
+
+  color_edge_annotator = sv.Color.from_hex("#e1e1e1")
+
+  color_vertex_annotator = sv.Color.from_hex('#ffc0cb')
+
+  key_points = sv.KeyPoints(
+      xy=xy, confidence=scores
+  )
+
+  edge_annotator = sv.EdgeAnnotator(
+      color=color_edge_annotator,
+      thickness=1
+  )
+  vertex_annotator = VertexAnnotatorHeart(
+      color=color_vertex_annotator,
+      radius=10
+  )
+  annotated_frame = edge_annotator.annotate(
+      scene=image.copy(),
+      key_points=key_points
+  )
+  annotated_frame = vertex_annotator.annotate(
+      scene=annotated_frame,
+      key_points=key_points
+  )
+
+  return annotated_frame
+

src/vit_pose/vertex_annotator_heart.py

@@ -0,0 +1,53 @@
+import cv2
+import numpy as np
+import supervision as sv
+from supervision.annotators.base import ImageType
+from supervision.keypoint.core import KeyPoints
+from supervision.utils.conversion import ensure_cv2_image_for_annotation
+
+
+class VertexAnnotatorHeart(sv.VertexAnnotator):
+    """
+    Extend the original class to add vertices with the shape of heart or star.
+    """
+
+    def __init__(self, color=sv.Color.RED, radius=5):
+        super().__init__(color=color, radius=radius)
+
+    @ensure_cv2_image_for_annotation
+    def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
+        assert isinstance(scene, np.ndarray)
+        if len(key_points) == 0:
+            return scene
+
+        for xy in key_points.xy:
+            for x, y in xy:
+                # Two circles (top lobes of heart)
+                cv2.circle(
+                    img=scene,
+                    center=(int(x) - self.radius // 2, int(y)),
+                    radius=self.radius // 2,
+                    color=self.color.as_bgr(),
+                    thickness=-1,
+                )
+                cv2.circle(
+                    img=scene,
+                    center=(int(x) + self.radius // 2, int(y)),
+                    radius=self.radius // 2,
+                    color=self.color.as_bgr(),
+                    thickness=-1,
+                )
+
+                # Triangle (bottom of heart)
+                pts = np.array(
+                    [
+                        [int(x) - self.radius, int(y)],
+                        [int(x) + self.radius, int(y)],
+                        [int(x), y + self.radius * 2],
+                    ],
+                    np.int32,
+                ).reshape((-1, 1, 2))
+
+                cv2.fillPoly(scene, [pts], self.color.as_bgr())
+
+        return scene

src/vit_pose/vertex_annotator_star.py

@@ -0,0 +1,54 @@
+import cv2
+import numpy as np
+import supervision as sv
+from supervision.annotators.base import ImageType
+from supervision.keypoint.core import KeyPoints
+from supervision.utils.conversion import ensure_cv2_image_for_annotation
+
+
+class VertexAnnotatorStar(sv.VertexAnnotator):
+    """
+    Extend the original class to add vertices with the shape of heart or star.
+    """
+
+    def __init__(self, color=sv.Color.RED, radius=5, thickness=-1):
+        super().__init__(color=color, radius=radius)
+        self.thickness = thickness
+
+    @ensure_cv2_image_for_annotation
+    def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
+        """Draw a 5-pointed star on an image."""
+        assert isinstance(scene, np.ndarray)
+        if len(key_points) == 0:
+            return scene
+
+        for xy in key_points.xy:
+            for x, y in xy:
+                pts = np.array(
+                    [
+                        [x, y - self.radius],
+                        [x + self.radius // 3, y - self.radius // 3],
+                        [x + self.radius, y - self.radius // 3],
+                        [x + self.radius // 2, y + self.radius // 6],
+                        [x + 2 * self.radius // 3, y + self.radius],
+                        [x, y + self.radius // 2],
+                        [x - 2 * self.radius // 3, y + self.radius],
+                        [x - self.radius // 2, y + self.radius // 6],
+                        [x - self.radius, y - self.radius // 3],
+                        [x - self.radius // 3, y - self.radius // 3],
+                    ],
+                    np.int32,
+                ).reshape((-1, 1, 2))
+
+                if self.thickness == -1:
+                    cv2.fillPoly(scene, [pts], self.color.as_bgr())
+                else:
+                    cv2.polylines(
+                        scene,
+                        [pts],
+                        isClosed=True,
+                        color=self.color.as_bgr(),
+                        thickness=self.thickness,
+                    )
+
+        return scene