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__":