app.py

import subprocess
subprocess.run(["sh", "tddfa/build.sh"]) 

import gradio as gr
from gradio.components import Dropdown

import cv2 as cv
import torch
from torchvision import transforms
from DeePixBiS.Model import DeePixBiS

import yaml
import numpy as np
import pandas as pd
from skimage.io import imread, imsave
# from tddfa.TDDFA import TDDFA
from tddfa.utils.depth import depth
from tddfa.TDDFA_ONNX import TDDFA_ONNX

import torch.optim as optim
from DSDG.DUM.models.CDCNs_u import Conv2d_cd, CDCN_u

import io
import uuid
import numpy as np
from PIL import Image
import boto3

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
os.environ['OMP_NUM_THREADS'] = '4'
os.environ['AWS_ACCESS_KEY_ID'] = 'AKIA3JAMX4K53MFDKMGJ'
os.environ['AWS_SECRET_ACCESS_KEY'] = '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'

app_version = 'dsdg_vid_1'

device = torch.device("cpu")
labels = ['Live', 'Spoof']
PIX_THRESHOLD = 0.45
DSDG_THRESHOLD = 0.5
MIN_FACE_WIDTH_THRESHOLD = 210
examples = [
    ['examples/1_1_21_2_33_scene_fake.jpg'],
    ['examples/frame150_real.jpg'],
    ['examples/1_2.avi_125_real.jpg'],
    ['examples/1_3.avi_25_fake.jpg']]
faceClassifier = cv.CascadeClassifier('./DeePixBiS/Classifiers/haarface.xml')
tfms = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
deepix_model = DeePixBiS(pretrained=False)
deepix_model.load_state_dict(torch.load('./DeePixBiS/DeePixBiS.pth'))
deepix_model.eval()


depth_config_path = 'tddfa/configs/mb1_120x120.yml'  # 'tddfa/configs/mb1_120x120.yml
cfg = yaml.load(open(depth_config_path), Loader=yaml.SafeLoader)
tddfa = TDDFA_ONNX(gpu_mode=False, **cfg)


cdcn_model = CDCN_u(basic_conv=Conv2d_cd, theta=0.7)
cdcn_model = cdcn_model.to(device)
weights = torch.load('./DSDG/DUM/checkpoint/CDCN_U_P1_updated.pkl', map_location=device)
cdcn_model.load_state_dict(weights)
optimizer = optim.Adam(cdcn_model.parameters(), lr=0.001, weight_decay=0.00005)
cdcn_model.eval()


class Normaliztion_valtest(object):
    """
        same as mxnet, normalize into [-1, 1]
        image = (image - 127.5)/128
    """
    def __call__(self, image_x):
        image_x = (image_x - 127.5) / 128  # [-1,1]
        return image_x


def prepare_data_dsdg(images, boxes, depths):
    transform = transforms.Compose([Normaliztion_valtest()])
    files_total = 1
    image_x = np.zeros((files_total, 256, 256, 3))
    depth_x = np.ones((files_total, 32, 32))

    for i, (image, bbox, depth_img) in enumerate(
            zip(images, boxes, depths)):
        x, y, x2, y2 = bbox
        depth_img = cv.cvtColor(depth_img, cv.COLOR_RGB2GRAY)
        image = image[y:y2, x:x2]
        depth_img = depth_img[y:y2, x:x2]

        image_x[i, :, :, :] = cv.resize(image, (256, 256))
        # transform to binary mask --> threshold = 0 
        depth_x[i, :, :] = cv.resize(depth_img, (32, 32))
    image_x = image_x.transpose((0, 3, 1, 2))
    image_x = transform(image_x)
    image_x = torch.from_numpy(image_x.astype(float)).float()
    depth_x = torch.from_numpy(depth_x.astype(float)).float()
    return image_x, depth_x


def find_largest_face(faces):
    # find the largest face in the list
    largest_face = None
    largest_area = 0
    for face in faces:
        x, y, w, h = face
        area = w * h
        if area > largest_area:
            largest_area = area
            largest_face = face
    return largest_face


