139 lines
5.1 KiB
Python
139 lines
5.1 KiB
Python
import os
|
|
import random
|
|
from ultralytics import YOLO
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
class PhotoManager:
|
|
|
|
def __init__(self):
|
|
self.parent_folder = "/home/luca/git_repos/telegram_amicobot/data/photos"
|
|
self.file_extension = "jpg"
|
|
|
|
def get_random_photo(self):
|
|
files = [f for f in os.listdir(self.parent_folder) if os.path.isfile(os.path.join(self.parent_folder, f))]
|
|
if not files:
|
|
raise FileNotFoundError("No files found in the specified folder")
|
|
|
|
return os.path.join(self.parent_folder, random.choice(files))
|
|
|
|
def replace_faces(self, photo_bytes):
|
|
|
|
# Convert the byte array to a NumPy array
|
|
nparr = np.frombuffer(photo_bytes, np.uint8)
|
|
|
|
# Decode the image
|
|
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
original_img = img
|
|
|
|
# Specify the desired padding size
|
|
padding_perc = 5
|
|
padding_size = int(img.shape[1] * (1+padding_perc/100))
|
|
|
|
# Calculate new dimensions
|
|
new_width = img.shape[1] + 2 * padding_size
|
|
new_height = img.shape[0] + 2 * padding_size
|
|
|
|
# Create a blank image with the new dimensions
|
|
padded_image = np.zeros((new_height, new_width, 3), dtype=np.uint8)
|
|
|
|
# Copy the original image into the center of the blank image
|
|
padded_image[padding_size:padding_size + img.shape[0], padding_size:padding_size + img.shape[1]] = img
|
|
img = padded_image
|
|
|
|
# Load a pre-trained YOLOv8n model
|
|
# Net downloaded from https://github.com/noorkhokhar99/face-detection-yolov8
|
|
model = YOLO('yolov8n-face.pt')
|
|
names = model.model.names
|
|
|
|
# Perform inference on 'bus.jpg' with specified parameters with conf=0.5
|
|
results = model.predict(img, verbose=False, conf=0.3, device='cpu')
|
|
|
|
# Process detections
|
|
boxes = results[0].boxes.xywh.cpu()
|
|
clss = results[0].boxes.cls.cpu().tolist()
|
|
confs = results[0].boxes.conf.float().cpu().tolist()
|
|
|
|
# for box, cls, conf in zip(boxes, clss, confs):
|
|
# print(f"Class Name: {names[int(cls)]}, Confidence Score: {conf}, Bounding Box: {box}")
|
|
|
|
for box, cls, conf in zip(boxes, clss, confs):
|
|
x, y, w, h = map(int, box)
|
|
class_name = names[int(cls)]
|
|
|
|
# Calculate the bottom-right corner coordinates
|
|
x1, y1 = x - w/2, y - h/2
|
|
x2, y2 = x + w/2, y + h/2
|
|
|
|
x1 = int(x1)
|
|
y1 = int(y1)
|
|
x2 = int(x2)
|
|
y2 = int(y2)
|
|
|
|
# # Draw bounding box
|
|
# cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
# # Display class name and confidence
|
|
# label = f"{class_name}: {conf:.2f}"
|
|
# cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
|
|
|
overlay = cv2.imread('/home/luca/git_repos/telegram_amicobot/data/faces/piovra.png', cv2.IMREAD_UNCHANGED) # Load with alpha channel
|
|
|
|
overlay_size_h = overlay.shape[0]
|
|
box_size_h = h
|
|
scale_h = (box_size_h/overlay_size_h) * 100
|
|
|
|
overlay_size_w = overlay.shape[1]
|
|
box_size_w = w
|
|
scale_w = (box_size_w/overlay_size_w) * 100
|
|
|
|
|
|
scale_percent = max(scale_h, scale_w) * 1.5 # percent of original size
|
|
width = int(overlay.shape[1] * scale_percent / 100)
|
|
height = int(overlay.shape[0] * scale_percent / 100)
|
|
dim = (width, height)
|
|
|
|
# resize image
|
|
overlay = cv2.resize(overlay, dim, interpolation = cv2.INTER_AREA)
|
|
|
|
h, w = img.shape[:2]
|
|
h1, w1 = overlay.shape[:2]
|
|
|
|
# let store center coordinate as cx, cy
|
|
cx, cy = x, y#(w - w1) // 2, (h - h1) // 2 # Note the order of width and height here
|
|
cx = cx - int(overlay.shape[1]/2)
|
|
cy = cy - int(overlay.shape[0]/2) - int(overlay.shape[0] * 0.15)
|
|
|
|
if(cx < 0):
|
|
cx = 0
|
|
if(cy < 0):
|
|
cy = 0
|
|
|
|
# Ensure the overlay image has an alpha channel
|
|
if overlay.shape[2] == 3:
|
|
overlay = cv2.cvtColor(overlay, cv2.COLOR_BGR2BGRA)
|
|
|
|
# Create masks for overlay and background
|
|
overlay_mask = overlay[:, :, 3] / 255.0
|
|
background_mask = 1.0 - overlay_mask
|
|
|
|
# Resize overlay image to fit within specified region
|
|
overlay_resized = cv2.resize(overlay, (w1, h1))
|
|
|
|
# Blend the images using alpha blending
|
|
for c in range(0, 3):
|
|
img[cy:cy + h1, cx:cx + w1, c] = (overlay_resized[:, :, c] * overlay_mask +
|
|
img[cy:cy + h1, cx:cx + w1, c] * background_mask)
|
|
|
|
|
|
img = img[padding_size:padding_size + original_img.shape[0], padding_size:padding_size + original_img.shape[1]]
|
|
# Convert OpenCV image to bytes
|
|
_, image_bytes = cv2.imencode('.jpg', img)
|
|
image_bytes = image_bytes.tobytes()
|
|
|
|
return image_bytes
|
|
|
|
# cv2.imshow('Image from Byte Array', img)
|
|
# cv2.waitKey(0)
|
|
# cv2.destroyAllWindows() |