gexu13's picture
Upload 16 files
f74ae4b verified
import gradio as gr
from PIL import Image
import os
from image_segmentation_mask_rcnn import segment_person
from insert_person_into_stereo import insert_person_from_combined_stereo
from create_anaglyph import create_anaglyph
# Predefined sample files (make sure these exist in your project directory)
DEFAULT_BACKGROUNDS = {
"Backyard": "Side_By_Side_3D_Images/sbs_backyard.png",
"Campus": "Side_By_Side_3D_Images/sbs_campus.png",
"Downtown": "Side_By_Side_3D_Images/sbs_downtown.png",
"NEU": "Side_By_Side_3D_Images/sbs_neu.png",
"STEAM_CLOCK": "Side_By_Side_3D_Images/sbs_steam_clock.png",
"Trail": "Side_By_Side_3D_Images/sbs_trail.png"
}
DEFAULT_PEOPLE = {
"PERSON1": "person/person1.jpg",
"PERSON2": "person/person2.png",
"PERSON3": "person/person3.png",
"PERSON4": "person/person4.png",
"PERSON5": "person/person5.png",
}
def pipeline(person_image, stereo_image, depth, x, y):
segmented = segment_person(person_image)
left_image, right_image, _ = insert_person_from_combined_stereo(
stereo_image=stereo_image,
segmented_person=segmented,
depth=depth,
position=(x, y)
)
anaglyph = create_anaglyph(left_image, right_image)
return anaglyph
def get_image_dimensions(stereo_image, person_image):
if stereo_image is None or person_image is None:
return gr.update(), gr.update()
w_bg, h_bg = stereo_image.size
w_p, h_p = person_image.size
max_x = max(10, w_bg // 2 - w_p // 2) # Ensure > 0
max_y = max(10, h_bg) # Ensure > 0
return gr.update(minimum=0, maximum=max_x, value=max_x // 2), gr.update(minimum=0, maximum=max_y, value=int(h_bg * 0.9))
def main():
with gr.Blocks() as demo:
gr.Markdown("# 🧍‍➡️ 3D Anaglyph Composer")
with gr.Row():
person_input = gr.Image(type="pil", label="Person Image")
stereo_input = gr.Image(type="pil", label="Stereo Background")
# Sample selectors
with gr.Row():
with gr.Column():
gr.Markdown("### Sample People")
for label, path in DEFAULT_PEOPLE.items():
with gr.Row():
preview = gr.Image(value=path, label=label, interactive=False, show_label=False, width=128, height=128)
use_btn = gr.Button(f"Use {label}")
use_btn.click(lambda p=path: Image.open(p), outputs=person_input)
with gr.Column():
gr.Markdown("### Sample Backgrounds")
for label, path in DEFAULT_BACKGROUNDS.items():
with gr.Row():
preview = gr.Image(value=path, label=label, interactive=False, show_label=False)
use_btn = gr.Button(f"Use {label}")
use_btn.click(lambda p=path: Image.open(p), outputs=stereo_input)
depth_input = gr.Dropdown(["close", "medium", "far"], value="medium", label="Depth")
x_slider = gr.Slider(0, 2000, value=1000, label="Person X Position")
y_slider = gr.Slider(0, 2000, value=500, label="Person Y Position")
generate_btn = gr.Button("Generate 3D Anaglyph")
output_img = gr.Image(type="pil", label="Anaglyph 3D Image")
# Dynamically update position sliders when images are uploaded
person_input.change(get_image_dimensions, inputs=[stereo_input, person_input], outputs=[x_slider, y_slider])
stereo_input.change(get_image_dimensions, inputs=[stereo_input, person_input], outputs=[x_slider, y_slider])
generate_btn.click(
fn=pipeline,
inputs=[person_input, stereo_input, depth_input, x_slider, y_slider],
outputs=output_img
)
demo.launch()
if __name__ == "__main__":
main()