|
|
import cv2 |
|
|
import os |
|
|
import argparse |
|
|
import numpy as np |
|
|
from ultralytics import YOLO |
|
|
|
|
|
# ===================== 模型路径配置 ===================== |
|
|
# 请确保此路径下有模型文件，如果没有会自动下载 |
|
|
YOLO_FACE_MODEL_PATH = "/root/code/lenovo_qinzhe_ssh/sam3/sam3_new/face/yolov8n-face.pt" |
|
|
# ======================================================= |
|
|
|
|
|
# 加载YOLO人脸检测模型 |
|
|
yolo_model = YOLO(YOLO_FACE_MODEL_PATH) |
|
|
|
|
|
# 支持的图片格式（包含大小写） |
|
|
SUPPORTED_FORMATS = { |
|
|
'jpg': 'jpg', |
|
|
'jpeg': 'jpeg', |
|
|
'png': 'png', |
|
|
'bmp': 'bmp', |
|
|
'JPG': 'jpg', |
|
|
'JPEG': 'jpeg', |
|
|
'PNG': 'png', |
|
|
'BMP': 'bmp' |
|
|
} |
|
|
|
|
|
def detect_face_bbox(image, conf_threshold=0.5): |
|
|
""" |
|
|
使用YOLO检测人脸边界框 |
|
|
:param image: 输入BGR图像 |
|
|
:param conf_threshold: 置信度阈值 |
|
|
:return: 人脸边界框列表 [(x1, y1, x2, y2), ...]，若无检测结果返回空列表 |
|
|
""" |
|
|
results = yolo_model(image, conf=conf_threshold) |
|
|
bboxes = [] |
|
|
if results and len(results) > 0: |
|
|
for box in results[0].boxes: |
|
|
if box.conf[0] >= conf_threshold: |
|
|
x1, y1, x2, y2 = map(int, box.xyxy[0]) |
|
|
bboxes.append((x1, y1, x2, y2)) |
|
|
return bboxes |
|
|
|
|
|
def is_supported_image(file_path): |
|
|
""" |
|
|
检查文件是否为支持的图片格式 |
|
|
""" |
|
|
if not os.path.isfile(file_path): |
|
|
return False |
|
|
file_ext = os.path.splitext(file_path)[1].lstrip('.') |
|
|
return file_ext in SUPPORTED_FORMATS |
|
|
|
|
|
def get_output_path(input_img_path, output_arg): |
|
|
""" |
|
|
根据输出参数确定最终的输出路径 |
|
|
:param input_img_path: 输入图片路径 |
|
|
:param output_arg: 命令行传入的-o参数（目录或文件路径） |
|
|
:return: 最终的输出图片路径 |
|
|
""" |
|
|
# 如果输出参数是目录 |
|
|
if os.path.isdir(output_arg): |
|
|
img_name = os.path.basename(input_img_path) |
|
|
return os.path.join(output_arg, img_name) |
|
|
|
|
|
# 如果输出参数是文件路径（且目录存在） |
|
|
output_dir = os.path.dirname(output_arg) |
|
|
if output_dir and not os.path.exists(output_dir): |
|
|
os.makedirs(output_dir, exist_ok=True) |
|
|
|
|
|
# 检查输出文件格式是否支持 |
|
|
output_ext = os.path.splitext(output_arg)[1].lstrip('.') |
|
|
if output_ext not in SUPPORTED_FORMATS: |
|
|
# 使用输入文件的格式 |
|
|
input_ext = os.path.splitext(input_img_path)[1].lstrip('.') |
|
|
input_format = SUPPORTED_FORMATS.get(input_ext, 'jpg') |
|
|
output_arg = f"{os.path.splitext(output_arg)[0]}.{input_format}" |
|
|
|
|
|
return output_arg |
|
|
|
|
|
def crop_face_center_1024(input_path, output_arg, conf_threshold=0.5): |
|
|
""" |
|
|
以检测到的人脸为中心裁剪出1024×1024的图片 |
|
|
:param input_path: 输入图片文件路径 |
|
|
:param output_arg: 输出图片路径或目录路径 |
|
|
:param conf_threshold: YOLO人脸检测置信度阈值 |
|
|
:return: 裁剪是否成功 |
|
|
""" |
|
|
# 读取图片（确保以彩色模式读取） |
|
|
image = cv2.imread(input_path, cv2.IMREAD_COLOR) |
|
|
if image is None: |
|
|
print(f"❌ 无法读取图片：{input_path}") |
|
|
return False |
|
|
|
|
|
h, w = image.shape[:2] |
|
|
|
|
|
# ========== YOLO人脸检测 ========== |
|
|
face_bboxes = detect_face_bbox(image, conf_threshold) |
|
|
if not face_bboxes: |
|
|
print(f"⚠️ 未检测到人脸：{input_path}") |
|
|
return False |
|
|
|
|
|
# 取第一个检测到的人脸 |
|
|
x1, y1, x2, y2 = face_bboxes[0] |
|
|
|
|
|
# 计算人脸中心坐标 |
|
|
face_center_x = (x1 + x2) // 2 |
|
|
face_center_y = (y1 + y2) // 2 |
|
|
|
|
|
# 计算1024×1024裁剪区域的边界（以人脸为中心） |
|
|
crop_size = 1024 |
|
|
half_size = crop_size // 2 |
|
|
|
|
|
# 计算裁剪区域的左上角和右下角坐标 |
|
|
crop_x1 = face_center_x - half_size |
|
|
crop_y1 = face_center_y - half_size |
|
|
crop_x2 = crop_x1 + crop_size |
|
|
crop_y2 = crop_y1 + crop_size |
|
|
|
|
|
# 处理边界越界问题（如果裁剪区域超出图片范围） |
|
|
pad_left = max(0, -crop_x1) |
|
|
pad_top = max(0, -crop_y1) |
|
|
pad_right = max(0, crop_x2 - w) |
|
|
pad_bottom = max(0, crop_y2 - h) |
|
|
|
|
|
# 调整裁剪坐标到图片范围内 |
|
|
crop_x1_clamped = max(0, crop_x1) |
|
|
crop_y1_clamped = max(0, crop_y1) |
|
|
crop_x2_clamped = min(w, crop_x2) |
|
|
crop_y2_clamped = min(h, crop_y2) |
|
|
|
|
|
# 裁剪图片 |
|
|
cropped_image = image[crop_y1_clamped:crop_y2_clamped, crop_x1_clamped:crop_x2_clamped] |
|
|
|
|
|
# 如果有越界，用黑色填充（保持图片维度） |
|
|
if pad_left > 0 or pad_top > 0 or pad_right > 0 or pad_bottom > 0: |
|
|
cropped_image = cv2.copyMakeBorder( |
|
|
cropped_image, |
|
|
pad_top, pad_bottom, pad_left, pad_right, |
|
|
cv2.BORDER_CONSTANT, |
|
|
value=[0, 0, 0] # 黑色填充（BGR格式） |
|
|
) |
|
|
|
|
|
# 强制调整为1024×1024（确保输出尺寸准确） |
|
|
cropped_image = cv2.resize(cropped_image, (1024, 1024), interpolation=cv2.INTER_AREA) |
|
|
|
|
|
# 确定最终输出路径 |
|
|
output_path = get_output_path(input_path, output_arg) |
|
|
|
|
|
# 保存图片（根据格式选择合适的编码） |
|
|
file_ext = os.path.splitext(output_path)[1].lstrip('.') |
|
|
file_format = SUPPORTED_FORMATS.get(file_ext, 'jpg') |
|
|
|
|
|
save_success = False |
|
|
if file_format == 'jpg' or file_format == 'jpeg': |
|
|
# JPG格式：95%质量保存 |
|
|
save_success = cv2.imwrite(output_path, cropped_image, [int(cv2.IMWRITE_JPEG_QUALITY), 95]) |
|
|
elif file_format == 'png': |
|
|
# PNG格式：无损压缩 |
|
|
save_success = cv2.imwrite(output_path, cropped_image, [int(cv2.IMWRITE_PNG_COMPRESSION), 0]) |
|
|
else: |
|
|
# 其他格式默认保存 |
|
|
save_success = cv2.imwrite(output_path, cropped_image) |
|
|
|
|
|
if save_success: |
|
|
print(f"✅ 裁剪完成：{output_path} (尺寸：1024×1024，人脸中心：({face_center_x}, {face_center_y}))") |
|
|
return True |
|
|
else: |
|
|
print(f"❌ 保存图片失败：{output_path}") |
|
|
return False |
|
|
|
|
|
def main(): |
|
|
# 创建参数解析器 |
|
|
parser = argparse.ArgumentParser(description='以人脸为中心裁剪1024×1024图片的工具') |
|
|
|
|
|
# 添加命令行参数 |
|
|
parser.add_argument('-i', '--input', required=True, |
|
|
help='输入图片文件路径（仅支持单张图片，格式：jpg/jpeg/png/bmp）') |
|
|
parser.add_argument('-o', '--output', required=True, |
|
|
help='输出图片路径（如：/path/to/output.jpg）或输出目录路径（如：/path/to/output/）') |
|
|
parser.add_argument('-c', '--conf', type=float, default=0.5, |
|
|
help='人脸检测置信度阈值（默认0.5，值越高检测越严格）') |
|
|
|
|
|
# 解析参数 |
|
|
args = parser.parse_args() |
|
|
|
|
|
# 验证输入文件 |
|
|
if not os.path.exists(args.input): |
|
|
print(f"❌ 错误：输入文件不存在：{args.input}") |
|
|
return |
|
|
|
|
|
if not is_supported_image(args.input): |
|
|
print(f"❌ 错误：输入文件不是支持的图片格式（支持：{', '.join(set(SUPPORTED_FORMATS.values()))}）") |
|
|
return |
|
|
|
|
|
# 执行裁剪 |
|
|
success = crop_face_center_1024(args.input, args.output, args.conf) |
|
|
|
|
|
if success: |
|
|
print("\n🎉 图片裁剪任务完成！") |
|
|
else: |
|
|
print("\n❌ 图片裁剪任务失败！") |
|
|
|
|
|
if __name__ == "__main__": |
|
|
main()
|
|
|
|
