torch.hub update

README.md

@@ -190,20 +190,19 @@ The test commands above (e.g. for getting the benchmark performance & running re
 If you want to make contributions to this repo, please make a pull request and add instructions in the following format.
 
 <details>
-<summary><b>Using torch hub to predict normal</b> (contribution by <a href="https://github.com/hugoycj" target="_blank">hugoycj</a>)</summary>
+<summary><b>Using torch hub to predict normal</b> (contribution by <a href="https://github.com/hugoycj" target="_blank">hugoycj</a>, updated by <a href="https://github.com/pierremerriaux-leddartech/DSINE" target="_blank">Pierre M.</a>)</summary>
+
 
-NOTE: the code below is deprecated and should be modified (as the folder structure has changed).
 
 ```
 import torch
 import cv2
 import numpy as np
+import sys
 
-# Load the normal predictor model from torch hub
-normal_predictor = torch.hub.load("hugoycj/DSINE-hub", "DSINE", trust_repo=True)
+normal_predictor = torch.hub.load("pierremerriaux-leddartech/DSINE", "DSINE", trust_repo=True, source='github')
 
-# Load the input image using OpenCV
-image = cv2.imread(args.input, cv2.IMREAD_COLOR)
+image = cv2.imread('projects/dsine/samples/img/office_01.png', cv2.IMREAD_COLOR)
 h, w = image.shape[:2]
 
 # Use the model to infer the normal map from the input image
@@ -216,13 +215,13 @@ normal = (normal * 255).cpu().numpy().astype(np.uint8).transpose(1, 2, 0)
 normal = cv2.cvtColor(normal, cv2.COLOR_RGB2BGR)
 
 # Save the output normal map to a file
-cv2.imwrite(args.output, normal)
+cv2.imwrite('projects/dsine/samples/img/office_01_result.png', normal)
 ```
 
 If the network is unavailable to retrieve weights, you can use local weights for torch hub as shown below:
 
 ```
-normal_predictor = torch.hub.load("hugoycj/DSINE-hub", "DSINE", local_file_path='./checkpoints/dsine.pt', trust_repo=True)
+normal_predictor = torch.hub.load("pierremerriaux-leddartech/DSINE", "DSINE", local_file_path='./checkpoints/dsine.pt', trust_repo=True)
 ```
 </details>
 

hubconf.py

@@ -20,7 +20,7 @@ def _load_state_dict(local_file_path: Optional[str] = None):
 
 class Predictor:
     def __init__(self, model) -> None:
-        from models.dsine import DSINE
+        from models.dsine import v02
         self.device = torch.device('cuda')
         self.model = model
         self.transform = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
@@ -32,33 +32,44 @@ def infer_cv2(self, image):
 
     def infer_pil(self, img, intrins=None):
         import utils.utils as utils
+        from utils.projection import intrins_from_fov, intrins_from_txt
         img = np.array(img).astype(np.float32) / 255.0
         img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to(self.device)
         _, _, orig_H, orig_W = img.shape
 
         # zero-pad the input image so that both the width and height are multiples of 32
-        l, r, t, b = utils.pad_input(orig_H, orig_W)
-        img = F.pad(img, (l, r, t, b), mode="constant", value=0.0)
+        # l, r, t, b = utils.pad_input(orig_H, orig_W)
+
+        lrtb = utils.get_padding(orig_H, orig_W)
+        img = F.pad(img, lrtb, mode="constant", value=0.0)
+
+
         img = self.transform(img)
 
         if intrins is None:
-            intrins = utils.get_intrins_from_fov(new_fov=60.0, H=orig_H, W=orig_W, device=self.device).unsqueeze(0)
+            intrins = intrins_from_fov(new_fov=60.0, H=orig_H, W=orig_W, device=self.device).unsqueeze(0)
 
-        intrins[:, 0, 2] += l
-        intrins[:, 1, 2] += t
+        intrins[:, 0, 2] += lrtb[0]
+        intrins[:, 1, 2] += lrtb[2]
 
         with torch.no_grad():
             pred_norm = self.model(img, intrins=intrins)[-1]
-            pred_norm = pred_norm[:, :, t:t+orig_H, l:l+orig_W]
+            pred_norm = pred_norm[:, :, lrtb[2]:lrtb[2]+orig_H, lrtb[0]:lrtb[0]+orig_W]
 
         # pred_norm_np = pred_norm.cpu().detach().numpy()[0,:,:,:].transpose(1, 2, 0) # (H, W, 3)
         return pred_norm
 
 def DSINE(local_file_path: Optional[str] = None):
-    from models import dsine
-
+    from models.dsine import v02
+    import projects.dsine.config as config
+    import sys
+    old_sys_argv = sys.argv
+    sys.argv = [old_sys_argv[0]] + [os.path.join(os.path.dirname(os.path.realpath(__file__)), 'projects/dsine/experiments/exp001_cvpr2024/dsine.txt')]
+
+    args = config.get_args(test=True)
+    #args.exp_root = os.path.join(os.path.split(os.path.dirname(os.path.realpath(__file__)))[0], 'experiments')
     state_dict = _load_state_dict(local_file_path)
-    model = dsine.DSINE()
+    model = v02.DSINE_v02(args)
     model.load_state_dict(state_dict, strict=True)
     model.eval()
     model = model.to(torch.device("cuda"))

projects/__init__.py

@@ -3,10 +3,11 @@
 import os
 import argparse
 
-DATASET_DIR = '/media/gwangbin/ssd/datasets/'   # where datasets are stored         NOTE: this should be changed!
-EXPERIMENT_DIR = '/home/gwangbin/experiments/'  # where to save the experiments     NOTE: this should be changed!
+#DATASET_DIR = '/media/gwangbin/ssd/datasets/'   # where datasets are stored         NOTE: this should be changed!
+#EXPERIMENT_DIR = '/home/gwangbin/experiments/'  # where to save the experiments     NOTE: this should be changed!
 PROJECT_DIR = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0]
-
+EXPERIMENT_DIR  = os.path.join(PROJECT_DIR, 'experiments')
+DATASET_DIR  = os.path.join(PROJECT_DIR, 'datasets')
 
 def get_default_parser():
     parser = argparse.ArgumentParser(fromfile_prefix_chars='@', conflict_handler='resolve')

test_conf_hub.py

@@ -0,0 +1,21 @@
+import torch
+import cv2
+import numpy as np
+import sys
+
+normal_predictor = torch.hub.load("pierremerriaux-leddartech/DSINE", "DSINE", trust_repo=True, source='github')
+
+image = cv2.imread('projects/dsine/samples/img/office_01.png', cv2.IMREAD_COLOR)
+h, w = image.shape[:2]
+
+# Use the model to infer the normal map from the input image
+with torch.inference_mode():
+    normal = normal_predictor.infer_cv2(image)[0]  # Output shape: (H, W, 3)
+    normal = (normal + 1) / 2  # Convert values to the range [0, 1]
+
+# Convert the normal map to a displayable format
+normal = (normal * 255).cpu().numpy().astype(np.uint8).transpose(1, 2, 0)
+normal = cv2.cvtColor(normal, cv2.COLOR_RGB2BGR)
+
+# Save the output normal map to a file
+cv2.imwrite('projects/dsine/samples/img/office_01_result.png', normal)