update to new codebase

app.py

@@ -1,26 +1,26 @@
 
 import os
-os.system(f"pip install -U openmim")
-os.system(f"mim install mmcv")
-os.system(f"pip install git+https://github.com/jinlinyi/PerspectiveFields.git@dev#egg=perspective2d")
+try:
+    import perspective2d
+except:
+    os.system(f"pip install git+https://github.com/jinlinyi/PerspectiveFields.git@v1.0.0")
 
 
 import gradio as gr
 import cv2
 import copy
-import torch
-from PIL import Image, ImageDraw
-from glob import glob
 import numpy as np
 import os.path as osp
-from detectron2.config import get_cfg
-from detectron2.data.detection_utils import read_image
-from perspective2d.utils.predictor import VisualizationDemo
-import perspective2d.modeling  # noqa
-from perspective2d.config import get_perspective2d_cfg_defaults
-from perspective2d.utils import draw_from_r_p_f_cx_cy
 from datetime import datetime
 
+import torch
+from PIL import Image, ImageDraw
+from glob import glob
+
+from perspective2d import PerspectiveFields
+from perspective2d.utils import draw_perspective_fields, draw_from_r_p_f_cx_cy
+from perspective2d.perspectivefields import model_zoo
+
 
 
 
@@ -51,24 +51,6 @@ article = """
 
 
 
-
-def setup_cfg(args):
-    cfgs = []
-    configs = args['config_file'].split('#')
-    weights_id = args['opts'].index('MODEL.WEIGHTS') + 1
-    weights = args['opts'][weights_id].split('#')
-    for i, conf in enumerate(configs):
-        if len(conf) != 0:
-            tmp_opts = copy.deepcopy(args['opts'])
-            tmp_opts[weights_id] = weights[i]
-            cfg = get_cfg()
-            get_perspective2d_cfg_defaults(cfg)
-            cfg.merge_from_file(conf)
-            cfg.merge_from_list(tmp_opts)
-            cfg.freeze()
-            cfgs.append(cfg)
-    return cfgs
-
 def resize_fix_aspect_ratio(img, field, target_width=None, target_height=None):
     height = img.shape[0]
     width = img.shape[1]
@@ -98,37 +80,26 @@ def resize_fix_aspect_ratio(img, field, target_width=None, target_height=None):
     return img, field
 
 
-def inference(img, model_type):
-    img_h = img.shape[0]
+def inference(img_rgb, model_type):
     if model_type is None:
         return None, ""
-    perspective_cfg_list = setup_cfg(model_zoo[model_type])
-    demo = VisualizationDemo(cfg_list=perspective_cfg_list)
-
-    # img = read_image(image_path, format="BGR")
-    img = img[..., ::-1] # rgb->bgr
-    pred = demo.run_on_image(img)
+    pf_model = PerspectiveFields(model_type).eval().to(device)
+    pred = pf_model.inference(img_bgr=img_rgb[...,::-1])
+    img_h = img_rgb.shape[0]
     field = {
         'up': pred['pred_gravity_original'].cpu().detach(),
         'lati': pred['pred_latitude_original'].cpu().detach(),
     }
-    img, field = resize_fix_aspect_ratio(img, field, 640)
+    img_rgb, field = resize_fix_aspect_ratio(img_rgb, field, 640)
     if not model_zoo[model_type]['param']:
-        pred_vis = demo.draw(
-            image=img,
-            latimap=field['lati'],
-            gravity=field['up'],
-            latimap_format=pred['pred_latitude_original_mode'],
-        ).get_image()
+        pred_vis = draw_perspective_fields(
+            img_rgb,
+            field['up'],
+            torch.deg2rad(field['lati']),
+            color=(0,1,0),
+        )
         param = "Not Implemented"
     else:
-        if 'pred_general_vfov' not in pred.keys():
-            pred['pred_general_vfov'] = pred['pred_vfov']
-        if 'pred_rel_cx' not in pred.keys():
-            pred['pred_rel_cx'] = torch.FloatTensor([0])
-        if 'pred_rel_cy' not in pred.keys():
-            pred['pred_rel_cy'] = torch.FloatTensor([0])
-            
         r_p_f_rad = np.radians(
             [
                 pred['pred_roll'].cpu().item(),
@@ -143,14 +114,14 @@ def inference(img, model_type):
         param = f"roll {pred['pred_roll'].cpu().item() :.2f}\npitch {pred['pred_pitch'].cpu().item() :.2f}\nvertical fov {pred['pred_general_vfov'].cpu().item() :.2f}\nfocal_length {pred['pred_rel_focal'].cpu().item()*img_h :.2f}\n"
         param += f"principal point {pred['pred_rel_cx'].cpu().item() :.2f} {pred['pred_rel_cy'].cpu().item() :.2f}"
         pred_vis = draw_from_r_p_f_cx_cy(
-            img[:,:,::-1], 
+            img_rgb, 
             *r_p_f_rad,
             *cx_cy,
             'rad',
             up_color=(0,1,0),
         )
     print(f"""time {datetime.now().strftime("%H:%M:%S")}
