Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python | |
| import os | |
| try: | |
| import detectron2 | |
| except: | |
| os.system('pip install git+https://github.com/facebookresearch/detectron2.git') | |
| import logging | |
| logging.disable(logging.CRITICAL) # comment out to enable verbose logging | |
| ######################################################### | |
| import pathlib | |
| import gradio as gr | |
| import numpy as np | |
| import PIL.Image as Image | |
| import os | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| import matplotlib.pyplot as plt | |
| from PIL import Image | |
| from collections import defaultdict | |
| from pathlib import Path | |
| from torch.utils.data import Dataset, DataLoader | |
| from torchvision import transforms as T | |
| from tqdm import tqdm | |
| from types import SimpleNamespace | |
| from detectron2.checkpoint import DetectionCheckpointer | |
| from detectron2.data import MetadataCatalog | |
| from detectron2.utils.visualizer import Visualizer | |
| import config | |
| import utils as ut | |
| from eval_utils import MaskMerger | |
| from mask_former_trainer import setup, Trainer | |
| def load_model_cfg(dataset=None): | |
| args = SimpleNamespace(config_file='configs/maskformer/maskformer_R50_bs16_160k_dino.yaml', opts=["MODEL.DEVICE", "cpu", "GWM.DATASET", dataset], wandb_sweep_mode=False, resume_path=str('checkpoints/checkpoint_best.pth'), eval_only=True) | |
| cfg = setup(args) | |
| random_state = ut.random_state.PytorchRNGState(seed=cfg.SEED).to(torch.device(cfg.MODEL.DEVICE)) | |
| model = Trainer.build_model(cfg) | |
| checkpointer = DetectionCheckpointer(model, | |
| random_state=random_state, | |
| save_dir=None) | |
| checkpoint_path = 'checkpoints/checkpoint_best.pth' | |
| checkpoint = checkpointer.resume_or_load(checkpoint_path, resume=False) | |
| model.eval() | |
| return model, cfg | |
| def edgeness(masks): | |
| em = torch.zeros(1, masks.shape[-2], masks.shape[-1], device=masks.device) | |
| lm = em.clone() | |
| lm[..., :2] = 1. | |
| rm = em.clone() | |
| rm[...,-2:] = 1. | |
| tm = em.clone() | |
| tm[..., :2, :] = 1. | |
| bm = em.clone() | |
| bm[..., -2:,:] = 1. | |
| one = torch.tensor(1.,dtype= masks.dtype, device=masks.device) | |
| l = (masks * lm).flatten(-2).sum(-1) / lm.sum() | |
| l = torch.where(l > 0.3, one, l) | |
| r = (masks * rm).flatten(-2).sum(-1) / rm.sum() | |
| r = torch.where(r > 0.3, one, r) | |
| t = (masks * tm).flatten(-2).sum(-1) / tm.sum() | |
| t = torch.where(t > 0.3, one, t) | |
| b = (masks * bm).flatten(-2).sum(-1) / bm.sum() | |
| b = torch.where(b > 0.3, one, b) | |
| return (l + r + t + b ) | |
| def expand2sizedivisible(pil_img, background_color, size_divisibility): | |
| width, height = pil_img.size | |
| if width % size_divisibility == 0 and height % size_divisibility == 0: | |
| return pil_img | |
| result = Image.new(pil_img.mode, (width + (size_divisibility - width%size_divisibility)%size_divisibility, height + (size_divisibility - height%size_divisibility)%size_divisibility), background_color) | |
| result.paste(pil_img, (((size_divisibility - width%size_divisibility)%size_divisibility) // 2, ((size_divisibility - height%size_divisibility)%size_divisibility) // 2)) | |
| return result | |
| def cropfromsizedivisible(img, size_divisibility, orig_size): | |
| height, width = img.shape[:2] | |
| owidth, oheight = orig_size | |
| result = img[(height-oheight)//2:oheight+(height-oheight)//2, (width-owidth)//2:owidth+(width-owidth)//2] | |
| return result | |
| def evaluate_image(image_path): | |
| binary_threshold = 0.5 | |
| metadata = MetadataCatalog.get("__unused") | |
| image_pil = Image.open(image_path).convert('RGB') | |
| image_pil.thumbnail((384, 384)) | |
| osize = image_pil.size | |
| if cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY > 0: | |
| image_pil = expand2sizedivisible(image_pil, 0, cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY) | |
| image = np.asarray(image_pil) | |
| image_pt = torch.from_numpy(np.array(image)).permute(2,0,1) | |
| with torch.no_grad(): | |
| sample = [{'rgb': image_pt}] | |
| preds = model.forward_base(sample, keys=['rgb'], get_eval=True) | |
| masks_raw = torch.stack([x['sem_seg'] for x in preds], 0) | |
| K = masks_raw.shape[1] | |
| if K > 2: | |
| masks_softmaxed = torch.softmax(masks_raw, dim=1) | |
| masks_dict = merger(sample, masks_softmaxed) | |
| K = 2 | |
| masks = masks_dict['cos'] | |
| else: | |
| print(K) | |
| masks = masks_raw.softmax(1) | |
| masks_raw = F.interpolate(masks, size=(image_pt.shape[-2], image_pt.shape[-1]), mode='bilinear') # t s 1 h w | |
| bg = edgeness(masks_raw)[0].argmax().item() | |
| masks = masks_raw[0] > binary_threshold | |
| frame_visualizer = Visualizer(image, metadata) | |
| out = frame_visualizer.overlay_instances( | |
| masks=masks[[int(bg==0)]], | |
| alpha=0.3, | |
| assigned_colors=[(1,0,1)] | |
| ).get_image() | |
| return cropfromsizedivisible(out, cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY, osize) | |
| paths = sorted(pathlib.Path('samples').glob('*.jpg')) | |
| css = ".component-1 {height: 256px !important;}" | |
| model, cfg = load_model_cfg("DAVIS") | |
| merger = MaskMerger(cfg, model, merger_model="dino_vitb8") | |
| demo = gr.Interface( | |
| fn=evaluate_image, | |
| inputs=gr.Image(label='Image', type='filepath'), | |
| outputs=gr.Image(label='Annotated Image', type='numpy'), | |
| examples=[[path.as_posix(), 0.15, 6] for path in paths], | |
| title="Guess What Moves", | |
| description="#### Unsupervised image segmentation mode of [Guess What Moves](https://www.robots.ox.ac.uk/~vgg/research/gwm/)", | |
| css=css) | |
| demo.queue().launch() | |