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.
37 lines
1.4 KiB
37 lines
1.4 KiB
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
|
|
|
from ultralytics.engine.model import Model
|
|
from ultralytics.models import yolo # noqa
|
|
from ultralytics.nn.tasks import ClassificationModel, DetectionModel, PoseModel, SegmentationModel
|
|
|
|
|
|
class YOLO(Model):
|
|
"""
|
|
YOLO (You Only Look Once) object detection model.
|
|
"""
|
|
|
|
@property
|
|
def task_map(self):
|
|
"""Map head to model, trainer, validator, and predictor classes"""
|
|
return {
|
|
'classify': {
|
|
'model': ClassificationModel,
|
|
'trainer': yolo.classify.ClassificationTrainer,
|
|
'validator': yolo.classify.ClassificationValidator,
|
|
'predictor': yolo.classify.ClassificationPredictor, },
|
|
'detect': {
|
|
'model': DetectionModel,
|
|
'trainer': yolo.detect.DetectionTrainer,
|
|
'validator': yolo.detect.DetectionValidator,
|
|
'predictor': yolo.detect.DetectionPredictor, },
|
|
'segment': {
|
|
'model': SegmentationModel,
|
|
'trainer': yolo.segment.SegmentationTrainer,
|
|
'validator': yolo.segment.SegmentationValidator,
|
|
'predictor': yolo.segment.SegmentationPredictor, },
|
|
'pose': {
|
|
'model': PoseModel,
|
|
'trainer': yolo.pose.PoseTrainer,
|
|
'validator': yolo.pose.PoseValidator,
|
|
'predictor': yolo.pose.PosePredictor, }, }
|