Update app.py

app.py

@@ -67,7 +67,30 @@ COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
                 [0.700, 0.300, 0.600],[0.000, 0.447, 0.741], [0.850, 0.325, 0.098]]
 
 
-
+def load_smallmodel():
+    cfg_r50 = get_cfg()
+    add_deeplab_config(cfg_r50)
+    add_glee_config(cfg_r50)
+    conf_files_r50 = 'GLEE/configs/R50.yaml'
+    checkpoints_r50 = torch.load('GLEE_R50_Scaleup10m.pth') 
+    cfg_r50.merge_from_file(conf_files_r50)
+    GLEEmodel_r50 = GLEE_Model(cfg_r50, None, device, None, True).to(device)
+    GLEEmodel_r50.load_state_dict(checkpoints_r50, strict=False)
+    GLEEmodel_r50.eval()
+    return GLEEmodel_r50
+def load_bigmodel():
+    cfg_swin = get_cfg()
+    add_deeplab_config(cfg_swin)
+    add_glee_config(cfg_swin)
+    conf_files_swin = 'GLEE/configs/SwinL.yaml'
+    checkpoints_swin = torch.load('GLEE_SwinL_Scaleup10m.pth') 
+    cfg_swin.merge_from_file(conf_files_swin)
+    GLEEmodel_swin = GLEE_Model(cfg_swin, None, device, None, True).to(device)
+    GLEEmodel_swin.load_state_dict(checkpoints_swin, strict=False)
+    GLEEmodel_swin.eval()
+    return GLEEmodel_swin
+GLEEmodel_swin = load_bigmodel()
+GLEEmodel_r50 = load_smallmodel()
 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']
 OBJ365_class_names = [cat['name'] for cat in OBJ365_CATEGORIESV2]
 class_agnostic_name = ['object']
@@ -79,27 +102,6 @@ else:
     print('use cpu')
     device='cpu'
 
-cfg_r50 = get_cfg()
-add_deeplab_config(cfg_r50)
-add_glee_config(cfg_r50)
-conf_files_r50 = 'GLEE/configs/R50.yaml'
-checkpoints_r50 = torch.load('GLEE_R50_Scaleup10m.pth') 
-cfg_r50.merge_from_file(conf_files_r50)
-GLEEmodel_r50 = GLEE_Model(cfg_r50, None, device, None, True).to(device)
-GLEEmodel_r50.load_state_dict(checkpoints_r50, strict=False)
-GLEEmodel_r50.eval()
-
-
-cfg_swin = get_cfg()
-add_deeplab_config(cfg_swin)
-add_glee_config(cfg_swin)
-conf_files_swin = 'GLEE/configs/SwinL.yaml'
-checkpoints_swin = torch.load('GLEE_SwinL_Scaleup10m.pth') 
-cfg_swin.merge_from_file(conf_files_swin)
-GLEEmodel_swin = GLEE_Model(cfg_swin, None, device, None, True).to(device)
-GLEEmodel_swin.load_state_dict(checkpoints_swin, strict=False)
-GLEEmodel_swin.eval()
-
 pixel_mean = torch.Tensor( [123.675, 116.28, 103.53]).to(device).view(3, 1, 1)
 pixel_std = torch.Tensor([58.395, 57.12, 57.375]).to(device).view(3, 1, 1)
 normalizer = lambda x: (x - pixel_mean) / pixel_std