# ssdlite.py
# Vojtěch Golda
from torchvision.models.detection import ssdlite320_mobilenet_v3_large, SSDLite320_MobileNet_V3_Large_Weights
from torchvision.models.detection import ssd300_vgg16, SSD300_VGG16_Weights 
from util.model_util import preprocess, get_data_bb, get_model, batch2boxes

from models.model_util import ModelUtil



# 140, 217
# 35
class SSDlite(ModelUtil):
    def __init__(self, device, filename="models/ssdlite.pth"):
        self.device = device
        self.model = ssdlite320_mobilenet_v3_large
        self.weights = SSDLite320_MobileNet_V3_Large_Weights
        self.target_size = (320, 320)
        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)

