add base file for api

__init__.py

@@ -0,0 +1,10 @@
+from .localize_with_landmark import *
+from .check_mask import *
+from .check_liveness import *
+from .estimate_headpose import *
+from .extract_facevector import *
+from .extract_agegender import *
+from .extract_emotion import *
+from .configs import TASK_CONFIG
+
+

base.py

@@ -0,0 +1,73 @@
+import onnxruntime
+import random
+import numpy as np
+from pathlib import Path
+from numpy.typing import NDArray
+from typing import Any, List
+from .utils import count_gpus, get_memory_free_MiB
+from abc import ABC, abstractclassmethod
+
+__dir__ = Path(__file__).parent
+
+class ONNXBaseTask(ABC):
+    num_gpus: int = count_gpus()
+
+    def __init__(self, weight: str) -> None:
+        self.session = self.initialize_session(weight)
+        self.input_metadata = self.session.get_inputs()[0]
+        self.prepare_input = self.setup_prepare_input()
+        
+        # warmup model
+        input_height, input_width = self.input_metadata.shape[-2:]
+        temp = np.zeros((1, 3, int(input_height) if int(input_height) > 0 else 320, int(input_width) if int(input_width) > 0 else 320), dtype=np.float32)
+        self.run_session(temp)
+
+    @abstractclassmethod
+    def process_output(self, raw_outputs: List[NDArray], **kwargs) -> Any:
+        pass
+
+    @abstractclassmethod
+    def setup_prepare_input(self):
+        pass
+
+    def call(self, image) -> Any:
+        input_height, input_width = self.input_metadata.shape[-2:]
+
+        # predict
+        input_value = self.prepare_input(image, height=input_height, width=input_width)
+        raw_outputs = self.run_session(input_value)
+
+        return self.process_output(raw_outputs)
+    
+    def run_session(self, input_value: NDArray) -> List[NDArray]:
+        input_dict = {self.input_metadata.name : input_value}
+
+        return self.session.run(None, input_dict)
+    
+    def initialize_session(self, weight: str):
+        # get avaiable runtime
+        providers=[]
+        if self.num_gpus == 0:
+            providers += [("CPUExecutionProvider", {})]
+        else:
+            providers += [(
+                "CUDAExecutionProvider", 
+                {
+                    "device_id": random.choice([i for i in range(self.num_gpus) if get_memory_free_MiB(i) >= 1000])
+                }
+            )]
+
+        # init session
+        return onnxruntime.InferenceSession(
+            str(__dir__.parent.parent.parent/weight),
+            providers=providers
+        )
+        
+
+
+
+
+
+
+
+

utils.py

@@ -0,0 +1,323 @@
+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