You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
95 lines
3.8 KiB
95 lines
3.8 KiB
""" |
|
Top-level YOLO model interface. First principle usage example - https://github.com/ultralytics/ultralytics/issues/13 |
|
""" |
|
import torch |
|
import yaml |
|
|
|
from ultralytics import yolo |
|
from ultralytics.yolo.utils import LOGGER |
|
from ultralytics.yolo.utils.checks import check_yaml |
|
from ultralytics.yolo.utils.modeling import attempt_load_weights |
|
from ultralytics.yolo.utils.modeling.tasks import ClassificationModel, DetectionModel, SegmentationModel |
|
|
|
# map head: [model, trainer] |
|
MODEL_MAP = { |
|
"classify": [ClassificationModel, 'yolo.VERSION.classify.ClassificationTrainer'], |
|
"detect": [DetectionModel, 'yolo.VERSION.detect.DetectionTrainer'], |
|
"segment": [SegmentationModel, 'yolo.VERSION.segment.SegmentationTrainer']} |
|
|
|
|
|
class YOLO: |
|
|
|
def __init__(self, version=8) -> None: |
|
self.version = version |
|
self.ModelClass = None |
|
self.TrainerClass = None |
|
self.model = None |
|
self.trainer = None |
|
self.task = None |
|
self.ckpt = None |
|
|
|
def new(self, cfg: str): |
|
cfg = check_yaml(cfg) # check YAML |
|
with open(cfg, encoding='ascii', errors='ignore') as f: |
|
cfg = yaml.safe_load(f) # model dict |
|
self.ModelClass, self.TrainerClass, self.task = self._guess_model_trainer_and_task(cfg["head"][-1][-2]) |
|
self.model = self.ModelClass(cfg) # initialize |
|
|
|
def load(self, weights): |
|
self.ckpt = torch.load(weights, map_location="cpu") |
|
self.task = self.ckpt["train_args"]["task"] |
|
_, trainer_class_literal = MODEL_MAP[self.task] |
|
self.TrainerClass = eval(trainer_class_literal.replace("VERSION", f"v{self.version}")) |
|
self.model = attempt_load_weights(weights) |
|
|
|
def reset(self): |
|
for m in self.model.modules(): |
|
if hasattr(m, 'reset_parameters'): |
|
m.reset_parameters() |
|
for p in self.model.parameters(): |
|
p.requires_grad = True |
|
|
|
def train(self, **kwargs): |
|
if 'data' not in kwargs: |
|
raise Exception("data is required to train") |
|
if not self.model and not self.ckpt: |
|
raise Exception("model not initialized. Use .new() or .load()") |
|
|
|
kwargs["task"] = self.task |
|
kwargs["mode"] = "train" |
|
self.trainer = self.TrainerClass(overrides=kwargs) |
|
# load pre-trained weights if found, else use the loaded model |
|
self.trainer.model = self.trainer.load_model(weights=self.ckpt) if self.ckpt else self.model |
|
self.trainer.train() |
|
|
|
def resume(self, task=None, model=None): |
|
if not task: |
|
raise Exception( |
|
"pass the task type and/or model(optional) from which you want to resume: `model.resume(task=" |
|
")`") |
|
if task.lower() not in MODEL_MAP: |
|
raise Exception(f"unrecognised task - {task}. Supported tasks are {MODEL_MAP.keys()}") |
|
_, trainer_class_literal = MODEL_MAP[task.lower()] |
|
self.TrainerClass = eval(trainer_class_literal.replace("VERSION", f"v{self.version}")) |
|
self.trainer = self.TrainerClass(overrides={"task": task.lower(), "resume": model if model else True}) |
|
self.trainer.train() |
|
|
|
def _guess_model_trainer_and_task(self, head): |
|
# TODO: warn |
|
task = None |
|
if head.lower() in ["classify", "classifier", "cls", "fc"]: |
|
task = "classify" |
|
if head.lower() in ["detect"]: |
|
task = "detect" |
|
if head.lower() in ["segment"]: |
|
task = "segment" |
|
model_class, trainer_class = MODEL_MAP[task] |
|
# warning: eval is unsafe. Use with caution |
|
trainer_class = eval(trainer_class.replace("VERSION", f"v{self.version}")) |
|
|
|
return model_class, trainer_class, task |
|
|
|
def __call__(self, imgs): |
|
if not self.model: |
|
LOGGER.info("model not initialized!") |
|
return self.model(imgs)
|
|
|