| | 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() |
