Upload trial_insert.py

src/trial_insert.py

@@ -0,0 +1,101 @@
+import cv2
+import numpy as np
+
+
+def insert_person_V2(
+    left_image_path, right_image_path, person_image_path, depth="medium"
+):
+    # Load left and right stereoscopic images
+    left_image = cv2.imread(left_image_path)
+    right_image = cv2.imread(right_image_path)
+
+    # Load the segmented person image with alpha channel (transparency)
+    person = cv2.imread(person_image_path, cv2.IMREAD_UNCHANGED)
+
+    # Define scaling and disparity values for each depth level
+    depth_settings = {
+        "close": {"scale": 1.2, "disparity": 15},
+        "medium": {"scale": 1.0, "disparity": 10},
+        "far": {"scale": 0.7, "disparity": 5},
+    }
+    scale_factor = depth_settings[depth]["scale"]
+    disparity = depth_settings[depth]["disparity"]
+
+    # Resize person image according to the scale factor
+    person_resized = cv2.resize(
+        person, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_AREA
+    )
+    (ph, pw) = person_resized.shape[:2]
+
+    # Extract color and alpha channels from person image
+    person_rgb = person_resized[:, :, :3]
+    person_alpha = person_resized[:, :, 3] / 255.0  # Normalize alpha channel to [0, 1]
+
+    # Match brightness and contrast of the person to the background
+    person_rgb = match_brightness_contrast(left_image, person_rgb, person_alpha)
+
+    # Color match person to the background
+    person_rgb = match_color_tone(left_image, person_rgb, person_alpha)
+
+    # Create a blended version of the person with soft edges
+    person_blended = soft_edge_blending(person_rgb, person_alpha)
+
+    # Determine insertion position in the left image
+    x_offset, y_offset = 100, left_image.shape[0] - ph - 20
+    left_image[y_offset : y_offset + ph, x_offset : x_offset + pw] = blend_images(
+        left_image[y_offset : y_offset + ph, x_offset : x_offset + pw],
+        person_blended,
+        person_alpha,
+    )
+
+    # Create stereoscopic effect by shifting the person image in the right image
+    x_offset_right = x_offset + disparity
+    right_image[y_offset : y_offset + ph, x_offset_right : x_offset_right + pw] = (
+        blend_images(
+            right_image[y_offset : y_offset + ph, x_offset_right : x_offset_right + pw],
+            person_blended,
+            person_alpha,
+        )
+    )
+
+    return left_image, right_image
+
+
+def match_brightness_contrast(background, person, alpha):
+    # Calculate the mean brightness of the background where the person will be placed
+    background_mean = np.mean(background, axis=(0, 1))
+    person_mean = np.mean(person, axis=(0, 1))
+    adjustment = background_mean / (person_mean + 1e-6)
+    return cv2.convertScaleAbs(person, alpha=adjustment[0], beta=adjustment[1])
+
+
+def match_color_tone(background, person, alpha):
+    # Adjust color tone to match background
+    bg_mean, bg_std = cv2.meanStdDev(background)
+    person_mean, person_std = cv2.meanStdDev(person)
+    scale = (bg_std + 1e-6) / (person_std + 1e-6)
+    person = cv2.convertScaleAbs(
+        person, alpha=scale[0][0], beta=(bg_mean - person_mean)[0][0]
+    )
+    return person
+
+
+def soft_edge_blending(person, alpha):
+    # Apply Gaussian blur to soften the edges for better blending
+    blurred_alpha = cv2.GaussianBlur(alpha, (15, 15), 0)
+    person = cv2.merge(
+        (
+            person[:, :, 0] * blurred_alpha,
+            person[:, :, 1] * blurred_alpha,
+            person[:, :, 2] * blurred_alpha,
+        )
+    )
+    return person
+
+
+def blend_images(background, person, alpha):
+    # Blend person into background using the alpha mask
+    blended = (alpha[..., None] * person + (1 - alpha[..., None]) * background).astype(
+        np.uint8
+    )
+    return blended