Merge pull request #4 from MedericFourmy/add_cpu_option

Add cpu option

Author: nim65s

cosypose/integrated/detector.py

@@ -11,20 +11,17 @@ def __init__(self, model):
         self.model = model
         self.config = model.config
         self.category_id_to_label = {v: k for k, v in self.config.label_to_category_id.items()}
-
-    def cast(self, obj):
-        return obj.cuda()
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
     @torch.no_grad()
     def get_detections(self, images, detection_th=None,
                        output_masks=False, mask_th=0.8,
                        one_instance_per_class=False):
-        images = self.cast(images).float()
+        images = images.to(self.device).float()
         if images.shape[-1] == 3:
             images = images.permute(0, 3, 1, 2)
         if images.max() > 1:
             images = images / 255.
-            images = images.float().cuda()
         outputs_ = self.model([image_n for image_n in images])
 
         infos = []
@@ -46,12 +43,12 @@ def get_detections(self, images, detection_th=None,
                 infos.append(info)
 
         if len(bboxes) > 0:
-            bboxes = torch.stack(bboxes).cuda().float()
-            masks = torch.stack(masks).cuda()
+            bboxes = torch.stack(bboxes).to(self.device).float()
+            masks = torch.stack(masks).to(self.device)
         else:
             infos = dict(score=[], label=[], batch_im_id=[])
-            bboxes = torch.empty(0, 4).cuda().float()
-            masks = torch.empty(0, images.shape[1], images.shape[2], dtype=torch.bool).cuda()
+            bboxes = torch.empty(0, 4).to(self.device).float()
+            masks = torch.empty(0, images.shape[1], images.shape[2], dtype=torch.bool).to(self.device)
 
         outputs = tc.PandasTensorCollection(
             infos=pd.DataFrame(infos),

cosypose/rendering/bullet_batch_renderer.py

@@ -6,15 +6,16 @@
 from .bullet_scene_renderer import BulletSceneRenderer
 
 
-def init_renderer(urdf_ds, preload=True):
+def init_renderer(urdf_ds, preload=True, gpu_renderer=True):
     renderer = BulletSceneRenderer(urdf_ds=urdf_ds,
                                    preload_cache=preload,
-                                   background_color=(0, 0, 0))
+                                   background_color=(0, 0, 0),
+                                   gpu_renderer=gpu_renderer)
     return renderer
 
 
-def worker_loop(worker_id, in_queue, out_queue, object_set, preload=True):
-    renderer = init_renderer(object_set, preload=preload)
+def worker_loop(worker_id, in_queue, out_queue, object_set, preload=True, gpu_renderer=True):
+    renderer = init_renderer(object_set, preload=preload, gpu_renderer=gpu_renderer)
     while True:
         kwargs = in_queue.get()
         if kwargs is None:
@@ -38,10 +39,11 @@ def worker_loop(worker_id, in_queue, out_queue, object_set, preload=True):
 
 
 class BulletBatchRenderer:
-    def __init__(self, object_set, n_workers=8, preload_cache=True):
+    def __init__(self, object_set, n_workers=8, preload_cache=True, gpu_renderer=True):
         self.object_set = object_set
         self.n_workers = n_workers
-        self.init_plotters(preload_cache)
+        self.init_plotters(preload_cache, gpu_renderer)
+        self.gpu_renderer = gpu_renderer
 
     def render(self, obj_infos, TCO, K, resolution=(240, 320), render_depth=False):
         TCO = torch.as_tensor(TCO).detach()
@@ -79,17 +81,23 @@ def render(self, obj_infos, TCO, K, resolution=(240, 320), render_depth=False):
             images[data_id] = im[0]
             if render_depth:
                 depths[data_id] = depth[0]
-        images = torch.as_tensor(np.stack(images, axis=0)).pin_memory().cuda(non_blocking=True)
+        if self.gpu_renderer:
+            images = torch.as_tensor(np.stack(images, axis=0)).pin_memory().cuda(non_blocking=True)
+        else:
+            images = torch.as_tensor(np.stack(images, axis=0))
         images = images.float().permute(0, 3, 1, 2) / 255
 
         if render_depth:
-            depths = torch.as_tensor(np.stack(depths, axis=0)).pin_memory().cuda(non_blocking=True)
+            if self.gpu_renderer:
+                depths = torch.as_tensor(np.stack(depths, axis=0)).pin_memory().cuda(non_blocking=True)
+            else:
+                depths = torch.as_tensor(np.stack(depths, axis=0))
             depths = depths.float()
             return images, depths
         else:
             return images
 
-    def init_plotters(self, preload_cache):
+    def init_plotters(self, preload_cache, gpu_renderer):
         self.plotters = []
         self.in_queue = multiprocessing.Queue()
         self.out_queue = multiprocessing.Queue()
@@ -100,12 +108,13 @@ def init_plotters(self, preload_cache):
                                                   kwargs=dict(worker_id=n,
                                                               in_queue=self.in_queue,
                                                               out_queue=self.out_queue,
+                                                              object_set=self.object_set,
                                                               preload=preload_cache,
-                                                              object_set=self.object_set))
+                                                              gpu_renderer=gpu_renderer))
                 plotter.start()
                 self.plotters.append(plotter)
         else:
-            self.plotters = [init_renderer(self.object_set, preload_cache)]
+            self.plotters = [init_renderer(self.object_set, preload_cache, gpu_renderer)]
 
     def stop(self):
         if self.n_workers > 0: