jinlinyi/PerspectiveFields

1import spaces2import os3try:4    import perspective2d5except:6    os.system(f"pip install git+https://github.com/jinlinyi/PerspectiveFields.git@hf-debug")7 8 9import gradio as gr10import cv211import copy12import numpy as np13import os.path as osp14from datetime import datetime15 16import torch17from PIL import Image, ImageDraw18from glob import glob19 20from perspective2d import PerspectiveFields21from perspective2d.utils import draw_perspective_fields, draw_from_r_p_f_cx_cy22from perspective2d.perspectivefields import model_zoo23 24 25 26 27title = "Perspective Fields Demo"28 29description = """30<p style="text-align: center">31    <a href="https://jinlinyi.github.io/PerspectiveFields/" target="_blank">Project Page</a> | 32    <a href="https://arxiv.org/abs/2212.03239" target="_blank">Paper</a> | 33    <a href="https://github.com/jinlinyi/PerspectiveFields" target="_blank">Code</a> | 34    <a href="https://www.youtube.com/watch?v=sN5B_ZvMva8&themeRefresh=1" target="_blank">Video</a>35</p>36<h2>Gradio Demo</h2>37<p>Try our Gradio demo for Perspective Fields for single image camera calibration. You can click on one of the provided examples or upload your own image.</p>38<h3>Available Models:</h3>39<ol>40    <li><span style="color:red;">[NEW!!!]</span><strong>Paramnet-360Cities-edina:</strong> Our latest model trained on <a href="https://www.360cities.net/">360cities</a> and <a href="https://github.com/tien-d/EgoDepthNormal/tree/main#egocentric-depth-on-everyday-indoor-activities-edina-dataset">EDINA</a> dataset.</li>41    <li><strong>PersNet-360Cities:</strong> PerspectiveNet trained on the 360Cities dataset. This model predicts perspective fields and is designed to be robust and generalize well to both indoor and outdoor images.</li>42    <li><strong>PersNet_Paramnet-GSV-uncentered:</strong> A combination of PerspectiveNet and ParamNet trained on the Google Street View (GSV) dataset. This model predicts camera Roll, Pitch, and Field of View (FoV), as well as the Principal Point location.</li>43    <li><strong>PersNet_Paramnet-GSV-centered:</strong> PerspectiveNet+ParamNet trained on the GSV dataset. This model assumes the principal point is at the center of the image and predicts camera Roll, Pitch, and FoV.</li>44</ol>45"""46 47 48article = """49<p style='text-align: center'><a href='https://arxiv.org/abs/2212.03239' target='_blank'>Perspective Fields for Single Image Camera Calibrations</a> | <a href='https://github.com/jinlinyi/PerspectiveFields' target='_blank'>Github Repo</a></p>50"""51 52 53 54def resize_fix_aspect_ratio(img, field, target_width=None, target_height=None):55    height = img.shape[0]56    width = img.shape[1]57    if target_height is None:58        factor = target_width / width59    elif target_width is None:60        factor = target_height / height61    else:62        factor = max(target_width / width, target_height / height)63    if factor == target_width / width:64        target_height = int(height * factor)65    else:66        target_width = int(width * factor)67        68    img = cv2.resize(img, (target_width, target_height))69    for key in field:70        if key not in ['up', 'lati']:71            continue72        tmp = field[key].numpy()73        transpose = len(tmp.shape) == 374        if transpose:75            tmp = tmp.transpose(1,2,0)76        tmp = cv2.resize(tmp, (target_width, target_height))77        if transpose:78            tmp = tmp.transpose(2,0,1)79        field[key] = torch.tensor(tmp)80    return img, field81 82@spaces.GPU83def inference(img_rgb, model_type):84    if model_type is None:85        return None, ""86    pf_model = PerspectiveFields(model_type).eval().to(device)87    pred = pf_model.inference(img_bgr=img_rgb[...,::-1])88    img_h = img_rgb.shape[0]89    field = {90        'up': pred['pred_gravity_original'].cpu().detach(),91        'lati': pred['pred_latitude_original'].cpu().detach(),92    }93    img_rgb, field = resize_fix_aspect_ratio(img_rgb, field, 640)94    if not model_zoo[model_type]['param']:95        pred_vis = draw_perspective_fields(96            img_rgb,97            field['up'],98            torch.deg2rad(field['lati']),99            color=(0,1,0),100        )101        param = "Not Implemented"102    else:103        r_p_f_rad = np.radians(104            [105                pred['pred_roll'].cpu().item(),106                pred['pred_pitch'].cpu().item(),107                pred['pred_general_vfov'].cpu().item(),108            ]109        )110        cx_cy = [111            pred['pred_rel_cx'].cpu().item(),112            pred['pred_rel_cy'].cpu().item(),113        ]114        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"115        param += f"principal point {pred['pred_rel_cx'].cpu().item() :.2f} {pred['pred_rel_cy'].cpu().item() :.2f}"116        pred_vis = draw_from_r_p_f_cx_cy(117            img_rgb, 118            *r_p_f_rad,119            *cx_cy,120            'rad',121            up_color=(0,1,0),122        )123    print(f"""time {datetime.now().strftime("%H:%M:%S")}124    img.shape {img_rgb.shape}125    model_type {model_type}126    param {param}127    """128    )129    return Image.fromarray(pred_vis), param130 131examples = []132for img_name in glob('assets/imgs/*.*g'):133    examples.append([img_name])134print(examples)135 136device = 'cuda' if torch.cuda.is_available() else 'cpu'137 138info = """Select model\n"""139gr.Interface(140    fn=inference,141    inputs=[142        "image", 143        gr.Radio(144            list(model_zoo.keys()), 145            value=list(sorted(model_zoo.keys()))[0], 146            label="Model", 147            info=info,148        ),149    ],150    outputs=[gr.Image(label='Perspective Fields'), gr.Textbox(label='Pred Camera Parameters')],151    title=title,152    description=description,153    article=article,154    examples=examples,155).launch()