import numpy as np import cv2 import subprocess import math from itertools import product as product from numpy.typing import NDArray from typing import List import argparse import pynvml from dataclasses import dataclass from skimage import transform def parse_args(): @dataclass class Argument: image_path: str weight_path: str # parse argument parser = argparse.ArgumentParser( prog="Run AI Tasks", description="call builded task belong to Face", ) parser.add_argument( "--image", type=str, default="samples/An_2000.jpg", help="path to tested image" ) parser.add_argument( "--weight", type=str, default="weights/retinaface_mobilev3.onnx", help="path to weight" ) args = parser.parse_args() return Argument( image_path=args.image, weight_path=args.weight ) def get_memory_free_MiB(gpu_index): pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(int(gpu_index)) mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle) return mem_info.free // 1024 ** 2 def count_gpus(): try: output = subprocess.check_output(['nvidia-smi', '--query-gpu=count', '--format=csv,noheader'], encoding='utf-8') num_gpus = int(output.strip().split('\n')[0]) except subprocess.CalledProcessError: num_gpus = 0 return num_gpus def prepare_input_wraper(inter=1, mean=None, std=None, channel_first=True, color_space="BGR", is_scale=False): ''' THIS PROCESS WAY WILL OPTIMIZE RUNTIME (scaling will bit slower than no scaling) ========================================================================== inter: resize type (0: Nearest, 1: Linear, 2: Cubic) is_scale: whether we scale image in range(0,1) to normalize or not NOTE: image normalize with scale DIFFERENT normalize no scale mean: expected value of distribution std: standard deviation of distribution channel_first: True is (c,h,w), False is (h,w,c) color_space: BGR (default of cv2), RGB ========================================================================== ''' if mean is not None and std is not None: mean = mean if isinstance(mean, list) or isinstance(mean, tuple) else [mean]*3 std = std if isinstance(std, list) or isinstance(std, tuple) else [std]*3 def call(img: NDArray, width: int, height: int): ''' weight: input width of input model height: input height of input model ''' if img.shape[0] != height or img.shape[1] != width: image = cv2.resize(img.copy(), (width, height), interpolation=inter) else: image = img.copy() image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) if color_space == "RGB" else image image = image.transpose((2,0,1)) if channel_first else image image = image.astype(np.float32) # scale image in range(0,1) if is_scale: image /= 255 if mean is not None and std is not None: if channel_first: image[0, :, :] -= mean[0]; image[1, :, :] -= mean[1]; image[2, :, :] -= mean[2] image[0, :, :] /= std[0] ; image[1, :, :] /= std[1] ; image[2, :, :] /= std[2] else: image[:, :, 0] -= mean[0]; image[:, :, 1] -= mean[1]; image[:, :, 2] -= mean[2] image[:, :, 0] /= std[0] ; image[:, :, 1] /= std[1] ; image[:, :, 2] /= std[2] return image[np.newaxis, :] return call # =============================External Process image def class_letterbox(im, new_shape=(640, 640), color=(0, 0, 0), scaleup=True): # Resize and pad image while meeting stride-multiple constraints shape = im.shape[:2] # current shape [height, width] if isinstance(new_shape, int): new_shape = (new_shape, new_shape) if im.shape[0] == new_shape[0] and im.shape[1] == new_shape[1]: return im # Scale ratio (new / old) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: # only scale down, do not scale up (for better val mAP) r = min(r, 1.0) # Compute padding # ratio = r, r # width, height ratios new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding dw /= 2 # divide padding into 2 sides dh /= 2 if shape[::-1] != new_unpad: # resize im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border return im def get_new_box(src_w: int, src_h: int, bbox: List[int], scale: float): x, y, xmax, ymax = bbox box_w = (xmax - x) box_h = (ymax - y) # Re-calculate scale ratio scale = min((src_h-1)/box_h, min((src_w-1)/box_w, scale)) # get new width and height with scale ratio new_width = box_w * scale new_height = box_h * scale center_x, center_y = box_w/2+x, box_h/2+y # calculate bbox with new width and height left_top_x = center_x-new_width/2 left_top_y = center_y-new_height/2 right_bottom_x = center_x+new_width/2 right_bottom_y = center_y+new_height/2 # bbox must be in image if left_top_x < 0: right_bottom_x -= left_top_x left_top_x = 0 if left_top_y < 0: right_bottom_y -= left_top_y left_top_y = 0 if right_bottom_x > src_w-1: left_top_x -= right_bottom_x-src_w+1 right_bottom_x = src_w-1 if right_bottom_y > src_h-1: left_top_y -= right_bottom_y-src_h+1 right_bottom_y = src_h-1 return int(left_top_x), int(left_top_y),\ int(right_bottom_x), int(right_bottom_y) def align_face(image: NDArray, bounding_box: List[int], landmark: List[int], use_bbox: int=True): src = np.array(landmark).reshape(-1, 2) if use_bbox: # crop face x1, y1, x2, y2 = bounding_box image = image[y1:y2+1, x1:x2+1] # align src -= np.array([x1, y1]) des = np.array( [ [38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [38.2946, 92.3655], [70.7299, 92.2041], ] ) trans = transform.SimilarityTransform() trans.estimate(src, des) return cv2.warpAffine(image, trans.params[:2, :], dsize=(112, 112)) # =============================DETECT def get_largest_bbox(bboxes: NDArray) -> NDArray: # compute bbox area hbbox, wbbox = ( bboxes[:, 3] - bboxes[:, 1], bboxes[:, 2] - bboxes[:, 0], ) area = hbbox*wbbox return np.argmax(area) def get_input_size(image_height: int, image_width: int, limit_side_len: int) -> List[int]: ''' image_size: [ImageHeight, ImageWidth] ''' if max(image_height, image_width) >= limit_side_len: ratio = ( float(limit_side_len) / image_height if image_height < image_width else float(limit_side_len) / image_width ) else: ratio = 1. input_height = int((ratio*image_height // 32) * 32) input_width = int((ratio*image_width // 32) * 32) return input_height, input_width def prior_box(width: int, height: int, steps: List[int], min_sizes: List[List[int]]) -> NDArray: anchors = [] feature_maps = [ [math.ceil(height / step), math.ceil(width / step)] for step in steps ] for k, f in enumerate(feature_maps): for i, j in product(range(f[0]), range(f[1])): for min_size in min_sizes[k]: s_kx = min_size / width s_ky = min_size / height dense_cx = [x * steps[k] / width for x in [j + 0.5]] dense_cy = [y * steps[k] / height for y in [i + 0.5]] for cy, cx in product(dense_cy, dense_cx): anchors += [cx, cy, s_kx, s_ky] return np.reshape(anchors, (-1, 4)) def decode_boxes(bboxes: NDArray, priors: NDArray, variances: List[float], scale_factor: List[float]) -> NDArray: bboxes = np.concatenate( ( priors[:, :2] + bboxes[:, :2] * variances[0] * priors[:, 2:], priors[:, 2:] * np.exp(bboxes[:, 2:] * variances[1]), ), axis=1, ) bboxes[:, :2] -= bboxes[:, 2:] / 2 bboxes[:, 2:] += bboxes[:, :2] return bboxes * np.array(scale_factor * 2) def decode_landmarks(landmarks: NDArray, priors: NDArray, variances: List[float], scale_factor: List[float]) -> NDArray: landmarks = np.concatenate( ( priors[:, :2] + landmarks[:, :2] * variances[0] * priors[:, 2:], priors[:, :2] + landmarks[:, 2:4] * variances[0] * priors[:, 2:], priors[:, :2] + landmarks[:, 4:6] * variances[0] * priors[:, 2:], priors[:, :2] + landmarks[:, 6:8] * variances[0] * priors[:, 2:], priors[:, :2] + landmarks[:, 8:10] * variances[0] * priors[:, 2:], ), axis=1, ) return landmarks * np.array(scale_factor * 5) def intersection_over_union(bbox: NDArray, bboxes: NDArray, mode="Union") -> NDArray: """ Caculate IoU between detect and ground truth boxes :param crop_box:numpy array (4, ) :param bboxes:numpy array (n, 4):x1, y1, x2, y2 :return: numpy array, shape (n, ) Iou """ bbox_area = (bbox[2] - bbox[0] + 1) * (bbox[3] - bbox[1] + 1) areas = (bboxes[:, 2] - bboxes[:, 0] + 1) * (bboxes[:, 3] - bboxes[:, 1] + 1) xx1 = np.maximum(bbox[0], bboxes[:, 0]) yy1 = np.maximum(bbox[1], bboxes[:, 1]) xx2 = np.minimum(bbox[2], bboxes[:, 2]) yy2 = np.minimum(bbox[3], bboxes[:, 3]) # compute the width and height of the bounding box w = np.maximum(0, xx2 - xx1 + 1) h = np.maximum(0, yy2 - yy1 + 1) inter = w * h if mode == "Union": over = inter / (bbox_area + areas - inter) elif mode == "Minimum": over = inter / np.minimum(bbox_area, areas) return over def non_max_suppression(bboxes: NDArray, scores: NDArray, thresh: float, keep_top_k:int=100, mode:str="Union") -> List[int]: """ Bước 1: Tính diện tích của từng bbox Bước 2: Sort score của từng bbox theo thứ tự giảm dần và lấy vị trí index của chúng Bước 3: Theo thứ tự giảm dần của score, ta lấy bbox này giao với các bbox còn lại, sau đó loại bỏ bớt các vị trí mà phần giao của 2 bbox lớn hơn THRESHOLD """ # Sắp xếp độ tư tin giảm giần (lấy index) order = scores.argsort()[::-1][:keep_top_k] # Duyệt qua từng bbox với độ tự tin giảm dần để loại bỏ những bbox trùng nhau keep = [] while order.size > 0: i = order[0] keep.append(i) iou = intersection_over_union(bboxes[i], bboxes[order[1:]], mode=mode) # keep (cập nhật lại order bằng những gì còn lại sau khi loại bỏ) inds = np.where(iou <= thresh)[0] # [1,2,3,6,45,....] order = order[inds + 1] return keep