利用Python给TC整容

Enceladus

期盼了很久的潜力股终究只是一团CCC？
好不容易卷出的风眼像一坨？
结构很好，但冷水切崩？
利用深度学习，以上统统不是问题——
1、        数据准备：首先，加载并预处理需要美颜的图像
2、        融合方法：准备另一组顶超图像。在训练过程中从中抓取三张图像，结合它们的信息生成融合图像
3、        模型训练：使用一个深度网络来学习热带气旋图像的特征，通过L1损失函数来优化输出图像

import os, cv2, random, torch, numpy as np
from PIL import Image
from torch import nn, optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.utils import save_image
# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 超参数
image_size, batch_size, epochs, lr = 512, 16, 15, 1e-3
sample_dir = "fusion_samples"
os.makedirs(sample_dir, exist_ok=True)
transform = transforms.Compose([
    transforms.Resize((image_size, image_size), interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.ToTensor()
])
# 数据集定义
class FusionDataset(Dataset):
    def __init__(self, cyclone_dir, tc_space_dir, fusion_dir=None, transform=None):
        self.cyclone = sorted([f for f in os.listdir(cyclone_dir) if f.lower().endswith(('.png', '.jpg'))])
        self.tc = sorted([f for f in os.listdir(tc_space_dir) if f.lower().endswith(('.png', '.jpg'))])
        self.fusion = sorted([f for f in os.listdir(fusion_dir)]) if fusion_dir else None
        self.cyclone_dir, self.tc_dir, self.fusion_dir = cyclone_dir, tc_space_dir, fusion_dir
        self.transform = transform
    def __getitem__(self, idx):
        def load_img(folder, file):
            img = Image.open(os.path.join(folder, file)).convert("L")
            return self.transform(img)
        cyclone = load_img(self.cyclone_dir, self.cyclone[idx])
        # 平均融合3张参考图像
        refs = [load_img(self.tc_dir, self.tc) for i in random.sample(range(len(self.tc)), 3)]
        reference = torch.stack(refs).mean(0)
        # 生成融合图像
        if self.fusion_dir:
            fusion = load_img(self.fusion_dir, self.fusion[idx])
        else:
            grad_c = torch.abs(cyclone[:, :, 1:] - cyclone[:, :, :-1])
            grad_t = torch.abs(reference[:, :, 1:] - reference[:, :, :-1])
            grad_c = nn.functional.pad(grad_c, (0,1))
            grad_t = nn.functional.pad(grad_t, (0,1))
            weight = (grad_t > grad_c).float() * 0.6 + 0.4
            fusion = weight * reference + (1 - weight) * cyclone
        return cyclone, reference, fusion
    def __len__(self):
        return min(len(self.cyclone), len(self.tc))
# UNet结构定义
class UNet(nn.Module):
    def __init__(self):
        super().__init__()
        def conv_block(in_c, out_c):
            return nn.Sequential(
                nn.Conv2d(in_c, out_c, 3, padding=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(out_c, out_c, 3, padding=1),
                nn.ReLU(inplace=True))
        self.enc = nn.ModuleList([conv_block(2, 64), conv_block(64, 128), conv_block(128, 256), conv_block(256, 512)])
        self.pool = nn.MaxPool2d(2)
        self.bottleneck = conv_block(512, 1024)
        self.up = nn.ModuleList([nn.ConvTranspose2d(1024, 512, 2, 2), nn.ConvTranspose2d(512, 256, 2, 2), nn.ConvTranspose2d(256, 128, 2, 2), nn.ConvTranspose2d(128, 64, 2, 2)])
        self.dec = nn.ModuleList([conv_block(1024, 512), conv_block(512, 256), conv_block(256, 128), conv_block(128, 64)])
        self.out_conv = nn.Conv2d(64, 1, 1)
    def forward(self, x1, x2):
        x = torch.cat([x1, x2], dim=1)
        encs = []
        for layer in self.enc:
            x = layer(x)
            encs.append(x)
            x = self.pool(x)
        x = self.bottleneck(x)
        for i in range(4):
            x = self.up(x)
            x = torch.cat([x, encs[3 - i]], dim=1)
            x = self.dec(x)
        return self.out_conv(x)
# 准备数据和模型
dataset = FusionDataset("tropical storms", "category 5s", fusion_dir=None, transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
model = UNet().to(device)
criterion = nn.L1Loss()
optimizer = optim.Adam(model.parameters(), lr=lr)
# 训练循环
for epoch in range(epochs):
    for i, (cyc, tc, target) in enumerate(dataloader):
        cyc, tc, target = cyc.to(device), tc.to(device), target.to(device)
        output = model(cyc, tc)
        loss = criterion(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print(f"Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(dataloader)}], Loss: {loss.item():.4f}")
    if epoch >= 4:
        with torch.no_grad():
            sample = model(cyc[:1], tc[:1])
            sample = (sample - sample.min()) / (sample.max() - sample.min() + 1e-5) * (255 + 25) - 25 #色彩调整
            save_image(sample / 255., f"{sample_dir}/epoch_{epoch+1}.jpg")

4、        结果展示（2024 MALIKSI、2024 PRAPIROON）：
5、        不要在正式场合使用整容图像
整容前


整容后

9916-Bart

太强了，伟大

Enceladus

帮助每一个水货圆梦
2406山神

Enceladus

2408悟空

GJ229

让台风“大开眼界”是吧

Enceladus

2409云雀