prototype.py

import cv2
import numpy as np
import os
import mediapipe

from utils import extract_point, compute_distance



class ImageProcessor:
    def __init__(self, picture, folder_path , model):
        self.image = picture
        self.height , self.width = self.image.shape[:2]
        self.folder_path = folder_path
        self.model = model  

    def detect_and_overlay(self, write = False, output = None):

        detections = self.model.process(self.image).detections

        if not detections: 
            self.image = cv2.putText(self.image, "Unable to detect faces :(", (int(self.width//10), int(self.height//2)), fontFace = cv2.FONT_HERSHEY_SIMPLEX, fontScale = 3, color = (0,0,0),thickness =  7)
            self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
            return self.image

        for i, elem in enumerate(detections):
            print(i)
            x, y, w, h = self.get_bounding_box(elem)
            gadget_path, nose_path = self.select_gadgets(i)


            if nose_path:
              nose = extract_point(self, elem.location_data.relative_keypoints[2])
              eye = extract_point(self, elem.location_data.relative_keypoints[0]) 
              cv2.circle(self.image,  (int(nose[0]), int(nose[1])), int(compute_distance(nose, eye)/2), (255,0,0), -1)

             
            try:
                roi_x1, roi_y1, roi_x2, roi_y2 = self.calculate_roi_head(x, y, w, h)
                gadget = self.read_and_resize_gadget(gadget_path, roi_x2 - roi_x1, roi_y2 - roi_y1)
                self.overlay_gadget(gadget, roi_x1, roi_y1, roi_x2, roi_y2)
            except: 
                continue

        # self.image =  cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)

        if write: 
            self.display_result(output)

        return self.image
        
        

    def get_bounding_box(self, elem):
        bbox = elem.location_data.relative_bounding_box
        x = int(bbox.xmin * self.width)
        y = int(bbox.ymin * self.height)
        w = int(bbox.width * self.width)
        h = int(bbox.height * self.height)
        return x, y, w, h

    def calculate_roi_head(self, x, y, w, h):
        roi_height = 60 #hard coded, need to make it fit the head
        roi_width = int(w * 2)
        roi_x1 = int(x + (w - roi_width) // 2)
        vertical_offset = 20  #hard coded, need to make it fit the head
        roi_y1 = int(max(y - roi_height // 2 - vertical_offset, 0))
        roi_y2 = roi_y1 + roi_height
        roi_x2 = roi_x1 + roi_width

        return roi_x1, roi_y1, roi_x2, roi_y2


    def select_gadgets(self, index):
        if index == 0:
            gadget =  "anklers.png"
            nose = True

        else:
            gadget = "hat.png"
            nose = False

        return gadget, nose

    def read_and_resize_gadget(self, gadget_path, width, height):
        gadget = cv2.imread(gadget_path, cv2.IMREAD_UNCHANGED)
        gadget_resized = cv2.resize(gadget, (width, height))
        return gadget_resized

    def overlay_gadget(self, gadget, x1, y1, x2, y2):
        alpha_gadget = gadget[:, :, 3] / 255.0
        alpha_gadget_resized = np.stack([alpha_gadget] * 3, axis=-1)
        gadget_bgr = gadget[:, :, :3]
        gadget_bgr = cv2.cvtColor(gadget_bgr, cv2.COLOR_BGR2RGB)

        roi = self.image[y1:y2, x1:x2]

        roi = cv2.resize(roi, (gadget.shape[1], gadget.shape[0]))

        overlay = (1 - alpha_gadget_resized) * roi + alpha_gadget_resized * gadget_bgr
        self.image[y1:y2, x1:x2] = overlay

    def display_result(self, output):
        if not output:
            output = "image"
        cv2.imwrite( os.path.join("results", "{}.png".format(output)), self.image )



def activate(image):
    folder_path = 'gadget_path' #replace with you gadgets folder
    model = mediapipe.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.8)
    processor = ImageProcessor(image, folder_path, model)
    return processor.detect_and_overlay()