Add 1024 face crop tool

face_crop_1024.py

@@ -0,0 +1,203 @@
+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()