Upload folder using huggingface_hub

app.py

@@ -0,0 +1,32 @@
+from PIL import Image
+from face_detection import detect_faces
+from heatmap import get_heatmap
+from pose import detect_pose, draw_pose
+
+def prepare_image(img_path):
+    # detect face
+    annotated_image, face_bboxes = detect_faces(img_path)
+    # detect pose + bounding_box
+    pose_results = detect_pose(img_path)
+    pose_result = pose_results[0]
+    body_bboxes, body_keypoints = pose_result.boxes, pose_result.keypoints 
+    # generate heatmap
+    heatmap = get_heatmap(img_path)
+    heatmap_pil = Image.fromarray(heatmap)
+    # pose on heatmap
+    heatmap_n_pose = draw_pose(heatmap_pil,body_keypoints) #PIL.Image
+    
+    # area in bounding_box below head (rectangular)
+    face_y1, face_w = face_bboxes[0][1] #xywh
+    face_center_y = (face_w - face_y1) // 2
+    body_box = body_bboxes.numpy().xywh[0]
+    body_box[1] = face_center_y
+    
+    # Calculate the coordinates for cropping
+    x, y, w, h = map(int, body_box)
+    cropped_region = heatmap_n_pose.crop((x, y, x + w, y + h))
+    original_image = Image.open(img_path)
+    original_image.paste(cropped_region, (x, y))
+    
+    # Return the modified image
+    return original_image

face_detection.py

@@ -0,0 +1,11 @@
+# YuNet-face-detection
+import cv2
+from yunet2 import YuNet2
+
+MODEL_PATH = "./yunet-ckpts/yunet_2023mar.onnx"
+yunet_model = YuNet2(modelPath=MODEL_PATH)
+
+def detect_faces(src_img):
+    cv2_img = cv2.imread(src_img)
+    annotated_image, bboxes = yunet_model.detect(cv2_img)
+    return annotated_image, bboxes

heatmap.py

@@ -0,0 +1,32 @@
+# intel-isl/MiDaS
+import cv2
+import torch
+import torch.nn.functional as F
+
+model_type = "DPT_Large"     # MiDaS v3 - Large     (highest accuracy, slowest inference speed)
+# model_type = "DPT_Hybrid"   # MiDaS v3 - Hybrid    (medium accuracy, medium inference speed)
+# model_type = "MiDaS_small"  # MiDaS v2.1 - Small   (lowest accuracy, highest inference speed)
+
+midas = torch.hub.load("intel-isl/MiDaS", model_type)
+midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
+
+if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
+    transform = midas_transforms.dpt_transform
+else:
+    transform = midas_transforms.small_transform
+
+def get_heatmap(src_img):
+    cv2image = cv2.imread(src_img)
+    img = cv2.cvtColor(cv2image, cv2.COLOR_BGR2RGB)
+    input_batch = transform(img)
+    with torch.inference_mode():
+        prediction = midas(input_batch)
+
+    prediction = F.interpolate(
+        prediction.unsqueeze(1),
+        size=img.shape[:2],
+        mode="bicubic",
+        align_corners=False,
+    ).squeeze()
+    output = prediction.cpu().numpy()
+    return output

pose.py

