import torch

class ImageObjectDetector(object):
    
    def __init__(self):
        custom_path = r"object_extractor.pt"
        self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=custom_path, force_reload=False, source = 'github')
        self.model.conf = 0.1
        self.classes = self.model.names
        self.total_time = 0
        self.total_img_processed = 0
        
    def detect(self, image_path : str,  save_crop_images : bool = False, save_path : str = None):
        pred = self.model(image_path)
        df = pred.pandas().xyxy[0]
        croped = pred.crop(save = save_crop_images, save_dir = save_path, exist_ok=True)
        return df, croped
    
class DocumentObjects(ImageObjectDetector):
    
    def __init__(self):
        super().__init__()
        self.items_to_extract = ['stamp']
        self.keymapper = {"sign" : "signature", "checked_item" : "checkedItem", "qr_code" : "qrCode", "bar_code" : "barCode"}
    
    def detect_objects(self, img_path : str):
        df, _ = self.detect(img_path, False, None)

        if df.shape[0] < 0:
            items = []
        else:
            items = df['name'].unique()
        response = {}
        for key in self.items_to_extract:
            temp = {}
            loc = []
            if key in items:
                temp['found'] = True
                for ind in df[df['name'] == key].index:
                    cord = [df.loc[ind, 'xmin'], df.loc[ind, 'ymin'], df.loc[ind, 'xmax'], df.loc[ind, 'ymax']]
                    loc.append(cord)
            else:
                temp['found'] = False
            temp['loc'] = loc
            response[self.keymapper.get(key,key)] = temp
        return response