File size: 3,890 Bytes
1bc2204 9aeeedd 1bc2204 9fc43dc 1bc2204 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
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
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' #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() |