add: adding nyu model

app.py

@@ -15,43 +15,46 @@ torch.set_grad_enabled(False)
 #     "anhquancao/monoscene_kitti", trust_remote_code=True, revision='bf033f87c2a86b60903ab811b790a1532c1ae313'
 # )#.cuda()
 model = MonoScene.load_from_checkpoint(
-        "monoscene_kitti.ckpt",
-        dataset="kitti",
-        n_classes=20,
-        feature = 64,
-        project_scale = 2,
-        full_scene_size = (256, 256, 32),
-    )
+    "monoscene_nyu.ckpt",
+    dataset="NYU",
+    feature=200,
+    project_scale=1,
+    full_scene_size=(60, 36, 60),
+)
 
-img_W, img_H = 1220, 370
+img_W, img_H = 640, 480
 
 
 def predict(img):
     img = np.array(img, dtype=np.float32, copy=False) / 255.0
 
     normalize_rgb = transforms.Compose(
-            [
-                transforms.ToTensor(),
-                transforms.Normalize(
-                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
-                ),
-            ]
-        )
+        [
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+            ),
+        ]
+    )
     img = normalize_rgb(img)
-   
+
     batch = get_projections(img_W, img_H)
     batch["img"] = img
     for k in batch:
-        batch[k] = batch[k].unsqueeze(0)#.cuda()
+        batch[k] = batch[k].unsqueeze(0)  # .cuda()
 
     pred = model(batch).squeeze()
-    # print(pred.shape)
-    pred = majority_pooling(pred, k_size=2)
-    fig = draw(pred, batch['fov_mask_2'])
-
+    y_pred = torch.softmax(pred["ssc_logit"], dim=1).detach().cpu().numpy()
+    cam_pose = np.asarray([[ 9.6699458e-01,  4.2662762e-02,  2.5120059e-01,  0.0000000e+00],
+       [-2.5147417e-01,  1.0867463e-03,  9.6786356e-01,  0.0000000e+00],
+       [ 4.1018680e-02, -9.9908894e-01,  1.1779292e-02,  1.1794727e+00],
+       [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  1.0000000e+00]])
+    vox_origin = np.array([-1.54591799,  0.8907361 , -0.05 ])
+    
+    fig = draw(y_pred.squeeze(),cam_pose, vox_origin)
 
     return fig
-   
+
 
 description = """
 MonoScene Demo on SemanticKITTI Validation Set (Sequence 08), which uses the <b>camera parameters of Sequence 08</b>.
@@ -66,7 +69,7 @@ The output is <b>downsampled by 2</b> for faster rendering. <b>Darker</b> colors
 </center>
 """
 title = "MonoScene: Monocular 3D Semantic Scene Completion"
-article="""
+article = """
 <center>
 We also released a <b>smaller</b> MonoScene model (Half resolution - w/o 3D CRP) at: <a href="https://huggingface.co/spaces/CVPR/monoscene_lite">https://huggingface.co/spaces/CVPR/monoscene_lite</a>
     <img src='https://visitor-badge.glitch.me/badge?page_id=anhquancao.MonoScene&left_color=darkmagenta&right_color=purple' alt='visitor badge'>
@@ -110,11 +113,10 @@ examples = [
 ]
 
 
-
 demo = gr.Interface(
-    predict, 
-    gr.Image(shape=(1220, 370)), 
-    gr.Plot(),  
+    predict,
+    gr.Image(shape=(1220, 370)),
+    gr.Plot(),
     article=article,
     title=title,
     enable_queue=True,
@@ -124,4 +126,4 @@ demo = gr.Interface(
     description=description)
 
 
-demo.launch(enable_queue=True, debug=False)
+demo.launch(enable_queue=True, debug=False)