Spaces:
Sleeping
Sleeping
| #file: creat_anaglyph.py | |
| # Description: This script creates a red-cyan anaglyph stereo image by inserting a person into a stereo image. | |
| from PIL import Image | |
| import numpy as np | |
| import torch | |
| from torchvision.transforms import ToPILImage | |
| # preprocess the human image to remove the black background | |
| def preprocess_person_image(person_image_path): | |
| # uploaded human image | |
| person_image = Image.open(person_image_path).convert('RGBA') | |
| data = np.array(person_image) | |
| # separate color channels | |
| r, g, b, a = data.T | |
| # define the threshold for black background | |
| black_threshold = 1 | |
| black_areas = (r < black_threshold) & (g < black_threshold) & (b < black_threshold) | |
| # set black background to transparent | |
| data[..., 3][black_areas.T] = 0 # only modify the alpha channel | |
| # create a new image | |
| transparent_image = Image.fromarray(data) | |
| return transparent_image | |
| # gradio compatible version of preprocess_person_image | |
| def preprocess_person_image_gradio(person_image): | |
| # ensure the image is in RGBA mode | |
| if person_image.mode != 'RGBA': | |
| person_image = person_image.convert('RGBA') | |
| # load the human image | |
| data = np.array(person_image) | |
| # separate color channels | |
| r, g, b, a = data.T | |
| # define the threshold for black background | |
| black_threshold = 1 | |
| black_areas = (r < black_threshold) & (g < black_threshold) & (b < black_threshold) | |
| # set black background to transparent | |
| data[..., 3][black_areas.T] = 0 # 只修改 alpha 通道 | |
| # create a new image | |
| transparent_image = Image.fromarray(data) | |
| return transparent_image | |
| def insert_person_to_stereo(stereo_image_path, person_image_path, depth_option): | |
| # load the stitched image | |
| stereo_image = Image.open(stereo_image_path).convert('RGB') | |
| width, height = stereo_image.size | |
| # assume the stitched image is symmetrical | |
| left_image = stereo_image.crop((0, 0, width // 2, height)) | |
| right_image = stereo_image.crop((width // 2, 0, width, height)) | |
| # preprocess the human image | |
| person_image = preprocess_person_image(person_image_path) | |
| person_width, person_height = person_image.size | |
| # define disparity options based on image width | |
| max_disparity = width // 20 | |
| disparity_options = { | |
| 'close': max_disparity// 5, | |
| 'medium': max_disparity // 15, | |
| 'far': max_disparity // 20 | |
| } | |
| # get the corresponding disparity value | |
| disparity = disparity_options.get(depth_option, max_disparity // 2) | |
| # calculate the insertion position to align the bottom of the human image with the bottom of the scene image and center horizontally | |
| x_position = (width // 4) - (person_width // 2) + disparity | |
| y_position = height - person_height | |
| # insert the human image into the left and right views | |
| left_image.paste(person_image, (x_position, y_position), person_image) | |
| right_image.paste(person_image, (x_position - disparity, y_position), person_image) | |
| # combine the left and right views into a red-cyan stereo image | |
| left_array = np.array(left_image) # convert the left image to an array | |
| right_array = np.array(right_image) # convert the right image to an array | |
| # create a red-cyan stereo image | |
| anaglyph = np.zeros_like(left_array) | |
| anaglyph[..., 0] = left_array[..., 0] # red channel from left image | |
| anaglyph[..., 1] = right_array[..., 1] # green channel from right image | |
| anaglyph[..., 2] = right_array[..., 2] # blue channel from right image | |
| # convert to an image and save | |
| anaglyph_image = Image.fromarray(anaglyph) # convert the array to an image | |
| anaglyph_image.save('anaglyph.png') # save the image | |
| # gradio compatible version of insert_person_to_stereo | |
| def insert_person_to_stereo_gradio(stereo_image, person_image, disparity): | |
| # load the stitched image | |
| # ensure left_image is in RGB mode | |
| if person_image.mode != "RGBA": | |
| masked_image_pil = person_image.convert("RGBA") | |
| if stereo_image.mode != 'RGB': | |
| stereo_image = stereo_image.convert('RGB') | |
| width, height = stereo_image.size | |
| # assume the stitched image is symmetrical | |
| left_image = stereo_image.crop((0, 0, width // 2, height)) | |
| right_image = stereo_image.crop((width // 2, 0, width, height)) | |
| # preprocess the human image | |
| person_image = preprocess_person_image_gradio(person_image) | |
| person_width, person_height = person_image.size | |
| # calculate the insertion position to align the bottom of the human image with the bottom of the scene image and center horizontally | |
| x_position = (width // 4) - (person_width // 2) + disparity | |
| y_position = height - person_height | |
| # let's paste the person image into the left and right views | |
| left_image.paste(person_image, (x_position, y_position), person_image) | |
| right_image.paste(person_image, (x_position - disparity, y_position), person_image) | |
| # combine the left and right views into a red-cyan stereo image | |
| left_array = np.array(left_image) | |
| right_array = np.array(right_image) | |
| # create a red-cyan stereo image | |
| anaglyph = np.zeros_like(left_array) | |
| anaglyph[..., 0] = left_array[..., 0] # red channel from left image | |
| anaglyph[..., 1] = right_array[..., 1] # green channel from right image | |
| anaglyph[..., 2] = right_array[..., 2] # blue channel from right image | |
| # convert to an image and return | |
| anaglyph_image = Image.fromarray(anaglyph) | |
| return anaglyph_image | |
| # Example | |
| insert_person_to_stereo('img/scenery.jpg', 'img/masked.png', 'far') |