@@ -0,0 +1,80 @@
+# YOLO
+import cv2
+import torch
+import numpy as np
+from ultralytics import YOLO
+from PIL import Image, ImageDraw
+
+model_name = 'yolo11s-pose.pt'
+model = YOLO(model_name)
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+half = device=='cuda'
+
+connection_color_map = {
+    (0, 1): 'lightgreen',   # Nose to Left Eye
+    (0, 2): 'lightgreen',   # Nose to Right Eye
+    (1, 3): 'lightgreen',   # Left Eye to Left Ear
+    (2, 4): 'lightgreen',   # Right Eye to Right Ear
+    (0, 5): 'lightblue',    # Nose to Left Shoulder
+    (0, 6): 'lightblue',    # Nose to Right Shoulder
+    (5, 7): 'blue',         # Left Shoulder to Left Elbow
+    (7, 9): 'blue',         # Left Elbow to Left Wrist
+    (6, 8): 'red',          # Right Shoulder to Right Elbow
+    (8, 10): 'red',         # Right Elbow to Right Wrist
+    (5, 11): 'purple',      # Left Shoulder to Left Hip
+    (6, 12): 'purple',      # Right Shoulder to Right Hip
+    (11, 13): 'purple',     # Left Hip to Left Knee
+    (13, 15): 'purple',     # Left Knee to Left Ankle
+    (12, 14): 'purple',     # Right Hip to Right Knee
+    (14, 16): 'purple',     # Right Knee to Right Ankle
+    (11, 12): 'orange'      # Left Hip to Right Hip
+}
+
+
+def detect_pose(src_img):
+    results = model.predict(
+            source=src_img,
+            # conf=conf_threshold,
+            # iou=iou_threshold,
+            max_det=1,
+            show_labels=True,
+            show_conf=True,
+            imgsz=32*5,
+            classes=[0],
+            # line_width=3,
+            half=half,
+            device=device,
+        )
+    return results
+
+def draw_pose(skeleton_image:Image.Image, keypoints, threshold=0.7, line_width = 3,circle_radius = 6):
+    coords = keypoints.numpy().xy[0]
+    scores = keypoints.data.squeeze()[:,2].numpy()
+    # Create a blank image with the same size as the input image
+    # skeleton_image = Image.new("RGB", im.size, "black")  # White background
+    draw = ImageDraw.Draw(skeleton_image)
+    
+    # Draw lines connecting the keypoints
+    for connection, color in connection_color_map.items():
+        try:
+            start_point = coords[connection[0]]
+            end_point = coords[connection[1]]
+            if scores[connection[0]] > threshold and scores[connection[1]] > threshold:
+                draw.line([tuple(start_point), tuple(end_point)], fill=color, width=line_width)
+                draw.ellipse(
+                    (start_point[0] - circle_radius, start_point[1] - circle_radius,
+                    start_point[0] + circle_radius, start_point[1] + circle_radius),
+                fill=color
+                )
+                draw.ellipse(
+                    (end_point[0] - circle_radius, end_point[1] - circle_radius,
+                    end_point[0] + circle_radius, end_point[1] + circle_radius),
+                fill=color
+                )
+                
+        except IndexError:
+            pass # Handle cases where keypoints might be missing
+
+    # Save the skeleton image
+    return skeleton_image

yunet-ckpts/yunet_2023mar.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4
+size 232589

yunet-ckpts/yunet_2023mar_int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:321aa5a6afabf7ecc46a3d06bfab2b579dc96eb5c3be7edd365fa04502ad9294
+size 100416

yunet.py

@@ -0,0 +1,55 @@
+# This file is part of OpenCV Zoo project.
+# It is subject to the license terms in the LICENSE file found in the same directory.
+#
+# Copyright (C) 2021, Shenzhen Institute of Artificial Intelligence and Robotics for Society, all rights reserved.
+# Third party copyrights are property of their respective owners.
+
+from itertools import product
+
+import numpy as np
+import cv2 as cv
+
+class YuNet:
+    def __init__(self, modelPath, inputSize=[320, 320], confThreshold=0.6, nmsThreshold=0.3, topK=5000, backendId=0, targetId=0):
+        self._modelPath = modelPath
+        self._inputSize = tuple(inputSize) # [w, h]
+        self._confThreshold = confThreshold
+        self._nmsThreshold = nmsThreshold
+        self._topK = topK
+        self._backendId = backendId
+        self._targetId = targetId
+
+        self._model = cv.FaceDetectorYN.create(
+            model=self._modelPath,
+            config="",
+            input_size=self._inputSize,
+            score_threshold=self._confThreshold,
+            nms_threshold=self._nmsThreshold,
+            top_k=self._topK,
+            backend_id=self._backendId,
+            target_id=self._targetId)
+
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    def setBackendAndTarget(self, backendId, targetId):
+        self._backendId = backendId
+        self._targetId = targetId
+        self._model = cv.FaceDetectorYN.create(
+            model=self._modelPath,
+            config="",
+            input_size=self._inputSize,
+            score_threshold=self._confThreshold,
+            nms_threshold=self._nmsThreshold,
+            top_k=self._topK,
+            backend_id=self._backendId,
+            target_id=self._targetId)
+
+    def setInputSize(self, input_size):
+        self._model.setInputSize(tuple(input_size))
+
+    def infer(self, image):
+        # Forward
+        faces = self._model.detect(image)
+        return np.array([]) if faces[1] is None else faces[1]

