new version by using mediapipe

app.py

@@ -1,117 +1,141 @@
 import streamlit as st
 import cv2 as cv
 import numpy as np
+import mediapipe as mp
 import tempfile
 
-# Paths to your files
-model_path = "graph_opt.pb"
+# Initialize MediaPipe Pose
+mp_pose = mp.solutions.pose
+pose = mp_pose.Pose()
 
-# Load the TensorFlow model
-net = cv.dnn.readNetFromTensorflow(model_path)
-if net.empty():
-    raise FileNotFoundError(f"Model file not found or cannot be loaded: {model_path}")
+# Function to rotate an image around its center
+def rotate_image(image, angle):
+    h, w = image.shape[:2]
+    center = (w // 2, h // 2)
+    rot_matrix = cv.getRotationMatrix2D(center, angle, 1.0)
+    rotated_image = cv.warpAffine(image, rot_matrix, (w, h), flags=cv.INTER_LINEAR)
+    return rotated_image
 
-BODY_PARTS = {"Nose": 0, "Neck": 1, "RShoulder": 2, "RElbow": 3, "RWrist": 4,
-              "LShoulder": 5, "LElbow": 6, "LWrist": 7, "RHip": 8, "RKnee": 9,
-              "RAnkle": 10, "LHip": 11, "LKnee": 12, "LAnkle": 13, "REye": 14,
-              "LEye": 15, "REar": 16, "LEar": 17, "Background": 18}
-
-# Adjust POSE_PAIRS to include only wrist to elbow parts
-POSE_PAIRS = [["RElbow", "RWrist"], ["LElbow", "LWrist"]]
+# Function to overlay the image with rotation following arm direction
+def overlay_image_alpha_rotated(img, img_overlay, pos, alpha_mask, angle):
+    x, y = pos
 
-width = 368
-height = 368
-inWidth = width
-inHeight = height
-thr = 0.2
+    # Rotate the overlay image in the opposite direction
+    img_overlay_rotated = rotate_image(img_overlay, -angle)  # Negate the angle
+    alpha_mask_rotated = rotate_image(alpha_mask, -angle)  # Negate the angle
 
-def overlay_image_alpha(img, img_overlay, pos, alpha_mask):
-    x, y = pos
+    # Calculate the new position after rotation
+    h, w = img_overlay_rotated.shape[:2]
+    new_pos = (x - w // 2, y - h // 2)
 
-    y1, y2 = max(0, y), min(img.shape[0], y + img_overlay.shape[0])
-    x1, x2 = max(0, x), min(img.shape[1], x + img_overlay.shape[1])
+    # Image ranges
+    y1, y2 = max(0, new_pos[1]), min(img.shape[0], new_pos[1] + h)
+    x1, x2 = max(0, new_pos[0]), min(img.shape[1], new_pos[0] + w)
 
-    y1o, y2o = max(0, -y), min(img_overlay.shape[0], img.shape[0] - y)
-    x1o, x2o = max(0, -x), min(img_overlay.shape[1], img.shape[1] - x)
+    # Overlay ranges
+    y1o, y2o = max(0, -new_pos[1]), min(h, img.shape[0] - new_pos[1])
+    x1o, x2o = max(0, -new_pos[0]), min(w, img.shape[1] - new_pos[0])
 
+    # Exit if nothing to overlay
     if y1 >= y2 or x1 >= x2 or y1o >= y2o or x1o >= x2o:
         return
 
+    # Blend overlay within the determined ranges
     img_crop = img[y1:y2, x1:x2]
-    img_overlay_crop = img_overlay[y1o:y2o, x1o:x2o]
-    alpha = alpha_mask[y1o:y2o, x1o:x2o, np.newaxis]
+    img_overlay_crop = img_overlay_rotated[y1o:y2o, x1o:x2o]
+    alpha = alpha_mask_rotated[y1o:y2o, x1o:x2o, np.newaxis]
     alpha_inv = 1.0 - alpha
 
     img_crop[:] = alpha * img_overlay_crop + alpha_inv * img_crop
 
-def poseDetector(frame, overlay_img):
-    frameWidth = frame.shape[1]
-    frameHeight = frame.shape[0]
-
-    net.setInput(cv.dnn.blobFromImage(frame, 1.0, (inWidth, inHeight), (127.5, 127.5, 127.5), swapRB=True, crop=False))
-    out = net.forward()
-    out = out[:, :19, :, :]
+# Function to handle JPG images without alpha channel
+def add_alpha_channel(image):
+    b_channel, g_channel, r_channel = cv.split(image)
+    alpha_channel = np.ones(b_channel.shape, dtype=b_channel.dtype) * 255  # Creating a fully opaque alpha channel
+    return cv.merge((b_channel, g_channel, r_channel, alpha_channel))
 
-    assert (len(BODY_PARTS) == out.shape[1])
-
-    points = []
-    for i in range(len(BODY_PARTS)):
-        heatMap = out[0, i, :, :]
+# Function to detect pose and overlay image following arm direction
+def poseDetector(frame, overlay_img):
+    # Convert the image to RGB as MediaPipe expects RGB images
+    frame_rgb = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
 
-        _, conf, _, point = cv.minMaxLoc(heatMap)
-        x = (frameWidth * point[0]) / out.shape[3]
-        y = (frameHeight * point[1]) / out.shape[2]
+    # Process the image and get the pose landmarks
+    results = pose.process(frame_rgb)
 
-        points.append((int(x), int(y)) if conf > thr else None)
+    if results.pose_landmarks:
+        # Extract relevant landmarks
+        right_wrist = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST]
+        left_wrist = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_WRIST]
+        right_elbow = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ELBOW]
+        left_elbow = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ELBOW]
 
