akhaliq HF staff commited on
Commit
1f65035
·
1 Parent(s): c9a7da4

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +69 -0
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
+