"""
Created By: ishwor subedi
Date: 2024-07-10
"""
import cv2
import numpy as np
from ultralytics import YOLO


class ObjectDetection:
    def __init__(self, device="cpu"):
        self.device = device
        self.hand_detector = YOLO(model="artifacts/detection/hand_yolov8s.pt")
        self.face_detector = YOLO(model="artifacts/detection/face_yolov8s.pt")

    def hand_detection(self, image_path: str):
        hand_detection_results = self.hand_detector.predict(image_path, show=False, device=self.device, iou=0.8)

        orig_image = hand_detection_results[0].orig_img

        boxes = hand_detection_results[0].boxes.xyxy  #

        return boxes, orig_image

    def face_detection(self, image_path: str):
        face_detection_results = self.face_detector.predict(image_path, show=False, device=self.device, iou=0.8)

        boxes = face_detection_results[0].boxes.xyxy  #
        orig_image = face_detection_results[0].orig_img

        return boxes, orig_image

    def detect(self, hand: bool, face: bool, image_path: str):

        if hand == True and face == True:
            hand_boxes, orig_image = self.hand_detection(image_path)
            face_boxes, orig_image = self.face_detection(image_path)

            mask = np.zeros_like(orig_image)

            for i in range(len(hand_boxes)):
                box = hand_boxes[i].tolist()
                x, y, x1, y1 = map(int, box[:4])
                cv2.rectangle(mask, (x, y), (x1, y1), (255, 255, 255), -1)
                print("hand")

            for j in range(len(face_boxes)):
                box = face_boxes[j].tolist()
                x, y, x1, y1 = map(int, box[:4])
                cv2.rectangle(mask, (x, y), (x1, y1), (255, 255, 255), -1)
                print("face")

            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)

            return mask



        elif hand == True and face == False:
            hand_boxes, orig_image = self.hand_detection(image_path)

            mask = np.zeros_like(orig_image)

            for i in range(len(hand_boxes)):
                box = hand_boxes[i].tolist()
                x, y, x1, y1 = map(int, box[:4])
                cv2.rectangle(mask, (x, y), (x1, y1), (255, 255, 255), -1)

            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)

            return mask
        elif hand == False and face == True:
            face_boxes, orig_image = self.face_detection(image_path)

            mask = np.zeros_like(orig_image)

            for j in range(len(face_boxes)):
                box = face_boxes[j].tolist()
                x, y, x1, y1 = map(int, box[:4])
                cv2.rectangle(mask, (x, y), (x1, y1), (255, 255, 255), -1)

            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)

            return mask
        else:
            return "none"