{"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': 3,  # 5         # Adjust based on time constraints\n    'lr': 1e-3,   # 1e-4\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 to start with identity mapping\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        # Output is the flow field (displacement)\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] (values are offsets in normalized coordinates [-1, 1])\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    # flow is usually predicted in range suitable for [-1, 1] if tanh used, \n    # but here we predicted raw logits. Let's clamp or tanh if needed, \n    # or just assume network learns the scale.\n    # To stabilize, we usually assume flow is small.\n    \n    sampling_grid = grid + flow\n    \n    # Permute grid to [B, H, W, 2] for grid_sample\n    sampling_grid = sampling_grid.permute(0, 2, 3, 1)\n    \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    \"\"\"\n    Apply the low-res flow predicted by model to the high-res original image.\n    Using OpenCV Remap for efficiency on CPU after prediction.\n    \"\"\"\n    # Load High Res Image\n    img_cv = cv2.imread(img_path)\n    h_orig, w_orig = img_cv.shape[:2]\n    \n    # Flow is [1, 2, Model_H, Model_W]\n    # We need to resize flow to [h_orig, w_orig]\n    flow_np = flow_tensor.detach().cpu().numpy()[0] # [2, H, W]\n    flow_np = np.transpose(flow_np, (1, 2, 0))      # [H, W, 2]\n    \n    # Resize flow to original resolution\n    # Note: The flow values are in Normalized Device Coordinates [-1, 1].\n    # We just need to resize the map; the values remain relative to image dimensions.\n    flow_resized = cv2.resize(flow_np, (w_orig, h_orig), interpolation=cv2.INTER_LINEAR)\n    \n    # Convert Normalized Flow [-1, 1] to Absolute Pixel Coordinates\n    # Grid Sample logic: -1 is left/top, 1 is right/bottom.\n    # pixel_x = (norm_x + 1) * 0.5 * (W-1)\n    \n    # Create absolute meshgrid\n    grid_x, grid_y = np.meshgrid(np.arange(w_orig), np.arange(h_orig))\n    \n    # Extract deltas\n    delta_x_norm = flow_resized[:, :, 0]\n    delta_y_norm = flow_resized[:, :, 1]\n    \n    # Convert normalized deltas to pixel deltas roughly\n    # Actually, grid_sample takes (grid + flow). \n    # So map_x = (grid_x_norm + flow_x) converted to pixels\n    \n    # Let's do it rigorously:\n    # 1. Normalize current pixel coords to [-1, 1]\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    # 2. Add flow\n    sample_x = norm_x + delta_x_norm\n    sample_y = norm_y + delta_y_norm\n    \n    # 3. Convert back to pixels\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    # Remap\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\n# ==========================================\ndef train():\n    dataset = LensDataset(CONFIG['train_dir'], CONFIG['img_size'], is_train=True)\n    dataloader = DataLoader(dataset, batch_size=CONFIG['batch_size'], shuffle=True, num_workers=CONFIG['num_workers'])\n    \n    model = LensUNet().to(CONFIG['device'])\n    optimizer = torch.optim.Adam(model.parameters(), lr=CONFIG['lr'])\n    criterion = nn.L1Loss() # L1 preserves edges better than MSE\n    \n    print(f\"Starting training on {len(dataset)} pairs for {CONFIG['epochs']} epochs...\")\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'])\n            img_gt = img_gt.to(CONFIG['device'])\n            \n            optimizer.zero_grad()\n            \n            # Predict Flow\n            flow = model(img_in)\n            \n            # Warp Input\n            warped = warp_image(img_in, flow)\n            \n            # Loss: Reconstruction + Smoothness (TV Loss)\n            recon_loss = criterion(warped, img_gt)\n            \n            # Total Variation loss on flow to prevent tearing/noise\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            loss.backward()\n            optimizer.step()\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 on resized image\n            flow = model(img_tensor)\n            \n            # Apply flow to original high-res image\n            # raw_path is a tuple in dataloader (batch size 1), extract string\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-22T15:42:34.999288Z","iopub.execute_input":"2026-02-22T15:42:34.999609Z"}},"outputs":[],"execution_count":null}]}