Spaces:
Runtime error
Runtime error
| import argparse | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import kornia as K | |
| from kornia.contrib import FaceDetector, FaceDetectorResult | |
| import gradio as gr | |
| import face_detection | |
| def detect_faces(img: np.ndarray, method:str): | |
| frame = np.array(img) | |
| kornia_detections = kornia_detect(frame) | |
| retina_detections = retina_detect(frame) | |
| retina_mobile_detections = retina_mobilenet_detect(frame) | |
| dsfd_detections = dsfd_detect(frame) | |
| # if method == "Kornia YuNet": | |
| # re_im = kornia_detect(frame) | |
| # elif method == "RetinaFace": | |
| # re_im = retina_detect(frame) | |
| return kornia_detections, retina_detections, retina_mobile_detections, dsfd_detections | |
| def scale_image(img: np.ndarray, size: int) -> np.ndarray: | |
| h, w = img.shape[:2] | |
| scale = 1.0 * size / w | |
| return cv2.resize(img, (int(w * scale), int(h * scale))) | |
| def base_detect(detector, img): | |
| img = scale_image(img, 400) | |
| detections = detector.detect(img) | |
| img_vis = img.copy() | |
| for box in detections: | |
| img_vis = cv2.rectangle(img_vis, | |
| box[:2].astype(int).tolist(), | |
| box[2:4].astype(int).tolist(), | |
| (0, 255, 0), 1) | |
| return img_vis | |
| def retina_detect(img): | |
| detector = face_detection.build_detector( | |
| "RetinaNetResNet50", confidence_threshold=.5, nms_iou_threshold=.3) | |
| img_vis = base_detect(detector, img) | |
| return img_vis | |
| def retina_mobilenet_detect(img): | |
| detector = face_detection.build_detector( | |
| "RetinaNetMobileNetV1", confidence_threshold=.5, nms_iou_threshold=.3) | |
| img_vis = base_detect(detector, img) | |
| return img_vis | |
| def dsfd_detect(img): | |
| detector = face_detection.build_detector( | |
| "DSFDDetector", confidence_threshold=.5, nms_iou_threshold=.3) | |
| img_vis = base_detect(detector, img) | |
| return img_vis | |
| def kornia_detect(img): | |
| # select the device | |
| device = torch.device('cpu') | |
| vis_threshold = 0.6 | |
| # load the image and scale | |
| img_raw = scale_image(img, 400) | |
| # preprocess | |
| img = K.image_to_tensor(img_raw, keepdim=False).to(device) | |
| img = K.color.bgr_to_rgb(img.float()) | |
| # create the detector and find the faces ! | |
| face_detection = FaceDetector().to(device) | |
| with torch.no_grad(): | |
| dets = face_detection(img) | |
| dets = [FaceDetectorResult(o) for o in dets] | |
| # show image | |
| img_vis = img_raw.copy() | |
| for b in dets: | |
| if b.score < vis_threshold: | |
| continue | |
| # draw face bounding box | |
| img_vis = cv2.rectangle(img_vis, | |
| b.top_left.int().tolist(), | |
| b.bottom_right.int().tolist(), | |
| (0, 255, 0), | |
| 1) | |
| return img_vis | |
| input_image = gr.components.Image() | |
| image_kornia = gr.components.Image(label="Kornia YuNet") | |
| image_retina = gr.components.Image(label="RetinaFace") | |
| image_retina_mobile = gr.components.Image(label="Retina Mobilenet") | |
| image_dsfd = gr.components.Image(label="DSFD") | |
| confidence_slider = gr.components.Slider(minimum=0.1, maximum=0.9, value=0.5, label="Confidence Threshold") | |
| nms_slider = gr.components.Slider(minimum=0.1, maximum=0.9, value=0.5, label="Min Number of Neighbours") | |
| # scale_slider = gr.components.Slider(minimum=1.1, maximum=2.0, value=1.3, step=0.1, label="Scale Factor") | |
| # classifier_radio = gr.components.Radio(s) | |
| methods_dropdown = gr.components.Dropdown(["Kornia YuNet", "RetinaFace", "RetinaMobile", "DSFD"], value="Kornia YuNet", label="Choose a method") | |
| description = """Face Detection""" | |
| Iface = gr.Interface( | |
| fn=detect_faces, | |
| inputs=[input_image, methods_dropdown],#, size_slider, neighbour_slider, scale_slider], | |
| outputs=[image_kornia, image_retina, image_retina_mobile, image_dsfd], | |
| examples=[["data/9_Press_Conference_Press_Conference_9_86.jpg"], ["data/12_Group_Group_12_Group_Group_12_39.jpg"], ["data/31_Waiter_Waitress_Waiter_Waitress_31_55.jpg"]], | |
| title="Face Detection", | |
| ).launch() |