Refactor predict_normal function and add DSINE demo

main.py

@@ -20,7 +20,7 @@ model = utils.load_checkpoint("./checkpoints/dsine.pt", model)
 model.eval()
 
 
-def predict_normal(img_np: np.ndarray):
+def predict_normal(img_np: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     # normalize
     normalize = transforms.Normalize(
         mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
@@ -28,7 +28,7 @@ def predict_normal(img_np: np.ndarray):
 
     with torch.no_grad():
         img = np.array(img_np).astype(np.float32) / 255.0
-        img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to("cpu")
+        img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to(device)
         _, _, orig_H, orig_W = img.shape
 
         # zero-pad the input image so that both the width and height are multiples of 32
@@ -39,7 +39,7 @@ def predict_normal(img_np: np.ndarray):
         # NOTE: if intrins is not given, we just assume that the principal point is at the center
         # and that the field-of-view is 60 degrees (feel free to modify this assumption)
         intrins = utils.get_intrins_from_fov(
-            new_fov=60.0, H=orig_H, W=orig_W, device="cpu"
+            new_fov=60.0, H=orig_H, W=orig_W, device=device
         ).unsqueeze(0)
 
         intrins[:, 0, 2] += l
@@ -48,7 +48,6 @@ def predict_normal(img_np: np.ndarray):
         pred_norm = model(img, intrins=intrins)[-1]
         pred_norm = pred_norm[:, :, t : t + orig_H, l : l + orig_W]
 
-        # save to output folder
         # NOTE: by saving the prediction as uint8 png format, you lose a lot of precision
         # if you want to use the predicted normals for downstream tasks, we recommend saving them as float32 NPY files
         pred_norm_np = (
@@ -60,6 +59,12 @@ def predict_normal(img_np: np.ndarray):
 
 
 with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # DSINE
+        Unofficial Gradio demo of [DSINE: Rethinking Inductive Biases for Surface Normal Estimation](https://github.com/baegwangbin/DSINE)
+        """
+    )
     with gr.Group():
         with gr.Row():
             input_img = gr.Image(label="Input image", image_mode="RGB")
@@ -71,9 +76,12 @@ with gr.Blocks() as demo:
 
     with Modal(visible=True, allow_user_close=False) as modal:
         gr.Markdown(
-            "To use this space, you must agree to the terms and conditions. found [here](https://github.com/baegwangbin/DSINE/blob/main/LICENSE)."
+            """
+            To use this space, you must agree to the terms and conditions.
+            Found [HERE](https://github.com/baegwangbin/DSINE/blob/main/LICENSE).
+            """,
         )
-        btn = gr.Button("I agree")
+        btn = gr.Button("I Agree to the Terms and Conditions")
         btn.click(lambda: Modal(visible=False), None, modal)
 
 if __name__ == "__main__":