221 lines
7.5 KiB
Python
221 lines
7.5 KiB
Python
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
import onnxruntime as ort
|
|||
|
|
|
|||
|
|
|
|||
|
|
class QualityOfClarity:
|
|||
|
|
def __init__(self, low_thresh=0.10, high_thresh=0.20):
|
|||
|
|
self.low_thresh = low_thresh
|
|||
|
|
self.high_thresh = high_thresh
|
|||
|
|
|
|||
|
|
def reblur(self, data, width, height):
|
|||
|
|
data = np.array(data, dtype=np.float32).reshape((height, width))
|
|||
|
|
|
|||
|
|
# 创建一维核
|
|||
|
|
kernel = np.ones((9,), np.float32) / 9.0
|
|||
|
|
|
|||
|
|
# 垂直方向模糊处理
|
|||
|
|
BVer = cv2.filter2D(data, -1, kernel.reshape(-1, 1), borderType=cv2.BORDER_REPLICATE)
|
|||
|
|
|
|||
|
|
# 水平方向模糊处理
|
|||
|
|
BHor = cv2.filter2D(data, -1, kernel.reshape(1, -1), borderType=cv2.BORDER_REPLICATE)
|
|||
|
|
|
|||
|
|
s_FVer, s_FHor, s_Vver, s_VHor = 0.0, 0.0, 0.0, 0.0
|
|||
|
|
|
|||
|
|
# 计算垂直方向的差分
|
|||
|
|
D_Fver = np.abs(data[1:, :] - data[:-1, :])
|
|||
|
|
D_BVer = np.abs(BVer[1:, :] - BVer[:-1, :])
|
|||
|
|
|
|||
|
|
# 计算垂直方向的累积
|
|||
|
|
s_FVer = np.sum(D_Fver)
|
|||
|
|
s_Vver = np.sum(np.maximum(0.0, D_Fver - D_BVer))
|
|||
|
|
|
|||
|
|
# 计算水平方向的差分
|
|||
|
|
D_FHor = np.abs(data[:, 1:] - data[:, :-1])
|
|||
|
|
D_BHor = np.abs(BHor[:, 1:] - BHor[:, :-1])
|
|||
|
|
|
|||
|
|
# 计算水平方向的累积
|
|||
|
|
s_FHor = np.sum(D_FHor)
|
|||
|
|
s_VHor = np.sum(np.maximum(0.0, D_FHor - D_BHor))
|
|||
|
|
|
|||
|
|
|
|||
|
|
b_FVer = (s_FVer - s_Vver) / s_FVer
|
|||
|
|
b_FHor = (s_FHor - s_VHor) / s_FHor
|
|||
|
|
blur_val = max(b_FVer, b_FHor)
|
|||
|
|
|
|||
|
|
return blur_val
|
|||
|
|
|
|||
|
|
def grid_max_reblur(self, img, rows, cols):
|
|||
|
|
height, width = img.shape
|
|||
|
|
row_height = height // rows
|
|||
|
|
col_width = width // cols
|
|||
|
|
blur_val = float('-inf')
|
|||
|
|
|
|||
|
|
for y in range(rows):
|
|||
|
|
for x in range(cols):
|
|||
|
|
grid = img[y * row_height: (y + 1) * row_height, x * col_width: (x + 1) * col_width]
|
|||
|
|
this_grad_blur_val = self.reblur(grid, col_width, row_height)
|
|||
|
|
if this_grad_blur_val > blur_val:
|
|||
|
|
blur_val = this_grad_blur_val
|
|||
|
|
|
|||
|
|
return max(blur_val, 0.0)
|
|||
|
|
|
|||
|
|
def clarity_estimate(self, image):
|
|||
|
|
# x, y, w, h = rect
|
|||
|
|
# if w < 9 or h < 9:
|
|||
|
|
# return 0.0
|
|||
|
|
|
|||
|
|
src_data = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|||
|
|
# src_data = gray_data[y:y+h, x:x+w]
|
|||
|
|
blur_val = self.grid_max_reblur(src_data, 2, 2)
|
|||
|
|
clarity = 1.0 - blur_val
|
|||
|
|
|
|||
|
|
T1, T2 = 0.0, 1.0
|
|||
|
|
if clarity <= T1:
|
|||
|
|
clarity = 0.0
|
|||
|
|
elif clarity >= T2:
|
|||
|
|
clarity = 1.0
|
|||
|
|
else:
|
|||
|
|
clarity = (clarity - T1) / (T2 - T1)
|
|||
|
|
|
|||
|
|
return clarity
|
|||
|
|
|
|||
|
|
def check(self, image):
|
|||
|
|
clarity = self.clarity_estimate(image)
|
|||
|
|
if clarity < self.low_thresh:
|
|||
|
|
level = "LOW"
|
|||
|
|
elif self.low_thresh <= clarity < self.high_thresh:
|
|||
|
|
level = "MEDIUM"
|
|||
|
|
else:
|
|||
|
|
level = "HIGH"
|
|||
|
|
return level != 'LOW'
|
|||
|
|
# return {'level': level, "score": clarity}
|
|||
|
|
|
|||
|
|
|
|||
|
|
class QualityChecker: # check resolution and clarity of image
|
|||
|
|
def __init__(self, v0=70.0, v1=100.0, v2=210.0, v3=230.0, hw = (112,112)):
|
|||
|
|
self.bright_thresh0 = v0
|
|||
|
|
self.bright_thresh1 = v1
|
|||
|
|
self.bright_thresh2 = v2
|
|||
|
|
self.bright_thresh3 = v3
|
|||
|
|
self.middle_thresh = (self.bright_thresh1 + self.bright_thresh2) / 2
|
|||
|
|
self.rolu_thrds = hw # （h，w）
|
|||
|
|
|
|||
|
|
def get_bright_score(self, bright):
|
|||
|
|
bright_score = 1.0 / (abs(bright - self.middle_thresh) + 1)
|
|||
|
|
return bright_score
|
|||
|
|
|
|||
|
|
def grid_max_bright(self, img, rows, cols):
|
|||
|
|
row_height = img.shape[0] // rows
|
|||
|
|
col_width = img.shape[1] // cols
|
|||
|
|
|
|||
|
|
# 使用列表生成式获取所有网格的平均亮度
|
|||
|
|
grid_means = [
|
|||
|
|
np.mean(img[y * row_height:(y + 1) * row_height, x * col_width:(x + 1) * col_width])
|
|||
|
|
for y in range(rows)
|
|||
|
|
for x in range(cols)
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# 获取最大亮度值
|
|||
|
|
bright_val = max(grid_means)
|
|||
|
|
return max(bright_val, 0)
|
|||
|
|
|
|||
|
|
def check_bright(self, face_image):
|
|||
|
|
gray = cv2.cvtColor(face_image, cv2.COLOR_BGR2GRAY) if len(face_image.shape) == 3 else face_image
|
|||
|
|
bright_value = self.grid_max_bright(gray, 3, 3)
|
|||
|
|
|
|||
|
|
if bright_value < self.bright_thresh0 or bright_value > self.bright_thresh3:
|
|||
|
|
level = "LOW"
|
|||
|
|
elif (self.bright_thresh0 <= bright_value < self.bright_thresh1) or (self.bright_thresh2 < bright_value <= self.bright_thresh3):
|
|||
|
|
level = "MEDIUM"
|
|||
|
|
else:
|
|||
|
|
level = "HIGH"
|
|||
|
|
|
|||
|
|
return level == "HIGH"
|
|||
|
|
|
|||
|
|
def check_resolution(self, face_image):
|
|||
|
|
H, W = face_image.shape[:2]
|
|||
|
|
if H < self.rolu_thrds[0] or W < self.rolu_thrds[1]:
|
|||
|
|
return False
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
|
|||
|
|
class QualityOfPose:
|
|||
|
|
def __init__(self, pad=0.3, yaw_thrd=30, pitch_thrd=25, var_onnx_path = './checkpoints/fsanet-var.onnx', conv_onnx_path='./checkpoints/fsanet-conv.onnx') -> None:
|
|||
|
|
self.pad = pad
|
|||
|
|
self.input_width = 64
|
|||
|
|
self.input_height = 64
|
|||
|
|
self.yaw_thrd = abs(yaw_thrd)
|
|||
|
|
self.pitch_thrd = abs(pitch_thrd)
|
|||
|
|
self.var_fsanet = ort.InferenceSession(var_onnx_path)
|
|||
|
|
self.conv_fsanet = ort.InferenceSession(conv_onnx_path)
|
|||
|
|
|
|||
|
|
self.var_input_names = [input_.name for input_ in self.var_fsanet.get_inputs()]
|
|||
|
|
self.var_output_names = [output.name for output in self.var_fsanet.get_outputs()]
|
|||
|
|
|
|||
|
|
self.conv_input_names = [input_.name for input_ in self.conv_fsanet.get_inputs()]
|
|||
|
|
self.conv_output_names = [output.name for output in self.conv_fsanet.get_outputs()]
|
|||
|
|
|
|||
|
|
def transform(self, mat):
|
|||
|
|
h, w = mat.shape[:2]
|
|||
|
|
nh = int(h + self.pad * h)
|
|||
|
|
nw = int(w + self.pad * w)
|
|||
|
|
nx1 = max(0, (nw - w) // 2)
|
|||
|
|
ny1 = max(0, (nh - h) // 2)
|
|||
|
|
|
|||
|
|
# Create a padded canvas and copy the image into the center
|
|||
|
|
canvas = np.zeros((nh, nw, 3), dtype=np.uint8)
|
|||
|
|
canvas[ny1:ny1 + h, nx1:nx1 + w] = mat
|
|||
|
|
|
|||
|
|
# Resize the image to the input dimensions
|
|||
|
|
canvas = cv2.resize(canvas, (self.input_width, self.input_height))
|
|||
|
|
|
|||
|
|
# Normalize the image in-place
|
|||
|
|
canvas = canvas.astype(np.float32)
|
|||
|
|
mean = 127.5
|
|||
|
|
scale = 1.0 / 127.5
|
|||
|
|
canvas = (canvas - mean) * scale
|
|||
|
|
|
|||
|
|
# Convert to CHW format
|
|||
|
|
canvas = np.transpose(canvas, (2, 0, 1))
|
|||
|
|
|
|||
|
|
# Create a tensor
|
|||
|
|
input_tensor = np.expand_dims(canvas, axis=0).astype(np.float32)
|
|||
|
|
return input_tensor
|
|||
|
|
|
|||
|
|
def detect_angle(self, img):
|
|||
|
|
input_tensor = self.transform(img)
|
|||
|
|
|
|||
|
|
var_output = self.var_fsanet.run(
|
|||
|
|
self.var_output_names, {self.var_input_names[0]: input_tensor}
|
|||
|
|
)[0]
|
|||
|
|
conv_output = self.conv_fsanet.run(
|
|||
|
|
self.conv_output_names, {self.conv_input_names[0]: input_tensor}
|
|||
|
|
)[0]
|
|||
|
|
yaw, pitch, roll = np.mean(np.vstack((var_output,conv_output)), axis=0)
|
|||
|
|
|
|||
|
|
return yaw, pitch, roll
|
|||
|
|
|
|||
|
|
def check(self, image):
|
|||
|
|
yaw, pitch, roll = self.detect_angle(image)
|
|||
|
|
|
|||
|
|
if abs(yaw) <= self.yaw_thrd and abs(pitch) <= self.pitch_thrd:
|
|||
|
|
return "frontFace"
|
|||
|
|
elif yaw < -1.0 * self.yaw_thrd:
|
|||
|
|
return "rightFace"
|
|||
|
|
elif yaw > self.yaw_thrd:
|
|||
|
|
return "leftFace"
|
|||
|
|
elif pitch > self.pitch_thrd:
|
|||
|
|
return "upFace"
|
|||
|
|
elif pitch < -1.0 * self.pitch_thrd:
|
|||
|
|
return "downFace"
|
|||
|
|
else:
|
|||
|
|
return "implementError"
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
qp = QualityOfPose()
|
|||
|
|
img = cv2.imread("/home/bns/liteAitoolkit/Demos_onnx/examples/det_oeyecmouth.jpg")
|
|||
|
|
angles = qp.check(img)
|
|||
|
|
print(f"ONNXRuntime Version yaw: {angles[0]} pitch: {angles[1]} roll: {angles[2]}")
|
|||
|
|
pass
|