app.py

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()