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.
96 lines
3.8 KiB
96 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)
|