diff --git a/image-inpainting/results/best_model.pt b/image-inpainting/results/best_model.pt new file mode 100644 index 0000000..a85a5e3 Binary files /dev/null and b/image-inpainting/results/best_model.pt differ diff --git a/image-inpainting/results/testset/my_submission_name.npz b/image-inpainting/results/testset/my_submission_name.npz new file mode 100644 index 0000000..47a1db9 Binary files /dev/null and b/image-inpainting/results/testset/my_submission_name.npz differ diff --git a/image-inpainting/src/__pycache__/architecture.cpython-314.pyc b/image-inpainting/src/__pycache__/architecture.cpython-314.pyc index eabdd34..5abafe7 100644 Binary files a/image-inpainting/src/__pycache__/architecture.cpython-314.pyc and b/image-inpainting/src/__pycache__/architecture.cpython-314.pyc differ diff --git a/image-inpainting/src/__pycache__/datasets.cpython-314.pyc b/image-inpainting/src/__pycache__/datasets.cpython-314.pyc index 26497b0..b6f25ff 100644 Binary files a/image-inpainting/src/__pycache__/datasets.cpython-314.pyc and b/image-inpainting/src/__pycache__/datasets.cpython-314.pyc differ diff --git a/image-inpainting/src/__pycache__/utils.cpython-314.pyc b/image-inpainting/src/__pycache__/utils.cpython-314.pyc index ee868f4..8911436 100644 Binary files a/image-inpainting/src/__pycache__/utils.cpython-314.pyc and b/image-inpainting/src/__pycache__/utils.cpython-314.pyc differ diff --git a/image-inpainting/src/architecture.py b/image-inpainting/src/architecture.py index ba1f092..e63224c 100644 --- a/image-inpainting/src/architecture.py +++ b/image-inpainting/src/architecture.py @@ -5,7 +5,109 @@ """ import torch +import torch.nn as nn -class MyModel(torch.nn.Module): - # TODO: Implement the model architecture. - pass \ No newline at end of file +class ConvBlock(nn.Module): + """Convolutional block with Conv2d -> BatchNorm -> ReLU""" + def __init__(self, in_channels, out_channels, kernel_size=3, padding=1): + super().__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding) + self.bn = nn.BatchNorm2d(out_channels) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.relu(self.bn(self.conv(x))) + +class DownBlock(nn.Module): + """Downsampling block with two conv blocks and max pooling""" + def __init__(self, in_channels, out_channels): + super().__init__() + self.conv1 = ConvBlock(in_channels, out_channels) + self.conv2 = ConvBlock(out_channels, out_channels) + self.pool = nn.MaxPool2d(2) + + def forward(self, x): + skip = self.conv2(self.conv1(x)) + return self.pool(skip), skip + +class UpBlock(nn.Module): + """Upsampling block with transposed conv and two conv blocks""" + def __init__(self, in_channels, out_channels): + super().__init__() + self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2) + self.conv1 = ConvBlock(in_channels, out_channels) # in_channels because of concatenation + self.conv2 = ConvBlock(out_channels, out_channels) + + def forward(self, x, skip): + x = self.up(x) + # Handle dimension mismatch by interpolating x to match skip's size + if x.shape[2:] != skip.shape[2:]: + x = nn.functional.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) + return x + +class MyModel(nn.Module): + """U-Net style architecture for image inpainting""" + def __init__(self, n_in_channels: int, base_channels: int = 64): + super().__init__() + + # Initial convolution + self.init_conv = ConvBlock(n_in_channels, base_channels) + + # Encoder (downsampling path) + self.down1 = DownBlock(base_channels, base_channels * 2) + self.down2 = DownBlock(base_channels * 2, base_channels * 4) + self.down3 = DownBlock(base_channels * 4, base_channels * 8) + + # Bottleneck + self.bottleneck1 = ConvBlock(base_channels * 8, base_channels * 16) + self.bottleneck2 = ConvBlock(base_channels * 16, base_channels * 16) + + # Decoder (upsampling path) + self.up1 = UpBlock(base_channels * 16, base_channels * 8) + self.up2 = UpBlock(base_channels * 8, base_channels * 4) + self.up3 = UpBlock(base_channels * 4, base_channels * 2) + + # Final upsampling and output + self.final_up = nn.ConvTranspose2d(base_channels * 2, base_channels, kernel_size=2, stride=2) + self.final_conv1 = ConvBlock(base_channels * 2, base_channels) + self.final_conv2 = ConvBlock(base_channels, base_channels) + + # Output layer + self.output = nn.Conv2d(base_channels, 3, kernel_size=1) + self.sigmoid = nn.Sigmoid() # To ensure output is in [0, 1] range + + def forward(self, x): + # Initial convolution + x0 = self.init_conv(x) + + # Encoder + x1, skip1 = self.down1(x0) + x2, skip2 = self.down2(x1) + x3, skip3 = self.down3(x2) + + # Bottleneck + x = self.bottleneck1(x3) + x = self.bottleneck2(x) + + # Decoder with skip connections + x = self.up1(x, skip3) + x = self.up2(x, skip2) + x = self.up3(x, skip1) + + # Final layers + x = self.final_up(x) + # Handle dimension mismatch for final concatenation + if x.shape[2:] != x0.shape[2:]: + x = nn.functional.interpolate(x, size=x0.shape[2:], mode='bilinear', align_corners=False) + x = torch.cat([x, x0], dim=1) + x = self.final_conv1(x) + x = self.final_conv2(x) + + # Output + x = self.output(x) + x = self.sigmoid(x) + + return x \ No newline at end of file diff --git a/image-inpainting/src/datasets.py b/image-inpainting/src/datasets.py index e809603..d5e74eb 100644 --- a/image-inpainting/src/datasets.py +++ b/image-inpainting/src/datasets.py @@ -4,6 +4,7 @@ datasets.py """ +from torchvision import transforms import torch import numpy as np import random @@ -15,16 +16,25 @@ IMAGE_DIMENSION = 100 def create_arrays_from_image(image_array: np.ndarray, offset: tuple, spacing: tuple) -> tuple[np.ndarray, np.ndarray]: - image_array, known_array = None, None + image_array = np.transpose(image_array, (2, 0, 1)) + known_array = np.zeros_like(image_array) - # TODO: Implement the logic to create input and known arrays based on offset and spacing + known_array[:, offset[1]::spacing[1], offset[0]::spacing[0]] = 1 + + image_array[known_array == 0] = 0 + known_array = known_array[0:1] return image_array, known_array def resize(img: Image): - pass + resize_transforms = transforms.Compose([ + transforms.Resize((IMAGE_DIMENSION, IMAGE_DIMENSION)), + transforms.CenterCrop((IMAGE_DIMENSION, IMAGE_DIMENSION)) + ]) + return resize_transforms(img) def preprocess(input_array: np.ndarray): - pass + input_array = np.asarray(input_array, dtype=np.float32) / 255.0 + return input_array class ImageDataset(torch.utils.data.Dataset): """ @@ -38,6 +48,20 @@ class ImageDataset(torch.utils.data.Dataset): return len(self.imagefiles) def __getitem__(self, idx:int): - pass + index = int(idx) - # TODO: Implement the __init__, __len__, and __getitem__ methods \ No newline at end of file + image = Image.open(self.imagefiles[index]) + image = np.asarray(resize(image)) + image = preprocess(image) + spacing_x = random.randint(2,6) + spacing_y = random.randint(2,6) + offset_x = random.randint(0,8) + offset_y = random.randint(0,8) + spacing = (spacing_x, spacing_y) + offset = (offset_x, offset_y) + input_array, known_array = create_arrays_from_image(image.copy(), offset, spacing) + target_image = torch.from_numpy(np.transpose(image, (2,0,1))) + input_array = torch.from_numpy(input_array) + known_array = torch.from_numpy(known_array) + input_array = torch.cat((input_array, known_array), dim=0) + return input_array, target_image \ No newline at end of file diff --git a/image-inpainting/src/main.py b/image-inpainting/src/main.py index 1316cae..967a616 100644 --- a/image-inpainting/src/main.py +++ b/image-inpainting/src/main.py @@ -17,30 +17,34 @@ if __name__ == '__main__': config_dict['seed'] = 42 config_dict['testset_ratio'] = 0.1 config_dict['validset_ratio'] = 0.1 - config_dict['results_path'] = os.path.join("results") - config_dict['data_path'] = os.path.join("data", "dataset") + # Get the absolute path based on the script's location + script_dir = os.path.dirname(os.path.abspath(__file__)) + project_root = os.path.dirname(script_dir) + 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'] = 1e-3 + config_dict['learningrate'] = 5e-4 # Slightly lower for more stable training config_dict['weight_decay'] = 1e-5 # default is 0 - config_dict['n_updates'] = 50000 - config_dict['batchsize'] = 32 - config_dict['early_stopping_patience'] = 3 + config_dict['n_updates'] = 200 + config_dict['batchsize'] = 16 # Reduced due to larger model + config_dict['early_stopping_patience'] = 5 # More patience for complex model config_dict['use_wandb'] = False config_dict['print_train_stats_at'] = 10 config_dict['print_stats_at'] = 100 - config_dict['plot_at'] = 100 + config_dict['plot_at'] = 10 config_dict['validate_at'] = 100 network_config = { - 'n_in_channels': 4 + 'n_in_channels': 4, + 'base_channels': 32 # Start with 32, can increase to 64 for even better results } config_dict['network_config'] = network_config train(**config_dict) - testset_path = os.path.join("data", "challenge_testset.npz") + testset_path = os.path.join(project_root, "data", "challenge_testset.npz") state_dict_path = os.path.join(config_dict['results_path'], "best_model.pt") save_path = os.path.join(config_dict['results_path'], "testset", "my_submission_name.npz") plot_path = os.path.join(config_dict['results_path'], "testset", "plots") diff --git a/image-inpainting/src/utils.py b/image-inpainting/src/utils.py index 27603fd..e8e91f5 100644 --- a/image-inpainting/src/utils.py +++ b/image-inpainting/src/utils.py @@ -18,7 +18,7 @@ def plot(inputs, targets, predictions, path, update): os.makedirs(path, exist_ok=True) fig, axes = plt.subplots(ncols=3, figsize=(15, 5)) - for i in range(len(inputs)): + for i in range(5): for ax, data, title in zip(axes, [inputs, targets, predictions], ["Input", "Target", "Prediction"]): ax.clear() ax.set_title(title)