|
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 |
|
|
|
|
|
|
|
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 |
|
roi_width = int(w * 2) |
|
roi_x1 = int(x + (w - roi_width) // 2) |
|
vertical_offset = 20 |
|
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' |
|
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() |
