File size: 3,490 Bytes
5ed15ed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import numpy as np
import cv2
import gradio as gr

from opencv_zoo.models.face_detection_yunet.yunet import YuNet
from opencv_zoo.models.license_plate_detection_yunet.lpd_yunet import LPD_YuNet

# Instantiate face detection YuNet
face_model = YuNet(
    modelPath="opencv_zoo/models/face_detection_yunet/face_detection_yunet_2023mar.onnx",
    inputSize=[320, 320],
    confThreshold=0.9,
    nmsThreshold=0.3,
    topK=5000,
    backendId=cv2.dnn.DNN_BACKEND_OPENCV,
    targetId=cv2.dnn.DNN_TARGET_CPU,
)

# Instantiate license plate detection YuNet
lpd_model = LPD_YuNet(
    modelPath="opencv_zoo/models/license_plate_detection_yunet/license_plate_detection_lpd_yunet_2023mar.onnx",
    confThreshold=0.9,
    nmsThreshold=0.3,
    topK=5000,
    keepTopK=750,
    backendId=cv2.dnn.DNN_BACKEND_OPENCV,
    targetId=cv2.dnn.DNN_TARGET_CPU,
)


def json_detections(face_results, lpd_results):
    json_result = {}
    json_result["faces"] = []
    json_result["license_plates"] = []

    for det in face_results if face_results is not None else []:
        bbox = det[0:4].astype(np.int32)
        json_result["faces"].append(
            {
                "xmin": int(bbox[0]),
                "ymin": int(bbox[1]),
                "xmax": int(bbox[0]) + int(bbox[2]),
                "ymax": int(bbox[1]) + int(bbox[3]),
            }
        )

    for det in lpd_results if lpd_results is not None else []:
        bbox = det[:-1].astype(np.int32)
        x1, y1, x2, y2, x3, y3, x4, y4 = bbox

        xmin = min(x1, x2, x3, x4)
        xmax = max(x1, x2, x3, x4)
        ymin = min(y1, y2, y3, y4)
        ymax = max(y1, y2, y3, y4)

        json_result["license_plates"].append(
            {
                "xmin": int(xmin),
                "ymin": int(ymin),
                "xmax": int(xmax),
                "ymax": int(ymax),
            }
        )

    return json_result


def overlay_results(image, face_results, lpd_results):
    # Draw face results on the input image
    for det in face_results if face_results is not None else []:
        bbox = det[0:4].astype(np.int32)
        cv2.rectangle(
            image,
            (bbox[0], bbox[1]),
            (bbox[0] + bbox[2], bbox[1] + bbox[3]),
            (0, 0, 0),
            -1,
        )

    # Draw lpd results on the input image
    for det in lpd_results:
        bbox = det[:-1].astype(np.int32)
        x1, y1, x2, y2, x3, y3, x4, y4 = bbox

        # The output of this is technically a parallelogram, but we will
        # just black out the rectangle
        xmin = min(x1, x2, x3, x4)
        xmax = max(x1, x2, x3, x4)
        ymin = min(y1, y2, y3, y4)
        ymax = max(y1, y2, y3, y4)

        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 0, 0), -1)

    return image


def predict(image):

    h, w, _ = image.shape

    # Inference
    face_model.setInputSize([w, h])
    lpd_model.setInputSize([w, h])
    infer_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    face_results = face_model.infer(infer_image)
    lpd_results = lpd_model.infer(infer_image)

    # Process output
    image = overlay_results(image, face_results, lpd_results)
    json = json_detections(face_results, lpd_results)

    return image, json


demo = gr.Interface(
    title="Face and License Plate Obfuscator - YuNet",
    fn=predict,
    inputs=gr.Image(type="numpy", label="Original Image"),
    outputs=[
        gr.Image(type="numpy", label="Output Image"),
        gr.JSON(),
    ],
)

demo.launch()