{"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":"none","dataSources":[{"sourceType":"competition","sourceId":130932,"databundleVersionId":15769099},{"sourceType":"datasetVersion","sourceId":14920175,"datasetId":9546753,"databundleVersionId":15786705}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport cv2\nimport json\nimport torch\nimport numpy as np\nfrom PIL import Image\nfrom scipy.optimize import minimize\nfrom torch.utils.data import Dataset\nfrom torchvision import transforms\nfrom torch.utils.data import DataLoader\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import MultiStepLR\nimport torchvision.models as models\nimport torch.nn as nn\nimport cv2\nfrom tqdm import tqdm\n\n\nclass KParameterPredictionDataset(Dataset):\n    def __init__(self, data_dir=\"/kaggle/input/automatic-lens-correction/lens-correction-train-cleaned\", labels_path=\"/kaggle/input/datasets/ramazanturann/k-labels/k_labels.json\"):\n        self.data_dir = data_dir\n        exclude_files = [\"training_metadata.json\", \"dataset-metadata.json\"]\n        images = os.listdir(data_dir)\n        images = [image for image in images if image not in exclude_files]\n        self.images = list(set([image.replace(\"generated.jpg\", \"\").replace(\"original.jpg\", \"\") for image in images]))\n\n        if os.path.exists(labels_path):\n            with open(labels_path, \"r\") as f:\n                self.labels = json.load(f)\n        else:\n            self.labels = self._precompute_labels(labels_path)\n\n        self.transform = transforms.Compose([\n            transforms.Resize((224, 224)),\n            transforms.RandomApply([transforms.GaussianBlur(3)], p=0.25),\n            transforms.RandomApply([transforms.ColorJitter(brightness=0.3)], p=0.25),\n            transforms.ToTensor(),\n            transforms.RandomErasing(p=0.25),\n            transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                                 std=[0.229, 0.224, 0.225])\n        ])\n\n    def _precompute_labels(self, labels_path):\n        from tqdm import tqdm\n        labels = {}\n        for img_id in tqdm(self.images, desc=\"Computing k params\"):\n            distorted_path = os.path.join(self.data_dir, img_id + \"original.jpg\")\n            corrected_path = os.path.join(self.data_dir, img_id + \"generated.jpg\")\n            distorted_img = Image.open(distorted_path).convert(\"RGB\")\n            corrected_img = Image.open(corrected_path).convert(\"RGB\")\n            k_params = self.compute_k_params(distorted_img, corrected_img)\n            labels[img_id] = k_params.tolist()\n        with open(labels_path, \"w\") as f:\n            json.dump(labels, f)\n        return labels\n\n    def compute_k_params(self, distorted, corrected):\n        distorted_np = np.array(distorted.resize((224, 224)))\n        corrected_np = np.array(corrected.resize((224, 224)))\n        h, w = distorted_np.shape[:2]\n        K = np.array([[w, 0, w/2],\n                      [0, w, h/2],\n                      [0, 0,   1]], dtype=np.float64)\n\n        def loss(params):\n            k1, k2 = params\n            dist_coeffs = np.array([k1, k2, 0, 0, 0])\n            undistorted = cv2.undistort(distorted_np, K, dist_coeffs)\n            return np.mean(np.abs(undistorted.astype(float) - corrected_np.astype(float)))\n\n        result = minimize(loss, [-0.2, 0.0], method='Nelder-Mead',\n                         options={'maxiter': 100, 'xatol': 1e-4})\n        return result.x\n\n    def get_edge_feature(self, img):\n        gray = cv2.cvtColor(np.array(img.resize((224, 224))), cv2.COLOR_RGB2GRAY)\n        edges = cv2.Canny(gray, 100, 200)\n        return transforms.ToTensor()(Image.fromarray(edges))\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, idx):\n        image = self.images[idx]\n        distorted_path = os.path.join(self.data_dir, image + \"original.jpg\")\n        corrected_path = os.path.join(self.data_dir, image + \"generated.jpg\")\n\n        distorted_img = Image.open(distorted_path).convert(\"RGB\")\n\n        orig_w, orig_h = distorted_img.size\n        img_size = torch.tensor([orig_w / 1920.0, orig_h / 1080.0], dtype=torch.float32)\n\n        k_params = torch.tensor(self.labels[image], dtype=torch.float32)\n\n        distorted_tensor = self.transform(distorted_img)\n\n        return distorted_tensor, img_size, k_params","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-02-22T17:23:07.983601Z","iopub.execute_input":"2026-02-22T17:23:07.984057Z","iopub.status.idle":"2026-02-22T17:23:08.011186Z","shell.execute_reply.started":"2026-02-22T17:23:07.984007Z","shell.execute_reply":"2026-02-22T17:23:08.009920Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\nimport os\nimport random\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom PIL import Image\n\n# Config\ndata_dir = \"/kaggle/input/automatic-lens-correction/lens-correction-train-cleaned\"\nlabels_path = \"/kaggle/input/datasets/ramazanturann/k-labels/k_labels.json\"\n\n# Load labels\nwith open(labels_path, \"r\") as f:\n    labels = json.load(f)\n\n# Get image IDs\nexclude_files = [\"training_metadata.json\", \"dataset-metadata.json\"]\nimages = os.listdir(data_dir)\nimages = [img for img in images if img not in exclude_files]\nimage_ids = list(set([img.replace(\"generated.jpg\", \"\").replace(\"original.jpg\", \"\") for img in images]))\n\n# Pick random samples\nsamples = random.sample(image_ids, min(4, len(image_ids)))\n\nfig, axes = plt.subplots(len(samples), 3, figsize=(15, 5 * len(samples)))\nfig.suptitle(\"Dataset Visualization\", fontsize=16, fontweight='bold')\n\nfor row, img_id in enumerate(samples):\n    distorted = Image.open(os.path.join(data_dir, img_id + \"original.jpg\")).convert(\"RGB\")\n    corrected = Image.open(os.path.join(data_dir, img_id + \"generated.jpg\")).convert(\"RGB\")\n    \n    k1, k2 = labels[img_id]\n    \n    # Distorted\n    axes[row, 0].imshow(distorted)\n    axes[row, 0].set_title(f\"Distorted (original)\\n{img_id[:30]}...\", fontsize=9)\n    axes[row, 0].axis(\"off\")\n    \n    # Corrected\n    axes[row, 1].imshow(corrected)\n    axes[row, 1].set_title(f\"Corrected (generated)\\nk1={k1:.4f}, k2={k2:.4f}\", fontsize=9)\n    axes[row, 1].axis(\"off\")\n    \n    # Diff map\n    diff = np.abs(np.array(distorted.resize((224, 224))).astype(float) - \n                  np.array(corrected.resize((224, 224))).astype(float))\n    diff_norm = (diff / diff.max() * 255).astype(np.uint8)\n    axes[row, 2].imshow(diff_norm)\n    axes[row, 2].set_title(f\"Pixel Diff\\nMAE={diff.mean():.2f}\", fontsize=9)\n    axes[row, 2].axis(\"off\")\n\nplt.tight_layout()\nplt.savefig(\"dataset_viz.png\", dpi=150, bbox_inches='tight')\nplt.show()\n\n# K params distribution\nk1s = [v[0] for v in labels.values()]\nk2s = [v[1] for v in labels.values()]\n\nfig, axes = plt.subplots(1, 2, figsize=(12, 4))\naxes[0].hist(k1s, bins=50, color='steelblue', edgecolor='white')\naxes[0].set_title(\"k1 Distribution\"); axes[0].set_xlabel(\"k1\")\naxes[1].hist(k2s, bins=50, color='salmon', edgecolor='white')\naxes[1].set_title(\"k2 Distribution\"); axes[1].set_xlabel(\"k2\")\nplt.tight_layout()\nplt.savefig(\"k_distribution.png\", dpi=150, bbox_inches='tight')\nplt.show()\n\nprint(f\"Total samples: {len(labels)}\")\nprint(f\"k1 → mean: {np.mean(k1s):.4f}, std: {np.std(k1s):.4f}, range: [{min(k1s):.4f}, {max(k1s):.4f}]\")\nprint(f\"k2 → mean: {np.mean(k2s):.4f}, std: {np.std(k2s):.4f}, range: [{min(k2s):.4f}, {max(k2s):.4f}]\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T17:23:08.013254Z","iopub.execute_input":"2026-02-22T17:23:08.013665Z","iopub.status.idle":"2026-02-22T17:23:18.666784Z","shell.execute_reply.started":"2026-02-22T17:23:08.013634Z","shell.execute_reply":"2026-02-22T17:23:18.665755Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#import json\n#from tqdm import tqdm\n\n#dataset = KParameterPredictionDataset()\n#labels = {}\n\n#for i in tqdm(range(len(dataset))):\n   # img_id = dataset.images[i]\n   # distorted_path = os.path.join(dataset.data_dir, img_id + \"original.jpg\")\n   # corrected_path = os.path.join(dataset.data_dir, img_id + \"generated.jpg\")\n    \n   # distorted_img = Image.open(distorted_path).convert(\"RGB\")\n   # corrected_img = Image.open(corrected_path).convert(\"RGB\")\n    \n   # k_params = dataset.compute_k_params(distorted_img, corrected_img)\n   # labels[img_id] = k_params.tolist()\n\n#with open(\"k_labels.json\", \"w\") as f:\n    #json.dump(labels, f)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T17:23:18.667927Z","iopub.execute_input":"2026-02-22T17:23:18.668399Z","iopub.status.idle":"2026-02-22T17:23:18.672904Z","shell.execute_reply.started":"2026-02-22T17:23:18.668372Z","shell.execute_reply":"2026-02-22T17:23:18.671824Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DBCv3(nn.Module):\n    def __init__(self, num_params=2):\n        super().__init__()\n        resnet = models.resnet50(pretrained=True)\n        self.backbone = nn.Sequential(*list(resnet.children())[:-1])\n\n        edge_resnet = models.resnet18(pretrained=False)\n        edge_resnet.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n        self.edge_backbone = nn.Sequential(*list(edge_resnet.children())[:-1])\n\n        self.fc = nn.Sequential(\n            nn.Linear(2048 + 2, 512),\n            nn.ReLU(),\n            nn.Linear(512, num_params)\n        )\n\n    def forward(self, x, img_size):\n        img_features = self.backbone(x).flatten(1)\n        combined = torch.cat([img_features, img_size], dim=1)\n        return self.fc(combined)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T17:23:18.674621Z","iopub.execute_input":"2026-02-22T17:23:18.675033Z","iopub.status.idle":"2026-02-22T17:23:18.693320Z","shell.execute_reply.started":"2026-02-22T17:23:18.674994Z","shell.execute_reply":"2026-02-22T17:23:18.691838Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ResolutionBatchSampler:\n    def __init__(self, dataset, batch_size):\n        self.groups = {}\n        for i, img_id in enumerate(dataset.images):\n            img = Image.open(os.path.join(dataset.data_dir, img_id + \"original.jpg\"))\n            size = img.size\n            if size not in self.groups:\n                self.groups[size] = []\n            self.groups[size].append(i)\n            img.close()\n        self.batch_size = batch_size\n\n    def __iter__(self):\n        for indices in self.groups.values():\n            for i in range(0, len(indices), self.batch_size):\n                yield indices[i:i + self.batch_size]\n\n    def __len__(self):\n        return sum(len(v) // self.batch_size for v in self.groups.values())\n\n\ndef undistort_and_crop(img, K, dist_coeffs):\n    h, w = img.shape[:2]\n    new_K, roi = cv2.getOptimalNewCameraMatrix(K, dist_coeffs, (w, h), 1, (w, h))\n    undistorted = cv2.undistort(img, K, dist_coeffs, None, new_K)\n    x, y, w, h = roi\n    return undistorted[y:y+h, x:x+w]\n\n\ndef loss_with_prior(preds, targets):\n    mse = nn.MSELoss()(preds, targets)\n    k1_penalty = torch.relu(preds[:, 0]).mean()\n    return mse + 0.01 * k1_penalty\n\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndataset = KParameterPredictionDataset()\n\ntrain_size = int(0.7 * len(dataset))\nval_size = int(0.15 * len(dataset))\ntest_size = len(dataset) - train_size - val_size\ntrain_set, val_set, test_set = torch.utils.data.random_split(dataset, [train_size, val_size, test_size])\n\ntrain_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=4)\nval_loader = DataLoader(val_set, batch_size=128, shuffle=False, num_workers=4)\n\nmodel = DBCv3(num_params=2).to(device)\noptimizer = AdamW([\n    {\"params\": model.backbone.parameters(), \"lr\": 1e-5},\n    {\"params\": model.edge_backbone.parameters(), \"lr\": 1e-4},\n    {\"params\": model.fc.parameters(), \"lr\": 1e-3}\n])\n\nscheduler = MultiStepLR(optimizer, milestones=[20, 40], gamma=0.1)\n\nbest_val_loss = float('inf')\npatience = 3\npatience_counter = 0\n\nfor epoch in range(10):\n    model.train()\n    train_loss = 0\n    for distorted, img_size, k_params in train_loader:\n        distorted = distorted.to(device)\n        img_size = img_size.to(device)\n        k_params = k_params.to(device)\n        optimizer.zero_grad()\n        preds = model(distorted, img_size)\n        loss = loss_with_prior(preds, k_params)\n        loss.backward()\n        optimizer.step()\n        train_loss += loss.item()\n\n    model.eval()\n    val_loss = 0\n    with torch.no_grad():\n        for distorted, edge, img_size, k_params in val_loader:\n            distorted = distorted.to(device)\n            edge = edge.to(device)\n            img_size = img_size.to(device)\n            k_params = k_params.to(device)\n            preds = model(distorted, edge, img_size)\n            val_loss += loss_with_prior(preds, k_params).item()\n\n    scheduler.step()\n    train_loss_avg = train_loss / len(train_loader)\n    val_loss_avg = val_loss / len(val_loader)\n    print(f\"Epoch {epoch+1}/60 | Train Loss: {train_loss_avg} | Val Loss: {val_loss_avg}\")\n\n    if val_loss_avg < best_val_loss:\n        best_val_loss = val_loss_avg\n        patience_counter = 0\n        torch.save(model.state_dict(), \"dbcv3_best.pth\")\n        print(f\"Model Saved. Best Val Loss: {best_val_loss}\")\n    else:\n        patience_counter += 1\n        print(f\"Patience: {patience_counter}/{patience}\")\n        if patience_counter >= patience:\n            print(\"Early stopping.\")\n            break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T17:23:18.694678Z","iopub.execute_input":"2026-02-22T17:23:18.695129Z","execution_failed":"2026-02-22T19:09:09.902Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def inference_and_save(model, test_dir, output_dir, device):\n    os.makedirs(output_dir, exist_ok=True)\n    model.eval()\n\n    transform = transforms.Compose([\n        transforms.Resize((224, 224)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                             std=[0.229, 0.224, 0.225])\n    ])\n\n    test_images = [f for f in os.listdir(test_dir) if f.endswith(\".jpg\")]\n\n    for img_name in tqdm(test_images, desc=\"Inference\"):\n        img_path = os.path.join(test_dir, img_name)\n        img = Image.open(img_path).convert(\"RGB\")\n\n        orig_w, orig_h = img.size\n        img_size = torch.tensor([orig_w / 1920.0, orig_h / 1080.0], dtype=torch.float32).unsqueeze(0).to(device)\n\n        gray = cv2.cvtColor(np.array(img.resize((224, 224))), cv2.COLOR_RGB2GRAY)\n        edges = cv2.Canny(gray, 100, 200)\n        edge_tensor = transforms.ToTensor()(Image.fromarray(edges)).unsqueeze(0).to(device)\n\n        img_tensor = transform(img).unsqueeze(0).to(device)\n\n        with torch.no_grad():\n            k_params = model(img_tensor, edge_tensor, img_size).cpu().numpy()[0]\n\n        k1, k2 = k_params[0], k_params[1]\n\n        img_np = np.array(img)\n        h, w = img_np.shape[:2]\n        K = np.array([[w,   0, w/2],\n                      [0,   w, h/2],\n                      [0,   0,   1]], dtype=np.float64)\n        dist_coeffs = np.array([k1, k2, 0, 0, 0])\n\n        corrected = undistort_and_crop(img_np, K, dist_coeffs)\n        corrected_resized = cv2.resize(corrected, (orig_w, orig_h))\n\n        output_path = os.path.join(output_dir, img_name)\n        cv2.imwrite(output_path, cv2.cvtColor(corrected_resized, cv2.COLOR_RGB2BGR))\n\nmodel = DBCv3(num_params=2).to(device)\nmodel.load_state_dict(torch.load(\"dbcv3_best.pth\"))\n\ninference_and_save(\n    model=model,\n    test_dir=\"/kaggle/input/automatic-lens-correction/test-originals\",\n    output_dir=\"/kaggle/working/corrected_images\",\n    device=device\n)\n\nimport zipfile\nwith zipfile.ZipFile(\"/kaggle/working/submission.zip\", \"w\") as zf:\n    for img_name in os.listdir(\"/kaggle/working/corrected_images\"):\n        zf.write(os.path.join(\"/kaggle/working/corrected_images\", img_name), img_name)\n","metadata":{"trusted":true,"execution":{"execution_failed":"2026-02-22T19:09:09.904Z"}},"outputs":[],"execution_count":null}]}