code release

modelhub/onnx/CenterFace/CenterFace.py

@@ -0,0 +1,111 @@
+from pathlib import Path
+from typing import List
+
+import numpy as np
+from xlib import math as lib_math
+from xlib.image import ImageProcessor
+from xlib.onnxruntime import (InferenceSession_with_device, ORTDeviceInfo,
+                              get_available_devices_info)
+
+
+class CenterFace:
+    """
+    CenterFace face detection model.
+
+    arguments
+
+     device_info    ORTDeviceInfo
+
+        use CenterFace.get_available_devices()
+        to determine a list of avaliable devices accepted by model
+
+    raises
+     Exception
+    """
+
+    @staticmethod
+    def get_available_devices() -> List[ORTDeviceInfo]:
+        # CenterFace ONNX model does not work correctly on CPU
+        # but it is much faster than Pytorch version
+        return get_available_devices_info(include_cpu=False)
+
+    def __init__(self, device_info : ORTDeviceInfo ):
+
+        if device_info not in CenterFace.get_available_devices():
+            raise Exception(f'device_info {device_info} is not in available devices for CenterFace')
+
+        path = Path(__file__).parent / 'CenterFace.onnx'
+        self._sess = sess = InferenceSession_with_device(str(path), device_info)
+        self._input_name = sess.get_inputs()[0].name
+
+    def extract(self, img, threshold : float = 0.5, fixed_window=0, min_face_size=40):
+        """
+        arguments
+
+         img    np.ndarray      ndim 2,3,4
+
+         fixed_window(0)    int  size
+                                 0 mean don't use
+                                 fit image in fixed window
+                                 downscale if bigger than window
+                                 pad if smaller than window
+                                 increases performance, but decreases accuracy
+
+        returns a list of [l,t,r,b] for every batch dimension of img
+        """
+
+        ip = ImageProcessor(img)
+        N,H,W,_ = ip.get_dims()
+
+        if fixed_window != 0:
+            fixed_window = max(64, max(1, fixed_window // 32) * 32 )
+            img_scale = ip.fit_in(fixed_window, fixed_window, pad_to_target=True, allow_upscale=False)
+        else:
+            ip.pad_to_next_divisor(64, 64)
+            img_scale = 1.0
+
+        img = ip.ch(3).swap_ch().to_uint8().as_float32().get_image('NCHW')
+
+        heatmaps, scales, offsets = self._sess.run(None, {self._input_name: img})
+        faces_per_batch = []
+
+        for heatmap, offset, scale in zip(heatmaps, offsets, scales):
+            faces = []
+            for face in self.refine(heatmap, offset, scale, H, W, threshold):
+                l,t,r,b,c = face
+
+                if img_scale != 1.0:
+                    l,t,r,b = l/img_scale, t/img_scale, r/img_scale, b/img_scale
+
+                bt = b-t
+                if min(r-l,bt) < min_face_size:
+                   continue
+                b += bt*0.1
+
+                faces.append( (l,t,r,b) )
+
+            faces_per_batch.append(faces)
+
+        return faces_per_batch
+
+    def refine(self, heatmap, offset, scale, h, w, threshold):
+        heatmap = heatmap[0]
+        scale0, scale1 = scale[0, :, :], scale[1, :, :]
+        offset0, offset1 = offset[0, :, :], offset[1, :, :]
+        c0, c1 = np.where(heatmap > threshold)
+        bboxlist = []
+        if len(c0) > 0:
+            for i in range(len(c0)):
+                s0, s1 = np.exp(scale0[c0[i], c1[i]]) * 4, np.exp(scale1[c0[i], c1[i]]) * 4
+                o0, o1 = offset0[c0[i], c1[i]], offset1[c0[i], c1[i]]
+                s = heatmap[c0[i], c1[i]]
+                x1, y1 = max(0, (c1[i] + o1 + 0.5) * 4 - s1 / 2), max(0, (c0[i] + o0 + 0.5) * 4 - s0 / 2)
+                x1, y1 = min(x1, w), min(y1, h)
+                bboxlist.append([x1, y1, min(x1 + s1, w), min(y1 + s0, h), s])
+
+            bboxlist = np.array(bboxlist, dtype=np.float32)
+
+            bboxlist = bboxlist[ lib_math.nms(bboxlist[:,0], bboxlist[:,1], bboxlist[:,2], bboxlist[:,3], bboxlist[:,4], 0.3), : ]
+            bboxlist = [x for x in bboxlist if x[-1] >= 0.5]
+
+        return bboxlist

modelhub/onnx/FaceMesh/FaceMesh.py

@@ -0,0 +1,58 @@
+from pathlib import Path
+from typing import List
+
+from xlib.image import ImageProcessor
+from xlib.onnxruntime import (InferenceSession_with_device, ORTDeviceInfo,
+                              get_available_devices_info)
+
+
+class FaceMesh:
+    """
+    Google FaceMesh detection model.
+    
+    arguments
+
+     device_info    ORTDeviceInfo
+
+        use FaceMesh.get_available_devices()
+        to determine a list of avaliable devices accepted by model
+
+    raises
+     Exception
+    """
+
+    @staticmethod
+    def get_available_devices() -> List[ORTDeviceInfo]:
+        return get_available_devices_info()
+
+    def __init__(self, device_info : ORTDeviceInfo):
+        if device_info not in FaceMesh.get_available_devices():
+            raise Exception(f'device_info {device_info} is not in available devices for FaceMesh')
+
+        path = Path(__file__).parent / 'FaceMesh.onnx'
+        self._sess = sess = InferenceSession_with_device(str(path), device_info)
+        self._input_name = sess.get_inputs()[0].name
+        self._input_width = 192
+        self._input_height = 192
+
+    def extract(self, img):
+        """
+        arguments
+
+         img    np.ndarray      HW,HWC,NHWC uint8/float32
+
+        returns (N,468,3)
+        """
+        ip = ImageProcessor(img)
+        N,H,W,_ = ip.get_dims()
+
+        h_scale = H / self._input_height
+        w_scale = W / self._input_width
+
+        feed_img = ip.resize( (self._input_width, self._input_height) ).to_ufloat32().ch(3).get_image('NHWC')
+
+        lmrks = self._sess.run(None, {self._input_name: feed_img})[0]
+        lmrks = lmrks.reshape( (N,468,3))
+        lmrks *= (w_scale, h_scale, 1)
+
+        return lmrks

modelhub/onnx/S3FD/S3FD.py

@@ -0,0 +1,89 @@
+from pathlib import Path
+from typing import List
+
+import numpy as np
+from xlib import math as lib_math
+from xlib.image import ImageProcessor
+from xlib.onnxruntime import (InferenceSession_with_device, ORTDeviceInfo,
+                              get_available_devices_info)
+
+
+class S3FD:
+
+    @staticmethod
+    def get_available_devices() -> List[ORTDeviceInfo]:
+        return get_available_devices_info()
+
+    def __init__(self, device_info : ORTDeviceInfo ):
+        if device_info not in S3FD.get_available_devices():
+            raise Exception(f'device_info {device_info} is not in available devices for S3FD')
+
+        path = Path(__file__).parent / 'S3FD.onnx'
+        self._sess = sess = InferenceSession_with_device(str(path), device_info)
+        self._input_name = sess.get_inputs()[0].name
+
+
+    def extract(self, img : np.ndarray, threshold=0.95, fixed_window=0, min_face_size=40):
+        """
+
+            img     HW,HWC,NHWC     [0..255]
+        """
+        ip = ImageProcessor(img)
+
+        if fixed_window != 0:
+            fixed_window = max(64, max(1, fixed_window // 32) * 32 )
+            img_scale = ip.fit_in(fixed_window, fixed_window, pad_to_target=True, allow_upscale=False)
+        else:
+            ip.pad_to_next_divisor(64, 64)
+            img_scale = 1.0
+
+        img = ip.ch(3).to_uint8().as_float32().apply( lambda img: img - [104,117,123]).get_image('NCHW')
+
+        batches_bbox = self._sess.run(None, {self._input_name: img})
+
+        faces_per_batch = []
+        for batch in range(img.shape[0]):
+            bbox = self.refine( [ x[batch] for x in batches_bbox ], threshold )
+
+            faces = []
+            for l,t,r,b,c in bbox:
+                if img_scale != 1.0:
+                    l,t,r,b = l/img_scale, t/img_scale, r/img_scale, b/img_scale
+
+                bt = b-t
+                if min(r-l,bt) < min_face_size:
+                    continue
+                b += bt*0.1
+
+                faces.append ( (l,t,r,b) )
+
+            faces_per_batch.append(faces)
+
+        return faces_per_batch
+
+
+    def refine(self, olist, threshold):
+        bboxlist = []
+        variances = [0.1, 0.2]
+        for i in range(len(olist) // 2):
+            ocls, oreg = olist[i * 2], olist[i * 2 + 1]
+
+            stride = 2**(i + 2)    # 4,8,16,32,64,128
+            for hindex, windex in [*zip(*np.where(ocls[1, :, :] > threshold))]:
+                axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride
+                score = ocls[1, hindex, windex]
+                loc = np.ascontiguousarray(oreg[:, hindex, windex]).reshape((1, 4))
+                priors = np.array([[axc, ayc, stride * 4, stride * 4]])
+                bbox = np.concatenate((priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+                                       priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])), 1)
+                bbox[:, :2] -= bbox[:, 2:] / 2
+                bbox[:, 2:] += bbox[:, :2]
+                x1, y1, x2, y2 = bbox[0]
+                bboxlist.append([x1, y1, x2, y2, score])
+
+        if len(bboxlist) != 0:
+            bboxlist = np.array(bboxlist)
+            bboxlist = bboxlist[ lib_math.nms(bboxlist[:,0], bboxlist[:,1], bboxlist[:,2], bboxlist[:,3], bboxlist[:,4], 0.3), : ]
+            bboxlist = [x for x in bboxlist if x[-1] >= 0.5]
+
+        return bboxlist

modelhub/onnx/YoloV5Face/YoloV5Face.py

@@ -0,0 +1,145 @@
+from pathlib import Path
+from typing import List
+import numpy as np
+from xlib import math as lib_math
+from xlib.image import ImageProcessor
+from xlib.onnxruntime import (InferenceSession_with_device, ORTDeviceInfo,
+                              get_available_devices_info)
+
+
+class YoloV5Face:
+    """
+    YoloV5Face face detection model.
+
+    arguments
+
+     device_info    ORTDeviceInfo
+
+        use YoloV5Face.get_available_devices()
+        to determine a list of avaliable devices accepted by model
+
+    raises
+     Exception
+    """
+
+    @staticmethod
+    def get_available_devices() -> List[ORTDeviceInfo]:
+        return get_available_devices_info()
+
+    def __init__(self, device_info : ORTDeviceInfo ):
+        if device_info not in YoloV5Face.get_available_devices():
+            raise Exception(f'device_info {device_info} is not in available devices for YoloV5Face')
+
+        path = Path(__file__).parent / 'YoloV5Face.onnx'
+        self._sess = sess = InferenceSession_with_device(str(path), device_info)
+        self._input_name = sess.get_inputs()[0].name
+
+    def extract(self, img, threshold : float = 0.3, fixed_window=0, min_face_size=8, augment=False):
+        """
+        arguments
+
+         img    np.ndarray      ndim 2,3,4
+
+         fixed_window(0)    int  size
+                                 0 mean don't use
+                                 fit image in fixed window
+                                 downscale if bigger than window
+                                 pad if smaller than window
+                                 increases performance, but decreases accuracy
+
+         min_face_size(8)
+
+         augment(False)     bool    augment image to increase accuracy
+                                    decreases performance
+
+        returns a list of [l,t,r,b] for every batch dimension of img
+        """
+
+        ip = ImageProcessor(img)
+        _,H,W,_ = ip.get_dims()
+        if H > 2048 or W > 2048:
+            fixed_window = 2048
+
+        if fixed_window != 0:
+            fixed_window = max(32, max(1, fixed_window // 32) * 32 )
+            img_scale = ip.fit_in(fixed_window, fixed_window, pad_to_target=True, allow_upscale=False)
+        else:
+            ip.pad_to_next_divisor(64, 64)
+            img_scale = 1.0
+
+        ip.ch(3).to_ufloat32()
+
+        _,H,W,_ = ip.get_dims()
+        
+        preds = self._get_preds(ip.get_image('NCHW'))
+
+        if augment:
+            rl_preds = self._get_preds( ip.flip_horizontal().get_image('NCHW') )
+            rl_preds[:,:,0] = W-rl_preds[:,:,0]
+            preds = np.concatenate([preds, rl_preds],1)
+
+        faces_per_batch = []
+        for pred in preds:
+            pred = pred[pred[...,4] >= threshold]
+
+            x,y,w,h,score = pred.T
+
+            l, t, r, b = x-w/2, y-h/2, x+w/2, y+h/2
+            keep = lib_math.nms(l,t,r,b, score, 0.5)
+            l, t, r, b = l[keep], t[keep], r[keep], b[keep]
+
+            faces = []
+            for l,t,r,b in np.stack([l, t, r, b], -1):
+                if img_scale != 1.0:
+                    l,t,r,b = l/img_scale, t/img_scale, r/img_scale, b/img_scale
+
+                if min(r-l,b-t) < min_face_size:
+                   continue
+                faces.append( (l,t,r,b) )
+
+            faces_per_batch.append(faces)
+
+        return faces_per_batch
+
+    def _get_preds(self, img):
+        N,C,H,W = img.shape
+        preds = self._sess.run(None, {self._input_name: img})
+        # YoloV5Face returns 3x [N,C*16,H,W].
+        # C = [cx,cy,w,h,thres, 5*x,y of landmarks, cls_id ]
+        # Transpose and cut first 5 channels.
+        pred0, pred1, pred2 = [pred.reshape( (N,C,16,pred.shape[-2], pred.shape[-1]) ).transpose(0,1,3,4,2)[...,0:5] for pred in preds]
+
+        pred0 = YoloV5Face.process_pred(pred0, W, H, anchor=[ [4,5],[8,10],[13,16] ]  ).reshape( (N, -1, 5) )
+        pred1 = YoloV5Face.process_pred(pred1, W, H, anchor=[ [23,29],[43,55],[73,105] ]  ).reshape( (N, -1, 5) )
+        pred2 = YoloV5Face.process_pred(pred2, W, H, anchor=[ [146,217],[231,300],[335,433] ]  ).reshape( (N, -1, 5) )
+
+        return np.concatenate( [pred0, pred1, pred2], 1 )[...,:5]
+
+    @staticmethod
+    def process_pred(pred, img_w, img_h, anchor):
+        pred_h = pred.shape[-3]
+        pred_w = pred.shape[-2]
+        anchor = np.float32(anchor)[None,:,None,None,:]
+
+        _xv, _yv,  = np.meshgrid(np.arange(pred_w), np.arange(pred_h), )
+        grid = np.stack((_xv, _yv), 2).reshape((1, 1, pred_h, pred_w, 2)).astype(np.float32)
+
+        stride = (img_w // pred_w, img_h // pred_h)
+
+        pred[..., [0,1,2,3,4] ] = YoloV5Face._np_sigmoid(pred[..., [0,1,2,3,4] ])
+
+        pred[..., 0:2] = (pred[..., 0:2]*2 - 0.5 + grid) * stride
+        pred[..., 2:4] = (pred[..., 2:4]*2)**2 * anchor
+        return pred
+
+    @staticmethod
+    def _np_sigmoid(x : np.ndarray):
+        """
+        sigmoid with safe check of overflow
+        """
+        x = -x
+        c = x > np.log( np.finfo(x.dtype).max )
+        x[c] = 0.0
+        result = 1 / (1+np.exp(x))
+        result[c] = 0.0
+        return result

modelhub/onnx/__init__.py

@@ -0,0 +1,4 @@
+from .CenterFace.CenterFace import CenterFace
+from .FaceMesh.FaceMesh import FaceMesh
+from .S3FD.S3FD import S3FD
+from .YoloV5Face.YoloV5Face import YoloV5Face