upload

app.py

@@ -80,7 +80,7 @@ def process_depth(
         batch_size=1 if processing_res == 0 else 0,
         show_progress_bar=False,
         mode='depth',
-        cmap='Spectral',
+        color_map='Spectral',
     )
 
     depth_pred = pipe_out.pred_np
@@ -252,7 +252,7 @@ def process_disparity(
         batch_size=1 if processing_res == 0 else 0,
         show_progress_bar=False,
         mode='disparity',
-        cmap='Spectral',
+        color_map='Spectral',
     )
 
     disparity_pred = pipe_out.pred_np

genpercept/genpercept_pipeline.py

@@ -41,6 +41,7 @@ from .util.image_util import (
 
 import matplotlib.pyplot as plt
 from genpercept.models.dpt_head import DPTNeckHeadForUnetAfterUpsampleIdentity
+from genpercept.util.image_util import process_normals
 
 
 class GenPerceptOutput(BaseOutput):
@@ -309,6 +310,9 @@ class GenPerceptPipeline(DiffusionPipeline):
         # Clip output range
         pipe_pred = pipe_pred.clip(0, 1)
 
+        if mode == 'normal':
+            pred_np = process_normals(torch.from_numpy(pred_np)[None])
+
         # Colorize
         if color_map is not None:
             assert self.mode == 'depth'

genpercept/util/image_util.py

@@ -25,6 +25,19 @@ import torch
 from torchvision.transforms import InterpolationMode
 from torchvision.transforms.functional import resize
 
+def process_normals(input_images:torch.Tensor):
+    normal_preds = input_images
+    bsz, d, h, w = normal_preds.shape
+    normal_preds = normal_preds / (torch.norm(normal_preds, p=2, dim=1).unsqueeze(1)+1e-5)
+    phi = torch.atan2(normal_preds[:,1,:,:], normal_preds[:,0,:,:]).mean(dim=0)
+    theta = torch.atan2(torch.norm(normal_preds[:,:2,:,:], p=2, dim=1), normal_preds[:,2,:,:]).mean(dim=0)
+    normal_pred = torch.zeros((d,h,w)).to(normal_preds)
+    normal_pred[0,:,:] = torch.sin(theta) * torch.cos(phi)
+    normal_pred[1,:,:] = torch.sin(theta) * torch.sin(phi)
+    normal_pred[2,:,:] = torch.cos(theta) 
+    angle_error = torch.acos(torch.clip(torch.cosine_similarity(normal_pred[None], normal_preds, dim=1),-0.999, 0.999))
+    normal_idx = torch.argmin(angle_error.reshape(bsz,-1).sum(-1))
+    return normal_preds[normal_idx]
 
 def colorize_depth_maps(
     depth_map, min_depth, max_depth, cmap="Spectral", valid_mask=None