Spaces:

onnx
/

ultraface

Runtime error

App Files Files Community

ultraface / app.py

akhaliq HF staff

Update app.py

a9f616c about 2 years ago

raw history blame contribute delete

No virus

2.87 kB

	# SPDX-License-Identifier: MIT


	import cv2
	import onnxruntime as ort
	import argparse
	import numpy as np
	from box_utils import predict
	import os
	import gradio as gr

	# ------------------------------------------------------------------------------------------------------------------------------------------------
	# 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"


	title="UltraFace"
	description="This model is a lightweight face detection model designed for edge computing devices."

	gr.Interface(inference,gr.inputs.Image(type="filepath",source="webcam"),gr.outputs.Image(type="file"),title=title,description=description).launch()