update to new version

app.py

@@ -3,24 +3,23 @@ import os
 try:
     import perspective2d
 except:
-    os.system(f"pip install git+https://github.com/jinlinyi/PerspectiveFields.git@dev#egg=perspective2d")
+    os.system(f"pip install git+https://github.com/jinlinyi/PerspectiveFields.git@v1.0.0")
 
 
 import gradio as gr
 import cv2
 import copy
+import numpy as np
+import os.path as osp
+from datetime import datetime
+
 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 perspective2d import PerspectiveFields
+from perspective2d.utils import draw_perspective_fields, draw_from_r_p_f_cx_cy
+from perspective2d.perspectivefields import model_zoo
 
 
 
@@ -32,25 +31,6 @@ description = ""
 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]
@@ -80,36 +60,26 @@ def resize_fix_aspect_ratio(img, field, target_width=None, target_height=None):
     return img, field
 
 
-def inference(img, model_type):
+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().cuda()
+    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(),
@@ -121,15 +91,21 @@ def inference(img, model_type):
             pred['pred_rel_cx'].cpu().item(),
             pred['pred_rel_cy'].cpu().item(),
         ]
-        param = f"roll {pred['pred_roll'].cpu().item() :.2f}\npitch {pred['pred_pitch'].cpu().item() :.2f}\nfov {pred['pred_general_vfov'].cpu().item() :.2f}\n"
+        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_rgb.shape}
+    model_type {model_type}
+    param {param}
+    """
+    )
     return Image.fromarray(pred_vis), param
 
 examples = []
@@ -138,46 +114,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(

requirements.txt

@@ -3,22 +3,17 @@
 
 torch==1.11.0+cu113
 torchvision==0.12.0+cu113
-mmcv==2.0.0rc4
 albumentations==1.3.0
-fvcore==0.1.5.post20221221
 gradio==3.29.0
-h5py==3.8.0
-imageio==2.27.0
-matplotlib==3.7.1
-numpy==1.24.2
-opencv_contrib_python==4.7.0.72
-opencv_python==4.7.0.72
-opencv_python_headless==4.7.0.72
-Pillow==9.5.0
-pycocotools==2.0.6
+albumentations
+matplotlib
+numpy
+omegaconf
+opencv-contrib-python
+pillow
 pyequilib==0.3.0
-scikit_learn==1.2.2
-scipy==1.9.1
-setuptools==59.5.0
-timm==0.6.13
-tqdm==4.65.0
+scikit-learn
+scipy
+setuptools
+timm
+yacs