Update image_segmenter.py

image_segmenter.py

@@ -7,9 +7,9 @@ import os
 import torch
 
 class ImageSegmenter:
-    def __init__(self, model_type="yolov8s-seg") -> None:
-        # Initialize parameters
-        self.model_type = model_type
+    def __init__(self, model_type="yolov8s-seg", device="cpu"):
+        self.device = device
+        self.model = YOLO(model_type).to(self.device)
         self.is_show_bounding_boxes = True
         self.is_show_segmentation_boundary = False
         self.is_show_segmentation = False