Add cpu option for segmentation model

segment_image.py

@@ -52,7 +52,8 @@ class SegmentImage:
                 region_overlap_threshold: float = 0.5,
                 class_id_region: int = 1,
                 class_id_line: int = 2,
-                min_polygon_points: int = 3):
+                min_polygon_points: int = 3,
+                device: str = 'cuda:0'):
 
         self.model_path = model_path 
         self.max_size = max_size 
@@ -66,6 +67,7 @@ class SegmentImage:
         self.class_id_region = class_id_region
         self.class_id_line = class_id_line
         self.min_polygon_points = min_polygon_points
+        self.device = device
     
         # Validate model path
         if not os.path.exists(self.model_path):
@@ -81,8 +83,12 @@ class SegmentImage:
             Exception: If model initialization fails
         """
         try:
-            self.model = RFDETRSeg2XLarge(pretrain_weights=self.model_path)
-            self.model.optimize_for_inference()
+            self.model = RFDETRSeg2XLarge(
+                pretrain_weights=self.model_path,
+                device=self.device
+            )
+            if self.device != 'cpu':
+                self.model.optimize_for_inference()
             print(f"✓ Segmentation model initialized successfully")
         except Exception as e:
             raise RuntimeError(f'Failed to initialize segmentation model: {e}')