Update segment_image.py

segment_image.py

@@ -58,7 +58,7 @@ class SegmentImage:
         """Function for initializing the line detection model."""
         try:
             # Load the trained line detection model
-            line_model = YOLO(self.line_model_path)
+            line_model = YOLO(self.line_model_path, task="instance-segmentation")
             return line_model
         except Exception as e:
             print('Failed to load the line detection model: %s' % e)
@@ -67,7 +67,7 @@ class SegmentImage:
         """Function for initializing the region detection model."""
         try:
             # Load the trained line detection model
-            region_model = YOLO(self.region_model_path)
+            region_model = YOLO(self.region_model_path, task="instance-segmentation")
             return region_model
         except Exception as e:
             print('Failed to load the region detection model: %s' % e)