模型

创建模型

mindyolo.models.model_factory.create_model(model_name, model_cfg=None, in_channels=3, num_classes=80, checkpoint_path='', **kwargs)

Source code in mindyolo/models/model_factory.py

def create_model(
    model_name: str,
    model_cfg: dict = None,
    in_channels: int = 3,
    num_classes: int = 80,
    checkpoint_path: str = "",
    **kwargs,
):
    model_args = dict(cfg=model_cfg, num_classes=num_classes, in_channels=in_channels)
    kwargs = {k: v for k, v in kwargs.items() if v is not None}

    if not is_model(model_name):
        raise RuntimeError(f"Unknown model {model_name}")

    create_fn = model_entrypoint(model_name)
    model = create_fn(**model_args, **kwargs)

    if checkpoint_path:
        assert os.path.isfile(checkpoint_path) and checkpoint_path.endswith(
            ".ckpt"
        ), f"[{checkpoint_path}] not a ckpt file."
        checkpoint_param = load_checkpoint(checkpoint_path)
        load_param_into_net(model, checkpoint_param)
        logger.info(f"Load checkpoint from [{checkpoint_path}] success.")

    return model

YOLOV3

mindyolo.models.yolov3(cfg, in_channels=3, num_classes=None, **kwargs)

Get yolov3 model.

Source code in mindyolo/models/yolov3.py

@register_model
def yolov3(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv3:
    """Get yolov3 model."""
    model = YOLOv3(cfg=cfg, in_channels=in_channels, num_classes=num_classes, **kwargs)
    return model

YOLOV4

mindyolo.models.yolov4(cfg, in_channels=3, num_classes=None, **kwargs)

Get yolov4 model.

Source code in mindyolo/models/yolov4.py

@register_model
def yolov4(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv4:
    """Get yolov4 model."""
    model = YOLOv4(cfg=cfg, in_channels=in_channels, num_classes=num_classes, **kwargs)
    return model

YOLOV5

mindyolo.models.yolov5(cfg, in_channels=3, num_classes=None, **kwargs)

Get yolov5 model.

Source code in mindyolo/models/yolov5.py

@register_model
def yolov5(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv5:
    """Get yolov5 model."""
    model = YOLOv5(cfg=cfg, in_channels=in_channels, num_classes=num_classes, **kwargs)
    return model

YOLOV7

mindyolo.models.yolov7(cfg, in_channels=3, num_classes=None, **kwargs)

Get yolov7 model.

Source code in mindyolo/models/yolov7.py

@register_model
def yolov7(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv7:
    """Get yolov7 model."""
    model = YOLOv7(cfg=cfg, in_channels=in_channels, num_classes=num_classes, **kwargs)
    return model

YOLOV8

mindyolo.models.yolov8(cfg, in_channels=3, num_classes=None, **kwargs)

Get yolov8 model.

Source code in mindyolo/models/yolov8.py

@register_model
def yolov8(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv8:
    """Get yolov8 model."""
    model = YOLOv8(cfg=cfg, in_channels=in_channels, num_classes=num_classes, **kwargs)
    return model

YOLOX

mindyolo.models.yolox(cfg, in_channels=3, num_classes=None, **kwargs)

Get yolox model.

Source code in mindyolo/models/yolox.py

@register_model
def yolox(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOX:
    """Get yolox model."""
    model = YOLOX(cfg, in_channels=in_channels, num_classes=num_classes, **kwargs)
    return model