-    for pair in POSE_PAIRS:
-        partFrom = pair[0]
-        partTo = pair[1]
+        # Convert landmark positions to pixel coordinates
+        h, w, _ = frame.shape
+        right_wrist_coord = (int(right_wrist.x * w), int(right_wrist.y * h))
+        left_wrist_coord = (int(left_wrist.x * w), int(left_wrist.y * h))
+        right_elbow_coord = (int(right_elbow.x * w), int(right_elbow.y * h))
+        left_elbow_coord = (int(left_elbow.x * w), int(left_elbow.y * h))
 
-        assert (partFrom in BODY_PARTS)
-        assert (partTo in BODY_PARTS)
+        # Calculate the length of the hand region (between wrist and elbow)
+        right_hand_length = int(np.sqrt((right_wrist_coord[0] - right_elbow_coord[0])**2 + (right_wrist_coord[1] - right_elbow_coord[1])**2) * 0.5)
+        left_hand_length = int(np.sqrt((left_wrist_coord[0] - left_elbow_coord[0])**2 + (left_wrist_coord[1] - left_elbow_coord[1])**2) * 0.5)
 
-        idFrom = BODY_PARTS[partFrom]
-        idTo = BODY_PARTS[partTo]
+        # Calculate the angle of the hand for rotation
+        right_angle = np.degrees(np.arctan2(right_wrist_coord[1] - right_elbow_coord[1], right_wrist_coord[0] - right_elbow_coord[0]))
+        left_angle = np.degrees(np.arctan2(left_wrist_coord[1] - left_elbow_coord[1], left_wrist_coord[0] - left_elbow_coord[0]))
 
-        if points[idFrom] and points[idTo]:
-            angle = np.degrees(np.arctan2(points[idTo][1] - points[idFrom][1], points[idTo][0] - points[idFrom][0]))
-            adjusted_angle = angle + 270
+        # Ensure the hand length is positive before resizing
+        if right_hand_length > 0:
+            # Resize overlay image to fit the hand length
+            right_overlay_resized = cv.resize(overlay_img, (right_hand_length, right_hand_length))
 
-            length = int(np.sqrt((points[idTo][0] - points[idFrom][0])**2 + (points[idTo][1] - points[idFrom][1])**2))
-            reduced_length = int(length * 0.5)
+            # Adjust position to overlay the image at the wrist
+            right_position = right_wrist_coord
 
