Spaces:
Runtime error
Runtime error
Create app.py
Browse files
app.py
ADDED
|
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: MIT
|
| 2 |
+
|
| 3 |
+
import cv2
|
| 4 |
+
import onnxruntime as ort
|
| 5 |
+
import argparse
|
| 6 |
+
import numpy as np
|
| 7 |
+
from dependencies.box_utils import predict
|
| 8 |
+
|
| 9 |
+
# ------------------------------------------------------------------------------------------------------------------------------------------------
|
| 10 |
+
# Face detection using UltraFace-320 onnx model
|
| 11 |
+
os.system("wget https://github.com/AK391/models/raw/main/vision/body_analysis/ultraface/models/version-RFB-320.onnx")
|
| 12 |
+
face_detector_onnx = "version-RFB-320.onnx"
|
| 13 |
+
|
| 14 |
+
# Start from ORT 1.10, ORT requires explicitly setting the providers parameter if you want to use execution providers
|
| 15 |
+
# other than the default CPU provider (as opposed to the previous behavior of providers getting set/registered by default
|
| 16 |
+
# based on the build flags) when instantiating InferenceSession.
|
| 17 |
+
# For example, if NVIDIA GPU is available and ORT Python package is built with CUDA, then call API as following:
|
| 18 |
+
# ort.InferenceSession(path/to/model, providers=['CUDAExecutionProvider'])
|
| 19 |
+
face_detector = ort.InferenceSession(face_detector_onnx)
|
| 20 |
+
|
| 21 |
+
# scale current rectangle to box
|
| 22 |
+
def scale(box):
|
| 23 |
+
width = box[2] - box[0]
|
| 24 |
+
height = box[3] - box[1]
|
| 25 |
+
maximum = max(width, height)
|
| 26 |
+
dx = int((maximum - width)/2)
|
| 27 |
+
dy = int((maximum - height)/2)
|
| 28 |
+
|
| 29 |
+
bboxes = [box[0] - dx, box[1] - dy, box[2] + dx, box[3] + dy]
|
| 30 |
+
return bboxes
|
| 31 |
+
|
| 32 |
+
# crop image
|
| 33 |
+
def cropImage(image, box):
|
| 34 |
+
num = image[box[1]:box[3], box[0]:box[2]]
|
| 35 |
+
return num
|
| 36 |
+
|
| 37 |
+
# face detection method
|
| 38 |
+
def faceDetector(orig_image, threshold = 0.7):
|
| 39 |
+
image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
|
| 40 |
+
image = cv2.resize(image, (320, 240))
|
| 41 |
+
image_mean = np.array([127, 127, 127])
|
| 42 |
+
image = (image - image_mean) / 128
|
| 43 |
+
image = np.transpose(image, [2, 0, 1])
|
| 44 |
+
image = np.expand_dims(image, axis=0)
|
| 45 |
+
image = image.astype(np.float32)
|
| 46 |
+
|
| 47 |
+
input_name = face_detector.get_inputs()[0].name
|
| 48 |
+
confidences, boxes = face_detector.run(None, {input_name: image})
|
| 49 |
+
boxes, labels, probs = predict(orig_image.shape[1], orig_image.shape[0], confidences, boxes, threshold)
|
| 50 |
+
return boxes, labels, probs
|
| 51 |
+
|
| 52 |
+
# ------------------------------------------------------------------------------------------------------------------------------------------------
|
| 53 |
+
# Main void
|
| 54 |
+
|
| 55 |
+
def inference(img):
|
| 56 |
+
color = (255, 128, 0)
|
| 57 |
+
|
| 58 |
+
orig_image = cv2.imread(img)
|
| 59 |
+
boxes, labels, probs = faceDetector(orig_image)
|
| 60 |
+
|
| 61 |
+
for i in range(boxes.shape[0]):
|
| 62 |
+
box = scale(boxes[i, :])
|
| 63 |
+
cv2.rectangle(orig_image, (box[0], box[1]), (box[2], box[3]), color, 4)
|
| 64 |
+
cv2.imwrite("out.png",orig_image)
|
| 65 |
+
return "out.png"
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
gr.Interface(inference,gr.inputs.Image(type="filepath"),gr.outputs.Image(type="file")).launch()
|
| 69 |
+
|