src/pixel3dmm/preprocessing/pipnet_utils.py

import importlib
import os
import torch.nn.parallel
import torch.utils.data
import torchvision.transforms as transforms


from pixel3dmm.preprocessing.PIPNet.FaceBoxesV2.faceboxes_detector import *
from pixel3dmm.preprocessing.PIPNet.lib.networks import *
from pixel3dmm.preprocessing.PIPNet.lib.functions import *
from pixel3dmm.preprocessing.PIPNet.lib.mobilenetv3 import mobilenetv3_large
from pixel3dmm import env_paths

def smooth(x, window_len=11, window='hanning'):
    """smooth the data using a window with requested size.

    This method is based on the convolution of a scaled window with the signal.
    The signal is prepared by introducing reflected copies of the signal
    (with the window size) in both ends so that transient parts are minimized
    in the begining and end part of the output signal.

    input:
        x: the input signal
        window_len: the dimension of the smoothing window; should be an odd integer
        window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
            flat window will produce a moving average smoothing.

    output:
        the smoothed signal

    example:

    t=linspace(-2,2,0.1)
    x=sin(t)+randn(len(t))*0.1
    y=smooth(x)

    see also:

    numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
    scipy.signal.lfilter

    TODO: the window parameter could be the window itself if an array instead of a string
    NOTE: length(output) != length(input), to correct this: return y[(window_len/2-1):-(window_len/2)] instead of just y.
    """

    if x.ndim != 1:
        raise ValueError("smooth only accepts 1 dimension arrays.")

    if x.size < window_len:
        raise ValueError( "Input vector needs to be bigger than window size.")

    if window_len < 3:
        return x

    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
        raise ValueError( "Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'")

    s = np.r_[x[window_len - 1:0:-1], x, x[-2:-window_len - 1:-1]]
    # print(len(s))
    if window == 'flat':  # moving average
        w = np.ones(window_len, 'd')
    else:
        w = eval('np.' + window + '(window_len)')

    y = np.convolve(w / w.sum(), s, mode='valid')
    return y

def get_cstm_crop(image, detections, detections_max, max_bbox : bool = False):
    #Image.fromarray(image).show()
    image_width = image.shape[1]
    image_height = image.shape[0]

    det_box_scale = 1.42 #2.0#1.42
    if detections[4]*1.42 * detections[5]*1.42 < detections_max[4] * 1.1 * detections_max[5] * 1.1:
        detections = detections_max
        det_box_scale = 1.1

    det_xmin = detections[2]
    det_ymin = detections[3]
    det_width = detections[4]
    det_height = detections[5]
    if det_width > det_height:
        det_ymin -= (det_width - det_height)//2
        det_height = det_width
    if det_width < det_height:
        det_xmin -= (det_height - det_width)//2
        det_width = det_height

    det_xmax = det_xmin + det_width - 1
    det_ymax = det_ymin + det_height - 1


    det_xmin -= int(det_width * (det_box_scale - 1) / 2)
    det_ymin -= int(det_height * (det_box_scale - 1) / 2)
    det_xmax += int(det_width * (det_box_scale - 1) / 2)
    det_ymax += int(det_height * (det_box_scale - 1) / 2)
    if det_xmin < 0 or det_ymin < 0:
        min_overflow = min(det_xmin, det_ymin)
        det_xmin += -min_overflow
        det_ymin += -min_overflow
    if det_xmax > image_width -1 or det_ymax > image_height - 1:
        max_overflow = max(det_xmax - image_width -1, det_ymax - image_height-1)
        det_xmax -= max_overflow
        det_ymax -= max_overflow

    det_width = det_xmax - det_xmin + 1
    det_height = det_ymax - det_ymin + 1
    det_crop = image[det_ymin:det_ymax, det_xmin:det_xmax, :]
    return det_crop, det_ymin, det_ymax, det_xmin, det_xmax
    #Image.fromarray(det_crop).show()
    #exit()


