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))}')
判断图片是否是人像全身图
最新推荐文章于 2024-11-19 10:55:40 发布