-    img.shape {img.shape}
+    img.shape {img_rgb.shape}
     model_type {model_type}
     param {param}
     """
@@ -163,46 +134,6 @@ for img_name in glob('assets/imgs/*.*g'):
 print(examples)
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model_zoo = {
-
-    'NEW:Paramnet-360Cities-edina-centered': {
-        'weights': ['https://www.dropbox.com/s/z2dja70bgy007su/paramnet_360cities_edina_rpf.pth'],
-        'opts': ['MODEL.WEIGHTS', 'models/paramnet_360cities_edina_rpf.pth', 'MODEL.DEVICE', device,],
-        'config_file': 'models/paramnet_360cities_edina_rpf.yaml',
-        'param': True,
-    },
-
-   'NEW:Paramnet-360Cities-edina-uncentered': {
-        'weights': ['https://www.dropbox.com/s/nt29e1pi83mm1va/paramnet_360cities_edina_rpfpp.pth'],
-        'opts': ['MODEL.WEIGHTS', 'models/paramnet_360cities_edina_rpfpp.pth', 'MODEL.DEVICE', device,],
-        'config_file': 'models/paramnet_360cities_edina_rpfpp.yaml',
-        'param': True,
-    },
-
-    'PersNet-360Cities': {
-        'weights': ['https://www.dropbox.com/s/czqrepqe7x70b7y/cvpr2023.pth'],
-        'opts': ['MODEL.WEIGHTS', 'models/cvpr2023.pth', 'MODEL.DEVICE', device,],
-        'config_file': 'models/cvpr2023.yaml',
-        'param': False,
-    },
-    'PersNet_Paramnet-GSV-uncentered': {
-        'weights': ['https://www.dropbox.com/s/ufdadxigewakzlz/paramnet_gsv_rpfpp.pth'],
-        'opts': ['MODEL.WEIGHTS', 'models/paramnet_gsv_rpfpp.pth', 'MODEL.DEVICE', device,],
-        'config_file': 'models/paramnet_gsv_rpfpp.yaml',
-        'param': True,
-    },
-    # trained on GSV dataset, predicts Perspective Fields + camera parameters (roll, pitch, fov), assuming centered principal point
-    'PersNet_Paramnet-GSV-centered': {
-        'weights': ['https://www.dropbox.com/s/g6xwbgnkggapyeu/paramnet_gsv_rpf.pth'],
-        'opts': ['MODEL.WEIGHTS', 'models/paramnet_gsv_rpf.pth', 'MODEL.DEVICE', device,],
-        'config_file': 'models/paramnet_gsv_rpf.yaml',
-        'param': True,
-    },
-}
-for model_id in model_zoo:
-    html = model_zoo[model_id]['weights'][0]
-    if not os.path.exists(os.path.join('models', html.split('/')[-1])):
-        os.system(f"wget -P models/ {html}")
 
 info = """Select model\n"""
 gr.Interface(