4. 纠正图像方向

任务说明

图像数据集中常存在因拍摄设备或传输过程导致的方向错乱问题。若直接用于模型训练,会引入错误的空间先验,降低模型收敛速度与最终精度。本案例实现一个自动纠正脚本:优先读取EXIF方向标签进行无损旋转;对无EXIF信息的图像,基于图像内容特征(如人脸、文字、天空位置)进行方向判别与纠正。

环境依赖

pip install Pillow numpy opencv-python-headless

核心实现

import os
from PIL import Image, ImageOps
import numpy as np
import cv2

# EXIF方向标签到旋转角度的映射
EXIF_ORIENTATION = {
    1: 0,    # 正常
    2: 0,    # 水平翻转(不旋转)
    3: 180,  # 旋转180度
    4: 180,  # 垂直翻转(等效180度)
    5: 90,   # 镜像+旋转90度(按90度处理)
    6: 90,   # 顺时针旋转90度
    7: 270,  # 镜像+旋转270度(按270度处理)
    8: 270   # 逆时针旋转90度(即顺时针270度)
}

def correct_by_exif(img):
    """基于EXIF方向信息纠正图像"""
    try:
        exif = img._getexif()
        if not exif:
            return img, False
        orientation = exif.get(274)  # 274 = Orientation标签
        if orientation in EXIF_ORIENTATION:
            angle = EXIF_ORIENTATION[orientation]
            if angle == 90:
                return img.transpose(Image.ROTATE_270), True
            elif angle == 180:
                return img.transpose(Image.ROTATE_180), True
            elif angle == 270:
                return img.transpose(Image.ROTATE_90), True
    except Exception:
        pass
    return img, False

def detect_rotation_by_content(img):
    """基于内容分析判断是否需要旋转(简化版:检测文字/人脸方向)"""
    # 转为OpenCV格式
    img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    h, w = img_cv.shape[:2]
    
    # 使用人脸检测作为方向线索(若有人脸,眼睛应在上方)
    face_cascade = cv2.CascadeClassifier(
        cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
    )
    gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.1, 5, minSize=(30, 30))
    
    if len(faces) > 0:
        # 取最大人脸
        largest = max(faces, key=lambda f: f[2] * f[3])
        x, y, fw, fh = largest
        # 人脸中心应在图像上半部分
        face_center_y = y + fh / 2
        if face_center_y > h * 0.6:  # 人脸偏下,说明图像倒置
            return img.rotate(180, expand=True), True
        elif face_center_y > h * 0.5:  # 人脸略偏下,旋转90度
            # 进一步判断:人脸宽度远大于高度则说明是横放
            if fw > fh * 1.5:
                return img.rotate(90, expand=True), True
        return img, False
    
    # 无检测结果时,使用图像边缘统计作为弱线索
    # 自然图像中,天空/亮区域通常在上方
    img_gray = np.array(img.convert('L'))
    top_half = img_gray[:h//2, :]
    bottom_half = img_gray[h//2:, :]
    if np.mean(bottom_half) > np.mean(top_half) * 1.15:  # 下方明显更亮
        return img.rotate(180, expand=True), True
    
    return img, False

def auto_correct_image(image_path, output_path=None):
    """自动纠正图像方向"""
    img = Image.open(image_path)
    original = img.copy()
    
    # 第一步:EXIF纠正
    img, corrected = correct_by_exif(img)
    
    # 第二步:若EXIF未纠正,尝试内容分析
    if not corrected:
        img, corrected = detect_rotation_by_content(img)
    
    # 保存结果
    if output_path is None:
        base, ext = os.path.splitext(image_path)
        output_path = f"{base}_corrected{ext}"
    
    # 清除EXIF方向标签,避免重复纠正
    if "exif" in img.info:
        exif_data = img.info["exif"]
        # 简单方式:移除方向标签
        img.save(output_path, exif=b"")
    else:
        img.save(output_path)
    
    return output_path, corrected

# 批量处理入口
def batch_correct(folder_path):
    """批量纠正文件夹内所有图片"""
    results = []
    for fname in os.listdir(folder_path):
        if fname.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            fpath = os.path.join(folder_path, fname)
            try:
                out_path, was_corrected = auto_correct_image(fpath)
                results.append((fname, was_corrected, out_path))
            except Exception as e:
                results.append((fname, False, f"错误: {e}"))
    return results

if __name__ == "__main__":
    # 单张测试
    test_img = "test_photo.jpg"
    if os.path.exists(test_img):
        out, corrected = auto_correct_image(test_img)
        print(f"纠正完成: {out}, 是否纠正: {corrected}")
    
    # 批量处理示例(取消注释使用)
    # results = batch_correct("./images_folder")
    # for name, corrected, out in results:
    #     print(f"{name}: {'已纠正' if corrected else '无需纠正'} -> {out}")

运行与验证

运行方式:

# 单张图片测试
python correct_orientation.py

# 或指定图片路径(修改main中的test_img变量)

验证方法:

  1. 准备测试图片:用手机拍摄一张照片(确保EXIF包含方向信息),旋转90度后保存为test_photo.jpg;再准备一张无EXIF信息的图片(可用截图工具生成)
  2. 运行脚本后,打开输出的_corrected后缀文件,目视检查方向是否正确
  3. 对于EXIF纠正路径:脚本会输出是否纠正: True,且输出图片方向正确
  4. 对于内容分析路径:可准备一张倒置的人脸照片(无EXIF),验证脚本能将其旋转180度

通过标准:测试图片经处理后方向正确;批量处理时无异常报错,且输出文件均能正常打开。

常见问题

  • 内容检测准确率有限:本案例的人脸检测方法仅覆盖常见场景,对无人物、无文字的抽象图片可能误判。实际工程中可结合场景分类模型(如Places365)或OCR文字方向检测提升鲁棒性。
  • EXIF标签清除:保存时清除EXIF方向标签是为了避免后续处理重复旋转。若需保留其他元数据,应使用piexif库精确移除Orientation字段。
  • 性能优化:批量处理大量图片时,可先读取EXIF快速过滤,仅对无方向标签的图片执行内容分析,减少计算开销。