import os 
import paddlehub as hub
import cv2

def draw_img(x,y,w,h,img,save_path):
    # # 判断人脸矩形宽高比是否为正脸
    if w / h < 1.2 and w / h > 0.8:  
        cv2.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), 2)  
        label = 'Frontal Face'
    else:  
        cv2.rectangle(img, (x, y), (x+w, y+h), (0, 0, 255), 2)  
        label = 'Not Frontal'
    # 在图片上打标签    
    cv2.putText(img, label, (x, y - 5),cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
    # cv2.imwrite(save_path, img)

# 文件夹路径
dir_path = '/home/dusr/data/xznsh/img/'  

# 文件保存路径
save_dir = '/home/dusr/data/xznsh/img_result'

# 获取folder_path下的所有文件
files = os.listdir(dir_path)  
# 遍历所有文件
for file in files:
    # 获取文件路径
    file_path = os.path.join(dir_path, file) 
    save_path = os.path.join(save_dir, "result_"+file)
    img = cv2.imread(file_path)
    # 判断是否为图片文件
    if file_path.endswith('.jpg') or file_path.endswith('.png'):

        # face_detector = hub.Module(name="pyramidbox_lite_server")
        face_detector = hub.Module(name="pyramidbox_face_detection")
        file_path = r'/home/dusr/data/xznsh/img/frame_1640.jpg'
        # result = face_detector.face_detection(images=[cv2.imread(file_path)])
        # result = face_detector.face_detection(paths=[file_path])
        result = face_detector.face_detection(paths=[file_path])
        print(file_path)
        print(result)
        # 获取图片路径 path的值
        path = result[0]['path']  
        # 获取data的值
        data = result[0]['data']
        if data != []:
            # 获取人脸置信度confidence的值
            confidence = data[0]['confidence']

            # 获取人脸左上角x坐标 left的值
            lx = data[0]['left']  

            # 获取人脸左上角y坐标 top的值
            ly = data[0]['top']

            # 获取人脸右下角x坐标 right的值 
            rx = data[0]['right']  

            # 获取人脸右下角y坐标 bottom的值
            ry = data[0]['bottom']
            w = rx -lx
            h = ry-ly
            x = lx
            y = ly

            draw_img(x,y,w,h,img,save_path)
        else:
            cv2.imwrite(save_path, img)
            print("No facial information was detected")