File size: 5,323 Bytes
c7f0cc1
 
 
 
 
 
 
 
 
 
 
6119602
0ba37f3
c7f0cc1
2a136c4
c7f0cc1
 
 
0ba37f3
f176972
c7f0cc1
 
1e03b30
c7f0cc1
 
2a136c4
 
c7f0cc1
f2f297f
c7f0cc1
f2f297f
c7f0cc1
 
2a136c4
 
 
c7f0cc1
 
678d96b
 
 
 
 
 
 
 
 
 
 
 
 
 
c7f0cc1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0ba37f3
 
 
 
 
c7f0cc1
0ba37f3
 
 
 
 
 
 
 
c7f0cc1
 
 
 
 
 
0ba37f3
 
 
c7f0cc1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0ba37f3
c7f0cc1
0ba37f3
c7f0cc1
 
 
 
 
 
 
 
 
 
 
 
 
1e03b30
c7f0cc1
 
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
153
154
155
from __future__ import annotations

import argparse
import os
import pathlib
import subprocess

if os.getenv('SYSTEM') == 'spaces':
    import mim

    mim.uninstall('mmcv-full', confirm_yes=True)
    mim.install('mmcv-full==1.5.2', is_yes=True)


    subprocess.call('pip install --upgrade pip'.split())
    subprocess.call('pip uninstall -y opencv-python'.split())
    subprocess.call('pip uninstall -y opencv-python-headless'.split())
    subprocess.call('pip install opencv-python-headless==4.5.5.64'.split())
    subprocess.call('pip install pycocotools'.split())
    subprocess.call("pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cpu/torch1.10/index.html".split())

import cv2
import gradio as gr
import numpy as np

print('60\% \imported')

from mmdet.apis import init_detector, inference_detector
print('70\% \imported')
from utils import show_result
print('80\% \imported')
from mmcv import Config

print('100\% \imported')


DESCRIPTION = '''# OpenPSG

This is an official demo for [OpenPSG](https://github.com/Jingkang50/OpenPSG). 

News: The PSG Challenge is NOW available on International Algorithm Case Competition and soon ECCV'22 SenseHuman Workshop! Prize pool πŸ€‘ US$150K πŸ€‘! 

Check out our [GitHub repo](https://github.com/Jingkang50/OpenPSG) and [official website](http://psgdataset.org/) for more details.

<div class="row">
  <div class="column">
    <img id="logo" src="https://camo.githubusercontent.com/880346b66831a8212074787ba9a2301b4d700bd8f765ca11e4845ac0ab34c230/68747470733a2f2f6c6976652e737461746963666c69636b722e636f6d2f36353533352f35323139333837393637375f373531613465306237395f6b2e6a7067" alt="logo" style="width:100%">
  </div>
  <div class="column">
    <img id="visualzation" src="https://github.com/Jingkang50/OpenPSG/blob/main/assets/psgtr_long.gif?raw=true" alt="visualzation" style="width:100%">
  </div>
</div>
'''
FOOTER = '<img id="visitor-badge" src="https://visitor-badge.glitch.me/badge?page_id=c-liangyu.openpsg" alt="visitor badge" />'


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', type=str, default='cpu')
    parser.add_argument('--theme', type=str)
    parser.add_argument('--share', action='store_true')
    parser.add_argument('--port', type=int)
    parser.add_argument('--disable-queue',
                        dest='enable_queue',
                        action='store_false')
    return parser.parse_args()


def update_input_image(image: np.ndarray) -> dict:
    if image is None:
        return gr.Image.update(value=None)
    scale = 1500 / max(image.shape[:2])
    if scale < 1:
        image = cv2.resize(image, None, fx=scale, fy=scale)
    return gr.Image.update(value=image)


def set_example_image(example: list) -> dict:
    return gr.Image.update(value=example[0])

class Model:
    def __init__(self, model_name, device='cpu'):
        model_ckt ='OpenPSG/checkpoints/epoch_60.pth'
        cfg = Config.fromfile('OpenPSG/configs/psgtr/psgtr_r50_psg_inference.py')
        self.model = init_detector(cfg, model_ckt, device=device)

    def infer(self, input_image, num_rel):
        result = inference_detector(self.model, input_image)
        return show_result(input_image,
                            result,
                            is_one_stage=True,
                            num_rel=num_rel,
                            show=True
                            )


def main():
    args = parse_args()

    with gr.Blocks(theme=args.theme, css='style.css') as demo:
        
        model = Model('psgtr', device=args.device)
        
        gr.Markdown(DESCRIPTION)

        with gr.Row():
            with gr.Column():
                with gr.Row():
                    input_image = gr.Image(label='Input Image', type='numpy')
                with gr.Group():
                    with gr.Row():
                        num_rel = gr.Slider(
                            5,
                            100,
                            step=5,
                            value=20,
                            label='Number of Relations')
                with gr.Row():
                    run_button = gr.Button(value='Run')
            with gr.Column():
                with gr.Row():
                    result = gr.Gallery(label='Result', type='numpy')

        with gr.Row():
            paths = sorted(pathlib.Path('images').rglob('*.jpg'))
            example_images = gr.Dataset(components=[input_image],
                                        samples=[[path.as_posix()]
                                                 for path in paths])

        gr.Markdown(FOOTER)

        input_image.change(fn=update_input_image,
                           inputs=input_image,
                           outputs=input_image)
        
        run_button.click(fn=model.infer,
                         inputs=[
                            input_image, num_rel
                         ],
                         outputs=result)

        example_images.click(fn=set_example_image,
                             inputs=example_images,
                             outputs=input_image)

    demo.launch(
        enable_queue=args.enable_queue,
        server_port=args.port,
        share=args.share,
    )


if __name__ == '__main__':
    main()