paulo061/codeformer

1import copy2import os3from pathlib import Path4 5import cv26import numpy as np7import torch8from torch import nn9 10from facelib.detection.yolov5face.models.common import Conv11from facelib.detection.yolov5face.models.yolo import Model12from facelib.detection.yolov5face.utils.datasets import letterbox13from facelib.detection.yolov5face.utils.general import (14    check_img_size,15    non_max_suppression_face,16    scale_coords,17    scale_coords_landmarks,18)19 20IS_HIGH_VERSION = tuple(map(int, torch.__version__.split('+')[0].split('.')[:3])) >= (1, 9, 0)21 22 23def isListempty(inList):24    if isinstance(inList, list): # Is a list25        return all(map(isListempty, inList))26    return False # Not a list27 28class YoloDetector:29    def __init__(30        self,31        config_name,32        min_face=10,33        target_size=None,34        device='cuda',35    ):36        """37        config_name: name of .yaml config with network configuration from models/ folder.38        min_face : minimal face size in pixels.39        target_size : target size of smaller image axis (choose lower for faster work). e.g. 480, 720, 1080.40                    None for original resolution.41        """42        self._class_path = Path(__file__).parent.absolute()43        self.target_size = target_size44        self.min_face = min_face45        self.detector = Model(cfg=config_name)46        self.device = device47 48 49    def _preprocess(self, imgs):50        """51        Preprocessing image before passing through the network. Resize and conversion to torch tensor.52        """53        pp_imgs = []54        for img in imgs:55            h0, w0 = img.shape[:2]  # orig hw56            if self.target_size:57                r = self.target_size / min(h0, w0)  # resize image to img_size58                if r < 1:59                    img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=cv2.INTER_LINEAR)60 61            imgsz = check_img_size(max(img.shape[:2]), s=self.detector.stride.max())  # check img_size62            img = letterbox(img, new_shape=imgsz)[0]63            pp_imgs.append(img)64        pp_imgs = np.array(pp_imgs)65        pp_imgs = pp_imgs.transpose(0, 3, 1, 2)66        pp_imgs = torch.from_numpy(pp_imgs).to(self.device)67        pp_imgs = pp_imgs.float()  # uint8 to fp16/3268        return pp_imgs / 255.0  # 0 - 255 to 0.0 - 1.069 70    def _postprocess(self, imgs, origimgs, pred, conf_thres, iou_thres):71        """72        Postprocessing of raw pytorch model output.73        Returns:74            bboxes: list of arrays with 4 coordinates of bounding boxes with format x1,y1,x2,y2.75            points: list of arrays with coordinates of 5 facial keypoints (eyes, nose, lips corners).76        """77        bboxes = [[] for _ in range(len(origimgs))]78        landmarks = [[] for _ in range(len(origimgs))]79 80        pred = non_max_suppression_face(pred, conf_thres, iou_thres)81 82        for image_id, origimg in enumerate(origimgs):83            img_shape = origimg.shape84            image_height, image_width = img_shape[:2]85            gn = torch.tensor(img_shape)[[1, 0, 1, 0]]  # normalization gain whwh86            gn_lks = torch.tensor(img_shape)[[1, 0, 1, 0, 1, 0, 1, 0, 1, 0]]  # normalization gain landmarks87            det = pred[image_id].cpu()88            scale_coords(imgs[image_id].shape[1:], det[:, :4], img_shape).round()89            scale_coords_landmarks(imgs[image_id].shape[1:], det[:, 5:15], img_shape).round()90 91            for j in range(det.size()[0]):92                box = (det[j, :4].view(1, 4) / gn).view(-1).tolist()93                box = list(94                    map(int, [box[0] * image_width, box[1] * image_height, box[2] * image_width, box[3] * image_height])95                )96                if box[3] - box[1] < self.min_face:97                    continue98                lm = (det[j, 5:15].view(1, 10) / gn_lks).view(-1).tolist()99                lm = list(map(int, [i * image_width if j % 2 == 0 else i * image_height for j, i in enumerate(lm)]))100                lm = [lm[i : i + 2] for i in range(0, len(lm), 2)]101                bboxes[image_id].append(box)102                landmarks[image_id].append(lm)103        return bboxes, landmarks104 105    def detect_faces(self, imgs, conf_thres=0.7, iou_thres=0.5):106        """107        Get bbox coordinates and keypoints of faces on original image.108        Params:109            imgs: image or list of images to detect faces on with BGR order (convert to RGB order for inference)110            conf_thres: confidence threshold for each prediction111            iou_thres: threshold for NMS (filter of intersecting bboxes)112        Returns:113            bboxes: list of arrays with 4 coordinates of bounding boxes with format x1,y1,x2,y2.114            points: list of arrays with coordinates of 5 facial keypoints (eyes, nose, lips corners).115        """116        # Pass input images through face detector117        images = imgs if isinstance(imgs, list) else [imgs]118        images = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in images]119        origimgs = copy.deepcopy(images)120 121        images = self._preprocess(images)122        123        if IS_HIGH_VERSION:124            with torch.inference_mode():  # for pytorch>=1.9 125                pred = self.detector(images)[0]126        else:127            with torch.no_grad():  # for pytorch<1.9128                pred = self.detector(images)[0]129 130        bboxes, points = self._postprocess(images, origimgs, pred, conf_thres, iou_thres)131 132        # return bboxes, points133        if not isListempty(points):134            bboxes = np.array(bboxes).reshape(-1,4)135            points = np.array(points).reshape(-1,10)136            padding = bboxes[:,0].reshape(-1,1)137            return np.concatenate((bboxes, padding, points), axis=1)138        else:139            return None140 141    def __call__(self, *args):142        return self.predict(*args)143