Roughed in Yolo class incl. training method

model_server/extensions/yolov8/models.py

+from typing import List
+
+from pydantic import BaseModel
+from ultralytics import YOLO
+
+from ...base.accessors import GenericImageDataAccessor
+from ...base.models import InstanceSegmentationModel
+from ...base.roiset import RoiSet
+
+class YoloV8Params(BaseModel):
+    pt_file: str
+    duplicate: bool = True
+
+class YoloModel(InstanceSegmentationModel):
+
+    def __init__(self, params: YoloV8Params, autoload=True):
+        # initialize from pretrained model
+        super().__init__(autoload, params)
+
+    def load(self):
+        self.yolo = YOLO(self.params.pt_file)
+
+    def _infer_yolo_seg(self, img:GenericImageDataAccessor) -> RoiSet:
+        # only populates bounding box info
+        pass
+
+    def _infer_yolo_det(self, img:GenericImageDataAccessor) -> RoiSet:
+        # also populates segmentation masks
+        pass
+
+    def label_instance_class(
+            self, img: GenericImageDataAccessor, mask: GenericImageDataAccessor, **kwargs
+    ) -> GenericImageDataAccessor:
+        roiset = self._infer_yolo_seg(img)
+        return roiset.acc_obj_ids
+
+    def export(self):
+        # export pretrained model
+        pass
+
+    def train(self, roisets: List[RoiSet]):
+        coco_list = [r.serialize_coco for r in roisets]
+        self.yolo.train(coco_list)
\ No newline at end of file