yunet2.py

@@ -0,0 +1,114 @@
+import cv2
+import cv2 as cv
+import numpy as np
+from yunet import YuNet
+
+
+# Valid combinations of backends and targets
+backend_target_pairs = [
+    [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_TARGET_CPU],
+    [cv.dnn.DNN_BACKEND_CUDA, cv.dnn.DNN_TARGET_CUDA],
+    [cv.dnn.DNN_BACKEND_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16],
+    [cv.dnn.DNN_BACKEND_TIMVX, cv.dnn.DNN_TARGET_NPU],
+    [cv.dnn.DNN_BACKEND_CANN, cv.dnn.DNN_TARGET_NPU],
+]
+
+
+class YuNet2:
+    def __init__(
+        self,
+        modelPath,
+        input_size=(320, 320),
+        conf_threshold=0.6,
+        nms_threshold=0.3,
+        top_k=5000,
+        backend_id=0,
+        target_id=0,
+    ):
+        self.model = YuNet(
+            modelPath=modelPath,
+            inputSize=input_size,
+            confThreshold=conf_threshold,
+            nmsThreshold=nms_threshold,
+            topK=top_k,
+            backendId=backend_id,
+            targetId=target_id,
+        )
+
+    def detect(self, image, num_faces=None):
+        # If input is an image
+        if image is not None:
+            h, w, _ = image.shape
+
+            # Inference
+            self.model.setInputSize([w, h])
+            results = self.model.infer(image)
+
+            faces = results[:num_faces] if num_faces else results
+
+            bboxes = []
+
+            for face in faces:
+                bbox = face[0:4].astype(np.int32)  # x,y,w,h
+                bboxes.append(bbox)
+                x, y, w, h = bbox
+                # draw
+                cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 255), 2)
+
+            return image, bboxes
+
+    def resize(self, image, target_size=512, above_head_ratio=0.5):
+        height, width, _c = image.shape
+        ar = width / height
+        # downscale the image
+        if not target_size:
+            target_size = 512
+        if ar > 1:
+            # Landscape
+            new_height = target_size
+            new_width = int(target_size * ar)
+        elif ar < 1:
+            # Portrait
+            new_width = target_size
+            new_height = int(target_size / ar)
+        else:
+            # Square
+            new_width = target_size
+            new_height = target_size
+
+        resized = cv2.resize(
+            image, (new_width, new_height), interpolation=cv2.INTER_AREA
+        )
+
+        # Perform object detection on the resized image
+        dt_image, bboxes = self.detect(resized.copy())
+
+        # crop around face
+        if len(bboxes) >= 1:
+            x, y, w, h = bboxes[0]
+        else:
+            x, y, w, h = 0, 0, target_size, target_size
+        # 20% of image height
+        above_head_max = int(target_size * above_head_ratio)
+        x_center = int((x + (x + w)) / 2)
+        y_center = int((y + (y + h)) / 2)
+        # Calculate cropping box
+        top = int(max(0, y_center - above_head_max))
+        bottom = int(min(top + target_size, resized.shape[0]))
+
+        left = int(max(0, x_center - target_size // 2))
+        right = int(min(x_center + target_size // 2, resized.shape[1]))
+
+        # adjust width if necessory
+        _w = right - left
+        if _w != target_size:
+            dx = (
+                target_size - _w
+            )  # difference between the target size and the current width
+            nl = max(0, left - dx)
+            dr = dx - nl  # remaining adjustment needed for the right coordinate
+            left = nl
+            right += dr
+
+        cropped_image = resized[top:bottom, left:right]
+        return dt_image, cropped_image