Allowing user to select the image size

src/app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import numpy as np
 from PIL import Image
-from utils import load_model, segment_person, resize_image, split_stereo_image
+from utils import load_model, segment_person, resize_image, split_stereo_image,resize_image_to_width, resize_mask, resize_images
 from testing import get_image_names
 # Load model and processor once
 processor, model = load_model()
@@ -10,32 +10,55 @@ processor, model = load_model()
 default_bg = Image.new("RGB", (512, 512), color=(95, 147, 89))
 
 
-
-
-
 def generate_3d_outputs(person_img, background_img=None, shift_pixels=10,  person_size=100):
     # Resize images to match
-    image = resize_image(person_img, person_size)
-    background_img = background_img if background_img is not None else default_bg
 
+    background_img = background_img if background_img is not None else default_bg
 
     # Split background image into left and right halves
     leftBackground, rightBackground = split_stereo_image(Image.fromarray(background_img))
 
-    # Resize image to match background dimensions
+    ## Match person image to background image width
+    image = resize_image_to_width(person_img, leftBackground)
 
-    
-    image = Image.fromarray(np.array(image)).resize((leftBackground.shape[1], leftBackground.shape[0]))
     # Step 1: Segment person
     mask = segment_person(image, processor, model)
 
-    image_np = np.array(image)
+    # Resize mask based on person_size percentage
+    mask = resize_mask(person_size, mask)
+
+    # Resize image based on person_size percentage
+    image_np = resize_images(image, person_size)
+
+    # Apply mask to image
+    person_only = image_np * mask
+    person_segmentation = np.clip(person_only, 0, 255).astype(np.uint8)
+
+    # Resize mask and person_only to match background dimensions while preserving content
+    target_height, target_width = leftBackground.shape[:2]
+    current_height, current_width = mask.shape[:2]
+
+    # Calculate padding
+    pad_top = max(0, (target_height - current_height) // 2)
+    pad_bottom = max(0, target_height - current_height - pad_top)
+    pad_left = max(0, (target_width - current_width) // 2)
+    pad_right = max(0, target_width - current_width - pad_left)
+
+    # Pad mask and person_only arrays
+    mask = np.pad(mask, ((pad_top, pad_bottom), (pad_left, pad_right), (0,0)), mode='constant')    
+    person_only = np.pad(person_segmentation, ((pad_top, pad_bottom), (pad_left, pad_right), (0,0)), mode='constant')
+
+
+    # CROP MASK TO MATCH BACKGROUND DIMENSIONS FROM CENTER OF BACKGROUND
+    if(mask.shape[0] > target_height or mask.shape[1] > target_width):
+        mask = mask[mask.shape[0]//2-target_height//2:mask.shape[0]//2+target_height//2, mask.shape[1]//2-target_width//2:mask.shape[1]//2+target_width//2, :]
+        person_only = person_only[person_only.shape[0]//2-target_height//2:person_only.shape[0]//2+target_height//2, person_only.shape[1]//2-target_width//2:person_only.shape[1]//2+target_width//2, :]
 
+    # Convert background images to numpy arrays
     leftBackground_np = np.array(leftBackground)
     rightBackground_np = np.array(rightBackground)
 
-
-    person_only = image_np * mask
+    # Apply mask to background images
     leftBackground_only = leftBackground_np * (1 - mask)
     rightBackground_only = rightBackground_np * (1 - mask)
 
@@ -43,10 +66,9 @@ def generate_3d_outputs(person_img, background_img=None, shift_pixels=10,  perso
     person_left = np.roll(person_only, shift=-shift_pixels, axis=1)
     person_right = np.roll(person_only, shift=shift_pixels, axis=1)
     
-
+    
     left_eye = np.clip(person_right + leftBackground_only, 0, 255).astype(np.uint8)
     right_eye = np.clip(person_left + rightBackground_only, 0, 255).astype(np.uint8)
-    person_segmentation = np.clip(person_only, 0, 255).astype(np.uint8)
 
     # --- Combine left and right images side by side ---
     stereo_pair = np.concatenate([left_eye, right_eye], axis=1)

src/testing.py

@@ -3,27 +3,33 @@ import os
 import random
 
 def get_image_names():
-    # Get background images
-    background_dir = "public/images/backgrounds"
-    background_images = [f for f in os.listdir(background_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]
-    
     # Get people images
     people_dir = "public/images/people" 
-    people_images = [f for f in os.listdir(people_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]
-    
-
-    generate_testing_list = []
+    # Get background images
+    background_dir = "public/images/backgrounds"
 
-    random.shuffle(background_images)
-    random.shuffle(people_images)
 
-    total_length = len(people_images)
+    data = [
+      ['e.jpg', 'e.png', 10, 60],
+      ['f.jpg', 'f.png', 10, 60],
+      ['g.jpeg', 'g.png', 10, 40],      
+      ['a.jpg', 'a.png', 10, 130],
+      ['b.jpg', 'b.png', 10, 40],
+      ['c.jpg', 'c.png', 10, 60],
+      ['d.png', 'd.png', 10, 50]
+    ]
 
-    # Get 10 random pairs
-    for background in background_images:
-        person = people_images[random.randint(0, total_length - 1)]
-        generate_testing_list.append(list(( os.path.join(people_dir, person), os.path.join(background_dir, background), 10, 100)))
+    generate_testing_list = []
 
+    for i in range(len(data)):
+      generate_testing_list.append(
+        list((
+          os.path.join(people_dir, data[i][0]), 
+          os.path.join(background_dir, data[i][1]), 
+          data[i][2], 
+          data[i][3]
+          ))
+        )
     
     return generate_testing_list
 

src/utils.py

@@ -94,3 +94,61 @@ def split_stereo_image(image):
     else:
         return image, resize_image(image, 99)
     
+def resize_image_to_width(person_img, background_img):
+       # Resize image to match background dimensions
+    if (background_img.shape[1] > background_img.shape[0]):
+        width = background_img.shape[1]
+        img_array = np.array(person_img)
+        height = int(width * img_array.shape[0] / img_array.shape[1])
+        person_img = Image.fromarray(img_array).resize((width, height))
+        person_img = np.array(person_img)
+        image = Image.fromarray(person_img)
+    else:
+        height = background_img.shape[0]
+        img_array = np.array(person_img)
+        width = int(height * img_array.shape[1] / img_array.shape[0])
+        person_img = Image.fromarray(img_array).resize((width, height))
+        person_img = np.array(person_img)
+        image = Image.fromarray(person_img)
+
+
+    return image
+
+def resize_mask(person_size, mask):
+        
+    scale_factor = person_size / 100.0
+    mask_height, mask_width = mask.shape[:2]
+    new_height = int(mask_height * scale_factor)
+    new_width = int(mask_width * scale_factor)
+    
+    # Convert mask to PIL Image for resizing
+    mask_image = Image.fromarray((mask * 255).astype(np.uint8))
+    resized_mask = mask_image.resize((new_width, new_height))
+    
+    # Convert back to numpy and normalize to 0-1
+    mask = np.array(resized_mask).astype(np.float32) / 255.0
+    
+    # Add third channel dimension back if needed
+    if len(mask.shape) == 2:
+        mask = np.stack([mask] * 3, axis=-1)
+
+    return mask
+
+def resize_images(image, person_size):
+    image_np = np.array(image)
+    # Resize image based on person_size percentage
+
+    scale_factor = person_size / 100.0
+    img_height, img_width = image_np.shape[:2]
+    new_height = int(img_height * scale_factor)
+    new_width = int(img_width * scale_factor)
+    
+    # Convert image to PIL Image for resizing
+    image_pil = Image.fromarray(image_np)
+    resized_image = image_pil.resize((new_width, new_height))
+    
+    # Convert back to numpy
+    image = resized_image
+    image_np = np.array(image)
+
+    return image_np