import cv2
from ultralytics import YOLO

X1_OFFSET, Y1_OFFSET, X2_OFFSET, Y2_OFFSET = 0, 0, 0, 0 # need to tinker with later
COLOR = (0, 255, 0)

class YOLOFace:
    
    def __init__(self):
        self.net = YOLO("models/best.pt")
    def get_patches(self, img):
        patches = []
        for (x1, y1), (x2, y2) in self.forward(img):
            cv2.rectangle(img, (x1, y1), (x2, y2), COLOR, 2)
            if (x1 == x2 or y1 == y2):
                continue
            patches.append(img[y1:y2, x1:x2])
        return patches
    def forward(self, img):
        boxes = self.net.predict(img, verbose=False)[0].boxes.xyxy.cpu().detach().numpy()
        for x1, y1, x2, y2 in boxes:
            yield (int(x1.item()+X1_OFFSET), int(y1.item()+Y1_OFFSET)), (int(x2.item()+X2_OFFSET), int(y2.item()+Y2_OFFSET))
    __call__ = get_patches
    
YOLO_FACE = YOLOFace()