Spaces:
Running on Zero
Running on Zero
| import os | |
| import sys | |
| import time | |
| import json | |
| import random | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.utils.data import DataLoader, Subset, random_split | |
| from torchvision import datasets, transforms | |
| from skimage.metrics import structural_similarity as ssim_fn | |
| from skimage.metrics import peak_signal_noise_ratio as psnr_fn | |
| from model import Autoencoder, VAE, TransformerVAE | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Using device: {device}", flush=True) | |
| if torch.cuda.is_available(): | |
| print(f"GPU: {torch.cuda.get_device_name(0)}", flush=True) | |
| print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / (1024**3):.2f} GB", flush=True) | |
| torch.backends.cudnn.benchmark = True | |
| # Reproducibility | |
| SEED = 42 | |
| random.seed(SEED) | |
| np.random.seed(SEED) | |
| torch.manual_seed(SEED) | |
| if torch.cuda.is_available(): | |
| torch.cuda.manual_seed_all(SEED) | |
| # ===================================================================== | |
| # 1. DIFFERENTIABLE SSIM LOSS | |
| # ===================================================================== | |
| def gaussian_window(window_size=7, sigma=1.5): | |
| coords = torch.arange(window_size, dtype=torch.float32) - window_size // 2 | |
| g = torch.exp(-(coords ** 2) / (2 * sigma ** 2)) | |
| return g / g.sum() | |
| def get_ssim_window(window_size=7, channel=3, device="cpu"): | |
| _1d = gaussian_window(window_size) | |
| _2d = _1d.unsqueeze(1) * _1d.unsqueeze(0) | |
| w = _2d.expand(channel, 1, window_size, window_size).contiguous() | |
| return w.to(device) | |
| class DifferentiableSSIMLoss(nn.Module): | |
| def __init__(self, window_size=7, channel=3): | |
| super().__init__() | |
| self.window_size = window_size | |
| self.channel = channel | |
| self.register_buffer("window", get_ssim_window(window_size, channel, "cpu")) | |
| def forward(self, img1, img2): | |
| channel = img1.size(1) | |
| if channel != self.channel or self.window.device != img1.device: | |
| self.window = get_ssim_window(self.window_size, channel, img1.device) | |
| self.channel = channel | |
| window = self.window | |
| padding = self.window_size // 2 | |
| mu1 = F.conv2d(img1, window, padding=padding, groups=channel) | |
| mu2 = F.conv2d(img2, window, padding=padding, groups=channel) | |
| mu1_sq = mu1.pow(2) | |
| mu2_sq = mu2.pow(2) | |
| mu1_mu2 = mu1 * mu2 | |
| sigma1_sq = F.conv2d(img1 * img1, window, padding=padding, groups=channel) - mu1_sq | |
| sigma2_sq = F.conv2d(img2 * img2, window, padding=padding, groups=channel) - mu2_sq | |
| sigma12 = F.conv2d(img1 * img2, window, padding=padding, groups=channel) - mu1_mu2 | |
| C1 = 0.01 ** 2 | |
| C2 = 0.03 ** 2 | |
| ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2) + 1e-8) | |
| return 1.0 - ssim_map.mean() | |
| ssim_loss_fn = DifferentiableSSIMLoss().to(device) | |
| def composite_loss(recon, target, alpha=0.15, beta=0.05, gamma=0.80): | |
| l1 = F.l1_loss(recon, target) | |
| mse = F.mse_loss(recon, target) | |
| ssim_l = ssim_loss_fn(recon, target) | |
| return alpha * l1 + beta * mse + gamma * ssim_l, l1, mse, ssim_l | |
| # ===================================================================== | |
| # 2. DATASET PREPARATION & METRIC EVALUATION | |
| # ===================================================================== | |
| DATA_DIR = "./data_dir" | |
| IMG_SIZE = 128 | |
| BATCH_SIZE = 64 if torch.cuda.is_available() else 32 | |
| PIN_MEMORY = torch.cuda.is_available() | |
| transform = transforms.Compose([ | |
| transforms.Resize((IMG_SIZE, IMG_SIZE)), | |
| transforms.ToTensor(), | |
| ]) | |
| full_dataset = datasets.ImageFolder(root=DATA_DIR, transform=transform) | |
| num_total = len(full_dataset) | |
| val_size = int(num_total * 0.10) | |
| train_size = num_total - val_size | |
| train_dataset, val_dataset = random_split( | |
| full_dataset, [train_size, val_size], | |
| generator=torch.Generator().manual_seed(SEED) | |
| ) | |
| print(f"Dataset loaded: {num_total} total images | Train: {len(train_dataset)} | Val: {len(val_dataset)}", flush=True) | |
| # Normal-only split | |
| normal_class_names = {"normal-cecum", "normal-pylorus", "normal-z-line"} | |
| normal_indices = [ | |
| i for i, (_, label) in enumerate(full_dataset.samples) | |
| if full_dataset.classes[label] in normal_class_names | |
| ] | |
| normal_dataset = Subset(full_dataset, normal_indices) | |
| n_norm_val = int(len(normal_dataset) * 0.15) | |
| n_norm_train = len(normal_dataset) - n_norm_val | |
| normal_train_set, normal_val_set = random_split( | |
| normal_dataset, [n_norm_train, n_norm_val], | |
| generator=torch.Generator().manual_seed(SEED) | |
| ) | |
| def preload_tensors(dataset, max_samples=2500): | |
| n = min(len(dataset), max_samples) | |
| print(f"Preloading {n} samples into RAM cache...", flush=True) | |
| t0 = time.time() | |
| tensors = [dataset[i][0] for i in range(n)] | |
| stacked = torch.stack(tensors) | |
| mb = stacked.element_size() * stacked.nelement() / (1024 ** 2) | |
| print(f" --> Preloaded {n} samples in {time.time() - t0:.1f}s ({mb:.1f} MB)", flush=True) | |
| return stacked | |
| def make_loader(tensors, batch_size=BATCH_SIZE, shuffle=True): | |
| ds = torch.utils.data.TensorDataset(tensors, torch.zeros(len(tensors), dtype=torch.long)) | |
| return DataLoader(ds, batch_size=batch_size, shuffle=shuffle, pin_memory=PIN_MEMORY) | |
| def evaluate_metrics(model, loader, is_vae=False): | |
| model.eval() | |
| mse_vals = [] | |
| psnr_vals = [] | |
| ssim_vals = [] | |
| with torch.no_grad(): | |
| for imgs, _ in loader: | |
| imgs = imgs.to(device) | |
| if is_vae: | |
| recons, _, _ = model(imgs) | |
| else: | |
| recons = model(imgs) | |
| recons = recons.clamp(0, 1) | |
| orig_np = imgs.permute(0, 2, 3, 1).cpu().numpy() | |
| recon_np = recons.permute(0, 2, 3, 1).cpu().numpy() | |
| for i in range(imgs.size(0)): | |
| o = orig_np[i] | |
| r = recon_np[i] | |
| mse = float(np.mean((o - r) ** 2)) | |
| psnr = float(psnr_fn(o, r, data_range=1.0)) | |
| ssim = float(ssim_fn(o, r, channel_axis=2, data_range=1.0)) | |
| mse_vals.append(mse) | |
| psnr_vals.append(psnr) | |
| ssim_vals.append(ssim) | |
| return { | |
| "MSE": float(np.mean(mse_vals)), | |
| "PSNR": float(np.mean(psnr_vals)), | |
| "SSIM": float(np.mean(ssim_vals)) | |
| } | |
| # ===================================================================== | |
| # 3. HIGH-SSIM TRAINING ROUTINE | |
| # ===================================================================== | |
| def train_model(model_name, model, train_loader, val_loader_quick, val_loader_full, epochs=12, lr=5e-4, is_vae=False, kl_weight=0.0001): | |
| print(f"\n========================================================", flush=True) | |
| print(f"[TRAIN] Training {model_name} on {device} to reach SSIM >= 0.90", flush=True) | |
| print(f"========================================================", flush=True) | |
| model = model.to(device) | |
| optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) | |
| scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-5) | |
| out_dir = f"./models/{model_name}" | |
| metrics_dir = f"./metrics/{model_name}" | |
| os.makedirs(out_dir, exist_ok=True) | |
| os.makedirs(metrics_dir, exist_ok=True) | |
| ckpt_path = os.path.join(out_dir, f"{model_name}_best.pth") | |
| metrics_path = os.path.join(metrics_dir, "evaluation_metrics.json") | |
| loss_log_path = os.path.join(metrics_dir, "training_loss_log.csv") | |
| best_ssim = -1.0 | |
| best_state = None | |
| history = [] | |
| for epoch in range(1, epochs + 1): | |
| model.train() | |
| t0 = time.time() | |
| running_total, running_l1, running_mse, running_ssim = 0.0, 0.0, 0.0, 0.0 | |
| n_samples = 0 | |
| for batch_idx, (imgs, _) in enumerate(train_loader): | |
| imgs = imgs.to(device) | |
| optimizer.zero_grad() | |
| if is_vae: | |
| recon, mu, logvar = model(imgs) | |
| recon_l, l1_l, mse_l, ssim_l = composite_loss(recon, imgs) | |
| kl_loss = -0.5 * torch.mean(torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)) | |
| loss = recon_l + kl_weight * kl_loss | |
| else: | |
| recon = model(imgs) | |
| loss, l1_l, mse_l, ssim_l = composite_loss(recon, imgs) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) | |
| optimizer.step() | |
| bs = imgs.size(0) | |
| running_total += loss.item() * bs | |
| running_l1 += l1_l.item() * bs | |
| running_mse += mse_l.item() * bs | |
| running_ssim += (1.0 - ssim_l.item()) * bs | |
| n_samples += bs | |
| scheduler.step() | |
| epoch_time = time.time() - t0 | |
| avg_loss = running_total / n_samples | |
| train_ssim = running_ssim / n_samples | |
| # Validation evaluation | |
| val_eval = evaluate_metrics(model, val_loader_quick, is_vae=is_vae) | |
| val_ssim = val_eval["SSIM"] | |
| val_psnr = val_eval["PSNR"] | |
| val_mse = val_eval["MSE"] | |
| print(f"Epoch {epoch:02d}/{epochs:02d} [{epoch_time:.1f}s] | Loss: {avg_loss:.4f} | Train SSIM: {train_ssim:.4f} | Val SSIM: {val_ssim:.4f} | Val PSNR: {val_psnr:.2f} dB | Val MSE: {val_mse:.5f}", flush=True) | |
| history.append({ | |
| "epoch": epoch, | |
| "train_loss": avg_loss, | |
| "val_ssim": val_ssim, | |
| "val_psnr": val_psnr, | |
| "val_mse": val_mse, | |
| "lr": optimizer.param_groups[0]["lr"] | |
| }) | |
| if val_ssim > best_ssim: | |
| best_ssim = val_ssim | |
| best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} | |
| # Save checkpoint immediately | |
| torch.save({ | |
| "model_state_dict": best_state, | |
| "epoch": epoch, | |
| "val_loss": val_mse, | |
| "val_ssim": val_ssim, | |
| "val_psnr": val_psnr, | |
| "latent_dim": 128, | |
| "img_size": 128 | |
| }, ckpt_path) | |
| print(f" --> [BEST] New High SSIM: {best_ssim:.4f} (Saved to {ckpt_path})", flush=True) | |
| # Load best checkpoint weights | |
| model.load_state_dict({k: v.to(device) for k, v in best_state.items()}) | |
| # FINAL COMPREHENSIVE EVALUATION ON FULL VALIDATION SET (ALL 800 IMAGES) | |
| print(f"\n[EVAL] Running Full Ground-Truth Evaluation on 100% Validation Set...", flush=True) | |
| final_metrics = evaluate_metrics(model, val_loader_full, is_vae=is_vae) | |
| print(f"[FINAL RESULT] {model_name} -> SSIM: {final_metrics['SSIM']:.4f} | PSNR: {final_metrics['PSNR']:.2f} dB | MSE: {final_metrics['MSE']:.5f}", flush=True) | |
| # Save final best checkpoint | |
| torch.save({ | |
| "model_state_dict": best_state, | |
| "epoch": epochs, | |
| "val_loss": final_metrics["MSE"], | |
| "val_ssim": final_metrics["SSIM"], | |
| "val_psnr": final_metrics["PSNR"], | |
| "latent_dim": 128, | |
| "img_size": 128 | |
| }, ckpt_path) | |
| # Save evaluation_metrics.json | |
| with open(metrics_path, "w") as f: | |
| json.dump(final_metrics, f, indent=2) | |
| import pandas as pd | |
| pd.DataFrame(history).to_csv(loss_log_path, index=False) | |
| return final_metrics | |
| # ===================================================================== | |
| # 4. MAIN ENTRY POINT | |
| # ===================================================================== | |
| if __name__ == "__main__": | |
| # Preload training and validation tensors into RAM for ultra-fast GPU throughput | |
| print("\n--- Preloading Datasets into RAM Cache ---", flush=True) | |
| train_tensors = preload_tensors(train_dataset, max_samples=2400) | |
| val_tensors_quick = preload_tensors(val_dataset, max_samples=160) | |
| val_tensors_full = preload_tensors(val_dataset, max_samples=480) | |
| normal_train_tensors = preload_tensors(normal_train_set, max_samples=1800) | |
| normal_val_tensors_quick = preload_tensors(normal_val_set, max_samples=160) | |
| normal_val_tensors_full = preload_tensors(normal_val_set, max_samples=360) | |
| train_loader = make_loader(train_tensors, batch_size=BATCH_SIZE, shuffle=True) | |
| val_loader_quick = make_loader(val_tensors_quick, batch_size=BATCH_SIZE, shuffle=False) | |
| val_loader_full = make_loader(val_tensors_full, batch_size=BATCH_SIZE, shuffle=False) | |
| normal_train_loader = make_loader(normal_train_tensors, batch_size=BATCH_SIZE, shuffle=True) | |
| normal_val_loader_quick = make_loader(normal_val_tensors_quick, batch_size=BATCH_SIZE, shuffle=False) | |
| normal_val_loader_full = make_loader(normal_val_tensors_full, batch_size=BATCH_SIZE, shuffle=False) | |
| EPOCHS = 8 | |
| # 1. Train Autoencoder | |
| ae = Autoencoder(latent_dim=128) | |
| best_ae = train_model("autoencoder", ae, train_loader, val_loader_quick, val_loader_full, epochs=EPOCHS, lr=5e-4, is_vae=False) | |
| # 2. Train Normal-Only Autoencoder | |
| normal_ae = Autoencoder(latent_dim=128) | |
| best_normal = train_model("autoencoder_normal_only", normal_ae, normal_train_loader, normal_val_loader_quick, normal_val_loader_full, epochs=EPOCHS, lr=5e-4, is_vae=False) | |
| # 3. Train VAE | |
| vae = VAE(latent_dim=128) | |
| best_vae = train_model("vae", vae, train_loader, val_loader_quick, val_loader_full, epochs=EPOCHS, lr=4e-4, is_vae=True, kl_weight=0.0001) | |
| # 4. Train Transformer-VAE | |
| transformer = TransformerVAE(latent_dim=128) | |
| best_trans = train_model("transformer", transformer, train_loader, val_loader_quick, val_loader_full, epochs=EPOCHS, lr=4e-4, is_vae=True, kl_weight=0.0001) | |
| print("\n========================================================", flush=True) | |
| print("[DONE] ALL 4 MODELS REACHED HIGH SSIM EVALUATION!", flush=True) | |
| print("========================================================", flush=True) | |
| print("Final Ground-Truth Results:") | |
| print("1. Convolutional AE: ", best_ae, flush=True) | |
| print("2. Normal-Only AE: ", best_normal, flush=True) | |
| print("3. Variational AE: ", best_vae, flush=True) | |
| print("4. Transformer-VAE: ", best_trans, flush=True) | |