Upload 41 files

annotator/normalbae/__init__.py

@@ -28,7 +28,8 @@ class NormalBaeDetector:
         args.importance_ratio = 0.7
         model = NNET(args)
         model = utils.load_checkpoint(modelpath, model)
-        model = model.cuda()
+#        model = model.cuda()
+        model = model.cpu()
         model.eval()
         self.model = model
         self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
@@ -37,7 +38,8 @@ class NormalBaeDetector:
         assert input_image.ndim == 3
         image_normal = input_image
         with torch.no_grad():
-            image_normal = torch.from_numpy(image_normal).float().cuda()
+#            image_normal = torch.from_numpy(image_normal).float().cuda()
+            image_normal = torch.from_numpy(image_normal).float().cpu()
             image_normal = image_normal / 255.0
             image_normal = rearrange(image_normal, 'h w c -> 1 c h w')
             image_normal = self.norm(image_normal)

annotator/normalbae/models/submodules/efficientnet_repo/geffnet/helpers.py

@@ -14,6 +14,7 @@ def load_checkpoint(model, checkpoint_path):
     if checkpoint_path and os.path.isfile(checkpoint_path):
         print("=> Loading checkpoint '{}'".format(checkpoint_path))
         checkpoint = torch.load(checkpoint_path)
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
         if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
             new_state_dict = OrderedDict()
             for k, v in checkpoint['state_dict'].items():