import cv2
from rockx import RockX
import numpy as np
# BODY_DETECT_MIN_SCORE = 0.4
FACE_DETECT_MIN_SCORE = 0.6
#存放了模型和数据
ROCKX_DATA_DIR = "/home/coolpi/RockX_SDK_V1.6.0_20220913/sdk/rockx-data-rk3588"
#实例化RockX类
facepath = 'datatest/face/'
handle_body = RockX(RockX.ROCKX_MODULE_PERSON_DETECTION, rockx_data_dir=ROCKX_DATA_DIR)
handle_face = RockX(RockX.ROCKX_MODULE_FACE_DETECTION, rockx_data_dir=ROCKX_DATA_DIR)
# process_map = {
# RockX.ROCKX_MODULE_PERSON_DETECTION : ("body", "rockx_person_detect", BODY_DETECT_MIN_SCORE),
# }
#调用RockX类的rockx_person_detect方法
process_face = getattr(handle_face, "rockx_face_detect")
process_body = getattr(handle_body, "rockx_person_detect")
if __name__=='main':
dataset = 'datatest/sit.mp4'
cap = cv2.VideoCapture(dataset)
flag, frame = cap.read()
i = 0
det_results = []
while flag:
# 读取每一帧，flag表示是否读取成功，frame为图片内容。
flag, frame = cap.read()
# frame = cv2.pyrDown(frame)
if frame is None:
# print("No video captured")
continue
else:
width = frame.shape[1]
height = frame.shape[0]
################
ret, results = process_body(frame, width, height,RockX.ROCKX_PIXEL_FORMAT_BGR888)
for result in results:
if result.score >FACE_DETECT_MIN_SCORE:
b = [result.box.top, result.box.right, result.box.bottom, result.box.left,result.score]
img = frame[b[3]:b[1],b[0]:b[2]]
w = img.shape[1]
h = img.shape[0]
ret, results = process_face(
img, w, h,RockX.ROCKX_PIXEL_FORMAT_BGR888)
# results = self._filter_results(results)
################
for result in results:
if result.score >FACE_DETECT_MIN_SCORE:
b = [result.box.top, result.box.right, result.box.bottom, result.box.left,result.score]
cv2.rectangle(img, (b[2],b[0]), (b[1],b[3]), (0, 0, 0), -1)
cv2.imwrite(facepath+str(i)+'.jpg', frame)
i=i+1
print(i)
##############################
# cv2.rectangle(img, (x1, y1), (x2, y2), (255,0,0), 2)
# x1,y1 ------
# | |
# | |
# | |
# --------x2,y2
