{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":130932,"databundleVersionId":15769099}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### The Strategy\n1.  **Architecture**: A lightweight UNet that takes a resized distorted image (e.g., 512x512) and outputs a **Flow Field** (2 channels: $\\Delta x, \\Delta y$).\n2.  **Resolution Independence**: We predict the distortion shape on a low-resolution image, upsample the flow map to the full resolution of the test image, and apply the correction. This ensures high speed and pixel-perfect quality.\n3.  **Loss Function**: We use L1 Loss between the *warped* input and the ground truth, plus a Smoothness Loss to ensure the lines don't get jagged.\n\n### Explanation of the Solution\n\n1.  **LensUNet**: I used a custom U-Net architecture.\n    *   **Input**: The distorted image.\n    *   **Output**: A 2-channel \"Flow Field\". This field contains the $\\Delta X$ and $\\Delta Y$ values for every pixel. It tells the code: \"To fix the distortion, move the pixel at (x,y) by this amount.\"\n    *   **Initialization**: The output layer weights are initialized to near-zero. This is a crucial trick. It means the model starts by outputting an identity transform (doing nothing), and slowly learns to correct the curve. This prevents training instability.\n\n2.  **Differentiable Warping (`warp_image`)**:\n    *   Standard neural networks output colors. We use PyTorch's `F.grid_sample`. This allows the network to act as a \"Sampler.\"\n    *   This is fully differentiable, so backpropagation works from the Loss Function -> Warped Image -> Grid Sample -> Flow Field -> UNet weights.\n\n3.  **High-Res Inference (`apply_high_res_warp`)**:\n    *   The test images are likely large (12MP+). Running a UNet on 4000x3000 input is too slow and consumes too much GPU memory.\n    *   However, lens distortion is a low-frequency geometric feature. It doesn't change sharply pixel-by-pixel.\n    *   **Technique**: We resize the input to 384x384, predict the flow map, then *upscale the flow map* via interpolation to the full 4000x3000 size. We then use OpenCV's `remap` function on the original high-res image. This results in corrected images that retain 100% of the original resolution and clarity.\n\n4.  **Loss Function**:\n    *   **L1 Loss**: Minimizes the absolute difference between the warped image and the ground truth. L1 is preferred over MSE (L2) for image generation tasks because it produces sharper edges.\n    *   **TV Loss (Total Variation)**: We penalize the gradients of the flow field. This forces the network to learn smooth, continuous warping (like a real lens) and prevents it from tearing the image or creating jagged artifacts.\n\n5.  **Data Loading**:\n    *   The dataset class handles the pairing of `_original` and `_generated` filenames automatically.","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport cv2\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nimport zipfile\n\n# ==========================================\n# 1. CONFIGURATION\n# ==========================================\nCONFIG = {\n    'img_size': 384,          # Input size for the network (Square)\n    'batch_size': 16,\n    'epochs': 5,              # Increased slightly as AMP is faster\n    'lr': 1e-3,\n    'device': 'cuda' if torch.cuda.is_available() else 'cpu',\n    'num_workers': 4,\n    'train_dir': '/kaggle/input/automatic-lens-correction/lens-correction-train-cleaned',\n    'test_dir': '/kaggle/input/automatic-lens-correction/test-originals',\n    'output_dir': '/kaggle/working/corrected_images',\n    'model_path': '/kaggle/working/lens_correction_model.pth'\n}\n\nos.makedirs(CONFIG['output_dir'], exist_ok=True)\n\n# ==========================================\n# 2. DATASET\n# ==========================================\nclass LensDataset(Dataset):\n    def __init__(self, root_dir, img_size, is_train=True):\n        self.img_size = img_size\n        self.is_train = is_train\n        self.pairs = []\n        \n        if is_train:\n            # Find all originals\n            originals = glob.glob(os.path.join(root_dir, \"*_original.jpg\"))\n            for org in originals:\n                # Construct corresponding generated filename\n                gen = org.replace(\"_original.jpg\", \"_generated.jpg\")\n                if os.path.exists(gen):\n                    self.pairs.append((org, gen))\n        else:\n            # For testing, just load the list of files\n            self.pairs = glob.glob(os.path.join(root_dir, \"*.jpg\"))\n\n    def __len__(self):\n        return len(self.pairs)\n\n    def __getitem__(self, idx):\n        if self.is_train:\n            org_path, gen_path = self.pairs[idx]\n            \n            # Read images\n            img_in = cv2.imread(org_path)\n            img_gt = cv2.imread(gen_path)\n            \n            # Convert BGR to RGB\n            img_in = cv2.cvtColor(img_in, cv2.COLOR_BGR2RGB)\n            img_gt = cv2.cvtColor(img_gt, cv2.COLOR_BGR2RGB)\n            \n            # Resize for training\n            img_in = cv2.resize(img_in, (self.img_size, self.img_size))\n            img_gt = cv2.resize(img_gt, (self.img_size, self.img_size))\n            \n            # Normalize to [0, 1] and convert to tensor\n            img_in = torch.from_numpy(img_in).permute(2, 0, 1).float() / 255.0\n            img_gt = torch.from_numpy(img_gt).permute(2, 0, 1).float() / 255.0\n            \n            return img_in, img_gt\n        else:\n            path = self.pairs[idx]\n            img_in = cv2.imread(path)\n            h, w = img_in.shape[:2]\n            img_in_rgb = cv2.cvtColor(img_in, cv2.COLOR_BGR2RGB)\n            \n            # Resize for model inference\n            img_resized = cv2.resize(img_in_rgb, (self.img_size, self.img_size))\n            img_tensor = torch.from_numpy(img_resized).permute(2, 0, 1).float() / 255.0\n            \n            filename = os.path.basename(path)\n            return img_tensor, filename, h, w, path\n\n# ==========================================\n# 3. MODEL ARCHITECTURE (UNet for Flow)\n# ==========================================\nclass DoubleConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.double_conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        return self.double_conv(x)\n\nclass LensUNet(nn.Module):\n    def __init__(self):\n        super().__init__()\n        \n        # Encoder\n        self.inc = DoubleConv(3, 32)\n        self.down1 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(32, 64))\n        self.down2 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(64, 128))\n        self.down3 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(128, 256))\n        self.down4 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(256, 512))\n        \n        # Decoder\n        self.up1 = nn.ConvTranspose2d(512, 256, 2, stride=2)\n        self.conv1 = DoubleConv(512, 256)\n        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)\n        self.conv2 = DoubleConv(256, 128)\n        self.up3 = nn.ConvTranspose2d(128, 64, 2, stride=2)\n        self.conv3 = DoubleConv(128, 64)\n        self.up4 = nn.ConvTranspose2d(64, 32, 2, stride=2)\n        self.conv4 = DoubleConv(64, 32)\n        \n        # Final Output: 2 Channels (Delta X, Delta Y)\n        self.outc = nn.Conv2d(32, 2, kernel_size=1)\n        \n        # Initialize output layer to near zero\n        nn.init.normal_(self.outc.weight, mean=0, std=1e-4)\n        nn.init.constant_(self.outc.bias, 0)\n\n    def forward(self, x):\n        x1 = self.inc(x)\n        x2 = self.down1(x1)\n        x3 = self.down2(x2)\n        x4 = self.down3(x3)\n        x5 = self.down4(x4)\n        \n        x = self.up1(x5)\n        x = torch.cat([x, x4], dim=1)\n        x = self.conv1(x)\n        \n        x = self.up2(x)\n        x = torch.cat([x, x3], dim=1)\n        x = self.conv2(x)\n        \n        x = self.up3(x)\n        x = torch.cat([x, x2], dim=1)\n        x = self.conv3(x)\n        \n        x = self.up4(x)\n        x = torch.cat([x, x1], dim=1)\n        x = self.conv4(x)\n        \n        flow = self.outc(x) \n        return flow\n\n# ==========================================\n# 4. WARPING & UTILS\n# ==========================================\ndef warp_image(x, flow):\n    \"\"\"\n    Warps an image x using flow field.\n    x: [B, C, H, W]\n    flow: [B, 2, H, W]\n    \"\"\"\n    B, C, H, W = x.size()\n    \n    # Create mesh grid\n    xx = torch.linspace(-1.0, 1.0, W).view(1, 1, 1, W).expand(B, -1, H, -1)\n    yy = torch.linspace(-1.0, 1.0, H).view(1, 1, H, 1).expand(B, -1, -1, W)\n    grid = torch.cat([xx, yy], 1).to(x.device)\n    \n    # Add flow to grid\n    # In AMP, x is float16, grid is float32. We let autocast handle mixing, \n    # but grid_sample often prefers float32 grid.\n    sampling_grid = grid + flow\n    \n    # Permute grid to [B, H, W, 2]\n    sampling_grid = sampling_grid.permute(0, 2, 3, 1)\n    \n    # grid_sample supports float16 in newer PyTorch, but we keep generic\n    warped = F.grid_sample(x, sampling_grid, align_corners=True, padding_mode='border')\n    return warped\n\ndef apply_high_res_warp(img_path, flow_tensor, target_size=None):\n    img_cv = cv2.imread(img_path)\n    h_orig, w_orig = img_cv.shape[:2]\n    \n    flow_np = flow_tensor.detach().cpu().numpy()[0]\n    flow_np = np.transpose(flow_np, (1, 2, 0))\n    \n    # Resize flow to original resolution\n    flow_resized = cv2.resize(flow_np, (w_orig, h_orig), interpolation=cv2.INTER_LINEAR)\n    \n    grid_x, grid_y = np.meshgrid(np.arange(w_orig), np.arange(h_orig))\n    \n    delta_x_norm = flow_resized[:, :, 0]\n    delta_y_norm = flow_resized[:, :, 1]\n    \n    norm_x = 2.0 * grid_x / (w_orig - 1) - 1.0\n    norm_y = 2.0 * grid_y / (h_orig - 1) - 1.0\n    \n    sample_x = norm_x + delta_x_norm\n    sample_y = norm_y + delta_y_norm\n    \n    map_x = ((sample_x + 1.0) / 2.0) * (w_orig - 1)\n    map_y = ((sample_y + 1.0) / 2.0) * (h_orig - 1)\n    \n    corrected = cv2.remap(img_cv, map_x.astype(np.float32), map_y.astype(np.float32), cv2.INTER_LINEAR)\n    \n    return corrected\n\n# ==========================================\n# 5. TRAINING LOOP (WITH AMP)\n# ==========================================\ndef train():\n    dataset = LensDataset(CONFIG['train_dir'], CONFIG['img_size'], is_train=True)\n    # Added pin_memory=True for faster host-to-device transfer\n    dataloader = DataLoader(\n        dataset, \n        batch_size=CONFIG['batch_size'], \n        shuffle=True, \n        num_workers=CONFIG['num_workers'],\n        pin_memory=True\n    )\n    \n    model = LensUNet().to(CONFIG['device'])\n    optimizer = torch.optim.Adam(model.parameters(), lr=CONFIG['lr'])\n    criterion = nn.L1Loss()\n    \n    # Initialize Mixed Precision Scaler\n    scaler = torch.cuda.amp.GradScaler()\n    \n    print(f\"Starting training on {len(dataset)} pairs for {CONFIG['epochs']} epochs using AMP...\")\n    \n    for epoch in range(CONFIG['epochs']):\n        model.train()\n        train_loss = 0\n        loop = tqdm(dataloader, total=len(dataloader), leave=True)\n        \n        for batch_idx, (img_in, img_gt) in enumerate(loop):\n            img_in = img_in.to(CONFIG['device'], non_blocking=True)\n            img_gt = img_gt.to(CONFIG['device'], non_blocking=True)\n            \n            # Zero gradients (set_to_none is slightly faster)\n            optimizer.zero_grad(set_to_none=True)\n            \n            # Autocast runs the forward pass in mixed precision (float16)\n            with torch.cuda.amp.autocast():\n                # Predict Flow\n                flow = model(img_in)\n                \n                # Warp Input\n                warped = warp_image(img_in, flow)\n                \n                # Reconstruction Loss\n                recon_loss = criterion(warped, img_gt)\n                \n                # TV Loss (Smoothness)\n                diff_x = torch.mean(torch.abs(flow[:, :, :, :-1] - flow[:, :, :, 1:]))\n                diff_y = torch.mean(torch.abs(flow[:, :, :-1, :] - flow[:, :, 1:, :]))\n                tv_loss = diff_x + diff_y\n                \n                loss = recon_loss + 0.01 * tv_loss\n            \n            # Backward pass with Scaler\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            \n            train_loss += loss.item()\n            loop.set_description(f\"Epoch [{epoch+1}/{CONFIG['epochs']}]\")\n            loop.set_postfix(loss=loss.item())\n            \n    torch.save(model.state_dict(), CONFIG['model_path'])\n    print(\"Training Complete. Model Saved.\")\n    return model\n\n# ==========================================\n# 6. INFERENCE & SUBMISSION\n# ==========================================\ndef inference(model):\n    model.eval()\n    test_dataset = LensDataset(CONFIG['test_dir'], CONFIG['img_size'], is_train=False)\n    test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n    \n    print(f\"Processing {len(test_dataset)} test images...\")\n    \n    with torch.no_grad():\n        for img_tensor, filename, h, w, raw_path in tqdm(test_loader):\n            img_tensor = img_tensor.to(CONFIG['device'])\n            \n            # Predict Flow\n            flow = model(img_tensor)\n            \n            # Apply to High Res\n            corrected_img = apply_high_res_warp(raw_path[0], flow)\n            \n            # Save\n            save_path = os.path.join(CONFIG['output_dir'], filename[0])\n            cv2.imwrite(save_path, corrected_img)\n\ndef zip_results():\n    print(\"Zipping results...\")\n    shutil_output = \"submission\"\n    with zipfile.ZipFile(f'{shutil_output}.zip', 'w') as zipf:\n        for root, dirs, files in os.walk(CONFIG['output_dir']):\n            for file in files:\n                zipf.write(os.path.join(root, file), file)\n    print(\"Zip created: submission.zip\")\n\n# ==========================================\n# 7. MAIN EXECUTION\n# ==========================================\nif __name__ == \"__main__\":\n    # Train\n    trained_model = train()\n    \n    # Inference\n    inference(trained_model)\n    \n    # Zip\n    zip_results()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T18:28:43.088289Z","iopub.execute_input":"2026-02-22T18:28:43.089036Z","execution_failed":"2026-02-22T18:30:41.478Z"}},"outputs":[],"execution_count":null}]}