1

vision_detect/vision_detect/VisionDetect/node/detect_node.py

@@ -199,20 +199,26 @@ class DetectNode(Node):
 
     def _init_model(self):
         """init model"""
-        device_model = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        try:
-            self.model = YOLO(self.checkpoint_path).to(device_model)
-        except Exception as e:
-            self.get_logger().error(f'Failed to load YOLO model: {e}')
-            raise 
-        self.get_logger().info(f'Loading checkpoint from: {self.checkpoint_path}')
-
-        if self.checkpoint_path.endswith('-seg.pt') or self.checkpoint_path.endswith('.onnx') or self.checkpoint_path.endswith('.engine'):
+        if self.checkpoint_path.endswith('-seg.pt'):
+            device_model = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+            try:
+                self.model = YOLO(self.checkpoint_path).to(device_model)
+            except Exception as e:
+                self.get_logger().error(f'Failed to load YOLO model: {e}')
+                raise 
             self.function = self._seg_image
+
+        elif self.checkpoint_path.endswith('.onnx') or self.checkpoint_path.endswith('.engine'):
+            self.model = YOLO(self.checkpoint_path)
+            self.function = self._seg_image
+
         else:
             self.function = None
             self.get_logger().fatal(f'Unknown checkpoint: {self.checkpoint_path}')
             raise ValueError(f'Unknown checkpoint: {self.checkpoint_path}')
+
+        self.get_logger().info(f'Loading checkpoint from: {self.checkpoint_path}')
+
         self.get_logger().info("Initialize model done")
 
     def _init_publisher(self):