# retina.py
# Vojtěch Golda
from torchvision.models.detection import retinanet_resnet50_fpn_v2, RetinaNet_ResNet50_FPN_V2_Weights
from util.model_util import preprocess, get_data_bb, get_model, batch2boxes
from models.model_util import ModelUtil

# 10, 10, 86
# 33



class RetinaNet(ModelUtil):
    def __init__(self, device, filename="models/retina.pth"):
        self.device = device
        self.model = retinanet_resnet50_fpn_v2
        self.weights = RetinaNet_ResNet50_FPN_V2_Weights
        self.target_size = (400, 400)
        self.filename = filename

    def preprocess(*args, **kwargs):
        return preprocess(*args, **kwargs)

    def get_data(*args, **kwargs):
        return get_data_bb(*args, **kwargs)

    def get_model(*args, **kwargs):
        return get_model(*args, **kwargs)

    def batch2boxes(*args, **kwargs):
        return batch2boxes(*args, **kwargs)

