add files

app.py

@@ -0,0 +1,120 @@
+import gradio as gr
+import torch
+from torchvision import models, transforms
+from PIL import Image
+import numpy as np
+import cv2
+
+# Load pre-trained segmentation model
+model = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()
+
+# Function to perform segmentation
+def segment_person(image):
+    try:
+        if not isinstance(image, Image.Image):
+            raise ValueError("Invalid image format. Please upload a valid image.")
+
+        preprocess = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+        input_tensor = preprocess(image).unsqueeze(0)
+
+        with torch.no_grad():
+            output = model(input_tensor)['out'][0]
+
+        output_predictions = output.argmax(0)
+        person_class = 15  # Person class in DeepLabV3
+        mask = output_predictions == person_class
+        mask_np = mask.byte().cpu().numpy()
+        mask_resized = cv2.resize(mask_np, (image.width, image.height), interpolation=cv2.INTER_NEAREST)
+
+        person = np.array(image)
+        alpha_channel = (mask_resized * 255).astype(np.uint8)
+        rgba_image = np.dstack((person, alpha_channel))
+        return Image.fromarray(rgba_image, 'RGBA')
+
+    except Exception as e:
+        raise RuntimeError(f"Error in segmenting person: {e}")
+
+# Function to create stereoscopic images with inserted person
+def create_stereo_images(stereo_image, person_image, depth_level, x_position, y_position):
+    try:
+        if not isinstance(stereo_image, np.ndarray):
+            raise ValueError("Invalid stereo image format. Please upload a valid image.")
+
+        height, width, _ = stereo_image.shape
+        midpoint = width // 2
+        left_image = stereo_image[:, :midpoint]
+        right_image = stereo_image[:, midpoint:]
+
+        left_image = Image.fromarray(cv2.cvtColor(left_image, cv2.COLOR_BGR2RGBA))
+        right_image = Image.fromarray(cv2.cvtColor(right_image, cv2.COLOR_BGR2RGBA))
+
+        disparity_map = {"close": 10, "medium": 1, "far": -10}
+        scaling_factors = {"close": 1.2, "medium": 1.0, "far": 0.8}
+
+        if depth_level not in disparity_map:
+            raise ValueError("Invalid depth level. Choose from 'close', 'medium', or 'far'.")
+
+        disparity = disparity_map[depth_level]
+        person_scale = scaling_factors[depth_level]
+        person_image = person_image.resize(
+            (int(person_image.width * person_scale), int(person_image.height * person_scale)),
+            Image.LANCZOS
+        )
+
+        left_x, left_y = x_position, y_position
+        right_x = left_x + disparity
+
+        left_image.paste(person_image, (left_x, left_y), person_image)
+        right_image.paste(person_image, (right_x, left_y), person_image)
+
+        return left_image, right_image
+
+    except Exception as e:
+        raise RuntimeError(f"Error in creating stereo images: {e}")
+
+# Function to create an anaglyph image
+def create_anaglyph(left_image, right_image):
+    try:
+        left_r, _, _ = left_image.convert("RGB").split()
+        _, right_g, right_b = right_image.convert("RGB").split()
+        anaglyph_image = Image.merge("RGB", (left_r, right_g, right_b))
+        return anaglyph_image
+    except Exception as e:
+        raise RuntimeError(f"Error in creating anaglyph image: {e}")
+
+# Gradio interface function
+def process_images(person_image, stereo_image, depth_level, x_position, y_position):
+    try:
+        person_image = segment_person(person_image)
+        stereo_image = np.array(stereo_image)
+
+        left_image, right_image = create_stereo_images(stereo_image, person_image, depth_level, x_position, y_position)
+        anaglyph_image = create_anaglyph(left_image, right_image)
+
+        return anaglyph_image
+    except Exception as e:
+        return f"An error occurred: {e}"
+
+# Gradio app setup
+with gr.Blocks() as demo:
+    gr.Markdown("# 3D Anaglyph Image Creator")
+    gr.Markdown("Upload a person image, select depth, adjust position, and create a 3D anaglyph image.")
+    
+    with gr.Row():
+        person_image = gr.Image(label="Upload Person Image", type="pil")
+        stereo_image = gr.Image(label="Upload Stereo Background Image", type="pil")
+    
+    depth_level = gr.Radio(["close", "medium", "far"], label="Select Depth Level")
+    x_position = gr.Slider(0, 2000, step=1, label="X Position")
+    y_position = gr.Slider(0, 2000, step=1, label="Y Position")
+    
+    output_image = gr.Image(label="Generated Anaglyph Image")
+    
+    generate_button = gr.Button("Generate Anaglyph Image")
+    generate_button.click(process_images, [person_image, stereo_image, depth_level, x_position, y_position], output_image)
+
+# Launch the app
+demo.launch()