Spaces:
Runtime error
Runtime error
| import io | |
| import requests | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from skimage.measure import block_reduce | |
| import gradio as gr | |
| from transformers import DetrFeatureExtractor, DetrForSegmentation, DetrConfig | |
| from transformers.models.detr.feature_extraction_detr import rgb_to_id | |
| from diffusers import StableDiffusionInpaintPipeline | |
| # TODO: maybe need to port to `Blocks` system | |
| # allegedly provides: | |
| # Have multi-step interfaces, in which the output of one model becomes the | |
| # input to the next model, or have more flexible data flows in general. | |
| # and: | |
| # Change a component’s properties (for example, the choices in a dropdown) or its visibility based on user input | |
| # https://huggingface.co/course/chapter9/7?fw=pt | |
| torch.inference_mode() | |
| torch.no_grad() | |
| def load_segmentation_models(model_name: str = 'facebook/detr-resnet-50-panoptic'): | |
| feature_extractor = DetrFeatureExtractor.from_pretrained(model_name) | |
| model = DetrForSegmentation.from_pretrained(model_name) | |
| cfg = DetrConfig.from_pretrained(model_name) | |
| return feature_extractor, model, cfg | |
| def load_diffusion_pipeline(model_name: str = 'runwayml/stable-diffusion-inpainting'): | |
| return StableDiffusionInpaintPipeline.from_pretrained( | |
| model_name, | |
| revision='fp16', | |
| torch_dtype=torch.float16 | |
| ) | |
| def get_device(try_cuda=True): | |
| return torch.device('cuda' if try_cuda and torch.cuda.is_available() else 'cpu') | |
| def min_pool(x: torch.Tensor, kernel_size: int): | |
| pad_size = (kernel_size - 1) // 2 | |
| return -torch.nn.functional.max_pool2d(-x, kernel_size, (1, 1), padding=pad_size) | |
| def max_pool(x: torch.Tensor, kernel_size: int): | |
| pad_size = (kernel_size - 1) // 2 | |
| return torch.nn.functional.max_pool2d(x, kernel_size, (1, 1), padding=pad_size) | |
| def clean_mask(mask, min_kernel: int = 5, max_kernel: int = 23): | |
| mask = torch.Tensor(mask[None, None]).float() | |
| mask = min_pool(mask, min_kernel) | |
| mask = max_pool(mask, max_kernel) | |
| mask = mask.bool().squeeze().numpy() | |
| return mask | |
| device = get_device() | |
| feature_extractor, segmentation_model, segmentation_cfg = load_segmentation_models() | |
| # segmentation_model = segmentation_model.to(device) | |
| pipe = load_diffusion_pipeline() | |
| pipe = pipe.to(device) | |
| def fn_segmentation(image, max_kernel, min_kernel): | |
| inputs = feature_extractor(images=image, return_tensors="pt") | |
| outputs = segmentation_model(**inputs) | |
| processed_sizes = torch.as_tensor(inputs["pixel_values"].shape[-2:]).unsqueeze(0) | |
| result = feature_extractor.post_process_panoptic(outputs, processed_sizes)[0] | |
| panoptic_seg = Image.open(io.BytesIO(result["png_string"])).resize((image.width, image.height)) | |
| panoptic_seg = np.array(panoptic_seg, dtype=np.uint8) | |
| panoptic_seg_id = rgb_to_id(panoptic_seg) | |
| raw_masks = [] | |
| for s in result['segments_info']: | |
| m = panoptic_seg_id == s['id'] | |
| raw_masks.append(m.astype(np.uint8) * 255) | |
| masks = fn_clean(raw_masks, max_kernel, min_kernel) | |
| return masks, raw_masks | |
| def fn_clean(masks, max_kernel, min_kernel): | |
| out = [] | |
| for m in masks: | |
| m = torch.FloatTensor(m)[None, None] | |
| m = min_pool(m, min_kernel) | |
| m = max_pool(m, max_kernel) | |
| m = m.squeeze().numpy().astype(np.uint8) | |
| out.append(m) | |
| return out | |
| def fn_mask(image, mask_enabled): | |
| if len(mask_enabled) == 0: | |
| return image | |
| def fn_segmentation_diffusion(prompt, mask_indices, image, max_kernel, min_kernel, num_diffusion_steps): | |
| mask_indices = [int(i) for i in mask_indices.split(',')] | |
| inputs = feature_extractor(images=image, return_tensors="pt") | |
| outputs = segmentation_model(**inputs) | |
| processed_sizes = torch.as_tensor(inputs["pixel_values"].shape[-2:]).unsqueeze(0) | |
| result = feature_extractor.post_process_panoptic(outputs, processed_sizes)[0] | |
| panoptic_seg = Image.open(io.BytesIO(result["png_string"])).resize((image.width, image.height)) | |
| panoptic_seg = np.array(panoptic_seg, dtype=np.uint8) | |
| class_str = '\n'.join(segmentation_cfg.id2label[s['category_id']] for s in result['segments_info']) | |
| panoptic_seg_id = rgb_to_id(panoptic_seg) | |
| if len(mask_indices) > 0: | |
| mask = (panoptic_seg_id == mask_indices[0]) | |
| for idx in mask_indices[1:]: | |
| mask = mask | (panoptic_seg_id == idx) | |
| mask = clean_mask(mask, min_kernel=min_kernel, max_kernel=max_kernel) | |
| masked_image = np.array(image).copy() | |
| masked_image[mask] = 0 | |
| masked_image = Image.fromarray(masked_image).resize(image.size) | |
| mask = Image.fromarray(mask.astype(np.uint8) * 255).resize(image.size) | |
| if num_diffusion_steps == 0: | |
| return masked_image, masked_image, class_str | |
| STABLE_DIFFUSION_SMALL_EDGE = 512 | |
| assert masked_image.size == mask.size | |
| w, h = masked_image.size | |
| is_width_larger = w > h | |
| resize_ratio = STABLE_DIFFUSION_SMALL_EDGE / (h if is_width_larger else w) | |
| new_width = int(w * resize_ratio) if is_width_larger else STABLE_DIFFUSION_SMALL_EDGE | |
| new_height = STABLE_DIFFUSION_SMALL_EDGE if is_width_larger else int(h * resize_ratio) | |
| new_width += 8 - (new_width % 8) if is_width_larger else 0 | |
| new_height += 0 if is_width_larger else 8 - (new_height % 8) | |
| mask = mask.convert("RGB").resize((new_width, new_height)) | |
| masked_image = masked_image.convert("RGB").resize((new_width, new_height)) | |
| inpainted_image = pipe( | |
| height=new_height, | |
| width=new_width, | |
| prompt=prompt, | |
| image=masked_image, | |
| mask_image=mask, | |
| num_inference_steps=num_diffusion_steps | |
| ).images[0] | |
| return masked_image, inpainted_image, class_str | |
| # iface_segmentation = gr.Interface( | |
| # fn=fn_segmentation, | |
| # inputs=[ | |
| # "text", | |
| # "text", | |
| # gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg"), | |
| # gr.Slider(minimum=1, maximum=99, value=23, step=2), | |
| # gr.Slider(minimum=1, maximum=99, value=5, step=2), | |
| # gr.Slider(minimum=0, maximum=100, value=50, step=1), | |
| # ], | |
| # outputs=["text", gr.Image(type="pil"), gr.Image(type="pil"), "number", "text"] | |
| # ) | |
| # iface_diffusion = gr.Interface( | |
| # fn=fn_diffusion, | |
| # inputs=["text", gr.Image(type='pil'), gr.Image(type='pil'), "number", "text"], | |
| # outputs=[gr.Image(), gr.Image(), gr.Textbox()] | |
| # ) | |
| # iface = gr.Series( | |
| # iface_segmentation, iface_diffusion, | |
| # iface = gr.Interface( | |
| # fn=fn_segmentation_diffusion, | |
| # inputs=[ | |
| # "text", | |
| # "text", | |
| # gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg", type='pil'), | |
| # gr.Slider(minimum=1, maximum=99, value=23, step=2), | |
| # gr.Slider(minimum=1, maximum=99, value=5, step=2), | |
| # gr.Slider(minimum=0, maximum=100, value=50, step=1), | |
| # ], | |
| # outputs=[gr.Image(), gr.Image(), gr.Textbox(interactive=False)] | |
| # ) | |
| # iface = gr.Interface( | |
| # fn=fn_segmentation, | |
| # inputs=[ | |
| # gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg", type='pil'), | |
| # gr.Slider(minimum=1, maximum=99, value=23, step=2), | |
| # gr.Slider(minimum=1, maximum=99, value=5, step=2), | |
| # ], | |
| # outputs=gr.Gallery() | |
| # ) | |
| # iface.launch() | |
| demo = gr.Blocks() | |
| with demo: | |
| input_image = gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg", type='pil') | |
| mask_gallery = gr.Gallery() | |
| mask_storage = gr.State() | |
| max_slider = gr.Slider(minimum=1, maximum=99, value=23, step=2) | |
| min_slider = gr.Slider(minimum=1, maximum=99, value=5, step=2) | |
| bt_masks = gr.Button("Compute Masks") | |
| bt_masks.click(fn_segmentation, inputs=[input_image, max_slider, min_slider], outputs=[mask_gallery, mask_storage]) | |
| max_slider.change(fn_clean, inputs=[mask_storage, max_slider, min_slider], outputs=mask_gallery) | |
| min_slider.change(fn_clean, inputs=[mask_storage, max_slider, min_slider], outputs=mask_gallery) | |
| demo.launch() | |