{"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":"gpu","dataSources":[{"sourceType":"competition","sourceId":130932,"databundleVersionId":15769099},{"sourceType":"datasetVersion","sourceId":14920014,"datasetId":9546635,"databundleVersionId":15786526}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# import os\n# import glob\n# import cv2\n# import math\n# import torch\n# import torch.nn as nn\n# import torch.nn.functional as F\n# import numpy as np\n# from torch.utils.data import Dataset, DataLoader\n# import albumentations as A\n# from albumentations.pytorch import ToTensorV2\n# import time\n\n# # ==========================================\n# # 1. DATASET CLASS (डेटा लोडर)\n# # ==========================================\n# class LensDistortionKaggleDataset(Dataset):\n#     def __init__(self, data_dir, image_size=256):\n#         self.data_dir = data_dir\n#         self.image_size = image_size\n#         self.original_files = glob.glob(os.path.join(data_dir, \"*_original.jpg\"))\n        \n#         print(f\"Total training pairs found: {len(self.original_files)}\")\n        \n#         self.transform = A.Compose([\n#             A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\n#             ToTensorV2()\n#         ], additional_targets={'target_image': 'image'})\n\n#     def __len__(self):\n#         return len(self.original_files)\n\n#     def __getitem__(self, idx):\n#         distorted_path = self.original_files[idx]\n#         corrected_path = distorted_path.replace(\"_original.jpg\", \"_generated.jpg\")\n        \n#         distorted_img = cv2.imread(distorted_path)\n#         corrected_img = cv2.imread(corrected_path)\n        \n#         if distorted_img is None or corrected_img is None:\n#             distorted_img = np.zeros((self.image_size, self.image_size, 3), dtype=np.uint8)\n#             corrected_img = np.zeros((self.image_size, self.image_size, 3), dtype=np.uint8)\n#         else:\n#             distorted_img = cv2.cvtColor(distorted_img, cv2.COLOR_BGR2RGB)\n#             corrected_img = cv2.cvtColor(corrected_img, cv2.COLOR_BGR2RGB)\n#             distorted_img = cv2.resize(distorted_img, (self.image_size, self.image_size))\n#             corrected_img = cv2.resize(corrected_img, (self.image_size, self.image_size))\n        \n#         transformed = self.transform(image=distorted_img, target_image=corrected_img)\n#         return transformed['image'], transformed['target_image']\n\n# # ==========================================\n# # 2. MODEL ARCHITECTURE (U-Net + STN)\n# # ==========================================\n# class DoubleConv(nn.Module):\n#     def __init__(self, in_channels, out_channels):\n#         super().__init__()\n#         self.conv = nn.Sequential(\n#             nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),\n#             nn.BatchNorm2d(out_channels),\n#             nn.ReLU(inplace=True),\n#             nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),\n#             nn.BatchNorm2d(out_channels),\n#             nn.ReLU(inplace=True)\n#         )\n#     def forward(self, x): return self.conv(x)\n\n# class DistortionCorrectionUNet(nn.Module):\n#     def __init__(self, in_channels=3, out_channels=2):\n#         super().__init__()\n#         self.down1 = DoubleConv(in_channels, 32)\n#         self.down2 = DoubleConv(32, 64)\n#         self.down3 = DoubleConv(64, 128)\n#         self.down4 = DoubleConv(128, 256)\n#         self.pool = nn.MaxPool2d(2)\n#         self.bottleneck = DoubleConv(256, 512)\n        \n#         self.up1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n#         self.conv_up1 = DoubleConv(512, 256)\n#         self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n#         self.conv_up2 = DoubleConv(256, 128)\n#         self.up3 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n#         self.conv_up3 = DoubleConv(128, 64)\n#         self.up4 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)\n#         self.conv_up4 = DoubleConv(64, 32)\n        \n#         self.out_conv = nn.Conv2d(32, out_channels, kernel_size=3, padding=1)\n#         nn.init.zeros_(self.out_conv.weight) \n#         nn.init.zeros_(self.out_conv.bias)\n\n#     def forward(self, x):\n#         d1 = self.down1(x)\n#         d2 = self.down2(self.pool(d1))\n#         d3 = self.down3(self.pool(d2))\n#         d4 = self.down4(self.pool(d3))\n#         b = self.bottleneck(self.pool(d4))\n        \n#         u1 = self.up1(b)\n#         u1 = self.conv_up1(torch.cat([u1, d4], dim=1))\n#         u2 = self.up2(u1)\n#         u2 = self.conv_up2(torch.cat([u2, d3], dim=1))\n#         u3 = self.up3(u2)\n#         u3 = self.conv_up3(torch.cat([u3, d2], dim=1))\n#         u4 = self.up4(u3)\n#         u4 = self.conv_up4(torch.cat([u4, d1], dim=1))\n        \n#         flow = self.out_conv(u4) \n        \n#         B, C, H, W = x.shape\n#         yy, xx = torch.meshgrid(torch.linspace(-1, 1, H, device=x.device), \n#                                 torch.linspace(-1, 1, W, device=x.device), indexing='ij')\n#         identity_grid = torch.stack([xx, yy], dim=-1).unsqueeze(0).repeat(B, 1, 1, 1)\n        \n#         flow = flow.permute(0, 2, 3, 1)\n#         warped_grid = identity_grid + flow\n        \n#         corrected_image = F.grid_sample(x, warped_grid, mode='bilinear', padding_mode='border', align_corners=True)\n#         return corrected_image, flow\n\n# # ==========================================\n# # 3. CUSTOM LOSS FUNCTIONS (ईमानदार मॉडल के लिए)\n# # ==========================================\n# class EdgeLoss(nn.Module):\n#     def __init__(self):\n#         super(EdgeLoss, self).__init__()\n#         self.sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3)\n#         self.sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).view(1, 1, 3, 3)\n\n#     def forward(self, pred, target):\n#         pred_gray = pred.mean(dim=1, keepdim=True)\n#         target_gray = target.mean(dim=1, keepdim=True)\n\n#         self.sobel_x = self.sobel_x.to(pred.device)\n#         self.sobel_y = self.sobel_y.to(pred.device)\n\n#         pred_edge_x = F.conv2d(pred_gray, self.sobel_x, padding=1)\n#         pred_edge_y = F.conv2d(pred_gray, self.sobel_y, padding=1)\n#         target_edge_x = F.conv2d(target_gray, self.sobel_x, padding=1)\n#         target_edge_y = F.conv2d(target_gray, self.sobel_y, padding=1)\n\n#         return F.l1_loss(pred_edge_x, target_edge_x) + F.l1_loss(pred_edge_y, target_edge_y)\n\n# class SSIMLoss(nn.Module):\n#     def __init__(self, window_size=11, channel=3):\n#         super(SSIMLoss, self).__init__()\n#         self.window_size = window_size\n#         self.channel = channel\n#         self.window = self.create_window(window_size, self.channel)\n\n#     def gaussian(self, window_size, sigma):\n#         gauss = torch.tensor([math.exp(-(x - window_size//2)**2 / float(2 * sigma**2)) for x in range(window_size)])\n#         return gauss / gauss.sum()\n\n#     def create_window(self, window_size, channel):\n#         _1D_window = self.gaussian(window_size, 1.5).unsqueeze(1)\n#         _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)\n#         return _2D_window.expand(channel, 1, window_size, window_size).contiguous()\n\n#     def forward(self, img1, img2):\n#         L = 2.0 \n#         C1 = (0.01 * L) ** 2\n#         C2 = (0.03 * L) ** 2\n\n#         window = self.window.to(img1.device)\n\n#         mu1 = F.conv2d(img1, window, padding=self.window_size//2, groups=self.channel)\n#         mu2 = F.conv2d(img2, window, padding=self.window_size//2, groups=self.channel)\n#         mu1_sq, mu2_sq, mu1_mu2 = mu1.pow(2), mu2.pow(2), mu1 * mu2\n\n#         sigma1_sq = F.conv2d(img1 * img1, window, padding=self.window_size//2, groups=self.channel) - mu1_sq\n#         sigma2_sq = F.conv2d(img2 * img2, window, padding=self.window_size//2, groups=self.channel) - mu2_sq\n#         sigma12 = F.conv2d(img1 * img2, window, padding=self.window_size//2, groups=self.channel) - mu1_mu2\n\n#         ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))\n#         return 1 - ssim_map.mean()\n\n# # ==========================================\n# # 4. TRAINING LOOP (ट्रेनिंग प्रक्रिया)\n# # ==========================================\n# def train_model():\n#     DATA_DIR = '/kaggle/input/automatic-lens-correction/lens-correction-train-cleaned' \n#     BATCH_SIZE = 8 \n#     EPOCHS = 8\n#     LEARNING_RATE = 1e-4\n#     IMAGE_SIZE = 256 \n    \n#     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n#     print(f\"Using device: {device}\")\n\n#     dataset = LensDistortionKaggleDataset(data_dir=DATA_DIR, image_size=IMAGE_SIZE)\n#     if len(dataset) == 0:\n#         print(f\"Error: No images found. Check folder path.\")\n#         return\n        \n#     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\n\n#     model = DistortionCorrectionUNet().to(device)\n    \n#     # तीनों लॉस फंक्शन्स को इनिशियलाइज़ करना\n#     l1_criterion = nn.L1Loss()\n#     edge_criterion = EdgeLoss()\n#     ssim_criterion = SSIMLoss()\n    \n#     optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)\n\n#     print(\"--- Advanced Training Started with Hybrid Loss (L1 + Edge + SSIM) ---\")\n#     for epoch in range(EPOCHS):\n#         model.train()\n#         running_loss = 0.0\n#         start_time = time.time()\n        \n#         for batch_idx, (distorted, ground_truth) in enumerate(dataloader):\n#             distorted = distorted.to(device)\n#             ground_truth = ground_truth.to(device)\n\n#             optimizer.zero_grad()\n            \n#             # फॉरवर्ड पास\n#             corrected_image, flow = model(distorted)\n            \n#             # 1. पिक्सेल लॉस (कलर के लिए)\n#             loss_pixel = l1_criterion(corrected_image, ground_truth)\n#             # 2. एज लॉस (सीधी रेखाओं के लिए)\n#             loss_edge = edge_criterion(corrected_image, ground_truth)\n#             # 3. SSIM लॉस (इमेज क्वालिटी/टेक्सचर के लिए)\n#             loss_ssim = ssim_criterion(corrected_image, ground_truth)\n\n#             # === हाइब्रिड लॉस (अल्टीमेट फॉर्मूला) ===\n#             loss = loss_pixel + (0.5 * loss_edge) + (0.5 * loss_ssim)\n            \n#             # बैकवर्ड पास\n#             loss.backward()\n#             optimizer.step()\n\n#             running_loss += loss.item()\n            \n#             if batch_idx % 50 == 0:\n#                 print(f\"Epoch [{epoch+1}/{EPOCHS}] | Batch [{batch_idx}/{len(dataloader)}]\")\n#                 print(f\"Total Loss: {loss.item():.4f} (Pixel: {loss_pixel.item():.4f} | Edge: {loss_edge.item():.4f} | SSIM: {loss_ssim.item():.4f})\")\n\n#         epoch_time = time.time() - start_time\n#         avg_loss = running_loss / len(dataloader)\n#         print(f\"==> Epoch {epoch+1} Completed | Average Loss: {avg_loss:.4f} | Time: {epoch_time:.2f} sec\")\n        \n#         torch.save(model.state_dict(), f\"lens_correction_model_epoch_{epoch+1}.pth\")\n#         print(\"Model saved!\\n\")\n\n# if __name__ == \"__main__\":\n#     train_model()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T09:22:24.335705Z","iopub.execute_input":"2026-02-22T09:22:24.336481Z","execution_failed":"2026-02-22T09:22:58.493Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport glob\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nimport zipfile\n\n# ==========================================\n# 1. MODEL ARCHITECTURE (जो ट्रेनिंग में यूज़ हुआ था)\n# ==========================================\nclass DoubleConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n    def forward(self, x): return self.conv(x)\n\nclass DistortionCorrectionUNet(nn.Module):\n    def __init__(self, in_channels=3, out_channels=2):\n        super().__init__()\n        self.down1 = DoubleConv(in_channels, 32)\n        self.down2 = DoubleConv(32, 64)\n        self.down3 = DoubleConv(64, 128)\n        self.down4 = DoubleConv(128, 256)\n        self.pool = nn.MaxPool2d(2)\n        self.bottleneck = DoubleConv(256, 512)\n        \n        self.up1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n        self.conv_up1 = DoubleConv(512, 256)\n        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n        self.conv_up2 = DoubleConv(256, 128)\n        self.up3 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n        self.conv_up3 = DoubleConv(128, 64)\n        self.up4 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)\n        self.conv_up4 = DoubleConv(64, 32)\n        \n        self.out_conv = nn.Conv2d(32, out_channels, kernel_size=3, padding=1)\n\n    def forward(self, x):\n        d1 = self.down1(x)\n        d2 = self.down2(self.pool(d1))\n        d3 = self.down3(self.pool(d2))\n        d4 = self.down4(self.pool(d3))\n        b = self.bottleneck(self.pool(d4))\n        \n        u1 = self.up1(b)\n        u1 = self.conv_up1(torch.cat([u1, d4], dim=1))\n        u2 = self.up2(u1)\n        u2 = self.conv_up2(torch.cat([u2, d3], dim=1))\n        u3 = self.up3(u2)\n        u3 = self.conv_up3(torch.cat([u3, d2], dim=1))\n        u4 = self.up4(u3)\n        u4 = self.conv_up4(torch.cat([u4, d1], dim=1))\n        \n        flow = self.out_conv(u4) \n        \n        B, C, H, W = x.shape\n        yy, xx = torch.meshgrid(torch.linspace(-1, 1, H, device=x.device), \n                                torch.linspace(-1, 1, W, device=x.device), indexing='ij')\n        identity_grid = torch.stack([xx, yy], dim=-1).unsqueeze(0).repeat(B, 1, 1, 1)\n        \n        flow = flow.permute(0, 2, 3, 1)\n        warped_grid = identity_grid + flow\n        \n        corrected_image = F.grid_sample(x, warped_grid, mode='bilinear', padding_mode='border', align_corners=True)\n        return corrected_image, flow\n\n# ==========================================\n# 2. INFERENCE & ZIP CREATION (टेस्ट करना और फाइल बनाना)\n# ==========================================\ndef generate_test_predictions():\n    # --- सेटिंग्स ---\n    TEST_DIR = '/kaggle/input/automatic-lens-correction/test-originals'       # 1000 टेस्ट इमेजेज वाला फोल्डर\n    OUTPUT_DIR = '/kaggle/working/test-corrected'     # जहाँ सही की गई इमेजेज सेव होंगी\n    MODEL_WEIGHTS = '/kaggle/input/datasets/dheerajkannaujiya/lens-correction-trained-model/lens_correction_model_epoch_8.pth' # सबसे बेस्ट/आखिरी मॉडल\n    ZIP_FILENAME = 'submission_images.zip'\n    IMAGE_SIZE = 256 # जो ट्रेनिंग में यूज़ किया था\n\n    # आउटपुट फोल्डर बनाना\n    os.makedirs(OUTPUT_DIR, exist_ok=True)\n    \n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    print(f\"Using device: {device}\")\n    \n    # मॉडल लोड करना\n    model = DistortionCorrectionUNet().to(device)\n    if os.path.exists(MODEL_WEIGHTS):\n        model.load_state_dict(torch.load(MODEL_WEIGHTS, map_location=device))\n        print(f\"✅ Model '{MODEL_WEIGHTS}' loaded successfully!\")\n    else:\n        print(f\"❌ Error: Model weights '{MODEL_WEIGHTS}' not found.\")\n        return\n\n    model.eval() # इवैल्यूएशन मोड\n    \n    test_files = glob.glob(os.path.join(TEST_DIR, \"*.jpg\"))\n    if len(test_files) == 0:\n        print(f\"❌ Error: No images found in '{TEST_DIR}' folder.\")\n        return\n        \n    print(f\"Found {len(test_files)} test images. Starting correction...\")\n\n    # ज़िप फाइल तैयार करना\n    with zipfile.ZipFile(ZIP_FILENAME, 'w') as zipf:\n        with torch.no_grad(): # टेस्टिंग में ग्रेडिएंट्स बंद रखते हैं\n            for i, img_path in enumerate(test_files):\n                filename = os.path.basename(img_path)\n                \n                # 1. इमेज लोड और रिसाइज़\n                img = cv2.imread(img_path)\n                original_h, original_w = img.shape[:2] # ऑरिजनल साइज़ सेव करें\n                \n                img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n                img_resized = cv2.resize(img_rgb, (IMAGE_SIZE, IMAGE_SIZE))\n                \n                # नॉर्मलाइज़ेशन ([-1, 1] रेंज)\n                img_normalized = (img_resized / 255.0 - 0.5) / 0.5\n                img_tensor = torch.tensor(img_normalized, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(device)\n                \n                # 2. मॉडल से प्रेडिक्शन लेना\n                corrected_tensor, _ = model(img_tensor)\n                \n                # 3. पोस्ट-प्रोसेसिंग\n                corrected_img = corrected_tensor.squeeze().permute(1, 2, 0).cpu().numpy()\n                corrected_img = (corrected_img * 0.5 + 0.5) * 255.0 # डी-नॉर्मलाइज़\n                corrected_img = np.clip(corrected_img, 0, 255).astype(np.uint8)\n                \n                # BGR में बदलना और ऑरिजनल रिज़ॉल्यूशन में वापस लाना\n                corrected_img_bgr = cv2.cvtColor(corrected_img, cv2.COLOR_RGB2BGR)\n                corrected_img_final = cv2.resize(corrected_img_bgr, (original_w, original_h))\n                \n                # 4. सेव करना और ज़िप में डालना\n                save_path = os.path.join(OUTPUT_DIR, filename)\n                cv2.imwrite(save_path, corrected_img_final)\n                zipf.write(save_path, arcname=filename)\n                \n                if (i + 1) % 100 == 0:\n                    print(f\"Processed {i + 1}/{len(test_files)} images...\")\n\n    print(f\"\\n🎉 Success! All images corrected and saved to '{OUTPUT_DIR}'.\")\n    print(f\"📦 Zip file created: '{ZIP_FILENAME}'.\")\n    print(f\"➡️ Now upload '{ZIP_FILENAME}' to bounty.autohdr.com\")\n\nif __name__ == \"__main__\":\n    generate_test_predictions()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T13:29:39.887434Z","iopub.execute_input":"2026-02-22T13:29:39.888187Z","iopub.status.idle":"2026-02-22T13:30:35.804865Z","shell.execute_reply.started":"2026-02-22T13:29:39.888156Z","shell.execute_reply":"2026-02-22T13:30:35.804105Z"}},"outputs":[],"execution_count":null}]}