IOAI 2025 (Beijing, China), At-Home Round
import numpy as np
import torch
import math
import matplotlib.pyplot as plt
import pandas as pd
import os
from pathlib import Path
import zipfile
from tqdm import tqdm
from datetime import datetime, timedelta
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
TEST_PATH = "./train_v2/"
DATASET_PATH = TEST_PATH + "dataset.npz"
METADATA_PATH = TEST_PATH + "metadata.csv"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BASELINE_MODEL = TEST_PATH + "model_weights.pth"
# Load from .npz file
loaded = np.load(DATASET_PATH)
X_train_128 = loaded['X_train_128']
X_train_256 = loaded['X_train_256']
Y_train_128 = loaded['Y_train_128']
Y_train_256 = loaded['Y_train_256']
X_test_128 = loaded['X_test_128']
X_test_256 = loaded['X_test_256']
Y_test_128 = loaded['Y_test_128']
Y_test_256 = loaded['Y_test_256']
df = pd.read_csv(METADATA_PATH)
df_train_128 = df[(df['size'] == 128) & (df['split'] == 'train')][['i', 'j', 'start_time']]
df_train_256 = df[(df['size'] == 256) & (df['split'] == 'train')][['i', 'j', 'start_time']]
df_test_128 = df[(df['size'] == 128) & (df['split'] == 'test')][['i', 'j', 'start_time']]
df_test_256 = df[(df['size'] == 256) & (df['split'] == 'test')][['i', 'j', 'start_time']]
def pixel_to_latlon(i, j):
"""
Converts pixel indices (row i, column j) to geographic latitude and longitude.
Args:
i (int): Row index (from 0 at the top to 3499 at the bottom).
j (int): Column index (from 0 at the left to 6999 at the right).
Returns:
tuple: (latitude, longitude) as floats.
This function is used to convert from image/pixel coordinates
(such as a satellite image) to actual map coordinates.
"""
lat = 60.0 - i * 0.01 # Each row down is 0.01 degree further south
lon = -130.0 + j * 0.01 # Each column right is 0.01 degree further east
return lat, lon
def solar_elevation(x, y, dt_utc):
"""
Calculates the sun's elevation angle above the horizon for a specific
location (pixel) and UTC time.
Args:
x (int): Row index in the image (vertical position).
y (int): Column index in the image (horizontal position).
dt_utc (datetime or str): The date and time in UTC (string or datetime).
Returns:
float: Solar elevation angle in degrees.
This function tells you "how high is the sun in the sky"
for a given place and time.
"""
# If time is given as string, convert to datetime
if isinstance(dt_utc, str):
dt_utc = datetime.strptime(dt_utc, '%Y-%m-%d %H:%M:%S.%f')
# Convert pixel indices to latitude and longitude
lat, lon = pixel_to_latlon(x, y)
# Estimate local time (in hours) by longitude (15 degrees = 1 hour)
timezone_offset = lon / 15.0
local_time = dt_utc.hour + dt_utc.minute / 60 + timezone_offset
# Day of year (1-365/366)
N = dt_utc.timetuple().tm_yday
# Solar declination: angle between sun's rays and Earth's equator
decl = 23.44 * math.sin(math.radians(360 / 365 * (N - 81)))
# Hour angle: how far in time from solar noon
H = 15 * (local_time - 12) # degrees
# Convert everything to radians for math functions
phi = math.radians(lat)
delta = math.radians(decl)
H = math.radians(H)
# Calculate elevation using spherical trigonometry
sin_h = math.sin(phi) * math.sin(delta) + math.cos(phi) * math.cos(delta) * math.cos(H)
return sin_h
def parse_goes_time(timestr):
"""
Converts a GOES satellite timestamp string into a Python datetime.
Args:
timestr (str): Timestamp string, e.g. 's20242891100205'.
Returns:
datetime: The parsed datetime.
Format explanation:
- 's20242891100205' means:
year = 2024,
day-of-year = 289,
hour = 11,
minute = 00,
second = 20,
tenths of a second = 5
"""
year = int(timestr[1:5])
doy = int(timestr[5:8])
hour = int(timestr[8:10])
minute = int(timestr[10:12])
second = int(timestr[12:14])
micro = int(timestr[14]) * 100000 # tenths of a second
return datetime(year, 1, 1) + timedelta(days=doy - 1, hours=hour, minutes=minute, seconds=second, microseconds=micro)
def df_to_list_solar_elevation(df):
hs = []
for i, j, t in zip(df['i'], df['j'], df['start_time']):
h = solar_elevation(i, j, t)
dt = parse_goes_time(t)
# 正弦编码月份 (1-12 -> 0-2π)
month_sin = math.sin(2 * math.pi * (dt.month - 1) / 12)
month_cos = math.cos(2 * math.pi * (dt.month - 1) / 12)
# 正弦编码一天内的时间 (小时 -> 0-2π)
hour_decimal = dt.hour + dt.minute / 60 + dt.second / 3600
time_sin = math.sin(2 * math.pi * hour_decimal / 24)
time_cos = math.cos(2 * math.pi * hour_decimal / 24)
# 正弦编码经纬度 (假设i和j是经纬度坐标)
# 假设i是纬度，j是经度，需要根据实际数据范围调整
i, j = pixel_to_latlon(i, j)
lat_sin = math.sin(2 * math.pi * i / 180) # 纬度范围 -90到90
lat_cos = math.cos(2 * math.pi * i / 180)
lon_sin = math.sin(2 * math.pi * j / 360) # 经度范围 -180到180
lon_cos = math.cos(2 * math.pi * j / 360)
# 组合所有编码特征
encoded_features = [h, month_sin, month_cos, time_sin, time_cos,
lat_sin, lat_cos, lon_sin, lon_cos]
hs.append(encoded_features)
return hs
angles_train_128 = df_to_list_solar_elevation(df_train_128)
angles_train_256 = df_to_list_solar_elevation(df_train_256)
angles_test_128 = df_to_list_solar_elevation(df_test_128)
angles_test_256 = df_to_list_solar_elevation(df_test_256)
class HybridSolarElevationDataset(Dataset):
def __init__(self, X128, Y128, X256, Y256, angles128, angles256, transform=None):
self.X128 = X128
self.Y128 = Y128
self.X256 = X256
self.Y256 = Y256
self.angles128 = angles128
self.angles256 = angles256
self.transform = transform
self.len_128 = len(X128)
self.len_256 = len(X256)
def __len__(self):
return self.len_128 + self.len_256
def normalize_channels(self, x):
"""
对X[N,C,H,W]的每一个channel进行normalization
"""
# 确保x是4D张量 [N,C,H,W]
if x.dim() == 3:
x = x.unsqueeze(0) # [C,H,W] -> [1,C,H,W]
# 对每个channel分别进行normalization
for c in range(x.shape[1]):
channel_data = x[:, c, :, :]
mean = channel_data.mean()
std = channel_data.std()
if std > 0:
x[:, c, :, :] = (channel_data - mean) / std
return x.squeeze(0) if x.shape[0] == 1 else x
def __getitem__(self, idx):
if idx < self.len_128:
# 128x128 数据，上采样到256x256
x = torch.tensor(self.X128[idx], dtype=torch.float32)
y = torch.tensor(self.Y128[idx], dtype=torch.float32).unsqueeze(0)
# 对输入数据进行channel-wise normalization
x = self.normalize_channels(x)
# 使用双线性插值将128x128上采样到256x256
x = F.interpolate(x.unsqueeze(0), size=(256, 256), mode='nearest', align_corners=False).squeeze(0)
y = F.interpolate(y.unsqueeze(0), size=(256, 256), mode='nearest', align_corners=False).squeeze(0)
angle = torch.tensor(self.angles128[idx], dtype=torch.float32).unsqueeze(-1)
size_flag = 0 # 标记为128尺寸（已上采样）
else:
# 256x256 数据
x = torch.tensor(self.X256[idx - self.len_128], dtype=torch.float32)
y = torch.tensor(self.Y256[idx - self.len_128], dtype=torch.float32).unsqueeze(0)
# 对输入数据进行channel-wise normalization
x = self.normalize_channels(x)
angle = torch.tensor(self.angles256[idx - self.len_128], dtype=torch.float32).unsqueeze(-1)
size_flag = 1 # 标记为256尺寸
if self.transform:
x = self.transform(x)
y = self.transform(y)
return x, y, angle, size_flag
# 创建混合数据集实例
train_dataset = HybridSolarElevationDataset(
X_train_128, Y_train_128, X_train_256, Y_train_256,
angles_train_128, angles_train_256
)
test_dataset = HybridSolarElevationDataset(
X_test_128, Y_test_128, X_test_256, Y_test_256,
angles_test_128, angles_test_256
)
# 创建DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
def dice_loss(pred, target, eps=1e-6):
"""
Computes Dice Loss, which measures overlap between prediction and target.
Args:
pred (torch.Tensor): Raw logits from the model, shape [B, 1, H, W]
target (torch.Tensor): Ground truth masks, shape [B, 1, H, W]
eps (float): Small number to avoid division by zero.
Returns:
torch.Tensor: Dice loss (scalar)
"""
# Apply sigmoid to logits to get probabilities between 0 and 1
pred = torch.sigmoid(pred)
# Intersection and union for each image in the batch
intersection = (pred * target).sum(dim=(1,2,3))
union = pred.sum(dim=(1,2,3)) + target.sum(dim=(1,2,3))
# Dice loss = 1 - Dice coefficient (higher is worse)
return 1 - ((2 * intersection + eps) / (union + eps)).mean()
def train_segmentation(model, train_loader, val_loader, epochs=10, lr=1e-3, device="cuda"):
"""
Trains the segmentation model using both BCEWithLogitsLoss and Dice loss.
Args:
model (nn.Module): The segmentation model (U-Net).
train_loader (DataLoader): DataLoader for training data.
val_loader (DataLoader): DataLoader for validation data.
epochs (int): Number of training epochs.
lr (float): Learning rate.
device (str): 'cuda' for GPU or 'cpu' for CPU.
computes the combined loss, and updates the model weights. After each epoch, validation is performed.
"""
bce_loss = nn.BCEWithLogitsLoss() # Binary cross-entropy loss for logits
model.to(device) # Move model to device (GPU/CPU)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
for epoch in range(epochs):
evaluate_on_test_custom(model, test_loader, device="cuda")
model.train() # Set model to training mode (affects layers like BatchNorm, Dropout)
total_loss = 0.0
print(f"Epoch {epoch + 1}/{epochs}")
# Progress bar for training
pbar = tqdm(train_loader, desc="Training", leave=False)
for batch_idx, (x, y, angle, size_flag) in enumerate(pbar):
# 如果size_flag是0，使用nearest插值对x,y进行下采样缩小
if size_flag == 0:
B, C, H, W = x.shape
new_H, new_W = H // 2, W // 2
x = F.interpolate(x, size=(new_H, new_W), mode='nearest')
y = F.interpolate(y, size=(new_H, new_W), mode='nearest')
x = x.to(device)
y = y.to(device)
angle = angle.to(device)
optimizer.zero_grad() # Clear previous gradients
logits = model(x, angle) # [B, 1, H, W]
logits_clamped = torch.clamp(logits, -20, 20) # Clamp values for stability
# Loss = BCE (pixelwise) + Dice (region overlap)
loss = bce_loss(logits_clamped, y) + dice_loss(logits_clamped, y)
if torch.isnan(loss):
print(f"❌ Batch {batch_idx}: loss is NaN, skipping")
else:
loss.backward() # Compute gradients
optimizer.step() # Update weights
total_loss += loss.item()
print(f" Avg epoch loss: {total_loss:.4f}")
class UNetSegmentation(nn.Module):
"""
U-Net neural network for image segmentation.
Args:
in_channels (int): Number of input channels (e.g., 16 for your patches)
The U-Net consists of an encoder (downsampling), a bottleneck (middle),
and a decoder (upsampling). It is widely used for segmenting images,
where every pixel needs to be classified (e.g., rain/no rain).
"""
def __init__(self, in_channels):
super().__init__()
# Encoder path ("contracting" path): extracts features, reduces size
self.encoder1 = self.conv_block(in_channels, 64)
self.encoder2 = self.conv_block(64, 128)
self.encoder3 = self.conv_block(128, 256)
self.encoder4 = self.conv_block(256, 512)
self.pool = nn.MaxPool2d(2) # Downsamples by factor of 2
# Bottleneck (middle part)
self.mid = self.conv_block(512, 1024)
# Decoder path ("expanding" path): upscales, combines with encoder outputs
self.up4 = nn.ConvTranspose2d(1024, 512, 2, stride=2) # Upsample
self.dec4 = self.conv_block(1024, 512)
self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
self.dec3 = self.conv_block(512, 256)
self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
self.dec2 = self.conv_block(256, 128)
self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
self.dec1 = self.conv_block(128, 64)
# Final layer: reduces to 1 output channel per pixel (for binary segmentation)
self.final = nn.Conv2d(64, 1, kernel_size=1) # Output: logits per pixel
def conv_block(self, in_ch, out_ch):
"""
Helper function to build a block of two convolutional layers,
each followed by BatchNorm and ReLU activation.
Args:
in_ch (int): Number of input channels.
out_ch (int): Number of output channels.
"""
return nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1), # Keeps spatial size the same
nn.BatchNorm2d(out_ch), # Helps training
nn.ReLU(inplace=True), # Non-linearity
nn.Conv2d(out_ch, out_ch, 3, padding=1), # Another conv layer
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True)
)
def forward(self, x, a):
"""
Forward pass through the network.
Args:
x (torch.Tensor): Input tensor of shape [batch, channels, height, width]
Returns:
torch.Tensor: Output logits, shape [batch, 1, height, width]
"""
# Encoder: save outputs for skip connections
e1 = self.encoder1(x)
e2 = self.encoder2(self.pool(e1))
e3 = self.encoder3(self.pool(e2))
e4 = self.encoder4(self.pool(e3))
m = self.mid(self.pool(e4))
m = m * a.view(-1, 1, 1, 1)
# Decoder: upsample, concatenate with encoder outputs (skip connections), then convolve
d4 = self.dec4(torch.cat([self.up4(m), e4], dim=1))
d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1))
d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
return self.final(d1) # Output shape: [batch, 1, H, W]
def evaluate_on_test_custom(model, test_loader, device="cuda", threshold=0.5):
"""
评估分割模型在测试数据上的性能并打印指标。
Args:
model (nn.Module): 训练好的分割模型
test_loader (DataLoader): 测试数据的DataLoader
device (str): 'cuda' 用于GPU或 'cpu' 用于CPU
threshold (float): 将概率转换为二值掩码的阈值
Prints:
- IoU (Intersection over Union)
- Dice coefficient
- Precision
- Recall
- 图像级准确率 (如果图像中任何像素被检测为"雨")
Returns:
float: Dice系数和图像级准确率的平均值
"""
model.eval() # 设置为评估模式 (不更新dropout/batchnorm)
model.to(device)
total_iou = 0.0 # IoU累加器
total_dice = 0.0 # Dice系数累加器
total_prec = 0.0 # 精确率累加器
total_recall = 0.0 # 召回率累加器
total_batches = 0
rain_y_true = [] # 真实图像级标签列表 (雨/无雨)
rain_y_pred = [] # 预测图像级标签列表
with torch.no_grad(): # 评估时不需要计算梯度
pbar = tqdm(test_loader, desc="Test", leave=False)
for batch_idx, (x, y, angle, size_flag) in enumerate(pbar):
if size_flag == 0:
B, C, H, W = x.shape
new_H, new_W = H // 2, W // 2
x = F.interpolate(x, size=(new_H, new_W), mode='nearest')
y = F.interpolate(y, size=(new_H, new_W), mode='nearest')
x = x.to(device)
y = y.to(device)
angle = angle.to(device)
# 按照训练代码的方式处理角度信息
x = x + angle.view(-1, 1, 1, 1)
logits = model(x) # [B, 1, H, W]
probs = torch.sigmoid(logits) # 概率值在 [0, 1] 之间
preds = (probs > threshold).float() # 二值预测
# 计算IoU的交集和并集
intersection = (preds * y).sum(dim=(1, 2, 3))
union = ((preds + y) > 0).float().sum(dim=(1, 2, 3))
iou = (intersection / (union + 1e-6)).mean().item()
# 计算Dice系数
dice = (2 * intersection / (preds.sum(dim=(1,2,3)) + y.sum(dim=(1,2,3)) + 1e-6)).mean().item()
# 所有像素的精确率和召回率
tp = (preds * y).sum().item() # 真正例
fp = (preds * (1 - y)).sum().item() # 假正例
fn = ((1 - preds) * y).sum().item() # 假负例
precision = tp / (tp + fp + 1e-6)
recall = tp / (tp + fn + 1e-6)
# 添加到总计中以计算平均值
total_iou += iou
total_dice += dice
total_prec += precision
total_recall += recall
total_batches += 1
# 图像级雨/无雨: 如果任何像素是"雨"则为True
rain_y_true += [(y > 0.5).any(dim=(1,2,3)).cpu()]
rain_y_pred += [(preds > 0.5).any(dim=(1,2,3)).cpu()]
if total_batches == 0:
print("⚠️ 没有有效的批次")
return
# 合并图像级标签以计算准确率
rain_y_true = torch.cat(rain_y_true)
rain_y_pred = torch.cat(rain_y_pred)
acc = (rain_y_true == rain_y_pred).float().mean().item()
dice_final = total_dice / total_batches
print(f"\n📊 测试指标:")
print(f" • IoU : {total_iou / total_batches:.4f}")
print(f" • Dice : {dice_final:.4f}")
print(f" • Prec : {total_prec / total_batches:.4f}")
print(f" • Recall: {total_recall / total_batches:.4f}")
print(f" • 图像级雨检测准确率: {acc:.4f}")
total_score = (dice_final + acc) / 2
print(f"最终得分: {total_score:.4f}")
# 返回排行榜的组合得分
return dice_final, acc, total_score
model = UNetSegmentation(in_channels=16)
model.load_state_dict(torch.load(BASELINE_MODEL, map_location=torch.device(DEVICE)))
train_segmentation(model, train_loader, test_loader, epochs=10, lr=1e-3, device="cuda")
evaluate_on_test_custom(model, test_loader, device="cuda")