File size: 2,783 Bytes
5395ccf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e4e47d6
5395ccf
 
 
 
 
900e3ff
5395ccf
 
900e3ff
5395ccf
 
 
 
 
 
 
 
4579715
6b8f616
 
5395ccf
 
6b8f616
5395ccf
 
 
 
 
 
e4e47d6
6b8f616
6e5b173
5395ccf
 
 
 
 
 
 
 
 
e324aa4
 
 
c0ca1d8
e4e47d6
5395ccf
c0ca1d8
 
 
 
5395ccf
 
 
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
#try:
#    import detectron2
#except:
import os 
os.system('pip install git+https://github.com/SysCV/transfiner.git')

from matplotlib.pyplot import axis
import gradio as gr
import requests
import numpy as np
from torch import nn
import requests

import torch

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog


model_name='./configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x_4gpu_transfiner.yaml'

# model = model_zoo.get(model_name, trained=True)

cfg = get_cfg()
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(model_name)
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as w ell
cfg.MODEL.WEIGHTS = './output_3x_transfiner_r50.pth'

if not torch.cuda.is_available():
    cfg.MODEL.DEVICE='cpu'

predictor = DefaultPredictor(cfg)


def inference(image):
    img = np.asarray(image)

    #img = np.array(image)
    outputs = predictor(img)

    v = Visualizer(img, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]))
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    
    return out.get_image()



title = "Mask Transfiner R50 model"
description = "Demo for <a href='https://arxiv.org/abs/2111.13673'>Mask Transfiner for High-Quality Instance Segmentation, CVPR 2022</a> based on R50-FPN. To use it, simply upload your image, or click one of the examples to load them. It runs in the cpu environment provided by Hugging Face. Read more at the links below."
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2111.13673'>Mask Transfiner for High-Quality Instance Segmentation, CVPR 2022</a> | <a href='https://github.com/SysCV/transfiner'>Mask Transfiner Github</a></p>"

gr.Interface(
    inference, 
    [gr.inputs.Image(type="pil", label="Input")], 
    gr.outputs.Image(type="numpy", label="Output"),
    title=title,
    description=description,
    article=article,
    examples=[
            ["demo/sample_imgs/000000131444.jpg"],
            ["demo/sample_imgs/000000157365.jpg"],
            ["demo/sample_imgs/000000286849.jpg"],
            ["demo/sample_imgs/000000176037.jpg"],
            ["demo/sample_imgs/000000018737.jpg"],
            ["demo/sample_imgs/000000224200.jpg"],
            ["demo/sample_imgs/000000558073.jpg"],
            ["demo/sample_imgs/000000404922.jpg"],
            ["demo/sample_imgs/000000252776.jpg"],
            ["demo/sample_imgs/000000482477.jpg"],
            ["demo/sample_imgs/000000344909.jpg"]
        ]).launch()