Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| import cv2 as cv | |
| import numpy as np | |
| import mediapipe as mp | |
| import tempfile | |
| # Initialize MediaPipe Pose | |
| mp_pose = mp.solutions.pose | |
| pose = mp_pose.Pose() | |
| # 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 | |
| # 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 | |
| # 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 | |
| # Calculate the new position after rotation | |
| h, w = img_overlay_rotated.shape[:2] | |
| new_pos = (x - w // 2, y - h // 2) | |
| # 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) | |
| # 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_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 | |
| # 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)) | |
| # 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) | |
| # Process the image and get the pose landmarks | |
| results = pose.process(frame_rgb) | |
| 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] | |
| # 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)) | |
| # 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) | |
| # 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])) | |
| # 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)) | |
| # Adjust position to overlay the image at the wrist | |
| right_position = right_wrist_coord | |
| # 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) | |
| 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)) | |
| # Adjust position to overlay the image at the wrist | |
| left_position = left_wrist_coord | |
| # 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) | |
| return frame | |
| # Streamlit interface | |
| def main(): | |
| 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", "jpg", "jpeg"]) | |
| uploaded_video = st.sidebar.file_uploader("Upload Video File", type=["mp4", "mov", "avi"]) | |
| if uploaded_tattoo_img and uploaded_video: | |
| # 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() | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| 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") | |
| cap.release() | |
| if __name__ == "__main__": | |
| main() | |