📚 课程目录第4章纠正图像方向

4. 纠正图像方向

任务说明

图像数据集中常存在因拍摄或扫描导致的旋转错误（0°/90°/180°/270°）。若直接用于训练，会引入方向偏差，降低模型对物体姿态的鲁棒性。本案例实现一个自动纠正图像方向的工具：优先读取EXIF中的Orientation字段进行无损旋转；对无EXIF信息的图像，使用一个轻量级CNN分类器预测旋转角度并纠正。最终输出统一为正确方向的图像集，供后续训练使用。

实现方案

整体流程分为两条路径：

EXIF路径：使用Pillow读取图像EXIF Orientation标签（274），根据其值（1-8）执行对应的转置/旋转操作，并清除该标签避免二次处理。
CNN路径：对无EXIF或EXIF无法解析的图像，送入一个在合成旋转数据上训练的小型CNN（4分类：0°/90°/180°/270°），预测最可能的旋转角度并逆旋转纠正。

CNN采用类似AlexNet的简化结构，输入为224×224 RGB图像，输出4类概率。训练数据由任意方向图像随机旋转生成，保证模型学习的是“内容无关的方向特征”。

代码实现

import os
import numpy as np
from PIL import Image, ImageOps
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import random

# ---------- 1. EXIF方向纠正 ----------
def correct_exif_orientation(img):
    """根据EXIF Orientation纠正图像方向，返回新图像及是否已纠正"""
    try:
        exif = img._getexif()
        if exif is None:
            return img, False
        orientation = exif.get(274)  # 274 = Orientation标签
        if orientation is None:
            return img, False
        # 处理1-8的EXIF方向值
        if orientation == 2:
            img = ImageOps.mirror(img)
        elif orientation == 3:
            img = img.rotate(180, expand=True)
        elif orientation == 4:
            img = ImageOps.flip(img)
        elif orientation == 5:
            img = ImageOps.mirror(img).rotate(270, expand=True)
        elif orientation == 6:
            img = img.rotate(270, expand=True)
        elif orientation == 7:
            img = ImageOps.mirror(img).rotate(90, expand=True)
        elif orientation == 8:
            img = img.rotate(90, expand=True)
        # 清除EXIF中的方向信息，避免后续重复处理
        if hasattr(img, '_getexif'):
            img.info.pop('exif', None)
        return img, True
    except Exception:
        return img, False

# ---------- 2. CNN方向分类器 ----------
class DirectionCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 28 * 28, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(256, 4)  # 0°, 90°, 180°, 270°
        )

    def forward(self, x):
        return self.classifier(self.features(x))

# 训练数据生成：随机旋转
class RotatedDataset(Dataset):
    def __init__(self, image_paths, transform=None):
        self.paths = image_paths
        self.transform = transform or transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    def __len__(self):
        return len(self.paths) * 4  # 每张图生成4个方向样本

    def __getitem__(self, idx):
        orig_idx = idx // 4
        angle_idx = idx % 4  # 0:0°, 1:90°, 2:180°, 3:270°
        img = Image.open(self.paths[orig_idx]).convert('RGB')
        angle = angle_idx * 90
        if angle > 0:
            img = img.rotate(-angle, expand=True)  # 负号=逆时针
        img = self.transform(img)
        return img, angle_idx

def train_cnn(image_paths, epochs=5, batch_size=16):
    """在给定图像列表上训练方向分类器"""
    dataset = RotatedDataset(image_paths)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    model = DirectionCNN()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-4)

    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for imgs, labels in loader:
            optimizer.zero_grad()
            outputs = model(imgs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(loader):.4f}")
    return model

def predict_rotation(model, img):
    """预测图像需要纠正的角度（0/90/180/270）"""
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    img_t = transform(img.convert('RGB')).unsqueeze(0)
    model.eval()
    with torch.no_grad():
        logits = model(img_t)
        pred = torch.argmax(logits, dim=1).item()
    # 预测类别对应需要逆旋转的角度
    correction_angle = (360 - pred * 90) % 360
    return correction_angle

# ---------- 3. 主流程：自动纠正 ----------
def auto_correct(image_path, model=None):
    """纠正单张图像方向，返回纠正后的PIL图像"""
    img = Image.open(image_path)
    # 先尝试EXIF
    corrected_img, used_exif = correct_exif_orientation(img)
    if used_exif:
        return corrected_img, "exif"
    # EXIF无效或无，使用CNN
    if model is None:
        raise ValueError("无EXIF信息且未提供CNN模型")
    angle = predict_rotation(model, corrected_img)
    if angle != 0:
        corrected_img = corrected_img.rotate(angle, expand=True)
    return corrected_img, f"cnn:{angle}°"

# ---------- 4. 使用示例 ----------
if __name__ == "__main__":
    # 假设有训练图像列表（此处用示例路径，实际需替换）
    train_images = ["sample1.jpg", "sample2.jpg", "sample3.jpg"]
    # 训练CNN（若已有训练好的模型可跳过）
    cnn_model = train_cnn(train_images, epochs=3)

    # 纠正测试图像
    test_file = "test_rotated.jpg"
    if os.path.exists(test_file):
        result_img, method = auto_correct(test_file, cnn_model)
        result_img.save("corrected_" + test_file)
        print(f"纠正完成，使用方式: {method}")
    else:
        print("请提供测试图像路径")

落地与验证

依赖安装：pip install pillow torch torchvision numpy

运行步骤：

准备若干正常方向图像（如JPG格式）作为训练集，替换代码中train_images列表路径。
将一张旋转90°的测试图命名为test_rotated.jpg放在同目录。
运行脚本，观察训练损失下降及最终输出corrected_test_rotated.jpg。

判定标准：

若测试图含EXIF方向信息，输出文件应被正确旋转且无EXIF方向标签残留。
若无EXIF，CNN预测角度应与实际旋转角度一致（误差为0°或360°倍数），输出图像视觉上恢复正立。
训练损失应随epoch下降，最终分类准确率（可在训练集上自测）应高于90%。