File size: 2,684 Bytes
1f65035 |
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 |
# SPDX-License-Identifier: MIT
import cv2
import onnxruntime as ort
import argparse
import numpy as np
from dependencies.box_utils import predict
# ------------------------------------------------------------------------------------------------------------------------------------------------
# Face detection using UltraFace-320 onnx model
os.system("wget https://github.com/AK391/models/raw/main/vision/body_analysis/ultraface/models/version-RFB-320.onnx")
face_detector_onnx = "version-RFB-320.onnx"
# Start from ORT 1.10, ORT requires explicitly setting the providers parameter if you want to use execution providers
# other than the default CPU provider (as opposed to the previous behavior of providers getting set/registered by default
# based on the build flags) when instantiating InferenceSession.
# For example, if NVIDIA GPU is available and ORT Python package is built with CUDA, then call API as following:
# ort.InferenceSession(path/to/model, providers=['CUDAExecutionProvider'])
face_detector = ort.InferenceSession(face_detector_onnx)
# scale current rectangle to box
def scale(box):
width = box[2] - box[0]
height = box[3] - box[1]
maximum = max(width, height)
dx = int((maximum - width)/2)
dy = int((maximum - height)/2)
bboxes = [box[0] - dx, box[1] - dy, box[2] + dx, box[3] + dy]
return bboxes
# crop image
def cropImage(image, box):
num = image[box[1]:box[3], box[0]:box[2]]
return num
# face detection method
def faceDetector(orig_image, threshold = 0.7):
image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (320, 240))
image_mean = np.array([127, 127, 127])
image = (image - image_mean) / 128
image = np.transpose(image, [2, 0, 1])
image = np.expand_dims(image, axis=0)
image = image.astype(np.float32)
input_name = face_detector.get_inputs()[0].name
confidences, boxes = face_detector.run(None, {input_name: image})
boxes, labels, probs = predict(orig_image.shape[1], orig_image.shape[0], confidences, boxes, threshold)
return boxes, labels, probs
# ------------------------------------------------------------------------------------------------------------------------------------------------
# Main void
def inference(img):
color = (255, 128, 0)
orig_image = cv2.imread(img)
boxes, labels, probs = faceDetector(orig_image)
for i in range(boxes.shape[0]):
box = scale(boxes[i, :])
cv2.rectangle(orig_image, (box[0], box[1]), (box[2], box[3]), color, 4)
cv2.imwrite("out.png",orig_image)
return "out.png"
gr.Interface(inference,gr.inputs.Image(type="filepath"),gr.outputs.Image(type="file")).launch()
|