from deeplite_torch_zoo import get_model, list_models

model = get_model(
    model_name='edgevit_xs',        # model names for imagenet available via `list_models('imagenet')`
    dataset_name='imagenet',        # dataset name, since resnet18 is different for e.g. imagenet and cifar100
    pretrained=False,               # if True, will try to load a pre-trained checkpoint
)

# creating a model with 42 classes for transfer learning:

model = get_model(
    model_name='fasternet_t0',        # model names for imagenet available via `list_models('imagenet')`
    num_classes=42,                   # number of classes for transfer learning
    dataset_name='imagenet',   # take weights from checkpoint pre-trained on this dataset
    pretrained=False,                 # if True, will try to load all weights with matching tensor shapes
)

from deeplite_torch_zoo import get_model

model = get_model(
    model_name='yolo4n',        # creates a YOLOv4n model on COCO
    dataset_name='coco',        # (`n` corresponds to width factor 0.25, depth factor 0.33)
    pretrained=False,           # if True, will try to load a pre-trained checkpoint
)

# one could create a YOLO model with timm backbone,
# PAN neck and YOLOv8 decoupled anchor-free head like this:

model = get_model(
    model_name='yolo_timm_fbnetv3_d',    # creates a YOLO with FBNetV3-d backbone from timm
    dataset_name='coco',                 #
    pretrained=False,                    # if True, will try to load a pre-trained checkpoint
    custom_head='yolo8',                 # will replace default detection head
                                         # with YOLOv8 detection head
)

from deeplite_torch_zoo import get_dataloaders

dataloaders = get_dataloaders(
    data_root='./',                      # folder with data, will be used for download
    dataset_name='imagewoof',            # datasets to if applicable,
    num_workers=8,                       # number of dataloader workers
    batch_size=64,                       # dataloader batch size (train and test)
)

# dataloaders['train'] -> train dataloader
# dataloaders['test'] -> test dataloader
#
# see below for the list of supported datasets

from deeplite_torch_zoo import get_eval_function

eval_function = get_eval_function(
    model_name='yolo8s',
    dataset_name='voc',
)

# required arg signature is fixed for all eval functions
metrics = eval_function(model, test_dataloader)

from deeplite_torch_zoo.trainer import Detector

model = Detector(model_name='yolo7n')        # will create a wrapper around YOLOv7n model
                                             # (YOLOv7n model with YOLOv8 detection head)
model.train(data='VOC.yaml', epochs=100)     # same arguments as Ultralytics trainer

# alternatively:

torch_model = get_model(
    model_name='yolo7n',
    dataset_name='coco',
    pretrained=False,
    custom_head='yolo8',
)
model = Detector(torch_model=torch_model)    # either `model_name` or `torch_model`
model.train(data='VOC.yaml', epochs=100)     # should be provided