Spaces:
Running
Running
| import cv2 | |
| import cv2 as cv | |
| import numpy as np | |
| from yunet import YuNet | |
| # Valid combinations of backends and targets | |
| backend_target_pairs = [ | |
| [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_TARGET_CPU], | |
| [cv.dnn.DNN_BACKEND_CUDA, cv.dnn.DNN_TARGET_CUDA], | |
| [cv.dnn.DNN_BACKEND_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16], | |
| [cv.dnn.DNN_BACKEND_TIMVX, cv.dnn.DNN_TARGET_NPU], | |
| [cv.dnn.DNN_BACKEND_CANN, cv.dnn.DNN_TARGET_NPU], | |
| ] | |
| class ImageResizer: | |
| def __init__( | |
| self, | |
| modelPath, | |
| input_size=(320, 320), | |
| conf_threshold=0.6, | |
| nms_threshold=0.3, | |
| top_k=5000, | |
| backend_id=0, | |
| target_id=0, | |
| ): | |
| self.model = YuNet( | |
| modelPath=modelPath, | |
| inputSize=input_size, | |
| confThreshold=conf_threshold, | |
| nmsThreshold=nms_threshold, | |
| topK=top_k, | |
| backendId=backend_id, | |
| targetId=target_id, | |
| ) | |
| def detect(self, image, num_faces=None): | |
| # If input is an image | |
| if image is not None: | |
| h, w, _ = image.shape | |
| # Inference | |
| self.model.setInputSize([w, h]) | |
| results = self.model.infer(image) | |
| faces = results[:num_faces] if num_faces else results | |
| bboxs = [] | |
| for face in faces: | |
| bbox = face[0:4].astype(np.int32) # x,y,w,h | |
| x, y, w, h = bbox | |
| # draw | |
| cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 255), 2) | |
| bboxs.append(bbox) | |
| return image, bboxs | |
| def resize(self, image, target_size=512, above_head_ratio=0.5): | |
| height, width, _c = image.shape | |
| ar = width / height | |
| # downscale the image | |
| if not target_size: | |
| target_size = 512 | |
| if ar > 1: | |
| # Landscape | |
| new_height = target_size | |
| new_width = int(target_size * ar) | |
| elif ar < 1: | |
| # Portrait | |
| new_width = target_size | |
| new_height = int(target_size / ar) | |
| else: | |
| # Square | |
| new_width = target_size | |
| new_height = target_size | |
| resized = cv2.resize( | |
| image, (new_width, new_height), interpolation=cv2.INTER_AREA | |
| ) | |
| # Perform object detection on the resized image | |
| dt_image, bboxes = self.detect(resized.copy()) | |
| # crop around face | |
| if len(bboxes) >= 1: | |
| x, y, w, h = bboxes[0] | |
| else: | |
| x, y, w, h = 0, 0, target_size, target_size | |
| # 20% of image height | |
| above_head_max = int(target_size * above_head_ratio) | |
| x_center = int((x + (x + w)) / 2) | |
| y_center = int((y + (y + h)) / 2) | |
| # Calculate cropping box | |
| top = int(max(0, y_center - above_head_max)) | |
| bottom = int(min(top + target_size, resized.shape[0])) | |
| left = int(max(0, x_center - target_size // 2)) | |
| right = int(min(x_center + target_size // 2, resized.shape[1])) | |
| # adjust width if necessory | |
| _w = right - left | |
| if _w != target_size: | |
| dx = ( | |
| target_size - _w | |
| ) # difference between the target size and the current width | |
| nl = max(0, left - dx) | |
| dr = dx - nl # remaining adjustment needed for the right coordinate | |
| left = nl | |
| right += dr | |
| cropped_image = resized[top:bottom, left:right] | |
| return dt_image, cropped_image | |