def demo_image(image_dir, pid, save_dir, preprocess, cfg, input_size, net_stride, num_nb, use_gpu, flip=False, start_frame=0,
               vertical_crop : bool = False,
               static_crop : bool = False,
               max_bbox : bool = False,
               disable_cropping : bool = False,
               ):

    if cfg.use_gpu:
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device("cpu")

    detector = FaceBoxesDetector('FaceBoxes', f'{env_paths.CODE_BASE}/src/pixel3dmm/preprocessing/PIPNet/FaceBoxesV2/weights/FaceBoxesV2.pth', use_gpu, device)
    my_thresh = 0.6
    det_box_scale = 1.2
    meanface_indices, reverse_index1, reverse_index2, max_len = get_meanface(
        os.path.join(f'{env_paths.CODE_BASE}/src/pixel3dmm/preprocessing/', 'PIPNet', 'data', cfg.data_name, 'meanface.txt'), cfg.num_nb)

    if cfg.backbone == 'resnet18':
        resnet18 = models.resnet18(pretrained=cfg.pretrained)
        net = Pip_resnet18(resnet18, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size,
                           net_stride=cfg.net_stride)
    elif cfg.backbone == 'resnet50':
        resnet50 = models.resnet50(pretrained=cfg.pretrained)
        net = Pip_resnet50(resnet50, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size,
                           net_stride=cfg.net_stride)
    elif cfg.backbone == 'resnet101':
        resnet101 = models.resnet101(pretrained=cfg.pretrained)
        net = Pip_resnet101(resnet101, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size,
                            net_stride=cfg.net_stride)
    elif cfg.backbone == 'mobilenet_v2':
        mbnet = models.mobilenet_v2(pretrained=cfg.pretrained)
        net = Pip_mbnetv2(mbnet, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
    elif cfg.backbone == 'mobilenet_v3':
        mbnet = mobilenetv3_large()
        if cfg.pretrained:
            mbnet.load_state_dict(torch.load('lib/mobilenetv3-large-1cd25616.pth'))
        net = Pip_mbnetv3(mbnet, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
    else:
        print('No such backbone!')
        exit(0)


    net = net.to(device)

    weight_file = os.path.join(save_dir, 'epoch%d.pth' % (cfg.num_epochs - 1))
    state_dict = torch.load(weight_file, map_location=device)
    net.load_state_dict(state_dict)
    net.eval()

    if start_frame > 0:
        files = [f for f in os.listdir(f'{image_dir}/') if f.endswith('.jpg') or f.endswith('.png') and (((int(f.split('_')[-1].split('.')[0])-start_frame) % 3 )== 0)]
    else:
        files = [f for f in os.listdir(f'{image_dir}/') if f.endswith('.jpg') or f.endswith('.png')]
    files.sort()

    if not vertical_crop:
        all_detections = []
        all_images = []
        #all_normals = []
        succ_files = []
        for file_name in files:
                image = cv2.imread(f'{image_dir}/{file_name}')
                #normals = cv2.imread(f'{image_dir}/../normals/{file_name[:-4]}.png')

                if len(image.shape) < 3 or image.shape[-1] != 3:
                    continue

                image_height, image_width, _ = image.shape



                detections, _ = detector.detect(image, my_thresh, 1)
                dets_filtered = [det for det in detections if det[0] == 'face']
                dets_filtered.sort(key=lambda x: -1 * x[1])
                detections = dets_filtered
                if detections[0][1] < 0.75:
                    raise ValueError("Found face with too low detections confidence as max confidence")
                all_detections.append(detections[0])
                all_images.append(image)
                #all_normals.append(normals)
                succ_files.append(file_name)

        assert static_crop, 'Other options currently not supported anymore'
        if static_crop:
            #if max_bbox:
            det1_max = np.min(np.array([x[2] for x in all_detections]), axis=0)
            det2_max = np.min(np.array([x[3] for x in all_detections]), axis=0)
            det3_max = np.max(np.array([x[4]+x[2]-det1_max for x in all_detections]), axis=0)
            det4_max = np.max(np.array([x[5]+x[3]-det2_max for x in all_detections]), axis=0)
            det1 = np.mean(np.array([x[2] for x in all_detections]), axis=0)
            det2 = np.mean(np.array([x[3] for x in all_detections]), axis=0)
            det3 = np.mean(np.array([x[4] for x in all_detections]), axis=0)
            det4 = np.mean(np.array([x[5] for x in all_detections]), axis=0)

            det_smoothed = np.stack([det1, det2, det3, det4], axis=0).astype(np.int32)
            det_smoothed_max = np.stack([det1_max, det2_max, det3_max, det4_max], axis=0).astype(np.int32)
            all_detections_smoothed = []  # = [[x[0], x[1], x_smoothed[0], x_smoothed[1], x_smoothed[2], x_smoothed[3]] for x, x_smoothed in zip()]
            all_detections_max_smoothed = []  # = [[x[0], x[1], x_smoothed[0], x_smoothed[1], x_smoothed[2], x_smoothed[3]] for x, x_smoothed in zip()]
            for i, det in enumerate(all_detections):
                all_detections_smoothed.append(
                    [det[0], det[1], det_smoothed[0], det_smoothed[1], det_smoothed[2], det_smoothed[3]])
                all_detections_max_smoothed.append(
                    [det[0], det[1], det_smoothed_max[0], det_smoothed_max[1], det_smoothed_max[2], det_smoothed_max[3]])
            all_detections = all_detections_smoothed
            all_detections_max = all_detections_max_smoothed
        else:
            if len(all_detections) > 11:
                WINDOW_LENGTH = 11
                det1 = smooth(np.array([x[2] for x in all_detections]), window_len=WINDOW_LENGTH)
                det2 = smooth(np.array([x[3] for x in all_detections]), window_len=WINDOW_LENGTH)
                det3 = smooth(np.array([x[4] for x in all_detections]), window_len=WINDOW_LENGTH)
                det4 = smooth(np.array([x[5] for x in all_detections]), window_len=WINDOW_LENGTH)
                det_smoothed = np.stack([det1, det2,det3,det4], axis=1).astype(np.int32)
                all_detections_smoothed = [] #= [[x[0], x[1], x_smoothed[0], x_smoothed[1], x_smoothed[2], x_smoothed[3]] for x, x_smoothed in zip()]
                for i, det in enumerate(all_detections):
                    all_detections_smoothed.append([det[0], det[1], det_smoothed[i, 0], det_smoothed[i, 1], det_smoothed[i, 2], det_smoothed[i, 3]])
                all_detections = all_detections_smoothed
        # TODO: smooth detections!!!
        for file_name, detection, detection_max, image in zip(succ_files, all_detections, all_detections_max, all_images):

                        if not disable_cropping:
                            img_crop, det_ymin, det_ymax, det_xmin, det_xmax = get_cstm_crop(image, detection, detection_max, max_bbox=max_bbox)
                            #n_crop = get_cstm_crop(normals, detection)
                            image = img_crop
                        # save cropped image
                        os.makedirs(f'{image_dir}/../cropped/', exist_ok=True)
                        #os.makedirs(f'{image_dir}/../cropped_normals/', exist_ok=True)
                        cv2.imwrite(f'{image_dir}/../cropped/{file_name}', cv2.resize(image, (512, 512)))
                        #cv2.imwrite(f'{image_dir}/../cropped_normals/{file_name[:-4]}.png', cv2.resize(n_crop, (512, 512)))

                        # store cropping information:
                        if not os.path.exists(f'{image_dir}/../crop_ymin_ymax_xmin_xmax.npy'):
                            np.save(f'{image_dir}/../crop_ymin_ymax_xmin_xmax.npy', np.array([det_ymin, det_ymax, det_xmin, det_xmax]))
    else:
        for file_name in files:
            image = cv2.imread(f'{image_dir}/{file_name}')
            if image.shape[0] != image.shape[1]:
                image = image[220:-220, 640:-640, :]
            os.makedirs(f'{image_dir}/../cropped/', exist_ok=True)
            cv2.imwrite(f'{image_dir}/../cropped/{file_name}', cv2.resize(image, (512, 512)))


    # run landmark detection
    lms = []
    image_dir = f'{image_dir}/../cropped/'
    for file_name in files:
                image = cv2.imread(f'{image_dir}/{file_name}')

                if len(image.shape) < 3 or image.shape[-1] != 3:
                    continue
                if flip:
                    image = cv2.transpose(image)

                image_height, image_width, _ = image.shape
                detections, _ = detector.detect(image, my_thresh, 1)
                pred_export = None
                dets_filtered = [det for det in detections if det[0] == 'face']
                dets_filtered.sort(key=lambda x: -1 * x[1])
                detections = dets_filtered


                print(detections)
                for i in range(min(1, len(detections))):
                    if detections[i][1] < 0.99:
                        continue
                    det_xmin = detections[i][2]
                    det_ymin = detections[i][3]
                    det_width = detections[i][4]
                    det_height = detections[i][5]
                    det_xmax = det_xmin + det_width - 1
                    det_ymax = det_ymin + det_height - 1


                    det_xmin -= int(det_width * (det_box_scale - 1) / 2)
                    # remove a part of top area for alignment, see paper for details
                    det_ymin += int(det_height * (det_box_scale - 1) / 2)
                    det_xmax += int(det_width * (det_box_scale - 1) / 2)
                    det_ymax += int(det_height * (det_box_scale - 1) / 2)
                    det_xmin = max(det_xmin, 0)
                    det_ymin = max(det_ymin, 0)
                    det_xmax = min(det_xmax, image_width - 1)
                    det_ymax = min(det_ymax, image_height - 1)
                    det_width = det_xmax - det_xmin + 1
                    det_height = det_ymax - det_ymin + 1
                    cv2.rectangle(image, (det_xmin, det_ymin), (det_xmax, det_ymax), (0, 0, 255), 2)
                    det_crop = image[det_ymin:det_ymax, det_xmin:det_xmax, :]
                    #np.save(f'{CROP_DIR}/{pid[:-4]}.npy', np.array([det_ymin, det_ymax, det_xmin, det_xmax]))
                    det_crop = cv2.resize(det_crop, (input_size, input_size))
                    inputs = Image.fromarray(det_crop[:, :, ::-1].astype('uint8'), 'RGB')
                    #inputs.show()
                    inputs = preprocess(inputs).unsqueeze(0)
                    inputs = inputs.to(device)
                    lms_pred_x, lms_pred_y, lms_pred_nb_x, lms_pred_nb_y, outputs_cls, max_cls = forward_pip(net,
                                                                                                             inputs,
                                                                                                             preprocess,
                                                                                                             input_size,
                                                                                                             net_stride,
                                                                                                             num_nb)
                    lms_pred = torch.cat((lms_pred_x, lms_pred_y), dim=1).flatten()
                    tmp_nb_x = lms_pred_nb_x[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
                    tmp_nb_y = lms_pred_nb_y[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
                    tmp_x = torch.mean(torch.cat((lms_pred_x, tmp_nb_x), dim=1), dim=1).view(-1, 1)
                    tmp_y = torch.mean(torch.cat((lms_pred_y, tmp_nb_y), dim=1), dim=1).view(-1, 1)
                    lms_pred_merge = torch.cat((tmp_x, tmp_y), dim=1).flatten()
                    lms_pred = lms_pred.cpu().numpy()
                    lms_pred_merge = lms_pred_merge.cpu().numpy()
                    pred_export = np.zeros([cfg.num_lms, 2])
                    for i in range(cfg.num_lms):
                        x_pred = lms_pred_merge[i * 2] * det_width
                        y_pred = lms_pred_merge[i * 2 + 1] * det_height
                        pred_export[i, 0] = (x_pred + det_xmin) / image_width
                        pred_export[i, 1] = (y_pred + det_ymin) / image_height
                        cv2.circle(image, (int(x_pred) + det_xmin, int(y_pred) + det_ymin), 1, (0, 0, 255), 2)
                        if i == 76:
                            cv2.circle(image, (int(x_pred) + det_xmin, int(y_pred) + det_ymin), 1, (255, 0, 0), 2)

                if pred_export is not None:
                    print('exporting stuff to ' + image_dir)
                    landmakr_dir =  f'{image_dir}/../PIPnet_landmarks/'
                    os.makedirs(landmakr_dir, exist_ok=True)
                    np.save(landmakr_dir + f'/{file_name[:-4]}.npy', pred_export)
                    lms.append(pred_export)
                    exp_dir = image_dir + '/../PIPnet_annotated_images/'
                    os.makedirs(exp_dir, exist_ok=True)
                    cv2.imwrite(exp_dir + f'/{file_name}', image)

                # cv2.imshow('1', image)
                # cv2.waitKey(0)

    lms = np.stack(lms, axis=0)
    os.makedirs(f'{image_dir}/../pipnet', exist_ok=True)
    np.save(f'{image_dir}/../pipnet/test.npy', lms)