wjf5203 commited on
Commit
317822d
1 Parent(s): 730a8cb

add video func support

Browse files
Files changed (5) hide show
  1. .gitignore +151 -0
  2. GLEE/configs/vos_v0.yml +72 -0
  3. GLEE_vos_r50.pth +3 -0
  4. app.py +1185 -492
  5. app_v1.py +492 -0
.gitignore ADDED
@@ -0,0 +1,151 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+ .DS_Store
6
+ */.DS_Store
7
+
8
+ # Distribution / packaging
9
+ .Python
10
+ build/
11
+ develop-eggs/
12
+ dist/
13
+ downloads/
14
+ eggs/
15
+ .eggs/
16
+ lib/
17
+ lib64/
18
+ parts/
19
+ sdist/
20
+ var/
21
+ wheels/
22
+ share/python-wheels/
23
+ *.egg-info/
24
+ .installed.cfg
25
+ *.egg
26
+ MANIFEST
27
+
28
+ # PyInstaller
29
+ # Usually these files are written by a python script from a template
30
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
31
+ *.manifest
32
+ *.spec
33
+
34
+ # Installer logs
35
+ pip-log.txt
36
+ pip-delete-this-directory.txt
37
+
38
+ # Unit test / coverage reports
39
+ htmlcov/
40
+ .tox/
41
+ .nox/
42
+ .coverage
43
+ .coverage.*
44
+ .cache
45
+ nosetests.xml
46
+ coverage.xml
47
+ *.cover
48
+ *.py,cover
49
+ .hypothesis/
50
+ .pytest_cache/
51
+ cover/
52
+
53
+ # Translations
54
+ *.mo
55
+ *.pot
56
+
57
+ # Django stuff:
58
+ *.log
59
+ local_settings.py
60
+ db.sqlite3
61
+ db.sqlite3-journal
62
+
63
+ # Flask stuff:
64
+ instance/
65
+ .webassets-cache
66
+
67
+ # Scrapy stuff:
68
+ .scrapy
69
+
70
+ # Sphinx documentation
71
+ docs/_build/
72
+
73
+ # PyBuilder
74
+ .pybuilder/
75
+ target/
76
+
77
+ # Jupyter Notebook
78
+ .ipynb_checkpoints
79
+
80
+ # IPython
81
+ profile_default/
82
+ ipython_config.py
83
+
84
+ # pyenv
85
+ # For a library or package, you might want to ignore these files since the code is
86
+ # intended to run in multiple environments; otherwise, check them in:
87
+ # .python-version
88
+
89
+ # pipenv
90
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
91
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
92
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
93
+ # install all needed dependencies.
94
+ #Pipfile.lock
95
+
96
+ # poetry
97
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
98
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
99
+ # commonly ignored for libraries.
100
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
101
+ #poetry.lock
102
+
103
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow
104
+ __pypackages__/
105
+
106
+ # Celery stuff
107
+ celerybeat-schedule
108
+ celerybeat.pid
109
+
110
+ # SageMath parsed files
111
+ *.sage.py
112
+
113
+ # Environments
114
+ .env
115
+ .venv
116
+ env/
117
+ venv/
118
+ ENV/
119
+ env.bak/
120
+ venv.bak/
121
+
122
+ # Spyder project settings
123
+ .spyderproject
124
+ .spyproject
125
+
126
+ # Rope project settings
127
+ .ropeproject
128
+
129
+ # mkdocs documentation
130
+ /site
131
+
132
+ # mypy
133
+ .mypy_cache/
134
+ .dmypy.json
135
+ dmypy.json
136
+
137
+ # Pyre type checker
138
+ .pyre/
139
+
140
+ # pytype static type analyzer
141
+ .pytype/
142
+
143
+ # Cython debug symbols
144
+ cython_debug/
145
+
146
+ # PyCharm
147
+ # JetBrains specific template is maintainted in a separate JetBrains.gitignore that can
148
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
149
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
150
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
151
+ #.idea/
GLEE/configs/vos_v0.yml ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MODEL:
2
+ META_ARCHITECTURE: "STAnything"
3
+ MASK_ON: True
4
+ VISUAL_PROMPT: True
5
+ BACKBONE:
6
+ FREEZE_AT: 0
7
+ NAME: "build_resnet_backbone"
8
+ WEIGHTS: "weights/bert_r50_coco.pth"
9
+ PIXEL_MEAN: [123.675, 116.280, 103.530]
10
+ PIXEL_STD: [58.395, 57.120, 57.375]
11
+ RESNETS:
12
+ DEPTH: 50
13
+ STEM_TYPE: "basic" # not used
14
+ STEM_OUT_CHANNELS: 64
15
+ STRIDE_IN_1X1: False
16
+ OUT_FEATURES: ["res2", "res3", "res4", "res5"]
17
+ # NORM: "SyncBN"
18
+ RES5_MULTI_GRID: [1, 1, 1] # not used
19
+ SEM_SEG_HEAD:
20
+ NAME: "MaskDINOHead"
21
+ IGNORE_VALUE: 255
22
+ NUM_CLASSES: 80
23
+ LOSS_WEIGHT: 1.0
24
+ CONVS_DIM: 256
25
+ MASK_DIM: 256
26
+ NORM: "GN"
27
+ # pixel decoder
28
+ PIXEL_DECODER_NAME: "MaskDINOEncoder"
29
+ DIM_FEEDFORWARD: 1024
30
+ NUM_FEATURE_LEVELS: 3
31
+ TOTAL_NUM_FEATURE_LEVELS: 3
32
+ IN_FEATURES: ["res2", "res3", "res4", "res5"]
33
+ DEFORMABLE_TRANSFORMER_ENCODER_IN_FEATURES: ["res3", "res4", "res5"]
34
+ COMMON_STRIDE: 4
35
+ TRANSFORMER_ENC_LAYERS: 6
36
+ MaskDINO:
37
+ TRANSFORMER_DECODER_NAME: "MaskDINODecoder"
38
+ DEEP_SUPERVISION: True
39
+ NO_OBJECT_WEIGHT: 0.1
40
+ CLASS_WEIGHT: 4.0
41
+ MASK_WEIGHT: 5.0
42
+ DICE_WEIGHT: 5.0
43
+ BOX_WEIGHT: 5.0
44
+ GIOU_WEIGHT: 2.0
45
+ HIDDEN_DIM: 256
46
+ NUM_OBJECT_QUERIES: 300
47
+ NHEADS: 8
48
+ DROPOUT: 0.0
49
+ DIM_FEEDFORWARD: 2048
50
+ ENC_LAYERS: 0
51
+ PRE_NORM: False
52
+ ENFORCE_INPUT_PROJ: False
53
+ SIZE_DIVISIBILITY: 32
54
+ DEC_LAYERS: 9 # 9+1, 9 decoder layers, add one for the loss on learnable query
55
+ TRAIN_NUM_POINTS: 12544
56
+ OVERSAMPLE_RATIO: 3.0
57
+ IMPORTANCE_SAMPLE_RATIO: 0.75
58
+ INITIAL_PRED: True
59
+ TWO_STAGE: True
60
+ DN: "standard"
61
+ DN_NUM: 100
62
+ INITIALIZE_BOX_TYPE: "no"
63
+ TEST:
64
+ SEMANTIC_ON: False
65
+ INSTANCE_ON: True
66
+ PANOPTIC_ON: False
67
+ OVERLAP_THRESHOLD: 0.8
68
+ OBJECT_MASK_THRESHOLD: 0.25
69
+ TEXT:
70
+ ARCH: clip_teacher
71
+ LANGUAGE_BACKBONE:
72
+ LANG_DIM: 512
GLEE_vos_r50.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:15d7b06917287d47ce8402d08f56826360a5deebd2e2f11bc31264e8084457b8
3
+ size 730158966
app.py CHANGED
@@ -1,492 +1,1185 @@
1
- try:
2
- import detectron2
3
- except:
4
- import os
5
- os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
6
- os.system('cd GLEE/glee/models/pixel_decoder/ops && sh mask.sh')
7
- # os.system('python -m pip install -e detectron2')
8
-
9
- import gradio as gr
10
- import numpy as np
11
- import cv2
12
- import torch
13
-
14
- from detectron2.config import get_cfg
15
- from GLEE.glee.models.glee_model import GLEE_Model
16
- from GLEE.glee.config_deeplab import add_deeplab_config
17
- from GLEE.glee.config import add_glee_config
18
- import torch.nn.functional as F
19
- import torchvision
20
- import math
21
- from obj365_name import categories as OBJ365_CATEGORIESV2
22
-
23
-
24
- print(f"Is CUDA available: {torch.cuda.is_available()}")
25
- # True
26
- if torch.cuda.is_available():
27
- print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
28
- # Tesla T4
29
-
30
- def box_cxcywh_to_xyxy(x):
31
- x_c, y_c, w, h = x.unbind(-1)
32
- b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
33
- (x_c + 0.5 * w), (y_c + 0.5 * h)]
34
- return torch.stack(b, dim=-1)
35
-
36
-
37
-
38
- def scribble2box(img):
39
- if img.max()==0:
40
- return None, None
41
- rows = np.any(img, axis=1)
42
- cols = np.any(img, axis=0)
43
- all = np.any(img,axis=2)
44
- R,G,B,A = img[np.where(all)[0][0],np.where(all)[1][0]].tolist() # get color
45
- ymin, ymax = np.where(rows)[0][[0, -1]]
46
- xmin, xmax = np.where(cols)[0][[0, -1]]
47
- return np.array([ xmin,ymin, xmax,ymax]), (R,G,B)
48
-
49
-
50
- def LSJ_box_postprocess( out_bbox, padding_size, crop_size, img_h, img_w):
51
- # postprocess box height and width
52
- boxes = box_cxcywh_to_xyxy(out_bbox)
53
- lsj_sclae = torch.tensor([padding_size[1], padding_size[0], padding_size[1], padding_size[0]]).to(out_bbox)
54
- crop_scale = torch.tensor([crop_size[1], crop_size[0], crop_size[1], crop_size[0]]).to(out_bbox)
55
- boxes = boxes * lsj_sclae
56
- boxes = boxes / crop_scale
57
- boxes = torch.clamp(boxes,0,1)
58
-
59
- scale_fct = torch.tensor([img_w, img_h, img_w, img_h])
60
- scale_fct = scale_fct.to(out_bbox)
61
- boxes = boxes * scale_fct
62
- return boxes
63
-
64
- COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
65
- [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933],
66
- [0.494, 0.000, 0.556], [0.494, 0.000, 0.000], [0.000, 0.745, 0.000],
67
- [0.700, 0.300, 0.600],[0.000, 0.447, 0.741], [0.850, 0.325, 0.098]]
68
-
69
-
70
-
71
- coco_class_name = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
72
- OBJ365_class_names = [cat['name'] for cat in OBJ365_CATEGORIESV2]
73
- class_agnostic_name = ['object']
74
-
75
- if torch.cuda.is_available():
76
- print('use cuda')
77
- device = 'cuda'
78
- else:
79
- print('use cpu')
80
- device='cpu'
81
-
82
- cfg_r50 = get_cfg()
83
- add_deeplab_config(cfg_r50)
84
- add_glee_config(cfg_r50)
85
- conf_files_r50 = 'GLEE/configs/R50.yaml'
86
- checkpoints_r50 = torch.load('GLEE_R50_Scaleup10m.pth')
87
- cfg_r50.merge_from_file(conf_files_r50)
88
- GLEEmodel_r50 = GLEE_Model(cfg_r50, None, device, None, True).to(device)
89
- GLEEmodel_r50.load_state_dict(checkpoints_r50, strict=False)
90
- GLEEmodel_r50.eval()
91
-
92
-
93
- cfg_swin = get_cfg()
94
- add_deeplab_config(cfg_swin)
95
- add_glee_config(cfg_swin)
96
- conf_files_swin = 'GLEE/configs/SwinL.yaml'
97
- checkpoints_swin = torch.load('GLEE_SwinL_Scaleup10m.pth')
98
- cfg_swin.merge_from_file(conf_files_swin)
99
- GLEEmodel_swin = GLEE_Model(cfg_swin, None, device, None, True).to(device)
100
- GLEEmodel_swin.load_state_dict(checkpoints_swin, strict=False)
101
- GLEEmodel_swin.eval()
102
-
103
- pixel_mean = torch.Tensor( [123.675, 116.28, 103.53]).to(device).view(3, 1, 1)
104
- pixel_std = torch.Tensor([58.395, 57.12, 57.375]).to(device).view(3, 1, 1)
105
- normalizer = lambda x: (x - pixel_mean) / pixel_std
106
- inference_size = 800
107
- inference_type = 'resize_shot' # or LSJ
108
- size_divisibility = 32
109
-
110
- FONT_SCALE = 1.5e-3
111
- THICKNESS_SCALE = 1e-3
112
- TEXT_Y_OFFSET_SCALE = 1e-2
113
-
114
-
115
- if inference_type != 'LSJ':
116
- resizer = torchvision.transforms.Resize(inference_size)
117
-
118
-
119
- def segment_image(img,prompt_mode, categoryname, custom_category, expressiong, results_select, num_inst_select, threshold_select, mask_image_mix_ration, model_selection):
120
- if model_selection == 'GLEE-Plus (SwinL)':
121
- GLEEmodel = GLEEmodel_swin
122
- print('use GLEE-Plus')
123
- else:
124
- GLEEmodel = GLEEmodel_r50
125
- print('use GLEE-Lite')
126
-
127
- copyed_img = img['background'][:,:,:3].copy()
128
-
129
- ori_image = torch.as_tensor(np.ascontiguousarray( copyed_img.transpose(2, 0, 1)))
130
- ori_image = normalizer(ori_image.to(device))[None,]
131
- _,_, ori_height, ori_width = ori_image.shape
132
-
133
- if inference_type == 'LSJ':
134
- infer_image = torch.zeros(1,3,1024,1024).to(ori_image)
135
- infer_image[:,:,:inference_size,:inference_size] = ori_image
136
- else:
137
- resize_image = resizer(ori_image)
138
- image_size = torch.as_tensor((resize_image.shape[-2],resize_image.shape[-1]))
139
- re_size = resize_image.shape[-2:]
140
- if size_divisibility > 1:
141
- stride = size_divisibility
142
- # the last two dims are H,W, both subject to divisibility requirement
143
- padding_size = ((image_size + (stride - 1)).div(stride, rounding_mode="floor") * stride).tolist()
144
- infer_image = torch.zeros(1,3,padding_size[0],padding_size[1]).to(resize_image)
145
- infer_image[0,:,:image_size[0],:image_size[1]] = resize_image
146
- # reversed_image = infer_image*pixel_std + pixel_mean
147
- # reversed_image = torch.clip(reversed_image,min=0,max=255)
148
- # reversed_image = reversed_image[0].permute(1,2,0)
149
- # reversed_image = reversed_image.int().cpu().numpy().copy()
150
- # cv2.imwrite('test.png',reversed_image[:,:,::-1])
151
-
152
-
153
- if prompt_mode == 'categories' or prompt_mode == 'expression':
154
- if len(results_select)==0:
155
- results_select=['box']
156
- if prompt_mode == 'categories':
157
- if categoryname =="COCO-80":
158
- batch_category_name = coco_class_name
159
- elif categoryname =="OBJ365":
160
- batch_category_name = OBJ365_class_names
161
- elif categoryname =="Custom-List":
162
- batch_category_name = custom_category.split(',')
163
- else:
164
- batch_category_name = class_agnostic_name
165
-
166
- # mask_ori = torch.from_numpy(np.load('03_moto_mask.npy'))[None,]
167
- # mask_ori = (F.interpolate(mask_ori, (height, width), mode='bilinear') > 0).to(device)
168
- # prompt_list = [mask_ori[0]]
169
- prompt_list = []
170
- with torch.no_grad():
171
- (outputs,_) = GLEEmodel(infer_image, prompt_list, task="coco", batch_name_list=batch_category_name, is_train=False)
172
- topK_instance = max(num_inst_select,1)
173
- else:
174
- topK_instance = 1
175
- prompt_list = {'grounding':[expressiong]}
176
- with torch.no_grad():
177
- (outputs,_) = GLEEmodel(infer_image, prompt_list, task="grounding", batch_name_list=[], is_train=False)
178
-
179
-
180
- mask_pred = outputs['pred_masks'][0]
181
- mask_cls = outputs['pred_logits'][0]
182
- boxes_pred = outputs['pred_boxes'][0]
183
-
184
- scores = mask_cls.sigmoid().max(-1)[0]
185
- scores_per_image, topk_indices = scores.topk(topK_instance, sorted=True)
186
- if prompt_mode == 'categories':
187
- valid = scores_per_image>threshold_select
188
- topk_indices = topk_indices[valid]
189
- scores_per_image = scores_per_image[valid]
190
-
191
- pred_class = mask_cls[topk_indices].max(-1)[1].tolist()
192
- pred_boxes = boxes_pred[topk_indices]
193
-
194
-
195
- boxes = LSJ_box_postprocess(pred_boxes,padding_size,re_size, ori_height,ori_width)
196
- mask_pred = mask_pred[topk_indices]
197
- pred_masks = F.interpolate( mask_pred[None,], size=(padding_size[0], padding_size[1]), mode="bilinear", align_corners=False )
198
- pred_masks = pred_masks[:,:,:re_size[0],:re_size[1]]
199
- pred_masks = F.interpolate( pred_masks, size=(ori_height,ori_width), mode="bilinear", align_corners=False )
200
- pred_masks = (pred_masks>0).detach().cpu().numpy()[0]
201
-
202
- if 'mask' in results_select:
203
-
204
- zero_mask = np.zeros_like(copyed_img)
205
- for nn, mask in enumerate(pred_masks):
206
- # mask = mask.numpy()
207
- mask = mask.reshape(mask.shape[0], mask.shape[1], 1)
208
-
209
- lar = np.concatenate((mask*COLORS[nn%12][2], mask*COLORS[nn%12][1], mask*COLORS[nn%12][0]), axis = 2)
210
- zero_mask = zero_mask+ lar
211
-
212
-
213
- lar_valid = zero_mask>0
214
- masked_image = lar_valid*copyed_img
215
- img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,1)*255*(1-mask_image_mix_ration)
216
- max_p = img_n.max()
217
- img_n = 255*img_n/max_p
218
- ret = (~lar_valid*copyed_img)*mask_image_mix_ration + img_n
219
- ret = ret.astype('uint8')
220
- else:
221
- ret = copyed_img
222
-
223
- if 'box' in results_select:
224
-
225
- line_width = max(ret.shape) /200
226
-
227
- for nn,(classid, box) in enumerate(zip(pred_class,boxes)):
228
- x1,y1,x2,y2 = box.long().tolist()
229
- RGB = (COLORS[nn%12][2]*255,COLORS[nn%12][1]*255,COLORS[nn%12][0]*255)
230
- cv2.rectangle(ret, (x1,y1), (x2,y2), RGB, math.ceil(line_width) )
231
- if prompt_mode == 'categories' or (prompt_mode == 'expression' and 'expression' in results_select ):
232
- if prompt_mode == 'categories':
233
- label = ''
234
- if 'name' in results_select:
235
- label += batch_category_name[classid]
236
- if 'score' in results_select:
237
- label += str(scores_per_image[nn].item())[:4]
238
- else:
239
- label = expressiong
240
-
241
- if len(label)==0:
242
- continue
243
- height, width, _ = ret.shape
244
- FONT = cv2.FONT_HERSHEY_COMPLEX
245
- label_width, label_height = cv2.getTextSize(label, FONT, min(width, height) * FONT_SCALE, math.ceil(min(width, height) * THICKNESS_SCALE))[0]
246
-
247
- cv2.rectangle(ret, (x1,y1), (x1+label_width,(y1 -label_height) - int(height * TEXT_Y_OFFSET_SCALE)), RGB, -1)
248
-
249
- cv2.putText(
250
- ret,
251
- label,
252
- (x1, y1 - int(height * TEXT_Y_OFFSET_SCALE)),
253
- fontFace=FONT,
254
- fontScale=min(width, height) * FONT_SCALE,
255
- thickness=math.ceil(min(width, height) * THICKNESS_SCALE),
256
- color=(255,255,255),
257
- )
258
-
259
-
260
-
261
- ret = ret.astype('uint8')
262
- return ret
263
-
264
-
265
- else: #visual prompt
266
- topK_instance = 1
267
- copyed_img = img['background'][:,:,:3].copy()
268
- # get bbox from scribbles in layers
269
- bbox_list = [scribble2box(layer) for layer in img['layers'] ]
270
- visual_prompt_list = []
271
- visual_prompt_RGB_list = []
272
-
273
- for mask, (box,RGB) in zip(img['layers'], bbox_list):
274
- if box is None:
275
- continue
276
- if prompt_mode=='box':
277
- fakemask = np.zeros_like(copyed_img[:,:,0])
278
- x1 ,y1 ,x2, y2 = box
279
- fakemask[ y1:y2, x1:x2 ] = 1
280
- fakemask = fakemask>0
281
- elif prompt_mode=='point':
282
- fakemask = np.zeros_like(copyed_img[:,:,0])
283
- H,W = fakemask.shape
284
- x1 ,y1 ,x2, y2 = box
285
- center_x, center_y = (x1+x2)//2, (y1+y2)//2
286
- fakemask[ center_y-H//40:center_y+H//40, center_x-W//40:center_x+W//40 ] = 1
287
- fakemask = fakemask>0
288
- elif prompt_mode=='scribble':
289
- fakemask = mask[:,:,-1]
290
- fakemask = fakemask>0
291
-
292
- fakemask = torch.from_numpy(fakemask).unsqueeze(0).to(ori_image)
293
- if inference_type == 'LSJ':
294
- infer_visual_prompt = torch.zeros(1,1024,1024).to(ori_image)
295
- infer_visual_prompt[:,:inference_size,:inference_size] = fakemask
296
- else:
297
- resize_fakemask = resizer(fakemask)
298
- if size_divisibility > 1:
299
- # the last two dims are H,W, both subject to divisibility requirement
300
- infer_visual_prompt = torch.zeros(1,padding_size[0],padding_size[1]).to(resize_fakemask)
301
- infer_visual_prompt[:,:image_size[0],:image_size[1]] = resize_fakemask
302
-
303
-
304
- visual_prompt_list.append( infer_visual_prompt>0 )
305
- visual_prompt_RGB_list.append(RGB)
306
-
307
-
308
- mask_results_list = []
309
- for visual_prompt in visual_prompt_list:
310
-
311
- prompt_list = {'spatial':[visual_prompt]}
312
-
313
- with torch.no_grad():
314
- (outputs,_) = GLEEmodel(infer_image, prompt_list, task="coco", batch_name_list=['object'], is_train=False, visual_prompt_type=prompt_mode )
315
-
316
- mask_pred = outputs['pred_masks'][0]
317
- mask_cls = outputs['pred_logits'][0]
318
- boxes_pred = outputs['pred_boxes'][0]
319
-
320
- scores = mask_cls.sigmoid().max(-1)[0]
321
- scores_per_image, topk_indices = scores.topk(topK_instance, sorted=True)
322
-
323
- pred_class = mask_cls[topk_indices].max(-1)[1].tolist()
324
- pred_boxes = boxes_pred[topk_indices]
325
-
326
-
327
- boxes = LSJ_box_postprocess(pred_boxes,padding_size,re_size, ori_height,ori_width)
328
- mask_pred = mask_pred[topk_indices]
329
- pred_masks = F.interpolate( mask_pred[None,], size=(padding_size[0], padding_size[1]), mode="bilinear", align_corners=False )
330
- pred_masks = pred_masks[:,:,:re_size[0],:re_size[1]]
331
- pred_masks = F.interpolate( pred_masks, size=(ori_height,ori_width), mode="bilinear", align_corners=False )
332
- pred_masks = (pred_masks>0).detach().cpu().numpy()[0]
333
- mask_results_list.append(pred_masks)
334
-
335
- zero_mask = np.zeros_like(copyed_img)
336
- for mask,RGB in zip(mask_results_list,visual_prompt_RGB_list):
337
- mask = mask.reshape(mask.shape[-2], mask.shape[-1], 1)
338
- lar = np.concatenate((mask*RGB[0], mask*RGB[1],mask*RGB[2]), axis = 2)
339
- zero_mask = zero_mask+ lar
340
- lar_valid = zero_mask>0
341
- masked_image = lar_valid*copyed_img
342
- img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,255)*(1-mask_image_mix_ration)
343
- max_p = img_n.max()
344
- img_n = 255*img_n/max_p
345
- ret = (~lar_valid*copyed_img)*mask_image_mix_ration + img_n
346
- ret = ret.astype('uint8')
347
- # cv2.imwrite('00020_inst.jpg', cv2.cvtColor(ret, cv2.COLOR_BGR2RGB))
348
-
349
- return ret
350
-
351
-
352
-
353
- # def get_select_coordinates(img):
354
-
355
- # # img{'background': (H,W,3)
356
- # # 'layers': list[ (H,W,4(RGBA)) ], draw map
357
- # # 'composite': (H,W,4(RGBA))} ori_img concat drow
358
-
359
- # ori_img = img['background'][:,:,:3].copy()
360
-
361
- # # get bbox from scribbles in layers
362
- # bbox_list = [scribble2box(layer) for layer in img['layers'] ]
363
- # for mask, (box,RGB) in zip(img['layers'], bbox_list):
364
- # if box is None:
365
- # continue
366
- # cv2.rectangle(ori_img, (box[0],box[1]), (box[2],box[3]),RGB, 3)
367
- # return ori_img
368
-
369
- def visual_prompt_preview(img, prompt_mode):
370
-
371
- copyed_img = img['background'][:,:,:3].copy()
372
-
373
- # get bbox from scribbles in layers
374
- bbox_list = [scribble2box(layer) for layer in img['layers'] ]
375
- zero_mask = np.zeros_like(copyed_img)
376
-
377
- for mask, (box,RGB) in zip(img['layers'], bbox_list):
378
- if box is None:
379
- continue
380
-
381
- if prompt_mode=='box':
382
- fakemask = np.zeros_like(copyed_img[:,:,0])
383
- x1 ,y1 ,x2, y2 = box
384
- fakemask[ y1:y2, x1:x2 ] = 1
385
- fakemask = fakemask>0
386
- elif prompt_mode=='point':
387
- fakemask = np.zeros_like(copyed_img[:,:,0])
388
- H,W = fakemask.shape
389
- x1 ,y1 ,x2, y2 = box
390
- center_x, center_y = (x1+x2)//2, (y1+y2)//2
391
- fakemask[ center_y-H//40:center_y+H//40, center_x-W//40:center_x+W//40 ] = 1
392
- fakemask = fakemask>0
393
- else:
394
- fakemask = mask[:,:,-1]
395
- fakemask = fakemask>0
396
-
397
- mask = fakemask.reshape(fakemask.shape[0], fakemask.shape[1], 1)
398
- lar = np.concatenate((mask*RGB[0], mask*RGB[1],mask*RGB[2]), axis = 2)
399
- zero_mask = zero_mask+ lar
400
-
401
-
402
- img_n = copyed_img + np.clip(zero_mask,0,255)
403
- max_p = img_n.max()
404
- ret = 255*img_n/max_p
405
- ret = ret.astype('uint8')
406
- return ret
407
-
408
-
409
-
410
- with gr.Blocks() as demo:
411
- gr.Markdown('# GLEE: General Object Foundation Model for Images and Videos at Scale')
412
- gr.Markdown('## [Paper](https://arxiv.org/abs/2312.09158) - [Project Page](https://glee-vision.github.io) - [Code](https://github.com/FoundationVision/GLEE) ')
413
-
414
- gr.Markdown(
415
- '**The functionality demonstration demo app of GLEE. Select a Tab for image or video tasks. Image tasks includes arbitrary vocabulary object detection&segmentation, any form of object name or object caption detection, referring expression comprehension, and interactive segmentation. Video tasks add object tracking functionality based on image tasks.**'
416
- )
417
-
418
-
419
- with gr.Tab("Image task"):
420
- with gr.Row():
421
- with gr.Column():
422
-
423
- img_input = gr.ImageEditor()
424
- model_select = gr.Dropdown(
425
- ["GLEE-Lite (R50)", "GLEE-Plus (SwinL)"], value = "GLEE-Plus (SwinL)" , multiselect=False, label="Model",
426
- )
427
- with gr.Row():
428
- with gr.Column():
429
- prompt_mode_select = gr.Radio(["point", "scribble", "box", "categories", "expression"], label="Prompt", value= "categories" , info="What kind of prompt do you want to use?")
430
- category_select = gr.Dropdown(
431
- ["COCO-80", "OBJ365", "Custom-List", "Class-Agnostic"], value = "COCO-80" , multiselect=False, label="Categories", info="Choose an existing category list or class-agnostic"
432
- )
433
- custom_category = gr.Textbox(
434
- label="Custom Category",
435
- info="Input custom category list, seperate by ',' ",
436
- lines=1,
437
- value="dog, cat, car, person",
438
- )
439
- input_expressiong = gr.Textbox(
440
- label="Expression",
441
- info="Input any description of an object in the image ",
442
- lines=2,
443
- value="the red car",
444
- )
445
- # with gr.Column():
446
- with gr.Group():
447
- with gr.Accordion("Text based detection usage",open=False):
448
- gr.Markdown(
449
- 'Press the "Detect & Segment" button directly to try the effect using the COCO category.<br />\
450
- GLEE supports three kind of object perception methods: category list, textual description, and class-agnostic.<br />\
451
- 1.Select an existing category list from the "Categories" dropdown, like COCO or OBJ365, or customize your own list.<br />\
452
- 2.Enter arbitrary object name in "Custom Category", or choose the expression model and describe the object in "Expression Textbox" for single object detection only.<br />\
453
- 3.For class-agnostic mode, choose "Class-Agnostic" from the "Categories" dropdown.'
454
- )
455
- with gr.Accordion("Interactive segmentation usage",open=False):
456
- gr.Markdown(
457
- 'For interactive segmentation:<br />\
458
- 1.Draw points, boxes, or scribbles on the canvas for multiclass segmentation; use separate layers for different objects, adding layers with a "+" sign.<br />\
459
- 2.Point mode accepts a single point only; multiple points default to the centroid, so use boxes or scribbles for larger objects.<br />\
460
- 3.After drawing, click green "" on the right side of the image to preview the prompt visualization; the segmentation mask follows the chosen prompt colors.'
461
- )
462
-
463
- img_showbox = gr.Image(label="visual prompt area preview")
464
-
465
-
466
-
467
-
468
- with gr.Column():
469
- image_segment = gr.Image(label="detection and segmentation results")
470
- with gr.Accordion("Try More Visualization Options"):
471
- results_select = gr.CheckboxGroup(["box", "mask", "name", "score", "expression"], value=["box", "mask", "name", "score"], label="Shown Results", info="The results shown on image")
472
- num_inst_select = gr.Slider(1, 50, value=15, step=1, label="Num of topK instances for category based detection", info="Choose between 1 and 50 for better visualization")
473
- threshold_select = gr.Slider(0, 1, value=0.2, label="Confidence Threshold", info="Choose threshold ")
474
- mask_image_mix_ration = gr.Slider(0, 1, value=0.45, label="Image Brightness Ratio", info="Brightness between image and colored masks ")
475
-
476
-
477
-
478
- image_button = gr.Button("Detect & Segment")
479
- img_input.change(visual_prompt_preview, inputs = [img_input,prompt_mode_select] , outputs = img_showbox)
480
- image_button.click(segment_image, inputs=[img_input, prompt_mode_select, category_select, custom_category,input_expressiong, results_select, num_inst_select, threshold_select, mask_image_mix_ration,model_select], outputs=image_segment)
481
-
482
-
483
- with gr.Tab("Video task"):
484
- with gr.Row():
485
- gr.Markdown(
486
- '# Due to computational resource limitations, support for video tasks is being processed and is expected to be available within a week.'
487
- )
488
- video_input = gr.Image()
489
- video_button = gr.Button("Segment&Track")
490
-
491
- if __name__ == '__main__':
492
- demo.launch(inbrowser=True,share=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ try:
2
+ import detectron2
3
+ except:
4
+ import os
5
+ os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
6
+ os.system('cd GLEE/glee/models/pixel_decoder/ops && sh mask.sh')
7
+ # os.system('python -m pip install -e detectron2')
8
+
9
+ import gradio as gr
10
+ import numpy as np
11
+ import cv2
12
+ import torch
13
+ from os import path
14
+ from detectron2.config import get_cfg
15
+ from GLEE.glee.models.glee_model import GLEE_Model
16
+ from GLEE.glee.config_deeplab import add_deeplab_config
17
+ from GLEE.glee.config import add_glee_config
18
+ import torch.nn.functional as F
19
+ import torchvision
20
+ import math
21
+ from scipy.optimize import linear_sum_assignment
22
+ from obj365_name import categories as OBJ365_CATEGORIESV2
23
+ import copy
24
+
25
+
26
+ this_dir = path.dirname(path.abspath(__file__))
27
+
28
+
29
+ print(f"Is CUDA available: {torch.cuda.is_available()}")
30
+ # True
31
+ if torch.cuda.is_available():
32
+ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
33
+ # Tesla T4
34
+
35
+ def box_cxcywh_to_xyxy(x):
36
+ x_c, y_c, w, h = x.unbind(-1)
37
+ b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
38
+ (x_c + 0.5 * w), (y_c + 0.5 * h)]
39
+ return torch.stack(b, dim=-1)
40
+
41
+
42
+
43
+ def scribble2box(img):
44
+ if img.max()==0:
45
+ return None, None
46
+ rows = np.any(img, axis=1)
47
+ cols = np.any(img, axis=0)
48
+ all = np.any(img,axis=2)
49
+ R,G,B,A = img[np.where(all)[0][0],np.where(all)[1][0]].tolist() # get color
50
+ ymin, ymax = np.where(rows)[0][[0, -1]]
51
+ xmin, xmax = np.where(cols)[0][[0, -1]]
52
+ return np.array([ xmin,ymin, xmax,ymax]), (R,G,B)
53
+
54
+
55
+ def LSJ_box_postprocess( out_bbox, padding_size, crop_size, img_h, img_w):
56
+ # postprocess box height and width
57
+ boxes = box_cxcywh_to_xyxy(out_bbox)
58
+ lsj_sclae = torch.tensor([padding_size[1], padding_size[0], padding_size[1], padding_size[0]]).to(out_bbox)
59
+ crop_scale = torch.tensor([crop_size[1], crop_size[0], crop_size[1], crop_size[0]]).to(out_bbox)
60
+ boxes = boxes * lsj_sclae
61
+ boxes = boxes / crop_scale
62
+ boxes = torch.clamp(boxes,0,1)
63
+
64
+ scale_fct = torch.tensor([img_w, img_h, img_w, img_h])
65
+ scale_fct = scale_fct.to(out_bbox)
66
+ boxes = boxes * scale_fct
67
+ return boxes
68
+
69
+ COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
70
+ [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933],
71
+ [0.494, 0.000, 0.556], [0.494, 0.000, 0.000], [0.000, 0.745, 0.000],
72
+ [0.700, 0.300, 0.600],[0.000, 0.447, 0.741], [0.850, 0.325, 0.098]]
73
+
74
+
75
+
76
+ coco_class_name = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
77
+ YTBVISOVIS_class_name = ['lizard', 'cat', 'horse', 'eagle', 'frog', 'Horse', 'monkey', 'bear', 'parrot', 'giant_panda', 'truck', 'zebra', 'rabbit', 'skateboard', 'tiger', 'shark', 'Person', 'Poultry', 'Zebra', 'Airplane', 'elephant', 'Elephant', 'Turtle', 'snake', 'train', 'Dog', 'snowboard', 'airplane', 'Lizard', 'dog', 'Cat', 'earless_seal', 'boat', 'Tiger', 'motorbike', 'duck', 'fox', 'Monkey', 'Bird', 'Bear', 'tennis_racket', 'Rabbit', 'Giraffe', 'Motorcycle', 'fish', 'Boat', 'deer', 'ape', 'Bicycle', 'Parrot', 'Cow', 'turtle', 'mouse', 'owl', 'Fish', 'surfboard', 'Giant_panda', 'Sheep', 'hand', 'Vehical', 'sedan', 'leopard', 'person', 'giraffe', 'cow']
78
+ OBJ365_class_names = [cat['name'] for cat in OBJ365_CATEGORIESV2]
79
+ class_agnostic_name = ['object']
80
+
81
+ if torch.cuda.is_available():
82
+ print('use cuda')
83
+ device = 'cuda'
84
+ else:
85
+ print('use cpu')
86
+ device='cpu'
87
+
88
+ cfg_r50 = get_cfg()
89
+ add_deeplab_config(cfg_r50)
90
+ add_glee_config(cfg_r50)
91
+ conf_files_r50 = 'GLEE/configs/R50.yaml'
92
+ checkpoints_r50 = torch.load('GLEE_R50_Scaleup10m.pth')
93
+ cfg_r50.merge_from_file(conf_files_r50)
94
+ GLEEmodel_r50 = GLEE_Model(cfg_r50, None, device, None, True).to(device)
95
+ GLEEmodel_r50.load_state_dict(checkpoints_r50, strict=False)
96
+ GLEEmodel_r50.eval()
97
+
98
+ cfg_vos = get_cfg()
99
+ add_deeplab_config(cfg_vos)
100
+ add_glee_config(cfg_vos)
101
+ conf_files_vos = 'GLEE/configs/vos_v0.yaml'
102
+ cfg_vos.merge_from_file(conf_files_vos)
103
+
104
+
105
+ cfg_swin = get_cfg()
106
+ add_deeplab_config(cfg_swin)
107
+ add_glee_config(cfg_swin)
108
+ conf_files_swin = 'GLEE/configs/SwinL.yaml'
109
+ checkpoints_swin = torch.load('GLEE_SwinL_Scaleup10m.pth')
110
+ cfg_swin.merge_from_file(conf_files_swin)
111
+ GLEEmodel_swin = GLEE_Model(cfg_swin, None, device, None, True).to(device)
112
+ GLEEmodel_swin.load_state_dict(checkpoints_swin, strict=False)
113
+ GLEEmodel_swin.eval()
114
+
115
+ pixel_mean = torch.Tensor( [123.675, 116.28, 103.53]).to(device).view(3, 1, 1)
116
+ pixel_std = torch.Tensor([58.395, 57.12, 57.375]).to(device).view(3, 1, 1)
117
+ normalizer = lambda x: (x - pixel_mean) / pixel_std
118
+ inference_size = 800
119
+ video_inference_size = 720
120
+ inference_type = 'resize_shot' # or LSJ
121
+ size_divisibility = 32
122
+
123
+ FONT_SCALE = 1.5e-3
124
+ THICKNESS_SCALE = 1e-3
125
+ TEXT_Y_OFFSET_SCALE = 1e-2
126
+
127
+
128
+ if inference_type != 'LSJ':
129
+ resizer = torchvision.transforms.Resize(inference_size,antialias=True)
130
+ videoresizer = torchvision.transforms.Resize(video_inference_size,antialias=True)
131
+
132
+
133
+ def segment_image(img, prompt_mode, categoryname, custom_category, expressiong, results_select, num_inst_select, threshold_select, mask_image_mix_ration, model_selection):
134
+ if model_selection == 'GLEE-Plus (SwinL)':
135
+ GLEEmodel = GLEEmodel_swin
136
+ print('use GLEE-Plus')
137
+ else:
138
+ GLEEmodel = GLEEmodel_r50
139
+ print('use GLEE-Lite')
140
+
141
+ copyed_img = img['background'][:,:,:3].copy()
142
+
143
+ ori_image = torch.as_tensor(np.ascontiguousarray( copyed_img.transpose(2, 0, 1)))
144
+ ori_image = normalizer(ori_image.to(device))[None,]
145
+ _,_, ori_height, ori_width = ori_image.shape
146
+
147
+ if inference_type == 'LSJ':
148
+ infer_image = torch.zeros(1,3,1024,1024).to(ori_image)
149
+ infer_image[:,:,:inference_size,:inference_size] = ori_image
150
+ else:
151
+ resize_image = resizer(ori_image)
152
+ image_size = torch.as_tensor((resize_image.shape[-2],resize_image.shape[-1]))
153
+ re_size = resize_image.shape[-2:]
154
+ if size_divisibility > 1:
155
+ stride = size_divisibility
156
+ # the last two dims are H,W, both subject to divisibility requirement
157
+ padding_size = ((image_size + (stride - 1)).div(stride, rounding_mode="floor") * stride).tolist()
158
+ infer_image = torch.zeros(1,3,padding_size[0],padding_size[1]).to(resize_image)
159
+ infer_image[0,:,:image_size[0],:image_size[1]] = resize_image
160
+ # reversed_image = infer_image*pixel_std + pixel_mean
161
+ # reversed_image = torch.clip(reversed_image,min=0,max=255)
162
+ # reversed_image = reversed_image[0].permute(1,2,0)
163
+ # reversed_image = reversed_image.int().cpu().numpy().copy()
164
+ # cv2.imwrite('test.png',reversed_image[:,:,::-1])
165
+
166
+
167
+ if prompt_mode == 'categories' or prompt_mode == 'expression':
168
+ if len(results_select)==0:
169
+ results_select=['box']
170
+ if prompt_mode == 'categories':
171
+ if categoryname =="COCO-80":
172
+ batch_category_name = coco_class_name
173
+ elif categoryname =="OBJ365":
174
+ batch_category_name = OBJ365_class_names
175
+ elif categoryname =="Custom-List":
176
+ batch_category_name = custom_category.split(',')
177
+ else:
178
+ batch_category_name = class_agnostic_name
179
+
180
+ # mask_ori = torch.from_numpy(np.load('03_moto_mask.npy'))[None,]
181
+ # mask_ori = (F.interpolate(mask_ori, (height, width), mode='bilinear') > 0).to(device)
182
+ # prompt_list = [mask_ori[0]]
183
+ prompt_list = []
184
+ with torch.no_grad():
185
+ (outputs,_) = GLEEmodel(infer_image, prompt_list, task="coco", batch_name_list=batch_category_name, is_train=False)
186
+
187
+ topK_instance = max(num_inst_select,1)
188
+ else:
189
+ topK_instance = 1
190
+ prompt_list = {'grounding':[expressiong]}
191
+ with torch.no_grad():
192
+ (outputs,_) = GLEEmodel(infer_image, prompt_list, task="grounding", batch_name_list=[], is_train=False)
193
+
194
+
195
+ mask_pred = outputs['pred_masks'][0]
196
+ mask_cls = outputs['pred_logits'][0]
197
+ boxes_pred = outputs['pred_boxes'][0]
198
+
199
+ scores = mask_cls.sigmoid().max(-1)[0]
200
+ scores_per_image, topk_indices = scores.topk(topK_instance, sorted=True)
201
+ if prompt_mode == 'categories':
202
+ valid = scores_per_image>threshold_select
203
+ topk_indices = topk_indices[valid]
204
+ scores_per_image = scores_per_image[valid]
205
+
206
+ pred_class = mask_cls[topk_indices].max(-1)[1].tolist()
207
+ pred_boxes = boxes_pred[topk_indices]
208
+
209
+
210
+ boxes = LSJ_box_postprocess(pred_boxes,padding_size,re_size, ori_height,ori_width)
211
+ mask_pred = mask_pred[topk_indices]
212
+ pred_masks = F.interpolate( mask_pred[None,], size=(padding_size[0], padding_size[1]), mode="bilinear", align_corners=False )
213
+ pred_masks = pred_masks[:,:,:re_size[0],:re_size[1]]
214
+ pred_masks = F.interpolate( pred_masks, size=(ori_height,ori_width), mode="bilinear", align_corners=False )
215
+ pred_masks = (pred_masks>0).detach().cpu().numpy()[0]
216
+
217
+ if 'mask' in results_select:
218
+
219
+ zero_mask = np.zeros_like(copyed_img)
220
+ for nn, mask in enumerate(pred_masks):
221
+ # mask = mask.numpy()
222
+ mask = mask.reshape(mask.shape[0], mask.shape[1], 1)
223
+
224
+ lar = np.concatenate((mask*COLORS[nn%12][2], mask*COLORS[nn%12][1], mask*COLORS[nn%12][0]), axis = 2)
225
+ zero_mask = zero_mask+ lar
226
+
227
+
228
+ lar_valid = zero_mask>0
229
+ masked_image = lar_valid*copyed_img
230
+ img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,1)*255*(1-mask_image_mix_ration)
231
+ max_p = img_n.max()
232
+ img_n = 255*img_n/max_p
233
+ ret = (~lar_valid*copyed_img)*mask_image_mix_ration + img_n
234
+ ret = ret.astype('uint8')
235
+ else:
236
+ ret = copyed_img
237
+
238
+ if 'box' in results_select:
239
+
240
+ line_width = max(ret.shape) /200
241
+
242
+ for nn,(classid, box) in enumerate(zip(pred_class,boxes)):
243
+ x1,y1,x2,y2 = box.long().tolist()
244
+ RGB = (COLORS[nn%12][2]*255,COLORS[nn%12][1]*255,COLORS[nn%12][0]*255)
245
+ cv2.rectangle(ret, (x1,y1), (x2,y2), RGB, math.ceil(line_width) )
246
+ if prompt_mode == 'categories' or (prompt_mode == 'expression' and 'expression' in results_select ):
247
+ if prompt_mode == 'categories':
248
+ label = ''
249
+ if 'name' in results_select:
250
+ label += batch_category_name[classid]
251
+ if 'score' in results_select:
252
+ label += str(scores_per_image[nn].item())[:4]
253
+ else:
254
+ label = expressiong
255
+
256
+ if len(label)==0:
257
+ continue
258
+ height, width, _ = ret.shape
259
+ FONT = cv2.FONT_HERSHEY_COMPLEX
260
+ label_width, label_height = cv2.getTextSize(label, FONT, min(width, height) * FONT_SCALE, math.ceil(min(width, height) * THICKNESS_SCALE))[0]
261
+
262
+ cv2.rectangle(ret, (x1,y1), (x1+label_width,(y1 -label_height) - int(height * TEXT_Y_OFFSET_SCALE)), RGB, -1)
263
+
264
+ cv2.putText(
265
+ ret,
266
+ label,
267
+ (x1, y1 - int(height * TEXT_Y_OFFSET_SCALE)),
268
+ fontFace=FONT,
269
+ fontScale=min(width, height) * FONT_SCALE,
270
+ thickness=math.ceil(min(width, height) * THICKNESS_SCALE),
271
+ color=(255,255,255),
272
+ )
273
+
274
+
275
+
276
+ ret = ret.astype('uint8')
277
+ return ret
278
+
279
+
280
+ else: #visual prompt
281
+ topK_instance = 1
282
+ copyed_img = img['background'][:,:,:3].copy()
283
+ # get bbox from scribbles in layers
284
+ bbox_list = [scribble2box(layer) for layer in img['layers'] ]
285
+ visual_prompt_list = []
286
+ visual_prompt_RGB_list = []
287
+
288
+ for mask, (box,RGB) in zip(img['layers'], bbox_list):
289
+ if box is None:
290
+ continue
291
+ if prompt_mode=='box':
292
+ fakemask = np.zeros_like(copyed_img[:,:,0])
293
+ x1 ,y1 ,x2, y2 = box
294
+ fakemask[ y1:y2, x1:x2 ] = 1
295
+ fakemask = fakemask>0
296
+ elif prompt_mode=='point':
297
+ fakemask = np.zeros_like(copyed_img[:,:,0])
298
+ H,W = fakemask.shape
299
+ x1 ,y1 ,x2, y2 = box
300
+ center_x, center_y = (x1+x2)//2, (y1+y2)//2
301
+ fakemask[ center_y-H//40:center_y+H//40, center_x-W//40:center_x+W//40 ] = 1
302
+ fakemask = fakemask>0
303
+ elif prompt_mode=='scribble':
304
+ fakemask = mask[:,:,-1]
305
+ fakemask = fakemask>0
306
+
307
+ fakemask = torch.from_numpy(fakemask).unsqueeze(0).to(ori_image)
308
+ if inference_type == 'LSJ':
309
+ infer_visual_prompt = torch.zeros(1,1024,1024).to(ori_image)
310
+ infer_visual_prompt[:,:inference_size,:inference_size] = fakemask
311
+ else:
312
+ resize_fakemask = resizer(fakemask)
313
+ if size_divisibility > 1:
314
+ # the last two dims are H,W, both subject to divisibility requirement
315
+ infer_visual_prompt = torch.zeros(1,padding_size[0],padding_size[1]).to(resize_fakemask)
316
+ infer_visual_prompt[:,:image_size[0],:image_size[1]] = resize_fakemask
317
+
318
+
319
+ visual_prompt_list.append( infer_visual_prompt>0 )
320
+ visual_prompt_RGB_list.append(RGB)
321
+
322
+
323
+ mask_results_list = []
324
+ for visual_prompt in visual_prompt_list:
325
+
326
+ prompt_list = {'spatial':[visual_prompt]}
327
+
328
+ with torch.no_grad():
329
+ (outputs,_) = GLEEmodel(infer_image, prompt_list, task="coco", batch_name_list=['object'], is_train=False, visual_prompt_type=prompt_mode )
330
+
331
+ mask_pred = outputs['pred_masks'][0]
332
+ mask_cls = outputs['pred_logits'][0]
333
+ boxes_pred = outputs['pred_boxes'][0]
334
+
335
+ scores = mask_cls.sigmoid().max(-1)[0]
336
+ scores_per_image, topk_indices = scores.topk(topK_instance, sorted=True)
337
+
338
+ pred_class = mask_cls[topk_indices].max(-1)[1].tolist()
339
+ pred_boxes = boxes_pred[topk_indices]
340
+
341
+
342
+ boxes = LSJ_box_postprocess(pred_boxes,padding_size,re_size, ori_height,ori_width)
343
+ mask_pred = mask_pred[topk_indices]
344
+ pred_masks = F.interpolate( mask_pred[None,], size=(padding_size[0], padding_size[1]), mode="bilinear", align_corners=False )
345
+ pred_masks = pred_masks[:,:,:re_size[0],:re_size[1]]
346
+ pred_masks = F.interpolate( pred_masks, size=(ori_height,ori_width), mode="bilinear", align_corners=False )
347
+ pred_masks = (pred_masks>0).detach().cpu().numpy()[0]
348
+ mask_results_list.append(pred_masks)
349
+
350
+ zero_mask = np.zeros_like(copyed_img)
351
+ for mask,RGB in zip(mask_results_list,visual_prompt_RGB_list):
352
+ mask = mask.reshape(mask.shape[-2], mask.shape[-1], 1)
353
+ lar = np.concatenate((mask*RGB[0], mask*RGB[1],mask*RGB[2]), axis = 2)
354
+ zero_mask = zero_mask+ lar
355
+ lar_valid = zero_mask>0
356
+ masked_image = lar_valid*copyed_img
357
+ img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,255)*(1-mask_image_mix_ration)
358
+ max_p = img_n.max()
359
+ img_n = 255*img_n/max_p
360
+ ret = (~lar_valid*copyed_img)*mask_image_mix_ration + img_n
361
+ ret = ret.astype('uint8')
362
+ # cv2.imwrite('00020_inst.jpg', cv2.cvtColor(ret, cv2.COLOR_BGR2RGB))
363
+
364
+ return ret
365
+
366
+
367
+
368
+
369
+ def process_frames(frame_list):
370
+ clip_images = [torch.as_tensor(np.ascontiguousarray( frame[:,:,::-1].transpose(2, 0, 1))) for frame in frame_list]
371
+ processed_frames = []
372
+ for ori_image in clip_images:
373
+ ori_image = normalizer(ori_image.to(device))[None,]
374
+ _,_, ori_height, ori_width = ori_image.shape
375
+
376
+ if inference_type == 'LSJ':
377
+ infer_image = torch.zeros(1,3,1024,1024).to(ori_image)
378
+ infer_image[:,:,:inference_size,:inference_size] = ori_image
379
+ else:
380
+ resize_image = videoresizer(ori_image)
381
+ image_size = torch.as_tensor((resize_image.shape[-2],resize_image.shape[-1]))
382
+ re_size = resize_image.shape[-2:]
383
+ if size_divisibility > 1:
384
+ stride = size_divisibility
385
+ # the last two dims are H,W, both subject to divisibility requirement
386
+ padding_size = ((image_size + (stride - 1)).div(stride, rounding_mode="floor") * stride).tolist()
387
+ infer_image = torch.zeros(1,3,padding_size[0],padding_size[1]).to(resize_image)
388
+ infer_image[0,:,:image_size[0],:image_size[1]] = resize_image
389
+ processed_frames.append(infer_image)
390
+ return torch.cat(processed_frames,dim=0), padding_size,re_size,ori_height, ori_width # [clip_lenth,3,h,w]
391
+
392
+
393
+ def match_from_embds(tgt_embds, cur_embds):
394
+ cur_embds = cur_embds / cur_embds.norm(dim=1)[:, None]
395
+ tgt_embds = tgt_embds / tgt_embds.norm(dim=1)[:, None]
396
+ cos_sim = torch.mm(cur_embds, tgt_embds.transpose(0,1))
397
+
398
+ cost_embd = 1 - cos_sim
399
+
400
+ C = 1.0 * cost_embd
401
+ C = C.cpu()
402
+
403
+ indices = linear_sum_assignment(C.transpose(0, 1)) # target x current
404
+ indices = indices[1] # permutation that makes current aligns to target
405
+
406
+ return indices
407
+
408
+
409
+
410
+
411
+ def segment_video(video, prompt_mode, categoryname, custom_category, expressiong, results_select, num_inst_select, threshold_select, mask_image_mix_ration, model_selection,video_frames_select, prompter):
412
+
413
+ ### model selection
414
+ if model_selection == 'GLEE-Plus (SwinL)':
415
+ GLEEmodel = GLEEmodel_swin
416
+ print('use GLEE-Plus')
417
+ clip_length = 4 #batchsize
418
+ else:
419
+ GLEEmodel = GLEEmodel_r50
420
+ print('use GLEE-Lite')
421
+ clip_length = 8 #batchsize
422
+
423
+ # read video and get sparse frames
424
+ cap = cv2.VideoCapture(video)
425
+ video_fps = cap.get(cv2.CAP_PROP_FPS )
426
+ print('video fps:', video_fps)
427
+ frame_list = []
428
+ frac = video_fps/30
429
+ frame_count = 0
430
+ read_fps = 10
431
+
432
+ interval = int( frac *(30 /read_fps) ) #interval frames
433
+ while cap.isOpened():
434
+ ret, frame = cap.read()
435
+ frame_count += 1
436
+ # if frame is read correctly ret is True
437
+ if not ret:
438
+ print("Can't receive frame (stream end?). Exiting ...")
439
+ break
440
+ if frame_count % int(interval) == 0:
441
+ frame_list.append(frame)
442
+ cap.release()
443
+ first_frame = frame_list[0]
444
+ frame_list = frame_list[:video_frames_select] # max num of frames
445
+ print('num frames:', len(frame_list))
446
+
447
+
448
+ video_len = len(frame_list)
449
+
450
+ if prompt_mode == 'categories' or prompt_mode == 'expression':
451
+ if len(results_select)==0:
452
+ results_select=['box']
453
+ if prompt_mode == 'categories':
454
+ if categoryname =="COCO-80":
455
+ batch_category_name = coco_class_name
456
+ elif categoryname =="YTBVIS&OVIS":
457
+ batch_category_name = YTBVISOVIS_class_name
458
+ elif categoryname =="OBJ365":
459
+ batch_category_name = OBJ365_class_names
460
+ elif categoryname =="Custom-List":
461
+ batch_category_name = custom_category.split(',')
462
+ else:
463
+ batch_category_name = class_agnostic_name
464
+ task = 'coco'
465
+ prompt_list = []
466
+ topK_instance = num_inst_select
467
+ prompt_mode = 'categories'
468
+ results_select = ['mask', 'score', 'box', 'name']
469
+ else:
470
+ topK_instance = 1
471
+ initprompt_list = {'grounding':[expressiong]}
472
+ task = 'grounding'
473
+ batch_category_name = []
474
+
475
+
476
+ #split long video into clips to form a batch input
477
+ num_clips = math.ceil(video_len/clip_length)
478
+ logits_list, boxes_list, embed_list, masks_list = [], [], [], []
479
+ for c in range(num_clips):
480
+ start_idx = c*clip_length
481
+ end_idx = (c+1)*clip_length
482
+ clip_inputs = frame_list[start_idx:end_idx]
483
+ clip_images, padding_size,re_size,ori_height, ori_width = process_frames(clip_inputs)
484
+
485
+ if task=='grounding':
486
+ prompt_list = {'grounding': initprompt_list['grounding']*len(clip_images)}
487
+ with torch.no_grad():
488
+ (clip_output,_) = GLEEmodel(clip_images, prompt_list, task=task, batch_name_list=batch_category_name, is_train=False)
489
+
490
+ logits_list.append(clip_output['pred_logits'].detach())
491
+ boxes_list.append(clip_output['pred_boxes'].detach())
492
+ embed_list.append(clip_output['pred_track_embed'].detach())
493
+ masks_list.append(clip_output['pred_masks'].detach()) #.to(self.merge_device)
494
+ del clip_output
495
+ torch.cuda.empty_cache()
496
+ outputs = {
497
+ 'pred_logits':torch.cat(logits_list,dim=0),
498
+ 'pred_track_embed':torch.cat(embed_list,dim=0),
499
+ 'pred_masks':torch.cat(masks_list,dim=0),
500
+ 'pred_boxes': torch.cat(boxes_list,dim=0),
501
+ }
502
+
503
+ pred_logits = list(torch.unbind(outputs['pred_logits']))
504
+ pred_masks = list(torch.unbind(outputs['pred_masks']))
505
+ pred_embds = list(torch.unbind(outputs['pred_track_embed']))
506
+ pred_boxes = list(torch.unbind(outputs['pred_boxes']))
507
+ del outputs
508
+
509
+ out_logits = []
510
+ out_masks = []
511
+ out_embds = []
512
+ out_boxes = []
513
+ out_logits.append(pred_logits[0])
514
+ out_masks.append(pred_masks[0])
515
+ out_embds.append(pred_embds[0])
516
+ out_boxes.append(pred_boxes[0])
517
+ memory_embedding = out_embds[-1]
518
+
519
+ for i in range(1, len(pred_logits)):
520
+ # indices = self.match_from_embds(memory_embedding, pred_embds[i])
521
+ MA_embedding = torch.stack(out_embds[-5:]).mean(0)
522
+ indices = match_from_embds(MA_embedding, pred_embds[i])
523
+ out_logits.append(pred_logits[i][indices, :])
524
+ out_masks.append(pred_masks[i][indices, :, :])
525
+ out_embds.append(pred_embds[i][indices, :])
526
+ out_boxes.append(pred_boxes[i][indices, :])
527
+ score_weights = pred_logits[i][indices, :].sigmoid().max(-1)[0][:,None]
528
+ memory_embedding = (memory_embedding+pred_embds[i][indices, :]*score_weights )/(1+score_weights)
529
+
530
+
531
+ mask_cls = sum(out_logits)/len(out_logits)
532
+
533
+ scores = mask_cls.sigmoid().max(-1)[0]
534
+ scores_per_image, topk_indices = scores.topk(topK_instance, sorted=True)
535
+ valid = scores_per_image>threshold_select
536
+ topk_indices = topk_indices[valid]
537
+ scores_per_image = scores_per_image[valid]
538
+
539
+ out_logits = torch.stack(out_logits, dim=1)[topk_indices] # q numc -> q t numc
540
+ mask_pred = torch.stack(out_masks, dim=1)[topk_indices] # q h w -> numinst t h w
541
+ pred_boxes = torch.stack(out_boxes, dim=1)[topk_indices] # q 4 -> numinst t 4
542
+ perframe_score = out_logits.sigmoid().max(-1)[0].cpu().numpy()
543
+
544
+ pred_class = mask_cls[topk_indices].max(-1)[1].tolist()
545
+ boxes = LSJ_box_postprocess(pred_boxes,padding_size,re_size, ori_height,ori_width)
546
+ pred_masks = F.interpolate( mask_pred, size=(padding_size[0], padding_size[1]), mode="bilinear", align_corners=False )
547
+ pred_masks = pred_masks[:,:,:re_size[0],:re_size[1]]
548
+ pred_masks = F.interpolate( pred_masks, size=(ori_height,ori_width), mode="bilinear", align_corners=False )
549
+ pred_masks = (pred_masks>0).detach().cpu().numpy() # [numinst,t,h,w]
550
+ ourput_frames = []
551
+ for frameidx, ori_frame in enumerate(frame_list):
552
+ copyed_img = ori_frame.copy()
553
+ if 'mask' in results_select:
554
+
555
+ zero_mask = np.zeros_like(copyed_img)
556
+ for nn, (mask,score) in enumerate(zip(pred_masks[:,frameidx],perframe_score[:,frameidx])):
557
+ # mask = mask.numpy()
558
+ if score<threshold_select:
559
+ continue
560
+ mask = mask.reshape(mask.shape[0], mask.shape[1], 1)
561
+ lar = np.concatenate((mask*COLORS[nn%12][0], mask*COLORS[nn%12][1], mask*COLORS[nn%12][2]), axis = 2)
562
+ zero_mask = zero_mask+ lar
563
+
564
+
565
+ lar_valid = zero_mask>0
566
+ masked_image = lar_valid*copyed_img
567
+ img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,1)*255*(1-mask_image_mix_ration)
568
+ max_p = img_n.max()
569
+ img_n = 255*img_n/max_p
570
+ ret = (~lar_valid*copyed_img)*mask_image_mix_ration + img_n
571
+ ret = ret.astype('uint8')
572
+ else:
573
+ ret = copyed_img
574
+ if 'box' in results_select:
575
+
576
+ line_width = max(ret.shape) /200
577
+
578
+ for nn,(classid, box, score) in enumerate(zip(pred_class,boxes[:,frameidx],perframe_score[:,frameidx])):
579
+ if score<threshold_select:
580
+ continue
581
+ x1,y1,x2,y2 = box.long().tolist()
582
+ RGB = (COLORS[nn%12][0]*255,COLORS[nn%12][1]*255,COLORS[nn%12][2]*255)
583
+ cv2.rectangle(ret, (x1,y1), (x2,y2), RGB, math.ceil(line_width) )
584
+ if prompt_mode == 'categories' or (prompt_mode == 'expression' and 'expression' in results_select ):
585
+ if prompt_mode == 'categories':
586
+ label = ''
587
+ if 'name' in results_select:
588
+ label += batch_category_name[classid]
589
+ if 'score' in results_select:
590
+ label += str(score.item())[:4]
591
+ else:
592
+ label = expressiong
593
+ if 'score' in results_select:
594
+ label += str(score.item())[:4]
595
+
596
+ if len(label)==0:
597
+ continue
598
+ height, width, _ = ret.shape
599
+ FONT = cv2.FONT_HERSHEY_COMPLEX
600
+ label_width, label_height = cv2.getTextSize(label, FONT, min(width, height) * FONT_SCALE, math.ceil(min(width, height) * THICKNESS_SCALE))[0]
601
+
602
+ cv2.rectangle(ret, (x1,y1), (x1+label_width,(y1 -label_height) - int(height * TEXT_Y_OFFSET_SCALE)), RGB, -1)
603
+
604
+ cv2.putText(
605
+ ret,
606
+ label,
607
+ (x1, y1 - int(height * TEXT_Y_OFFSET_SCALE)),
608
+ fontFace=FONT,
609
+ fontScale=min(width, height) * FONT_SCALE,
610
+ thickness=math.ceil(min(width, height) * THICKNESS_SCALE),
611
+ color=(255,255,255),
612
+ )
613
+
614
+ ourput_frames.append(ret)
615
+ # ret = ret.astype('uint8')
616
+ size = (ori_width,ori_height)
617
+ output_file = "test.mp4"
618
+ out = cv2.VideoWriter(output_file,cv2.VideoWriter_fourcc(*'avc1'), read_fps, size)
619
+ for i in range(len(ourput_frames)):
620
+ out.write(ourput_frames[i])
621
+ out.release()
622
+ del out_logits, out_masks, out_embds, out_boxes, pred_masks
623
+ torch.cuda.empty_cache()
624
+ return output_file
625
+
626
+ else: # visual prompt vos
627
+
628
+ # image prompt segmentation
629
+ topK_instance = 1
630
+ copyed_img = prompter['background'][:,:,:3].copy()
631
+
632
+ ori_image = torch.as_tensor(np.ascontiguousarray( copyed_img.transpose(2, 0, 1)))
633
+ ori_image = normalizer(ori_image.to(device))[None,]
634
+ _,_, ori_height, ori_width = ori_image.shape
635
+
636
+ resize_image = videoresizer(ori_image)
637
+ image_size = torch.as_tensor((resize_image.shape[-2],resize_image.shape[-1]))
638
+ re_size = resize_image.shape[-2:]
639
+ if size_divisibility > 1:
640
+ stride = size_divisibility
641
+ # the last two dims are H,W, both subject to divisibility requirement
642
+ padding_size = ((image_size + (stride - 1)).div(stride, rounding_mode="floor") * stride).tolist()
643
+ infer_image = torch.zeros(1,3,padding_size[0],padding_size[1]).to(resize_image)
644
+ infer_image[0,:,:image_size[0],:image_size[1]] = resize_image
645
+ prompter['layers'] = prompter['layers'][:1] #only keep 1 prompt for VOS as model can only segment one object once infer
646
+ bbox_list = [scribble2box(layer) for layer in prompter['layers'] ]
647
+ visual_prompt_list = []
648
+ visual_prompt_RGB_list = []
649
+
650
+ for mask, (box,RGB) in zip(prompter['layers'], bbox_list):
651
+ if box is None:
652
+ continue
653
+ if prompt_mode=='box':
654
+ fakemask = np.zeros_like(copyed_img[:,:,0])
655
+ x1 ,y1 ,x2, y2 = box
656
+ fakemask[ y1:y2, x1:x2 ] = 1
657
+ fakemask = fakemask>0
658
+ elif prompt_mode=='point':
659
+ fakemask = np.zeros_like(copyed_img[:,:,0])
660
+ H,W = fakemask.shape
661
+ x1 ,y1 ,x2, y2 = box
662
+ center_x, center_y = (x1+x2)//2, (y1+y2)//2
663
+ fakemask[ center_y-H//40:center_y+H//40, center_x-W//40:center_x+W//40 ] = 1
664
+ fakemask = fakemask>0
665
+ elif prompt_mode=='scribble':
666
+ fakemask = mask[:,:,-1]
667
+ fakemask = fakemask>0
668
+
669
+ fakemask = torch.from_numpy(fakemask).unsqueeze(0).to(ori_image)
670
+ if inference_type == 'LSJ':
671
+ infer_visual_prompt = torch.zeros(1,1024,1024).to(ori_image)
672
+ infer_visual_prompt[:,:inference_size,:inference_size] = fakemask
673
+ else:
674
+ resize_fakemask = videoresizer(fakemask)
675
+ if size_divisibility > 1:
676
+ # the last two dims are H,W, both subject to divisibility requirement
677
+ infer_visual_prompt = torch.zeros(1,padding_size[0],padding_size[1]).to(resize_fakemask)
678
+ infer_visual_prompt[:,:image_size[0],:image_size[1]] = resize_fakemask
679
+
680
+
681
+ visual_prompt_list.append( infer_visual_prompt>0 )
682
+ visual_prompt_RGB_list.append(RGB)
683
+
684
+
685
+ mask_results_list = []
686
+ for visual_prompt in visual_prompt_list:
687
+
688
+ prompt_list = {'spatial':[visual_prompt]}
689
+
690
+ with torch.no_grad():
691
+ (outputs,_) = GLEEmodel(infer_image, prompt_list, task="coco", batch_name_list=['object'], is_train=False, visual_prompt_type=prompt_mode )
692
+
693
+ mask_pred = outputs['pred_masks'][0]
694
+ mask_cls = outputs['pred_logits'][0]
695
+ boxes_pred = outputs['pred_boxes'][0]
696
+
697
+ scores = mask_cls.sigmoid().max(-1)[0]
698
+ scores_per_image, topk_indices = scores.topk(topK_instance, sorted=True)
699
+
700
+ pred_class = mask_cls[topk_indices].max(-1)[1].tolist()
701
+ pred_boxes = boxes_pred[topk_indices]
702
+
703
+
704
+ boxes = LSJ_box_postprocess(pred_boxes,padding_size,re_size, ori_height,ori_width)
705
+ mask_pred = mask_pred[topk_indices]
706
+ pred_masks = F.interpolate( mask_pred[None,], size=(padding_size[0], padding_size[1]), mode="bilinear", align_corners=False )
707
+ first_frame_mask_padding = copy.deepcopy(pred_masks.detach())
708
+ pred_masks = pred_masks[:,:,:re_size[0],:re_size[1]]
709
+ pred_masks = F.interpolate( pred_masks, size=(ori_height,ori_width), mode="bilinear", align_corners=False )
710
+ pred_masks = (pred_masks>0).detach().cpu().numpy()[0]
711
+ mask_results_list.append(pred_masks)
712
+
713
+ zero_mask = np.zeros_like(copyed_img)
714
+ for mask,RGB in zip(mask_results_list,visual_prompt_RGB_list):
715
+ mask = mask.reshape(mask.shape[-2], mask.shape[-1], 1)
716
+ lar = np.concatenate((mask*RGB[0], mask*RGB[1],mask*RGB[2]), axis = 2)
717
+ zero_mask = zero_mask+ lar
718
+ lar_valid = zero_mask>0
719
+ masked_image = lar_valid*copyed_img
720
+ img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,255)*(1-mask_image_mix_ration)
721
+ max_p = img_n.max()
722
+ img_n = 255*img_n/max_p
723
+ ret = (~lar_valid*copyed_img)*mask_image_mix_ration + img_n
724
+ ret = ret.astype('uint8')
725
+ # import pdb;pdb.set_trace()
726
+ # cv2.imwrite('00020_inst.jpg', cv2.cvtColor(ret, cv2.COLOR_BGR2RGB))
727
+
728
+ output_vos_results = []
729
+ output_vos_results.append(ret[:,:,::-1])
730
+
731
+
732
+
733
+ #### vos process
734
+ checkpoints_VOS = torch.load('GLEE_vos_r50.pth')
735
+ GLEEmodel_VOS = GLEE_Model(cfg_vos, None, device, None, True).to(device)
736
+ GLEEmodel_VOS.load_state_dict(checkpoints_VOS, strict=False)
737
+ GLEEmodel_VOS.eval()
738
+
739
+
740
+ exist_obj_dict = {}
741
+ language_dict_features_dict_init = {}
742
+ language_dict_features_dict_prev = {}
743
+ point_sample_extra = {}
744
+
745
+
746
+
747
+
748
+
749
+ for frame_idx in range(video_len):
750
+
751
+ score_dict = {}
752
+
753
+ if frame_idx==0:
754
+
755
+ exist_obj_dict.update({1:first_frame_mask_padding[0]>0 })
756
+ prompt_list["spatial"] = [first_frame_mask_padding[0]>0]
757
+ frame_image, padding_size,re_size,ori_height, ori_width = process_frames(frame_list[frame_idx:frame_idx+1])
758
+ with torch.no_grad():
759
+ language_dict_features_dict_init[1], point_sample_extra[1] = \
760
+ GLEEmodel_VOS.vos_step1(frame_image, prompt_list, 'ytbvos', batch_name_list=['object'], is_train= False)
761
+ language_dict_features_dict_prev[1] = copy.deepcopy(language_dict_features_dict_init[1])
762
+ score_dict[1] = 1.0
763
+
764
+
765
+ if frame_idx>0:
766
+ cur_obj_id=1
767
+ frame_image, padding_size,re_size,ori_height, ori_width = process_frames(frame_list[frame_idx:frame_idx+1])
768
+ prompt_list["spatial"] = [exist_obj_dict[cur_obj_id]]
769
+ # import pdb;pdb.set_trace()
770
+ language_dict_features_init = copy.deepcopy(language_dict_features_dict_init[cur_obj_id]) # Important
771
+ language_dict_features_prev = copy.deepcopy(language_dict_features_dict_prev[cur_obj_id]) # Important
772
+ language_dict_features_cur = {}
773
+ language_dict_features_cur["hidden"] = torch.cat([language_dict_features_init["hidden"], language_dict_features_prev["hidden"]], dim=1)
774
+ language_dict_features_cur["masks"] = torch.cat([language_dict_features_init["masks"], language_dict_features_prev["masks"]], dim=1)
775
+
776
+ # concat initial prompt and last frame prompt for early fusion,but only use last frame point sampled feature for decocer self attention
777
+ with torch.no_grad():
778
+ frame_output,_ = GLEEmodel_VOS.vos_step2(frame_image, task='ytbvos', language_dict_features = language_dict_features_cur, \
779
+ last_extra = point_sample_extra[cur_obj_id], batch_name_list=['object'], is_train= False)
780
+
781
+ logits = frame_output['pred_scores'][0]
782
+
783
+ top_k_propose = 1
784
+ topk_values, topk_indexes = torch.topk(logits.sigmoid(), top_k_propose, dim=0)
785
+ mask_pred_result = frame_output['pred_masks'][0,topk_indexes] #[nk,1,H,W]
786
+ # pred_embeddings = frame_output['pred_track_embed'][0,topk_indexes.squeeze()] #[nk,256]
787
+
788
+ score_dict[cur_obj_id] = topk_values.item()
789
+
790
+
791
+ if score_dict[cur_obj_id] > 0.3:
792
+
793
+
794
+
795
+ mask_pred_result = F.interpolate(
796
+ mask_pred_result,
797
+ size=(padding_size[0], padding_size[1]),
798
+ mode="bilinear",
799
+ align_corners=False,
800
+ )
801
+ exist_obj_dict[cur_obj_id] = mask_pred_result[0,0]>0
802
+ mask_pred_result = mask_pred_result[:,:,:re_size[0],:re_size[1]]
803
+ mask_pred_result = F.interpolate( mask_pred_result, size=(ori_height,ori_width), mode="bilinear", align_corners=True )[0]
804
+ final_mask = mask_pred_result[0]>0
805
+ final_mask = final_mask.cpu().numpy()
806
+
807
+
808
+ copyed_img = frame_list[frame_idx]
809
+ zero_mask = np.zeros_like(copyed_img)
810
+ RGB = visual_prompt_RGB_list[0]
811
+ mask = final_mask.reshape(final_mask.shape[0], final_mask.shape[1], 1)
812
+ lar = np.concatenate((mask*RGB[2], mask*RGB[1],mask*RGB[0]), axis = 2)
813
+ zero_mask = zero_mask+ lar
814
+ lar_valid = zero_mask>0
815
+ masked_image = lar_valid*copyed_img
816
+ img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,255)*(1-mask_image_mix_ration)
817
+ max_p = img_n.max()
818
+ img_n = 255*img_n/max_p
819
+ ret = (~lar_valid*copyed_img)*mask_image_mix_ration + img_n
820
+ ret = ret.astype('uint8')
821
+ output_vos_results.append(ret)
822
+
823
+
824
+ if score_dict[cur_obj_id]>0.5: # update memory
825
+ prompt_list["spatial"] = [exist_obj_dict[cur_obj_id].unsqueeze(0)]
826
+ assert cur_obj_id in language_dict_features_dict_prev
827
+ with torch.no_grad():
828
+ language_dict_features_dict_prev[cur_obj_id], point_sample_extra[cur_obj_id] = \
829
+ GLEEmodel_VOS.vos_step1(frame_image, prompt_list, 'ytbvos', batch_name_list=['object'], is_train= False)
830
+
831
+ else: # add zero as mask
832
+ copyed_img = frame_list[frame_idx]
833
+ ret = copyed_img*mask_image_mix_ration
834
+ ret = ret.astype('uint8')
835
+ output_vos_results.append(ret)
836
+
837
+ size = (ori_width,ori_height)
838
+ output_file = "test.mp4"
839
+ out = cv2.VideoWriter(output_file,cv2.VideoWriter_fourcc(*'avc1'), read_fps, size)
840
+ for i in range(len(output_vos_results)):
841
+ out.write(output_vos_results[i])
842
+ out.release()
843
+ torch.cuda.empty_cache()
844
+ return output_file
845
+
846
+
847
+ def visual_prompt_preview(img, prompt_mode):
848
+ copyed_img = img['background'][:,:,:3].copy()
849
+ import pdb;pdb.set_trace()
850
+ # get bbox from scribbles in layers
851
+ bbox_list = [scribble2box(layer) for layer in img['layers'] ]
852
+ zero_mask = np.zeros_like(copyed_img)
853
+
854
+ for mask, (box,RGB) in zip(img['layers'], bbox_list):
855
+ if box is None:
856
+ continue
857
+
858
+ if prompt_mode=='box':
859
+ fakemask = np.zeros_like(copyed_img[:,:,0])
860
+ x1 ,y1 ,x2, y2 = box
861
+ fakemask[ y1:y2, x1:x2 ] = 1
862
+ fakemask = fakemask>0
863
+ elif prompt_mode=='point':
864
+ fakemask = np.zeros_like(copyed_img[:,:,0])
865
+ H,W = fakemask.shape
866
+ x1 ,y1 ,x2, y2 = box
867
+ center_x, center_y = (x1+x2)//2, (y1+y2)//2
868
+ fakemask[ center_y-H//40:center_y+H//40, center_x-W//40:center_x+W//40 ] = 1
869
+ fakemask = fakemask>0
870
+ else:
871
+ fakemask = mask[:,:,-1]
872
+ fakemask = fakemask>0
873
+
874
+ mask = fakemask.reshape(fakemask.shape[0], fakemask.shape[1], 1)
875
+ lar = np.concatenate((mask*RGB[0], mask*RGB[1],mask*RGB[2]), axis = 2)
876
+ zero_mask = zero_mask+ lar
877
+
878
+
879
+ img_n = copyed_img + np.clip(zero_mask,0,255)
880
+ max_p = img_n.max()
881
+ ret = 255*img_n/max_p
882
+ ret = ret.astype('uint8')
883
+ return ret
884
+
885
+ with gr.Blocks(theme=gr.themes.Default()) as demo:
886
+ gr.Markdown('# GLEE: General Object Foundation Model for Images and Videos at Scale')
887
+
888
+ # gr.HTML("<p> <img src='/file=GLEE_logo.png' aligh='center' style='float:left' width='6%' > <h1 class='title is-1 publication-title'> <p style='margin-left: 20px'> GLEE: General Object Foundation Model for Images and Videos at Scale </h1> ")
889
+
890
+ gr.Markdown(' [Paper](https://arxiv.org/abs/2312.09158) —— [Project Page](https://glee-vision.github.io) —— [Code](https://github.com/FoundationVision/GLEE) ')
891
+ # gr.HTML(“img src=“image link” alt=“A beautiful landscape”)
892
+ gr.Markdown(
893
+ 'The functionality demonstration demo app of GLEE. \
894
+ Image tasks includes **arbitrary vocabulary** object detection&segmentation, \
895
+ **any form of object name**, object caption detection, \
896
+ referring expression comprehension, and interactive segmentation. \
897
+ Video tasks add object tracking based on image tasks.'
898
+ )
899
+
900
+
901
+ with gr.Tab("Image task"):
902
+ with gr.Row():
903
+ with gr.Column():
904
+
905
+ img_input = gr.ImageEditor()
906
+ model_select = gr.Dropdown(
907
+ ["GLEE-Lite (R50)", "GLEE-Plus (SwinL)"], value = "GLEE-Plus (SwinL)" , multiselect=False, label="Model",
908
+ )
909
+ with gr.Row():
910
+ with gr.Column():
911
+ prompt_mode_select = gr.Radio([ "categories", "expression", "point", "scribble", "box"], label="Prompt", value= "categories" , info="What kind of prompt do you want to use?")
912
+ category_select = gr.Dropdown(
913
+ ["COCO-80", "OBJ365", "Custom-List", "Class-Agnostic"], visible=True, value = "COCO-80" , multiselect=False, label="Categories", info="Choose an existing category list or class-agnostic"
914
+ )
915
+ custom_category = gr.Textbox(
916
+ label="Custom Category",
917
+ info="Input custom category list, seperate by ',' ",
918
+ lines=1,
919
+ visible=False,
920
+ value="dog, cat, car, person",
921
+ )
922
+ input_expressiong = gr.Textbox(
923
+ label="Expression",
924
+ info="Input any description of an object in the image ",
925
+ lines=2,
926
+ visible=False,
927
+ value="the red car",
928
+ )
929
+
930
+ with gr.Accordion("Text based detection usage",open=False, visible=False) as textusage:
931
+ gr.Markdown(
932
+ 'GLEE supports three kind of object perception methods: category list, textual description, and class-agnostic.<br />\
933
+ 1.Select an existing category list from the "Categories" dropdown, like COCO or OBJ365, or customize your own list.<br />\
934
+ 2.Enter arbitrary object name in "Custom Category", or choose the expression model and describe the object in "Expression Textbox" for single object detection only.<br />\
935
+ 3.For class-agnostic mode, choose "Class-Agnostic" from the "Categories" dropdown.'
936
+ )
937
+
938
+
939
+ with gr.Group(visible=False,) as promptshow:
940
+
941
+ with gr.Accordion("Interactive segmentation usage",open=False):
942
+ gr.Markdown(
943
+ 'For interactive segmentation:<br />\
944
+ 1.Draw points, boxes, or scribbles on the canvas for multiclass segmentation; use separate layers for different objects, adding layers with a "+" sign.<br />\
945
+ 2.Point mode accepts a single point only; multiple points default to the centroid, so use boxes or scribbles for larger objects.<br />\
946
+ 3.After drawing, click green "√" to preview the prompt visualization; the segmentation mask follows the chosen prompt colors.'
947
+ )
948
+
949
+ img_showbox = gr.Image(label="visual prompt area preview")
950
+
951
+
952
+
953
+ def update_component_visible(prompt,category):
954
+ if prompt in ['point', 'scribble', 'box']:
955
+ return {
956
+ category_select:gr.Dropdown(visible=False),
957
+ custom_category:gr.Textbox(visible=False),
958
+ input_expressiong: gr.Textbox(visible=False),
959
+ promptshow:gr.Group(visible=True),
960
+ textusage:gr.Accordion(visible=False),
961
+ }
962
+ elif prompt == 'categories':
963
+ if category == "Custom-List":
964
+ return {
965
+ category_select:gr.Dropdown(visible=True),
966
+ custom_category:gr.Textbox(visible=True),
967
+ input_expressiong: gr.Textbox(visible=False),
968
+ promptshow:gr.Group(visible=False),
969
+ textusage:gr.Accordion(visible=True),
970
+ }
971
+ return {
972
+ category_select:gr.Dropdown(visible=True),
973
+ custom_category:gr.Textbox(visible=False),
974
+ input_expressiong: gr.Textbox(visible=False),
975
+ promptshow:gr.Group(visible=False),
976
+ textusage:gr.Accordion(visible=True),
977
+ }
978
+ else:
979
+ return {
980
+ category_select:gr.Dropdown(visible=False),
981
+ custom_category:gr.Textbox(visible=False),
982
+ input_expressiong: gr.Textbox(visible=True),
983
+ promptshow:gr.Group(visible=False),
984
+ textusage:gr.Accordion(visible=True),
985
+ }
986
+ def update_category_showcase(category):
987
+ if category == "Custom-List":
988
+ return {
989
+ category_select:gr.Dropdown(visible=True),
990
+ custom_category:gr.Textbox(visible=True),
991
+ input_expressiong: gr.Textbox(visible=False),
992
+ promptshow:gr.Group(visible=False),
993
+ textusage:gr.Accordion(visible=True),
994
+ }
995
+ else:
996
+ return {
997
+ category_select:gr.Dropdown(visible=True),
998
+ custom_category:gr.Textbox(visible=False),
999
+ input_expressiong: gr.Textbox(visible=False),
1000
+ promptshow:gr.Group(visible=False),
1001
+ textusage:gr.Accordion(visible=True),
1002
+ }
1003
+
1004
+ prompt_mode_select.input(update_component_visible,
1005
+ [prompt_mode_select,category_select],
1006
+ [category_select,custom_category,input_expressiong,promptshow,textusage])
1007
+ category_select.input(update_category_showcase,
1008
+ [category_select],
1009
+ [category_select,custom_category,input_expressiong,promptshow,textusage])
1010
+ # with gr.Column():
1011
+
1012
+
1013
+ with gr.Column():
1014
+ image_segment = gr.Image(label="detection and segmentation results")
1015
+ with gr.Accordion("Try More Visualization Options"):
1016
+ results_select = gr.CheckboxGroup(["box", "mask", "name", "score", "expression"], value=["box", "mask", "name", "score"], label="Shown Results", info="The results shown on image")
1017
+ num_inst_select = gr.Slider(1, 50, value=15, step=1, label="Num of topK instances for category based detection", info="Choose between 1 and 50 for better visualization")
1018
+ threshold_select = gr.Slider(0, 1, value=0.2, label="Confidence Threshold", info="Choose threshold ")
1019
+ mask_image_mix_ration = gr.Slider(0, 1, value=0.65, label="Image Brightness Ratio", info="Brightness between image and colored masks ")
1020
+
1021
+
1022
+
1023
+
1024
+ image_button = gr.Button("Detect & Segment")
1025
+ img_input.change(visual_prompt_preview, inputs = [img_input,prompt_mode_select] , outputs = img_showbox)
1026
+ image_button.click(segment_image, inputs=[img_input, prompt_mode_select, category_select, custom_category,input_expressiong, results_select, num_inst_select, threshold_select, mask_image_mix_ration,model_select], outputs=image_segment)
1027
+
1028
+
1029
+ with gr.Tab("Video task"):
1030
+ gr.Markdown(
1031
+ '#### Gradio only support .mp4 for HTML display. \
1032
+ Due to computing resource restrictions, we sample and play the input video in 10 fps, and single video is limited (or cropped) to 10 seconds'
1033
+ )
1034
+ with gr.Row():
1035
+
1036
+ with gr.Column(): # video input face
1037
+
1038
+ video_input = gr.Video(label="Input Video", interactive=True, sources=['upload'])
1039
+ video_model_select = gr.Dropdown(
1040
+ ["GLEE-Lite (R50)", "GLEE-Plus (SwinL)"], value = "GLEE-Lite (R50)" , multiselect=False, label="Model",
1041
+ )
1042
+ with gr.Row():
1043
+ with gr.Column():
1044
+ video_prompt_mode_select = gr.Radio([ "categories", "expression", "point", "scribble", "box"], label="Prompt", value= "categories" , info="What kind of prompt do you want to use?")
1045
+ video_category_select = gr.Dropdown(
1046
+ ["YTBVIS&OVIS", "COCO-80", "OBJ365", "Custom-List", "Class-Agnostic"], visible=True, value = "COCO-80" , multiselect=False, label="Categories", info="Choose an existing category list or class-agnostic"
1047
+ )
1048
+ video_custom_category = gr.Textbox(
1049
+ label="Custom Category",
1050
+ info="Input custom category list, seperate by ',' ",
1051
+ lines=1,
1052
+ visible=False,
1053
+ value="dog, cat, car, person",
1054
+ )
1055
+ video_input_expressiong = gr.Textbox(
1056
+ label="Expression",
1057
+ info="Input any description of an object in the image ",
1058
+ lines=2,
1059
+ visible=False,
1060
+ value="the red car",
1061
+ )
1062
+
1063
+ with gr.Accordion("Text based detection usage",open=False, visible=False) as video_textusage:
1064
+ gr.Markdown(
1065
+ 'GLEE supports three kind of object perception methods: category list, textual description, and class-agnostic.<br />\
1066
+ 1.Select an existing category list from the "Categories" dropdown, like COCO or OBJ365, or customize your own list.<br />\
1067
+ 2.Enter arbitrary object name in "Custom Category", or choose the expression model and describe the object in "Expression Textbox" for single object detection only.<br />\
1068
+ 3.For class-agnostic mode, choose "Class-Agnostic" from the "Categories" dropdown.'
1069
+ )
1070
+
1071
+
1072
+ with gr.Group(visible=False,) as video_promptshow:
1073
+
1074
+ with gr.Accordion("Interactive segmentation usage",open=False):
1075
+ gr.Markdown(
1076
+ 'For video interactive segmentation, draw a prompt on the first frame:<br />\
1077
+ 1.Draw points, boxes, or scribbles on the canvas for multiclass segmentation; only support one object tracking in interactive mode\
1078
+ 2.Point mode accepts a single point only; multiple points default to the centroid, so use boxes or scribbles for larger objects.<br />\
1079
+ 3.After drawing, click "Preview" to preview the prompt visualization; the segmentation mask follows the chosen prompt colors.'
1080
+ )
1081
+ with gr.Row():
1082
+ video_visual_prompter = gr.ImageEditor(label="visual prompter", show_label=True ,sources=['clipboard'])
1083
+ video_img_showbox = gr.Image(label="visual prompt area preview")
1084
+ video_prompt_preview = gr.Button("Preview")
1085
+ def update_video_component_visible(prompt,category, video):
1086
+ if prompt in ['point', 'scribble', 'box']:
1087
+ if video is None:
1088
+ return {
1089
+ video_category_select:gr.Dropdown(visible=False),
1090
+ video_custom_category:gr.Textbox(visible=False),
1091
+ video_input_expressiong: gr.Textbox(visible=False),
1092
+ video_promptshow:gr.Group(visible=True),
1093
+ video_textusage:gr.Accordion(visible=False),}
1094
+ else:
1095
+ cap = cv2.VideoCapture(video)
1096
+ ret, frame = cap.read()
1097
+ frame = frame[:,:,::-1].astype('uint8')
1098
+ zerolayers = np.zeros((frame.shape[0],frame.shape[1],1)).astype('uint8')
1099
+ alpha = 255+zerolayers
1100
+ newframe = np.concatenate((frame,alpha),axis=2)
1101
+ cap.release()
1102
+ return {
1103
+ video_category_select:gr.Dropdown(visible=False),
1104
+ video_custom_category:gr.Textbox(visible=False),
1105
+ video_input_expressiong: gr.Textbox(visible=False),
1106
+ video_promptshow:gr.Group(visible=True),
1107
+ video_textusage:gr.Accordion(visible=False),
1108
+ video_visual_prompter:gr.ImageEditor(value= {
1109
+ 'background':newframe,
1110
+ 'layers':[ ],
1111
+ 'composite':newframe }),
1112
+ }
1113
+ elif prompt == 'categories':
1114
+ if category == "Custom-List":
1115
+ return {
1116
+ video_category_select:gr.Dropdown(visible=True),
1117
+ video_custom_category:gr.Textbox(visible=True),
1118
+ video_input_expressiong: gr.Textbox(visible=False),
1119
+ video_promptshow:gr.Group(visible=False),
1120
+ video_textusage:gr.Accordion(visible=True),
1121
+ }
1122
+ return {
1123
+ video_category_select:gr.Dropdown(visible=True),
1124
+ video_custom_category:gr.Textbox(visible=False),
1125
+ video_input_expressiong: gr.Textbox(visible=False),
1126
+ video_promptshow:gr.Group(visible=False),
1127
+ video_textusage:gr.Accordion(visible=True),
1128
+ }
1129
+ else:
1130
+ return {
1131
+ video_category_select:gr.Dropdown(visible=False),
1132
+ video_custom_category:gr.Textbox(visible=False),
1133
+ video_input_expressiong: gr.Textbox(visible=True),
1134
+ video_promptshow:gr.Group(visible=False),
1135
+ video_textusage:gr.Accordion(visible=True),
1136
+ }
1137
+ def update_video_category_showcase(category):
1138
+ if category == "Custom-List":
1139
+ return {
1140
+ video_category_select:gr.Dropdown(visible=True),
1141
+ video_custom_category:gr.Textbox(visible=True),
1142
+ video_input_expressiong: gr.Textbox(visible=False),
1143
+ video_promptshow:gr.Group(visible=False),
1144
+ video_textusage:gr.Accordion(visible=True),
1145
+ }
1146
+ else:
1147
+ return {
1148
+ video_category_select:gr.Dropdown(visible=True),
1149
+ video_custom_category:gr.Textbox(visible=False),
1150
+ video_input_expressiong: gr.Textbox(visible=False),
1151
+ video_promptshow:gr.Group(visible=False),
1152
+ video_textusage:gr.Accordion(visible=True),
1153
+ }
1154
+
1155
+ video_prompt_mode_select.input(update_video_component_visible,
1156
+ [video_prompt_mode_select,video_category_select,video_input],
1157
+ [video_category_select,video_custom_category,video_input_expressiong,video_promptshow,video_textusage,video_visual_prompter])
1158
+ video_category_select.input(update_video_category_showcase,
1159
+ [video_category_select],
1160
+ [video_category_select,video_custom_category,video_input_expressiong,video_promptshow,video_textusage])
1161
+ video_input.change(update_video_component_visible,
1162
+ [video_prompt_mode_select,video_category_select,video_input],
1163
+ [video_category_select,video_custom_category,video_input_expressiong,video_promptshow,video_textusage,video_visual_prompter])
1164
+ with gr.Column():
1165
+ video_output = gr.Video(label="Video Results")
1166
+ with gr.Accordion("Try More Visualization Options"):
1167
+ video_frames_select = gr.Slider(1, 200, value=32, step=1, label="Max frames", info="The max length for video frames, you can select fewer frames reduce the waiting time to check the effect quickly")
1168
+ video_results_select = gr.CheckboxGroup(["box", "mask", "name", "score", "expression"], value=["box", "mask", "name", "score", "expression"], label="Shown Results", info="The results shown on image")
1169
+ video_num_inst_select = gr.Slider(1, 30, value=10, step=1, label="Num of topK instances for category based detection", info="Choose between 1 and 50 for better visualization")
1170
+ video_threshold_select = gr.Slider(0, 1, value=0.2, label="Confidence Threshold", info="Choose threshold ")
1171
+ video_mask_image_mix_ration = gr.Slider(0, 1, value=0.65, label="Image Brightness Ratio", info="Brightness between image and colored masks ")
1172
+
1173
+
1174
+
1175
+
1176
+ video_prompt_preview.click(visual_prompt_preview, inputs = [video_visual_prompter,video_prompt_mode_select] , outputs = video_img_showbox)
1177
+ video_button = gr.Button("Segment&Track")
1178
+ video_button.click(segment_video, inputs=[video_input, video_prompt_mode_select, video_category_select, video_custom_category, video_input_expressiong, video_results_select, video_num_inst_select, video_threshold_select, video_mask_image_mix_ration, video_model_select, video_frames_select, video_visual_prompter], outputs=video_output)
1179
+
1180
+
1181
+
1182
+
1183
+
1184
+ if __name__ == '__main__':
1185
+ demo.launch(inbrowser=True, allowed_paths=["./"])
app_v1.py ADDED
@@ -0,0 +1,492 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ try:
2
+ import detectron2
3
+ except:
4
+ import os
5
+ os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
6
+ os.system('cd GLEE/glee/models/pixel_decoder/ops && sh mask.sh')
7
+ # os.system('python -m pip install -e detectron2')
8
+
9
+ import gradio as gr
10
+ import numpy as np
11
+ import cv2
12
+ import torch
13
+
14
+ from detectron2.config import get_cfg
15
+ from GLEE.glee.models.glee_model import GLEE_Model
16
+ from GLEE.glee.config_deeplab import add_deeplab_config
17
+ from GLEE.glee.config import add_glee_config
18
+ import torch.nn.functional as F
19
+ import torchvision
20
+ import math
21
+ from obj365_name import categories as OBJ365_CATEGORIESV2
22
+
23
+
24
+ print(f"Is CUDA available: {torch.cuda.is_available()}")
25
+ # True
26
+ if torch.cuda.is_available():
27
+ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
28
+ # Tesla T4
29
+
30
+ def box_cxcywh_to_xyxy(x):
31
+ x_c, y_c, w, h = x.unbind(-1)
32
+ b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
33
+ (x_c + 0.5 * w), (y_c + 0.5 * h)]
34
+ return torch.stack(b, dim=-1)
35
+
36
+
37
+
38
+ def scribble2box(img):
39
+ if img.max()==0:
40
+ return None, None
41
+ rows = np.any(img, axis=1)
42
+ cols = np.any(img, axis=0)
43
+ all = np.any(img,axis=2)
44
+ R,G,B,A = img[np.where(all)[0][0],np.where(all)[1][0]].tolist() # get color
45
+ ymin, ymax = np.where(rows)[0][[0, -1]]
46
+ xmin, xmax = np.where(cols)[0][[0, -1]]
47
+ return np.array([ xmin,ymin, xmax,ymax]), (R,G,B)
48
+
49
+
50
+ def LSJ_box_postprocess( out_bbox, padding_size, crop_size, img_h, img_w):
51
+ # postprocess box height and width
52
+ boxes = box_cxcywh_to_xyxy(out_bbox)
53
+ lsj_sclae = torch.tensor([padding_size[1], padding_size[0], padding_size[1], padding_size[0]]).to(out_bbox)
54
+ crop_scale = torch.tensor([crop_size[1], crop_size[0], crop_size[1], crop_size[0]]).to(out_bbox)
55
+ boxes = boxes * lsj_sclae
56
+ boxes = boxes / crop_scale
57
+ boxes = torch.clamp(boxes,0,1)
58
+
59
+ scale_fct = torch.tensor([img_w, img_h, img_w, img_h])
60
+ scale_fct = scale_fct.to(out_bbox)
61
+ boxes = boxes * scale_fct
62
+ return boxes
63
+
64
+ COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
65
+ [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933],
66
+ [0.494, 0.000, 0.556], [0.494, 0.000, 0.000], [0.000, 0.745, 0.000],
67
+ [0.700, 0.300, 0.600],[0.000, 0.447, 0.741], [0.850, 0.325, 0.098]]
68
+
69
+
70
+
71
+ coco_class_name = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
72
+ OBJ365_class_names = [cat['name'] for cat in OBJ365_CATEGORIESV2]
73
+ class_agnostic_name = ['object']
74
+
75
+ if torch.cuda.is_available():
76
+ print('use cuda')
77
+ device = 'cuda'
78
+ else:
79
+ print('use cpu')
80
+ device='cpu'
81
+
82
+ cfg_r50 = get_cfg()
83
+ add_deeplab_config(cfg_r50)
84
+ add_glee_config(cfg_r50)
85
+ conf_files_r50 = 'GLEE/configs/R50.yaml'
86
+ checkpoints_r50 = torch.load('GLEE_R50_Scaleup10m.pth')
87
+ cfg_r50.merge_from_file(conf_files_r50)
88
+ GLEEmodel_r50 = GLEE_Model(cfg_r50, None, device, None, True).to(device)
89
+ GLEEmodel_r50.load_state_dict(checkpoints_r50, strict=False)
90
+ GLEEmodel_r50.eval()
91
+
92
+
93
+ cfg_swin = get_cfg()
94
+ add_deeplab_config(cfg_swin)
95
+ add_glee_config(cfg_swin)
96
+ conf_files_swin = 'GLEE/configs/SwinL.yaml'
97
+ checkpoints_swin = torch.load('GLEE_SwinL_Scaleup10m.pth')
98
+ cfg_swin.merge_from_file(conf_files_swin)
99
+ GLEEmodel_swin = GLEE_Model(cfg_swin, None, device, None, True).to(device)
100
+ GLEEmodel_swin.load_state_dict(checkpoints_swin, strict=False)
101
+ GLEEmodel_swin.eval()
102
+
103
+ pixel_mean = torch.Tensor( [123.675, 116.28, 103.53]).to(device).view(3, 1, 1)
104
+ pixel_std = torch.Tensor([58.395, 57.12, 57.375]).to(device).view(3, 1, 1)
105
+ normalizer = lambda x: (x - pixel_mean) / pixel_std
106
+ inference_size = 800
107
+ inference_type = 'resize_shot' # or LSJ
108
+ size_divisibility = 32
109
+
110
+ FONT_SCALE = 1.5e-3
111
+ THICKNESS_SCALE = 1e-3
112
+ TEXT_Y_OFFSET_SCALE = 1e-2
113
+
114
+
115
+ if inference_type != 'LSJ':
116
+ resizer = torchvision.transforms.Resize(inference_size)
117
+
118
+
119
+ def segment_image(img,prompt_mode, categoryname, custom_category, expressiong, results_select, num_inst_select, threshold_select, mask_image_mix_ration, model_selection):
120
+ if model_selection == 'GLEE-Plus (SwinL)':
121
+ GLEEmodel = GLEEmodel_swin
122
+ print('use GLEE-Plus')
123
+ else:
124
+ GLEEmodel = GLEEmodel_r50
125
+ print('use GLEE-Lite')
126
+
127
+ copyed_img = img['background'][:,:,:3].copy()
128
+
129
+ ori_image = torch.as_tensor(np.ascontiguousarray( copyed_img.transpose(2, 0, 1)))
130
+ ori_image = normalizer(ori_image.to(device))[None,]
131
+ _,_, ori_height, ori_width = ori_image.shape
132
+
133
+ if inference_type == 'LSJ':
134
+ infer_image = torch.zeros(1,3,1024,1024).to(ori_image)
135
+ infer_image[:,:,:inference_size,:inference_size] = ori_image
136
+ else:
137
+ resize_image = resizer(ori_image)
138
+ image_size = torch.as_tensor((resize_image.shape[-2],resize_image.shape[-1]))
139
+ re_size = resize_image.shape[-2:]
140
+ if size_divisibility > 1:
141
+ stride = size_divisibility
142
+ # the last two dims are H,W, both subject to divisibility requirement
143
+ padding_size = ((image_size + (stride - 1)).div(stride, rounding_mode="floor") * stride).tolist()
144
+ infer_image = torch.zeros(1,3,padding_size[0],padding_size[1]).to(resize_image)
145
+ infer_image[0,:,:image_size[0],:image_size[1]] = resize_image
146
+ # reversed_image = infer_image*pixel_std + pixel_mean
147
+ # reversed_image = torch.clip(reversed_image,min=0,max=255)
148
+ # reversed_image = reversed_image[0].permute(1,2,0)
149
+ # reversed_image = reversed_image.int().cpu().numpy().copy()
150
+ # cv2.imwrite('test.png',reversed_image[:,:,::-1])
151
+
152
+
153
+ if prompt_mode == 'categories' or prompt_mode == 'expression':
154
+ if len(results_select)==0:
155
+ results_select=['box']
156
+ if prompt_mode == 'categories':
157
+ if categoryname =="COCO-80":
158
+ batch_category_name = coco_class_name
159
+ elif categoryname =="OBJ365":
160
+ batch_category_name = OBJ365_class_names
161
+ elif categoryname =="Custom-List":
162
+ batch_category_name = custom_category.split(',')
163
+ else:
164
+ batch_category_name = class_agnostic_name
165
+
166
+ # mask_ori = torch.from_numpy(np.load('03_moto_mask.npy'))[None,]
167
+ # mask_ori = (F.interpolate(mask_ori, (height, width), mode='bilinear') > 0).to(device)
168
+ # prompt_list = [mask_ori[0]]
169
+ prompt_list = []
170
+ with torch.no_grad():
171
+ (outputs,_) = GLEEmodel(infer_image, prompt_list, task="coco", batch_name_list=batch_category_name, is_train=False)
172
+ topK_instance = max(num_inst_select,1)
173
+ else:
174
+ topK_instance = 1
175
+ prompt_list = {'grounding':[expressiong]}
176
+ with torch.no_grad():
177
+ (outputs,_) = GLEEmodel(infer_image, prompt_list, task="grounding", batch_name_list=[], is_train=False)
178
+
179
+
180
+ mask_pred = outputs['pred_masks'][0]
181
+ mask_cls = outputs['pred_logits'][0]
182
+ boxes_pred = outputs['pred_boxes'][0]
183
+
184
+ scores = mask_cls.sigmoid().max(-1)[0]
185
+ scores_per_image, topk_indices = scores.topk(topK_instance, sorted=True)
186
+ if prompt_mode == 'categories':
187
+ valid = scores_per_image>threshold_select
188
+ topk_indices = topk_indices[valid]
189
+ scores_per_image = scores_per_image[valid]
190
+
191
+ pred_class = mask_cls[topk_indices].max(-1)[1].tolist()
192
+ pred_boxes = boxes_pred[topk_indices]
193
+
194
+
195
+ boxes = LSJ_box_postprocess(pred_boxes,padding_size,re_size, ori_height,ori_width)
196
+ mask_pred = mask_pred[topk_indices]
197
+ pred_masks = F.interpolate( mask_pred[None,], size=(padding_size[0], padding_size[1]), mode="bilinear", align_corners=False )
198
+ pred_masks = pred_masks[:,:,:re_size[0],:re_size[1]]
199
+ pred_masks = F.interpolate( pred_masks, size=(ori_height,ori_width), mode="bilinear", align_corners=False )
200
+ pred_masks = (pred_masks>0).detach().cpu().numpy()[0]
201
+
202
+ if 'mask' in results_select:
203
+
204
+ zero_mask = np.zeros_like(copyed_img)
205
+ for nn, mask in enumerate(pred_masks):
206
+ # mask = mask.numpy()
207
+ mask = mask.reshape(mask.shape[0], mask.shape[1], 1)
208
+
209
+ lar = np.concatenate((mask*COLORS[nn%12][2], mask*COLORS[nn%12][1], mask*COLORS[nn%12][0]), axis = 2)
210
+ zero_mask = zero_mask+ lar
211
+
212
+
213
+ lar_valid = zero_mask>0
214
+ masked_image = lar_valid*copyed_img
215
+ img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,1)*255*(1-mask_image_mix_ration)
216
+ max_p = img_n.max()
217
+ img_n = 255*img_n/max_p
218
+ ret = (~lar_valid*copyed_img)*mask_image_mix_ration + img_n
219
+ ret = ret.astype('uint8')
220
+ else:
221
+ ret = copyed_img
222
+
223
+ if 'box' in results_select:
224
+
225
+ line_width = max(ret.shape) /200
226
+
227
+ for nn,(classid, box) in enumerate(zip(pred_class,boxes)):
228
+ x1,y1,x2,y2 = box.long().tolist()
229
+ RGB = (COLORS[nn%12][2]*255,COLORS[nn%12][1]*255,COLORS[nn%12][0]*255)
230
+ cv2.rectangle(ret, (x1,y1), (x2,y2), RGB, math.ceil(line_width) )
231
+ if prompt_mode == 'categories' or (prompt_mode == 'expression' and 'expression' in results_select ):
232
+ if prompt_mode == 'categories':
233
+ label = ''
234
+ if 'name' in results_select:
235
+ label += batch_category_name[classid]
236
+ if 'score' in results_select:
237
+ label += str(scores_per_image[nn].item())[:4]
238
+ else:
239
+ label = expressiong
240
+
241
+ if len(label)==0:
242
+ continue
243
+ height, width, _ = ret.shape
244
+ FONT = cv2.FONT_HERSHEY_COMPLEX
245
+ label_width, label_height = cv2.getTextSize(label, FONT, min(width, height) * FONT_SCALE, math.ceil(min(width, height) * THICKNESS_SCALE))[0]
246
+
247
+ cv2.rectangle(ret, (x1,y1), (x1+label_width,(y1 -label_height) - int(height * TEXT_Y_OFFSET_SCALE)), RGB, -1)
248
+
249
+ cv2.putText(
250
+ ret,
251
+ label,
252
+ (x1, y1 - int(height * TEXT_Y_OFFSET_SCALE)),
253
+ fontFace=FONT,
254
+ fontScale=min(width, height) * FONT_SCALE,
255
+ thickness=math.ceil(min(width, height) * THICKNESS_SCALE),
256
+ color=(255,255,255),
257
+ )
258
+
259
+
260
+
261
+ ret = ret.astype('uint8')
262
+ return ret
263
+
264
+
265
+ else: #visual prompt
266
+ topK_instance = 1
267
+ copyed_img = img['background'][:,:,:3].copy()
268
+ # get bbox from scribbles in layers
269
+ bbox_list = [scribble2box(layer) for layer in img['layers'] ]
270
+ visual_prompt_list = []
271
+ visual_prompt_RGB_list = []
272
+
273
+ for mask, (box,RGB) in zip(img['layers'], bbox_list):
274
+ if box is None:
275
+ continue
276
+ if prompt_mode=='box':
277
+ fakemask = np.zeros_like(copyed_img[:,:,0])
278
+ x1 ,y1 ,x2, y2 = box
279
+ fakemask[ y1:y2, x1:x2 ] = 1
280
+ fakemask = fakemask>0
281
+ elif prompt_mode=='point':
282
+ fakemask = np.zeros_like(copyed_img[:,:,0])
283
+ H,W = fakemask.shape
284
+ x1 ,y1 ,x2, y2 = box
285
+ center_x, center_y = (x1+x2)//2, (y1+y2)//2
286
+ fakemask[ center_y-H//40:center_y+H//40, center_x-W//40:center_x+W//40 ] = 1
287
+ fakemask = fakemask>0
288
+ elif prompt_mode=='scribble':
289
+ fakemask = mask[:,:,-1]
290
+ fakemask = fakemask>0
291
+
292
+ fakemask = torch.from_numpy(fakemask).unsqueeze(0).to(ori_image)
293
+ if inference_type == 'LSJ':
294
+ infer_visual_prompt = torch.zeros(1,1024,1024).to(ori_image)
295
+ infer_visual_prompt[:,:inference_size,:inference_size] = fakemask
296
+ else:
297
+ resize_fakemask = resizer(fakemask)
298
+ if size_divisibility > 1:
299
+ # the last two dims are H,W, both subject to divisibility requirement
300
+ infer_visual_prompt = torch.zeros(1,padding_size[0],padding_size[1]).to(resize_fakemask)
301
+ infer_visual_prompt[:,:image_size[0],:image_size[1]] = resize_fakemask
302
+
303
+
304
+ visual_prompt_list.append( infer_visual_prompt>0 )
305
+ visual_prompt_RGB_list.append(RGB)
306
+
307
+
308
+ mask_results_list = []
309
+ for visual_prompt in visual_prompt_list:
310
+
311
+ prompt_list = {'spatial':[visual_prompt]}
312
+
313
+ with torch.no_grad():
314
+ (outputs,_) = GLEEmodel(infer_image, prompt_list, task="coco", batch_name_list=['object'], is_train=False, visual_prompt_type=prompt_mode )
315
+
316
+ mask_pred = outputs['pred_masks'][0]
317
+ mask_cls = outputs['pred_logits'][0]
318
+ boxes_pred = outputs['pred_boxes'][0]
319
+
320
+ scores = mask_cls.sigmoid().max(-1)[0]
321
+ scores_per_image, topk_indices = scores.topk(topK_instance, sorted=True)
322
+
323
+ pred_class = mask_cls[topk_indices].max(-1)[1].tolist()
324
+ pred_boxes = boxes_pred[topk_indices]
325
+
326
+
327
+ boxes = LSJ_box_postprocess(pred_boxes,padding_size,re_size, ori_height,ori_width)
328
+ mask_pred = mask_pred[topk_indices]
329
+ pred_masks = F.interpolate( mask_pred[None,], size=(padding_size[0], padding_size[1]), mode="bilinear", align_corners=False )
330
+ pred_masks = pred_masks[:,:,:re_size[0],:re_size[1]]
331
+ pred_masks = F.interpolate( pred_masks, size=(ori_height,ori_width), mode="bilinear", align_corners=False )
332
+ pred_masks = (pred_masks>0).detach().cpu().numpy()[0]
333
+ mask_results_list.append(pred_masks)
334
+
335
+ zero_mask = np.zeros_like(copyed_img)
336
+ for mask,RGB in zip(mask_results_list,visual_prompt_RGB_list):
337
+ mask = mask.reshape(mask.shape[-2], mask.shape[-1], 1)
338
+ lar = np.concatenate((mask*RGB[0], mask*RGB[1],mask*RGB[2]), axis = 2)
339
+ zero_mask = zero_mask+ lar
340
+ lar_valid = zero_mask>0
341
+ masked_image = lar_valid*copyed_img
342
+ img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,255)*(1-mask_image_mix_ration)
343
+ max_p = img_n.max()
344
+ img_n = 255*img_n/max_p
345
+ ret = (~lar_valid*copyed_img)*mask_image_mix_ration + img_n
346
+ ret = ret.astype('uint8')
347
+ # cv2.imwrite('00020_inst.jpg', cv2.cvtColor(ret, cv2.COLOR_BGR2RGB))
348
+
349
+ return ret
350
+
351
+
352
+
353
+ # def get_select_coordinates(img):
354
+
355
+ # # img{'background': (H,W,3)
356
+ # # 'layers': list[ (H,W,4(RGBA)) ], draw map
357
+ # # 'composite': (H,W,4(RGBA))} ori_img concat drow
358
+
359
+ # ori_img = img['background'][:,:,:3].copy()
360
+
361
+ # # get bbox from scribbles in layers
362
+ # bbox_list = [scribble2box(layer) for layer in img['layers'] ]
363
+ # for mask, (box,RGB) in zip(img['layers'], bbox_list):
364
+ # if box is None:
365
+ # continue
366
+ # cv2.rectangle(ori_img, (box[0],box[1]), (box[2],box[3]),RGB, 3)
367
+ # return ori_img
368
+
369
+ def visual_prompt_preview(img, prompt_mode):
370
+
371
+ copyed_img = img['background'][:,:,:3].copy()
372
+
373
+ # get bbox from scribbles in layers
374
+ bbox_list = [scribble2box(layer) for layer in img['layers'] ]
375
+ zero_mask = np.zeros_like(copyed_img)
376
+
377
+ for mask, (box,RGB) in zip(img['layers'], bbox_list):
378
+ if box is None:
379
+ continue
380
+
381
+ if prompt_mode=='box':
382
+ fakemask = np.zeros_like(copyed_img[:,:,0])
383
+ x1 ,y1 ,x2, y2 = box
384
+ fakemask[ y1:y2, x1:x2 ] = 1
385
+ fakemask = fakemask>0
386
+ elif prompt_mode=='point':
387
+ fakemask = np.zeros_like(copyed_img[:,:,0])
388
+ H,W = fakemask.shape
389
+ x1 ,y1 ,x2, y2 = box
390
+ center_x, center_y = (x1+x2)//2, (y1+y2)//2
391
+ fakemask[ center_y-H//40:center_y+H//40, center_x-W//40:center_x+W//40 ] = 1
392
+ fakemask = fakemask>0
393
+ else:
394
+ fakemask = mask[:,:,-1]
395
+ fakemask = fakemask>0
396
+
397
+ mask = fakemask.reshape(fakemask.shape[0], fakemask.shape[1], 1)
398
+ lar = np.concatenate((mask*RGB[0], mask*RGB[1],mask*RGB[2]), axis = 2)
399
+ zero_mask = zero_mask+ lar
400
+
401
+
402
+ img_n = copyed_img + np.clip(zero_mask,0,255)
403
+ max_p = img_n.max()
404
+ ret = 255*img_n/max_p
405
+ ret = ret.astype('uint8')
406
+ return ret
407
+
408
+
409
+
410
+ with gr.Blocks() as demo:
411
+ gr.Markdown('# GLEE: General Object Foundation Model for Images and Videos at Scale')
412
+ gr.Markdown('## [Paper](https://arxiv.org/abs/2312.09158) - [Project Page](https://glee-vision.github.io) - [Code](https://github.com/FoundationVision/GLEE) ')
413
+
414
+ gr.Markdown(
415
+ '**The functionality demonstration demo app of GLEE. Select a Tab for image or video tasks. Image tasks includes arbitrary vocabulary object detection&segmentation, any form of object name or object caption detection, referring expression comprehension, and interactive segmentation. Video tasks add object tracking functionality based on image tasks.**'
416
+ )
417
+
418
+
419
+ with gr.Tab("Image task"):
420
+ with gr.Row():
421
+ with gr.Column():
422
+
423
+ img_input = gr.ImageEditor()
424
+ model_select = gr.Dropdown(
425
+ ["GLEE-Lite (R50)", "GLEE-Plus (SwinL)"], value = "GLEE-Plus (SwinL)" , multiselect=False, label="Model",
426
+ )
427
+ with gr.Row():
428
+ with gr.Column():
429
+ prompt_mode_select = gr.Radio(["point", "scribble", "box", "categories", "expression"], label="Prompt", value= "categories" , info="What kind of prompt do you want to use?")
430
+ category_select = gr.Dropdown(
431
+ ["COCO-80", "OBJ365", "Custom-List", "Class-Agnostic"], value = "COCO-80" , multiselect=False, label="Categories", info="Choose an existing category list or class-agnostic"
432
+ )
433
+ custom_category = gr.Textbox(
434
+ label="Custom Category",
435
+ info="Input custom category list, seperate by ',' ",
436
+ lines=1,
437
+ value="dog, cat, car, person",
438
+ )
439
+ input_expressiong = gr.Textbox(
440
+ label="Expression",
441
+ info="Input any description of an object in the image ",
442
+ lines=2,
443
+ value="the red car",
444
+ )
445
+ # with gr.Column():
446
+ with gr.Group():
447
+ with gr.Accordion("Text based detection usage",open=False):
448
+ gr.Markdown(
449
+ 'Press the "Detect & Segment" button directly to try the effect using the COCO category.<br />\
450
+ GLEE supports three kind of object perception methods: category list, textual description, and class-agnostic.<br />\
451
+ 1.Select an existing category list from the "Categories" dropdown, like COCO or OBJ365, or customize your own list.<br />\
452
+ 2.Enter arbitrary object name in "Custom Category", or choose the expression model and describe the object in "Expression Textbox" for single object detection only.<br />\
453
+ 3.For class-agnostic mode, choose "Class-Agnostic" from the "Categories" dropdown.'
454
+ )
455
+ with gr.Accordion("Interactive segmentation usage",open=False):
456
+ gr.Markdown(
457
+ 'For interactive segmentation:<br />\
458
+ 1.Draw points, boxes, or scribbles on the canvas for multiclass segmentation; use separate layers for different objects, adding layers with a "+" sign.<br />\
459
+ 2.Point mode accepts a single point only; multiple points default to the centroid, so use boxes or scribbles for larger objects.<br />\
460
+ 3.After drawing, click green "√" on the right side of the image to preview the prompt visualization; the segmentation mask follows the chosen prompt colors.'
461
+ )
462
+
463
+ img_showbox = gr.Image(label="visual prompt area preview")
464
+
465
+
466
+
467
+
468
+ with gr.Column():
469
+ image_segment = gr.Image(label="detection and segmentation results")
470
+ with gr.Accordion("Try More Visualization Options"):
471
+ results_select = gr.CheckboxGroup(["box", "mask", "name", "score", "expression"], value=["box", "mask", "name", "score"], label="Shown Results", info="The results shown on image")
472
+ num_inst_select = gr.Slider(1, 50, value=15, step=1, label="Num of topK instances for category based detection", info="Choose between 1 and 50 for better visualization")
473
+ threshold_select = gr.Slider(0, 1, value=0.2, label="Confidence Threshold", info="Choose threshold ")
474
+ mask_image_mix_ration = gr.Slider(0, 1, value=0.45, label="Image Brightness Ratio", info="Brightness between image and colored masks ")
475
+
476
+
477
+
478
+ image_button = gr.Button("Detect & Segment")
479
+ img_input.change(visual_prompt_preview, inputs = [img_input,prompt_mode_select] , outputs = img_showbox)
480
+ image_button.click(segment_image, inputs=[img_input, prompt_mode_select, category_select, custom_category,input_expressiong, results_select, num_inst_select, threshold_select, mask_image_mix_ration,model_select], outputs=image_segment)
481
+
482
+
483
+ with gr.Tab("Video task"):
484
+ with gr.Row():
485
+ gr.Markdown(
486
+ '# Due to computational resource limitations, support for video tasks is being processed and is expected to be available within a week.'
487
+ )
488
+ video_input = gr.Image()
489
+ video_button = gr.Button("Segment&Track")
490
+
491
+ if __name__ == '__main__':
492
+ demo.launch(inbrowser=True,share=True)