-            if reduced_length > 0 and overlay_img.shape[0] * 0.5 > 0:
-                overlay_resized = cv.resize(overlay_img, (reduced_length, int(overlay_img.shape[0] * 0.5)))
+            # Overlay the image on the right hand with rotation following the arm direction
+            alpha_mask = right_overlay_resized[:, :, 3] / 255.0
+            overlay_image_alpha_rotated(frame, right_overlay_resized[:, :, :3], right_position, alpha_mask, right_angle)
 
-                M = cv.getRotationMatrix2D((overlay_resized.shape[1] / 2, overlay_resized.shape[0] / 2), adjusted_angle, 1)
-                overlay_rotated = cv.warpAffine(overlay_resized, M, (overlay_resized.shape[1], overlay_resized.shape[0]), flags=cv.INTER_LINEAR, borderMode=cv.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))
+        if left_hand_length > 0:
+            # Resize overlay image to fit the hand length
+            left_overlay_resized = cv.resize(overlay_img, (left_hand_length, left_hand_length))
 
-                wrist_position = points[idTo]
-                position = (int(wrist_position[0] - overlay_rotated.shape[1] / 2), int(wrist_position[1] - overlay_rotated.shape[0] / 2))
+            # Adjust position to overlay the image at the wrist
+            left_position = left_wrist_coord
 
-                alpha_mask = overlay_rotated[:, :, 3] / 255.0
-                overlay_image_alpha(frame, overlay_rotated[:, :, :3], position, alpha_mask)
+            # Overlay the image on the left hand with rotation following the arm direction
+            alpha_mask = left_overlay_resized[:, :, 3] / 255.0
+            overlay_image_alpha_rotated(frame, left_overlay_resized[:, :, :3], left_position, alpha_mask, left_angle)
 
-    t, _ = net.getPerfProfile()
     return frame
 
+# Streamlit interface
 def main():
-    st.title("Video File Pose Detection and Tattoo Overlay")
+    st.title("Webcam Pose Detection with Tattoo Overlay")
 
+    # Sidebar for file upload
     st.sidebar.header("Upload Tattoo Image and Video")
-    uploaded_tattoo_img = st.sidebar.file_uploader("Upload Tattoo Image", type=["png"])
+    uploaded_tattoo_img = st.sidebar.file_uploader("Upload Tattoo Image", type=["png", "jpg", "jpeg"])
     uploaded_video = st.sidebar.file_uploader("Upload Video File", type=["mp4", "mov", "avi"])
 
     if uploaded_tattoo_img and uploaded_video:
-        tattoo_img = cv.imdecode(np.frombuffer(uploaded_tattoo_img.read(), np.uint8), cv.IMREAD_UNCHANGED)
+        # Load the uploaded tattoo image
+        file_bytes = np.frombuffer(uploaded_tattoo_img.read(), np.uint8)
+        tattoo_img = cv.imdecode(file_bytes, cv.IMREAD_UNCHANGED)
 
+        # If the image is in JPG format (no alpha channel), add an alpha channel
+        if tattoo_img.shape[2] == 3:  # Check if image has only 3 channels (BGR)
+            tattoo_img = add_alpha_channel(tattoo_img)
+
+        # Checkbox to start the processing
         run = st.checkbox('Run')
 
         if run:
+            # Temporary file to store the uploaded video
             tfile = tempfile.NamedTemporaryFile(delete=False)
             tfile.write(uploaded_video.read())
 
+            # Video capture
             cap = cv.VideoCapture(tfile.name)
 
             stframe = st.empty()
@@ -121,7 +145,10 @@ def main():
                 if not ret:
                     break
 
+                # Perform pose detection and overlay
                 frame = poseDetector(frame, tattoo_img)
+
+                # Convert the frame to RGB for display in Streamlit
                 frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
                 stframe.image(frame, channels="RGB")
 

graph_opt.pb

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:51f84ea82b3d0143dd4cf362e018c125e832c64771908ee8766fbd8b0328008d
-size 7804434