Shopify/background-replacement

1import warnings2warnings.filterwarnings("ignore", category=FutureWarning)  # nopep83warnings.filterwarnings("ignore", category=UserWarning)  # nopep84import os5import math6from tqdm import tqdm7import torch8from PIL import Image, ImageFilter9from scipy.ndimage import binary_dilation10import numpy as np11 12from captioner import init as init_captioner, derive_caption13from upscaler import init as init_upscaler14from segmenter import init as init_segmenter, segment15from depth_estimator import init as init_depth_estimator, get_depth_map16from pipeline import init as init_pipeline, run_pipeline17from image_utils import ensure_resolution, crop_centered18 19developer_mode = os.getenv('DEV_MODE', False)20 21# You must uncomment this initialization block!22# init_captioner()23# init_upscaler()24# init_segmenter()25# init_depth_estimator()26# init_pipeline()27 28# torch.cuda.empty_cache()29 30POSITIVE_PROMPT_SUFFIX = "commercial product photography, 24mm lens f/8"31NEGATIVE_PROMPT_SUFFIX = "cartoon, drawing, anime, semi-realistic, illustration, painting, art, text, greyscale, (black and white), lens flare, watermark, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, floating, levitating"32 33MEGAPIXELS = 1.034 35 36def replace_background(37    original,38    positive_prompt,39    negative_prompt,40    options,41):42    pbar = tqdm(total=7)43 44    print("Original size:", original.size)45 46    print("Captioning...")47    caption = derive_caption(original)48    pbar.update(1)49 50    print("Caption:", caption)51 52    torch.cuda.empty_cache()53 54    print(f"Ensuring resolution ({MEGAPIXELS}MP)...")55    resized = ensure_resolution(original, megapixels=MEGAPIXELS)56    pbar.update(1)57 58    print("Resized size:", resized.size)59 60    torch.cuda.empty_cache()61 62    print("Segmenting...")63    [cropped, crop_mask] = segment(resized)64    pbar.update(1)65 66    torch.cuda.empty_cache()67 68    print("Depth mapping...")69    depth_map = get_depth_map(resized)70    pbar.update(1)71 72    torch.cuda.empty_cache()73 74    print("Feathering the depth map...")75 76    # Convert crop mask to grayscale and to numpy array77    crop_mask_np = np.array(crop_mask.convert('L'))78 79    # Convert to binary and dilate (grow) the edges80    # adjust threshold as needed81    crop_mask_binary = crop_mask_np > options.get(82        'depth_map_feather_threshold')83    # adjust iterations as needed84    dilated_mask = binary_dilation(85        crop_mask_binary, iterations=options.get('depth_map_dilation_iterations'))86 87    # Convert back to PIL Image88    dilated_mask = Image.fromarray((dilated_mask * 255).astype(np.uint8))89 90    # Apply Gaussian blur and normalize91    dilated_mask_blurred = dilated_mask.filter(92        ImageFilter.GaussianBlur(radius=options.get('depth_map_blur_radius')))93    dilated_mask_blurred_np = np.array(dilated_mask_blurred) / 255.094 95    # Normalize depth map, apply blurred, dilated mask, and scale back96    depth_map_np = np.array(depth_map.convert('L')) / 255.097    masked_depth_map_np = depth_map_np * dilated_mask_blurred_np98    masked_depth_map_np = (masked_depth_map_np * 255).astype(np.uint8)99 100    # Convert back to PIL Image101    masked_depth_map = Image.fromarray(masked_depth_map_np).convert('RGB')102 103    pbar.update(1)104 105    final_positive_prompt = f"{caption}, {positive_prompt}, {POSITIVE_PROMPT_SUFFIX}"106    final_negative_prompt = f"{negative_prompt}, {NEGATIVE_PROMPT_SUFFIX}"107 108    print("Final positive prompt:", final_positive_prompt)109    print("Final negative prompt:", final_negative_prompt)110 111    print("Generating...")112 113    generated_images = run_pipeline(114        positive_prompt=final_positive_prompt,115        negative_prompt=final_negative_prompt,116        image=[masked_depth_map],117        seed=options.get('seed')118    )119    pbar.update(1)120 121    torch.cuda.empty_cache()122 123    print("Compositing...")124 125    composited_images = [126        Image.alpha_composite(127            generated_image.convert('RGBA'),128            crop_centered(cropped, generated_image.size)129        ) for generated_image in generated_images130    ]131    pbar.update(1)132    pbar.close()133 134    print("Done!")135 136    if developer_mode:137        pre_processing_images = [138            [resized, "Resized"],139            [crop_mask, "Crop mask"],140            [cropped, "Cropped"],141            [depth_map, "Depth map"],142            [dilated_mask, "Dilated mask"],143            [dilated_mask_blurred, "Dilated mask blurred"],144            [masked_depth_map, "Masked depth map"]145        ]146        return [147            composited_images,148            generated_images,149            pre_processing_images,150            caption,151        ]152    else:153        return [composited_images, None, None, None]154