判断图片是否是人像全身图

本文介绍了一个Python脚本,利用Mediapipe库进行人体姿态估计和关键点检测,对给定文件夹中的图像逐个检查并保存具有有效关键点的图片。脚本还跟踪处理进度并记录相关信息。
摘要由CSDN通过智能技术生成
import cv2
import os
import shutil
import time
import argparse
import numpy as np
import mediapipe as mp

parser = argparse.ArgumentParser()
parser.add_argument('--image_folder', type=str, required=True)
args = parser.parse_args()
image_ori = args.image_folder


def get_pure_pose(img_bgr, results):  # 坐标转换
    """
    mediapipe转换取坐标点
    """
    pose_landmarks = results.pose_landmarks
    if pose_landmarks is None:
        return None
    frame_height, frame_width = img_bgr.shape[0], img_bgr.shape[1]
    pose_landmarks = np.array(
        [[lmk.x * frame_width, lmk.y * frame_height]
         for lmk in pose_landmarks.landmark], dtype=int)
    del results
    return pose_landmarks


def get_keypoints(img):
    mp_pose = mp.solutions.pose
    # 进行关键点检测
    with mp_pose.Pose(
            static_image_mode=True,
            model_complexity=1,
            enable_segmentation=False,
            min_detection_confidence=0.8) as pose:#这个可以调整图像关键点的置信度
        results = pose.process(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        body_pose = get_pure_pose(img, results)
        pose.reset()

    # 判断是否存在点
    if body_pose is not None:
        # 遍历所有点
        for i in range(0, len(body_pose)):
            point = body_pose[i]
            if point[0] > img.shape[1] or point[1] > img.shape[0] or point[0] < 0 or point[1] < 0:
                return False
        return True
    else:
        return False


def save_progress(progress_path, image_cloth, item, image_name_index, good_num, all_num):
    with open(progress_path, 'w') as f:
        f.write(f"{image_cloth}\n{item}\n{image_name_index}\n{good_num}\n{all_num}")


def load_progress(progress_path):
    try:
        with open(progress_path, 'r') as f:
            lines = f.readlines()
            image_cloth = lines[0].strip()
            item = lines[1].strip()
            image_name_index = int(lines[2].strip())
            good_num = int(lines[3].strip())
            all_num = int(lines[4])
            return image_cloth, item, image_name_index, good_num, all_num
    except FileNotFoundError:
        return None


def main(image_ori):
    zip_name = image_ori.split('/')[-1]
    save_path = image_ori.split(f'{zip_name}')[0]
    folder_name_num = image_ori.split('/')[-1].split('ori')[0]

    result_folder = f"{save_path}/whole_body_result_{folder_name_num}"
    progress_folder = f"{save_path}/process_log"
    os.makedirs(result_folder, exist_ok=True)
    os.makedirs(progress_folder, exist_ok=True)

    progress_file = 'progress.txt'
    progress_path = progress_folder + f'/{folder_name_num}-{progress_file}'

    progress = load_progress(progress_path)
    if progress is None:
        print('暂无上次处理记录')
        image_cloth_start, item_start, image_name_index_start, good_num, all_num = '', '', 0, 0, 0
    else:
        print('从上次处理记录开始接着处理')
        image_cloth_start, item_start, image_name_index_start, good_num, all_num = progress

    for image_cloth_index, image_cloth in enumerate(os.listdir(image_ori)):
        if image_cloth < image_cloth_start:
            continue

        class_folder = os.path.join(image_ori, image_cloth)
        for item_index, item in enumerate(os.listdir(class_folder)):
            if image_cloth == image_cloth_start and item < item_start:
                continue

            item_folder = os.path.join(class_folder, item)
            image_names = os.listdir(item_folder)
            try:
                for image_name_index, image_name in enumerate(image_names):
                    if image_cloth == image_cloth_start and item == item_start and image_name_index < image_name_index_start:
                        continue

                    image_path = os.path.join(item_folder, image_name)
                    output_path = os.path.join(result_folder, image_name)

                    try:
                        img = cv2.imread(image_path)
                        flag = get_keypoints(img)
                        all_num += len(image_names)
                        if flag:
                            good_num += 1
                            shutil.copy2(image_path, output_path)
                        print(f'[{good_num}/{all_num}] ==> {image_cloth} ==> {item} ==> 款式:[{item_index + 1}/{len(os.listdir(class_folder))}] ==> 图片:[{image_name_index + 1}/{len(image_names)}] ==> {flag}')
                    except Exception as e:
                        print(e)
                        print(f"图片处理失败: {image_path}")

                    save_progress(progress_path, image_cloth, item, image_name_index, good_num, all_num)
            except Exception as e:
                print(e)
                print(f"数据层级结构不一致: {image_names}")

            print(f"finish 款式: {item}\n")

        print(f"finish 种类 {image_cloth}\n\n")

    print(f"finish 压缩包 {image_ori}")


if __name__ == '__main__':
    start = time.time()
    main(image_ori)
    end = time.time()
    print(f'Consume time: {time.strftime("%H:%M:%S", time.gmtime(end - start))}')

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值