diff --git a/image-inpainting/results/testset/tikaiz-29.3653.npz b/image-inpainting/results/testset/tikaiz-18.0253.npz similarity index 69% rename from image-inpainting/results/testset/tikaiz-29.3653.npz rename to image-inpainting/results/testset/tikaiz-18.0253.npz index 2fc85a0..5ce37e0 100644 Binary files a/image-inpainting/results/testset/tikaiz-29.3653.npz and b/image-inpainting/results/testset/tikaiz-18.0253.npz differ diff --git a/image-inpainting/src/__pycache__/architecture.cpython-313.pyc b/image-inpainting/src/__pycache__/architecture.cpython-313.pyc index 5295f71..cac7a2a 100644 Binary files a/image-inpainting/src/__pycache__/architecture.cpython-313.pyc and b/image-inpainting/src/__pycache__/architecture.cpython-313.pyc differ diff --git a/image-inpainting/src/__pycache__/datasets.cpython-313.pyc b/image-inpainting/src/__pycache__/datasets.cpython-313.pyc index c101f32..ded87ad 100644 Binary files a/image-inpainting/src/__pycache__/datasets.cpython-313.pyc and b/image-inpainting/src/__pycache__/datasets.cpython-313.pyc differ diff --git a/image-inpainting/src/__pycache__/train.cpython-313.pyc b/image-inpainting/src/__pycache__/train.cpython-313.pyc index 3b0020b..baa1f65 100644 Binary files a/image-inpainting/src/__pycache__/train.cpython-313.pyc and b/image-inpainting/src/__pycache__/train.cpython-313.pyc differ diff --git a/image-inpainting/src/__pycache__/utils.cpython-313.pyc b/image-inpainting/src/__pycache__/utils.cpython-313.pyc index 251d5ec..dc2f6aa 100644 Binary files a/image-inpainting/src/__pycache__/utils.cpython-313.pyc and b/image-inpainting/src/__pycache__/utils.cpython-313.pyc differ diff --git a/image-inpainting/src/architecture.py b/image-inpainting/src/architecture.py index 76a6e1e..aa6c8c1 100644 --- a/image-inpainting/src/architecture.py +++ b/image-inpainting/src/architecture.py @@ -70,9 +70,9 @@ class CBAM(nn.Module): class ConvBlock(nn.Module): """Convolutional block with Conv2d -> BatchNorm -> LeakyReLU""" - def __init__(self, in_channels, out_channels, kernel_size=3, padding=1, dropout=0.0): + def __init__(self, in_channels, out_channels, kernel_size=3, padding=1, dilation=1, dropout=0.0): super().__init__() - self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding) + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding, dilation=dilation) self.bn = nn.BatchNorm2d(out_channels) self.relu = nn.LeakyReLU(0.1, inplace=True) self.dropout = nn.Dropout2d(dropout) if dropout > 0 else nn.Identity() @@ -101,13 +101,13 @@ class ResidualConvBlock(nn.Module): class DownBlock(nn.Module): - """Downsampling block with conv blocks, residual connection, attention, and max pooling""" - def __init__(self, in_channels, out_channels, dropout=0.1): + """Simplified downsampling block with conv blocks, residual connection, and max pooling""" + def __init__(self, in_channels, out_channels, dropout=0.1, use_attention=True): super().__init__() self.conv1 = ConvBlock(in_channels, out_channels, dropout=dropout) self.conv2 = ConvBlock(out_channels, out_channels, dropout=dropout) self.residual = ResidualConvBlock(out_channels, dropout=dropout) - self.attention = CBAM(out_channels) + self.attention = CBAM(out_channels) if use_attention else nn.Identity() self.pool = nn.MaxPool2d(2) def forward(self, x): @@ -118,15 +118,14 @@ class DownBlock(nn.Module): return self.pool(skip), skip class UpBlock(nn.Module): - """Upsampling block with transposed conv, residual connection, attention, and conv blocks""" - def __init__(self, in_channels, out_channels, dropout=0.1): + """Simplified upsampling block with transposed conv, residual connection, and conv blocks""" + def __init__(self, in_channels, out_channels, dropout=0.1, use_attention=True): super().__init__() self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2) # After concat: out_channels (from upconv) + in_channels (from skip) self.conv1 = ConvBlock(out_channels + in_channels, out_channels, dropout=dropout) - self.conv2 = ConvBlock(out_channels, out_channels, dropout=dropout) self.residual = ResidualConvBlock(out_channels, dropout=dropout) - self.attention = CBAM(out_channels) + self.attention = CBAM(out_channels) if use_attention else nn.Identity() def forward(self, x, skip): x = self.up(x) @@ -135,7 +134,6 @@ class UpBlock(nn.Module): x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, skip], dim=1) x = self.conv1(x) - x = self.conv2(x) x = self.residual(x) x = self.attention(x) return x @@ -145,38 +143,35 @@ class MyModel(nn.Module): def __init__(self, n_in_channels: int, base_channels: int = 64, dropout: float = 0.1): super().__init__() - # Initial convolution with larger receptive field + # Initial convolution - simplified self.init_conv = nn.Sequential( - ConvBlock(n_in_channels, base_channels, kernel_size=7, padding=3), - ConvBlock(base_channels, base_channels), - ResidualConvBlock(base_channels) + ConvBlock(n_in_channels, base_channels, kernel_size=5, padding=2), + ConvBlock(base_channels, base_channels) ) - # Encoder (downsampling path) - self.down1 = DownBlock(base_channels, base_channels * 2, dropout=dropout) - self.down2 = DownBlock(base_channels * 2, base_channels * 4, dropout=dropout) - self.down3 = DownBlock(base_channels * 4, base_channels * 8, dropout=dropout) - self.down4 = DownBlock(base_channels * 8, base_channels * 16, dropout=dropout) + # Encoder (downsampling path) - attention only on deeper layers + self.down1 = DownBlock(base_channels, base_channels * 2, dropout=dropout, use_attention=False) + self.down2 = DownBlock(base_channels * 2, base_channels * 4, dropout=dropout, use_attention=False) + self.down3 = DownBlock(base_channels * 4, base_channels * 8, dropout=dropout, use_attention=True) + self.down4 = DownBlock(base_channels * 8, base_channels * 16, dropout=dropout, use_attention=True) - # Bottleneck with multiple residual blocks + # Simplified bottleneck with dilated convolutions self.bottleneck = nn.Sequential( ConvBlock(base_channels * 16, base_channels * 16, dropout=dropout), - ResidualConvBlock(base_channels * 16, dropout=dropout), - ResidualConvBlock(base_channels * 16, dropout=dropout), + ConvBlock(base_channels * 16, base_channels * 16, dilation=2, padding=2, dropout=dropout), ResidualConvBlock(base_channels * 16, dropout=dropout), CBAM(base_channels * 16) ) - # Decoder (upsampling path) - self.up1 = UpBlock(base_channels * 16, base_channels * 8, dropout=dropout) - self.up2 = UpBlock(base_channels * 8, base_channels * 4, dropout=dropout) - self.up3 = UpBlock(base_channels * 4, base_channels * 2, dropout=dropout) - self.up4 = UpBlock(base_channels * 2, base_channels, dropout=dropout) + # Decoder (upsampling path) - attention only on deeper layers + self.up1 = UpBlock(base_channels * 16, base_channels * 8, dropout=dropout, use_attention=True) + self.up2 = UpBlock(base_channels * 8, base_channels * 4, dropout=dropout, use_attention=True) + self.up3 = UpBlock(base_channels * 4, base_channels * 2, dropout=dropout, use_attention=False) + self.up4 = UpBlock(base_channels * 2, base_channels, dropout=dropout, use_attention=False) - # Final refinement layers + # Simplified final refinement layers self.final_conv = nn.Sequential( ConvBlock(base_channels * 2, base_channels), - ResidualConvBlock(base_channels), ConvBlock(base_channels, base_channels) ) diff --git a/image-inpainting/src/datasets.py b/image-inpainting/src/datasets.py index d5e74eb..9ad7db6 100644 --- a/image-inpainting/src/datasets.py +++ b/image-inpainting/src/datasets.py @@ -10,7 +10,7 @@ import numpy as np import random import glob import os -from PIL import Image +from PIL import Image, ImageEnhance IMAGE_DIMENSION = 100 @@ -32,17 +32,44 @@ def resize(img: Image): transforms.CenterCrop((IMAGE_DIMENSION, IMAGE_DIMENSION)) ]) return resize_transforms(img) + def preprocess(input_array: np.ndarray): input_array = np.asarray(input_array, dtype=np.float32) / 255.0 return input_array +def augment_image(img: Image, strength: float = 0.5) -> Image: + """Apply fast data augmentation with controlled strength""" + # Random horizontal flip + if random.random() > 0.5: + img = img.transpose(Image.FLIP_LEFT_RIGHT) + + # Random rotation (90, 180, 270 degrees) - less frequent for speed + if random.random() > 0.6: + angle = random.choice([90, 180, 270]) + img = img.rotate(angle) + + # Simplified color jitter - only one transformation per image for speed + rand = random.random() + if rand > 0.66: + # Brightness + factor = 1.0 + random.uniform(-0.15, 0.15) * strength + img = ImageEnhance.Brightness(img).enhance(factor) + elif rand > 0.33: + # Contrast + factor = 1.0 + random.uniform(-0.15, 0.15) * strength + img = ImageEnhance.Contrast(img).enhance(factor) + + return img + class ImageDataset(torch.utils.data.Dataset): """ - Dataset class for loading images from a folder + Dataset class for loading images from a folder with augmentation support """ - def __init__(self, datafolder: str): + def __init__(self, datafolder: str, augment: bool = True, augment_strength: float = 0.5): self.imagefiles = sorted(glob.glob(os.path.join(datafolder,"**","*.jpg"),recursive=True)) + self.augment = augment + self.augment_strength = augment_strength def __len__(self): return len(self.imagefiles) @@ -51,7 +78,13 @@ class ImageDataset(torch.utils.data.Dataset): index = int(idx) image = Image.open(self.imagefiles[index]) - image = np.asarray(resize(image)) + image = resize(image) + + # Apply augmentation + if self.augment: + image = augment_image(image, self.augment_strength) + + image = np.asarray(image) image = preprocess(image) spacing_x = random.randint(2,6) spacing_y = random.randint(2,6) diff --git a/image-inpainting/src/main.py b/image-inpainting/src/main.py index 9f61cd6..da50793 100644 --- a/image-inpainting/src/main.py +++ b/image-inpainting/src/main.py @@ -24,22 +24,22 @@ if __name__ == '__main__': config_dict['results_path'] = os.path.join(project_root, "results") config_dict['data_path'] = os.path.join(project_root, "data", "dataset") config_dict['device'] = None - config_dict['learningrate'] = 3e-4 # Optimal learning rate for AdamW - config_dict['weight_decay'] = 1e-4 # Slightly higher for better regularization - config_dict['n_updates'] = 5000 # More updates for better convergence - config_dict['batchsize'] = 8 # Smaller batch for better gradient estimates - config_dict['early_stopping_patience'] = 10 # More patience for complex model + config_dict['learningrate'] = 8e-4 # Higher max LR for OneCycleLR + config_dict['weight_decay'] = 1e-5 # Lower for faster learning + config_dict['n_updates'] = 3500 # Reduced for fast training + config_dict['batchsize'] = 16 # Larger batch for speed + config_dict['early_stopping_patience'] = 12 # Adjusted patience config_dict['use_wandb'] = False config_dict['print_train_stats_at'] = 10 config_dict['print_stats_at'] = 100 - config_dict['plot_at'] = 300 - config_dict['validate_at'] = 300 # Validate more frequently + config_dict['plot_at'] = 500 + config_dict['validate_at'] = 150 # More frequent validation network_config = { 'n_in_channels': 4, - 'base_channels': 48, # Good balance between capacity and memory - 'dropout': 0.1 # Regularization + 'base_channels': 40, # Reduced for lower complexity + 'dropout': 0.05 # Lower dropout for faster convergence } config_dict['network_config'] = network_config diff --git a/image-inpainting/src/train.py b/image-inpainting/src/train.py index 10bf917..0ffaaa4 100644 --- a/image-inpainting/src/train.py +++ b/image-inpainting/src/train.py @@ -15,44 +15,21 @@ import os from torch.utils.data import DataLoader from torch.utils.data import Subset -from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts +from torch.optim.lr_scheduler import OneCycleLR import wandb -class CombinedLoss(nn.Module): - """Combined loss: MSE + L1 + SSIM-like perceptual component""" - def __init__(self, mse_weight=1.0, l1_weight=0.5, edge_weight=0.1): +class RMSELoss(nn.Module): + """RMSE loss for direct optimization of evaluation metric""" + def __init__(self): super().__init__() - self.mse_weight = mse_weight - self.l1_weight = l1_weight - self.edge_weight = edge_weight self.mse = nn.MSELoss() - self.l1 = nn.L1Loss() - - # Sobel filters for edge detection - sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3) - sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).view(1, 1, 3, 3) - self.register_buffer('sobel_x', sobel_x.repeat(3, 1, 1, 1)) - self.register_buffer('sobel_y', sobel_y.repeat(3, 1, 1, 1)) - - def edge_loss(self, pred, target): - """Compute edge-aware loss using Sobel filters""" - pred_edge_x = torch.nn.functional.conv2d(pred, self.sobel_x, padding=1, groups=3) - pred_edge_y = torch.nn.functional.conv2d(pred, self.sobel_y, padding=1, groups=3) - target_edge_x = torch.nn.functional.conv2d(target, self.sobel_x, padding=1, groups=3) - target_edge_y = torch.nn.functional.conv2d(target, self.sobel_y, padding=1, groups=3) - - edge_loss = self.l1(pred_edge_x, target_edge_x) + self.l1(pred_edge_y, target_edge_y) - return edge_loss def forward(self, pred, target): - mse_loss = self.mse(pred, target) - l1_loss = self.l1(pred, target) - edge_loss = self.edge_loss(pred, target) - - total_loss = self.mse_weight * mse_loss + self.l1_weight * l1_loss + self.edge_weight * edge_loss - return total_loss + mse = self.mse(pred, target) + rmse = torch.sqrt(mse + 1e-8) # Add epsilon for numerical stability + return rmse def train(seed, testset_ratio, validset_ratio, data_path, results_path, early_stopping_patience, device, learningrate, @@ -111,15 +88,16 @@ def train(seed, testset_ratio, validset_ratio, data_path, results_path, early_st network.to(device) network.train() - # defining the loss - combined loss for better reconstruction - combined_loss = CombinedLoss(mse_weight=1.0, l1_weight=0.5, edge_weight=0.1).to(device) + # defining the loss - RMSE for direct optimization of evaluation metric + rmse_loss = RMSELoss().to(device) mse_loss = torch.nn.MSELoss() # Keep for evaluation # defining the optimizer with AdamW for better weight decay handling - optimizer = torch.optim.AdamW(network.parameters(), lr=learningrate, weight_decay=weight_decay) + optimizer = torch.optim.AdamW(network.parameters(), lr=learningrate, weight_decay=weight_decay, betas=(0.9, 0.99)) - # Learning rate scheduler for better convergence - scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=50, T_mult=2, eta_min=1e-6) + # OneCycleLR for fast convergence - ramps up then down over entire training + scheduler = OneCycleLR(optimizer, max_lr=learningrate, total_steps=n_updates, + pct_start=0.3, anneal_strategy='cos', div_factor=25.0, final_div_factor=1e4) if use_wandb: wandb.watch(network, mse_loss, log="all", log_freq=10) @@ -146,7 +124,7 @@ def train(seed, testset_ratio, validset_ratio, data_path, results_path, early_st output = network(input) - loss = combined_loss(output, target) + loss = rmse_loss(output, target) loss.backward() @@ -154,7 +132,7 @@ def train(seed, testset_ratio, validset_ratio, data_path, results_path, early_st torch.nn.utils.clip_grad_norm_(network.parameters(), max_norm=1.0) optimizer.step() - scheduler.step(i + len(loss_list) / len(dataloader_train)) + scheduler.step() # OneCycleLR steps once per optimizer step loss_list.append(loss.item())