added columns, improved depth, added error flags

app.py

@@ -15,7 +15,7 @@ import pandas as pd
 from skimage.io import imread, imsave
 # from tddfa.TDDFA import TDDFA
 from tddfa.utils.depth import depth
-from tddfa.TDDFA import TDDFA
+from tddfa.TDDFA_ONNX import TDDFA_ONNX
 
 import torch.optim as optim
 from DSDG.DUM.models.CDCNs_u import Conv2d_cd, CDCN_u
@@ -29,13 +29,15 @@ 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 = 'ddn1'
 
 device = torch.device("cpu")
 labels = ['Live', 'Spoof']
 pix_threshhold = 0.45
-dsdg_threshold = 0.003
+dsdg_threshold = 0.0015
 examples = [
     ['examples/1_1_21_2_33_scene_fake.jpg'],
     ['examples/frame150_real.jpg'],
@@ -53,9 +55,9 @@ deepix_model.load_state_dict(torch.load('./DeePixBiS/DeePixBiS.pth'))
 deepix_model.eval()
 
 
-depth_config_path = 'tddfa/configs/mb05_120x120.yml'  # 'tddfa/configs/mb1_120x120.yml
+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(gpu_mode=False, **cfg)
+tddfa = TDDFA_ONNX(gpu_mode=False, **cfg)
 
 
 cdcn_model = CDCN_u(basic_conv=Conv2d_cd, theta=0.7)
@@ -112,15 +114,18 @@ def find_largest_face(faces):
 
 
 def inference(img):
-    grey = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
+    if img is None:
+        return None, {}, None, None, {}, None, None
+    grey = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
     faces = faceClassifier.detectMultiScale(
         grey, scaleFactor=1.1, minNeighbors=4)
     face = find_largest_face(faces)
     
     if face is not None:
         x, y, w, h = face
-        faceRegion = img[y:y + h, x:x + w]
-        faceRegion = cv.cvtColor(faceRegion, cv.COLOR_BGR2RGB)
+        x2 = x + w
+        y2 = y + h
+        faceRegion = img[y:y2, x:x2]
         faceRegion = tfms(faceRegion)
         faceRegion = faceRegion.unsqueeze(0)
 
@@ -129,21 +134,19 @@ def inference(img):
         res_deepix = torch.mean(mask).item()
         cls_deepix = 'Real' if res_deepix >= pix_threshhold else 'Spoof'
         
-        label_deepix = f'{cls_deepix} {res_deepix:.2f}'
-        confidences_deepix = {label_deepix: res_deepix}
+        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), (x + w, y + h), color_deepix, 2)
-        cv.putText(img_deepix, label_deepix, (x, y + h + 30),
+        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)
 
         # else:
         dense_flag = True
-        boxes = list(face)
-        boxes.append(1)
-        param_lst, roi_box_lst = tddfa(img, [boxes])
+        box = [x, y, x2, y2, 1]
+        param_lst, roi_box_lst = tddfa(img, [box])
         
         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)
+        depth_img = depth(img, ver_lst, tddfa.tri, with_bg_flag=True)
         with torch.no_grad():
             map_score_list = []
             image_x, map_x = prepare_data([img], [list(face)], [depth_img])
@@ -167,13 +170,12 @@ def inference(img):
         if res_dsdg > 10:
             res_dsdg = 0.0
         cls_dsdg = 'Real' if res_dsdg >= dsdg_threshold else 'Spoof'
-        res_dsdg = res_dsdg * 100
+        res_dsdg = res_dsdg * 300
 
-        label_dsdg = f'{cls_dsdg} {res_dsdg:.2f}'
-        confidences_dsdg = {label_dsdg: res_deepix}
+        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), (x + w, y + h), color_dsdg, 2)
-        cv.putText(img_dsdg, label_dsdg, (x, y + h + 30),
+        img_dsdg = cv.rectangle(img.copy(), (x, y), (x2, y2), color_dsdg, 2)
+        cv.putText(img_dsdg, cls_dsdg, (x, y2 + 30),
                     cv.FONT_HERSHEY_COMPLEX, 1, color_dsdg)
 
         cls_deepix, cls_dsdg = [1 if cls_ == 'Real' else 0 for cls_ in [cls_deepix, cls_dsdg]]         
@@ -186,6 +188,12 @@ def inference(img):
 def upload_to_s3(image_array, app_version, *labels):
     folder = 'demo'
     bucket_name = 'livenessng'
+    if image_array is None:
+        return 'Error. Take a photo first.'
+    elif labels[-2] == -1:
+        return 'Error. Run the detection first.'
+    elif labels[0] is None:
+        return 'Error. Select the true label first.'
 
     # Initialize S3 client
     s3 = boto3.client('s3')
@@ -212,25 +220,27 @@ def upload_to_s3(image_array, app_version, *labels):
 demo = gr.Blocks()
 
 with demo:
-    input_img = gr.Image(source='webcam', shape=None, type='numpy')
-    btn_run = gr.Button(value="Run")
-    with gr.Column():
-        outputs=[
-            gr.Image(label='DeePixBiS', type='numpy'),
-            gr.Label(num_top_classes=2, label='DeePixBiS'),
-            gr.Image(label='DSDG', type='numpy'),
-            gr.Label(num_top_classes=2, label='DSDG')]
-    labels = [gr.Number(visible=False), gr.Number(visible=False)]
+    with gr.Row():
+        with gr.Column():
+            input_img = gr.Image(source='webcam', shape=None, type='numpy')
+            btn_run = gr.Button(value="Run")
+        with gr.Column():
+            outputs=[
+                gr.Image(label='DeePixBiS', type='numpy'),
+                gr.Label(num_top_classes=2, label='DeePixBiS'),
+                gr.Image(label='DSDG', type='numpy'),
+                gr.Label(num_top_classes=2, label='DSDG')]
+        with gr.Column():
+            radio = gr.Radio(
+                ["Real", "Spoof", "None"], label="True label", type='index')
+            flag = gr.Button(value="Flag")
+            status = gr.Textbox()
+            example_block = gr.Examples(examples, [input_img], outputs+labels)
+
+    labels = [gr.Number(visible=False, value=-1), gr.Number(visible=False, value=-1)]
     btn_run.click(inference, [input_img], outputs+labels)
-
     app_version_block = gr.Textbox(value=app_version, visible=False)
-    with gr.Column():
-        radio = gr.Radio(
-            ["Real", "Spoof", "None"], label="True label", type='index'
-        )
-        flag = gr.Button(value="Flag")
-        status = gr.Textbox()
-        flag.click(upload_to_s3, [input_img, app_version_block, radio]+labels, [status], show_progress=True)
+    flag.click(upload_to_s3, [input_img, app_version_block, radio]+labels, [status], show_progress=True)
 
 
 if __name__ == '__main__':