Spaces:
Runtime error
Runtime error
File size: 5,535 Bytes
5e88f62 d08ffbd 5e88f62 0ade859 6f13b9a 5e88f62 7151f8e 5e88f62 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 |
#!/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()
|