def extract_face(img):
    face = None
    if img is None:
        return face
    grey = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
    faces = faceClassifier.detectMultiScale(
        grey, scaleFactor=1.1, minNeighbors=4)
    if len(faces):
        face = find_largest_face(faces)
    return face


def deepix_model_inference(img, bbox):
    x, y, x2, y2 = bbox
    faceRegion = img[y:y2, x:x2]
    faceRegion = tfms(faceRegion)
    faceRegion = faceRegion.unsqueeze(0)
    mask, binary = deepix_model.forward(faceRegion)
    res_deepix = torch.mean(mask).item()
    cls_deepix = 'Real' if res_deepix >= PIX_THRESHOLD else 'Spoof'
    confidences_deepix = {'Real confidence': res_deepix}
    color_deepix = (0, 255, 0) if cls_deepix == 'Real' else (255, 0, 0)
    img_deepix = cv.rectangle(img.copy(), (x, y), (x2, y2), color_deepix, 2)
    cv.putText(img_deepix, cls_deepix, (x, y2 + 30),
                cv.FONT_HERSHEY_COMPLEX, 1, color_deepix)
    cls_deepix = 1 if cls_deepix == 'Real' else 0
    return img_deepix, confidences_deepix, cls_deepix


def dsdg_model_inference(img, bbox, dsdg_thresh):
    dsdg_thresh = dsdg_thresh / 10000
    dense_flag = True
    x, y, x2, y2 = bbox
    w = x2 - x
    h = y2 - y
    if w < MIN_FACE_WIDTH_THRESHOLD:
        color_dsdg = (0, 0, 0)
        text = f'Small res ({w}*{h})'
        img_dsdg = cv.rectangle(img.copy(), (x, y), (x2, y2), color_dsdg, 2)
        cv.putText(img_dsdg, text, (x, y2 + 30),
                   cv.FONT_HERSHEY_COMPLEX, 1, color_dsdg)
        cls_dsdg = -1
        return img_dsdg, {}, cls_dsdg
    bbox_conf = list(bbox)
    bbox_conf.append(1)
    param_lst, roi_box_lst = tddfa(img, [bbox_conf])
    ver_lst = tddfa.recon_vers(param_lst, roi_box_lst, dense_flag=dense_flag)
    depth_img = depth(img, ver_lst, tddfa.tri, with_bg_flag=False)
    with torch.no_grad():
        map_score_list = []
        image_x, map_x = prepare_data_dsdg([img], [list(bbox)], [depth_img])
        # get the inputs
        image_x = image_x.unsqueeze(0)
        map_x = map_x.unsqueeze(0)
        inputs = image_x.to(device)
        test_maps = map_x.to(device)
        optimizer.zero_grad()
        map_score = 0.0
        for frame_t in range(inputs.shape[1]):
            mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = cdcn_model(inputs[:, frame_t, :, :, :])
            score_norm = torch.sum(mu) / torch.sum(test_maps[:, frame_t, :, :])
            map_score += score_norm
        map_score = map_score / inputs.shape[1]
        map_score_list.append(map_score)
    res_dsdg = map_score_list[0].item()
    if res_dsdg > 10:
        res_dsdg = 0.0
    cls_dsdg = 'Real' if res_dsdg >= dsdg_thresh else 'Spoof'
    text = f'{cls_dsdg} {w}*{h}'
    confidences_dsdg = {'Real confidence': res_dsdg}
    color_dsdg = (0, 255, 0) if cls_dsdg == 'Real' else (255, 0, 0)
    img_dsdg = cv.rectangle(img.copy(), (x, y), (x2, y2), color_dsdg, 2)
    cv.putText(img_dsdg, text, (x, y2 + 30),
                cv.FONT_HERSHEY_COMPLEX, 1, color_dsdg)
    res_dsdg = res_dsdg * 1000000
    # cls_dsdg = 1 if cls_dsdg == 'Real' else 0
    return img_dsdg, confidences_dsdg, res_dsdg


def inference(img, dsdg_thresh):
    face = extract_face(img)
    if face is not None:
        x, y, w, h = face
        x2 = x + w
        y2 = y + h
        bbox = (x, y, x2, y2)
        # img_deepix, confidences_deepix, cls_deepix = deepix_model_inference(img, bbox)
        img_dsdg, confidences_dsdg, cls_dsdg = dsdg_model_inference(img, bbox, dsdg_thresh)
        return img, {}, 2, img_dsdg, confidences_dsdg, cls_dsdg
    else:
        return img, {}, None, img, {}, None


def process_video(vid_path, dsdg_thresh):
    cap = cv.VideoCapture(vid_path)

    input_width = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
    input_height = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))

    # Set video codec and create VideoWriter object to save the output video
    fourcc = cv.VideoWriter_fourcc(*'mp4v')
    output_vid_path = 'output_dsdg.mp4'
    out_dsdg = cv.VideoWriter(output_vid_path, fourcc, 20.0, (input_width, input_height))

    frame_counter = 0
    confidences_arr = []
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        # Process only every 5th frame
        if frame_counter % 5 == 0:
            # Run inference on the current frame
            _, _, _, img_dsdg, confidences_dsdg, res_dsdg = inference(frame, dsdg_thresh)
            if res_dsdg == -1:
                continue
            confidences_arr.append(res_dsdg)
            # Write the DSDG frame to the output video
            out_dsdg.write(img_dsdg)
        frame_counter += 1
    # Release resources
    cap.release()
    out_dsdg.release()
    if not confidences_arr:
        return vid_path, {'Not supported right now': 0}, -1, vid_path, 'Faces too small or not found', -1
    avg_conf = sum(confidences_arr) / len(confidences_arr)
    text_dsdg = f'Average real confidence: {avg_conf}\nFrames used: {len(confidences_arr)}\nConfidences:{confidences_arr}'
    return vid_path, {'Not supported right now': 0}, -1, output_vid_path, text_dsdg, avg_conf


def upload_to_s3(vid_path, app_version, *labels):
    folder = 'demo'
    bucket_name = 'livenessng'

    if vid_path is None:
        return 'Error. Take a photo first.'
    elif labels[-2] == -2:
        return 'Error. Run the detection first.'
    elif labels[0] is None:
        return 'Error. Select the true label first.'
    elif labels[0] == 2:
        labels[0] = -1

    # Initialize S3 client
    s3 = boto3.client('s3')

    # Encode labels and app version in video file name
    encoded_labels = '_'.join([str(int(label)) for label in labels])
    random_string = str(uuid.uuid4()).split('-')[-1]
    video_name = f"{folder}/{app_version}/{encoded_labels}_{random_string}.mp4"

    # Upload video to S3
    with open(vid_path, 'rb') as video_file:
        res = s3.upload_fileobj(video_file, bucket_name, video_name)

    # Return the S3 URL of the uploaded video
    status = 'Successfully uploaded'
    return status


demo = gr.Blocks()

with demo:
    with gr.Row():
        with gr.Column():
            input_vid = gr.Video(format='mp4', source='webcam')
            dsdg_thresh = gr.Slider(value=DSDG_THRESHOLD, label='DSDG threshold', maximum=3.0, step=0.05)
            btn_run = gr.Button(value="Run")
        with gr.Column():
            outputs=[
                gr.Video(label='DeePixBiS', format='mp4'),
                gr.Label(num_top_classes=2, label='DeePixBiS'),
                gr.Number(visible=False, value=-2),
                gr.Video(label='DSDG', format='mp4'),
                gr.Textbox(label='DSDG'),
                gr.Number(visible=False, value=-2)]
        with gr.Column():
            radio = gr.Radio(
                ["Spoof", "Real", "None"], label="True label", type='index')
            flag = gr.Button(value="Flag")
            status = gr.Textbox()
            # example_block = gr.Examples(examples, [input_vid], outputs)

    btn_run.click(process_video, [input_vid, dsdg_thresh], outputs)
    app_version_block = gr.Textbox(value=app_version, visible=False)
    flag.click(
        upload_to_s3,
        [input_vid, app_version_block, radio]+[outputs[2], outputs[5]],
        [status], show_progress=True)


if __name__ == '__main__':
    demo.queue(concurrency_count=2)
    demo.launch(share=False)