Skip to content
Snippets Groups Projects
Commit ef341002 authored by Christopher Randolph Rhodes's avatar Christopher Randolph Rhodes
Browse files

Removed roughed-in COCO models

parent 71ce1f57
No related branches found
No related tags found
No related merge requests found
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
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment