{"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}],"dockerImageVersionId":31286,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import cv2\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom pathlib import Path\nfrom scipy.optimize import minimize_scalar\nfrom tqdm import tqdm\n\nTRAIN_DIR = Path(\"/kaggle/input/automatic-lens-correction/lens-correction-train-cleaned\")\nTEST_DIR  = Path(\"/kaggle/input/automatic-lens-correction/test-originals\")\nOUT_DIR   = Path(\"/kaggle/working/corrected\"); OUT_DIR.mkdir(exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T11:43:31.859489Z","iopub.execute_input":"2026-02-21T11:43:31.859867Z","iopub.status.idle":"2026-02-21T11:43:33.207743Z","shell.execute_reply.started":"2026-02-21T11:43:31.859834Z","shell.execute_reply":"2026-02-21T11:43:33.206760Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def undistort(img: np.ndarray, k1: float) -> np.ndarray:\n    h, w = img.shape[:2]\n    cx, cy = w / 2, h / 2\n    ys, xs = np.mgrid[0:h, 0:w]\n    xn = (xs - cx) / cx\n    yn = (ys - cy) / cy\n    r2 = xn**2 + yn**2\n    factor = 1 + k1 * r2\n    xs_src = (xn * factor * cx + cx).astype(np.float32)\n    ys_src = (yn * factor * cy + cy).astype(np.float32)\n    return cv2.remap(img, xs_src, ys_src,\n                     interpolation=cv2.INTER_LINEAR,\n                     borderMode=cv2.BORDER_REPLICATE)\n\n\ndef line_straightness_score(img: np.ndarray) -> float:\n    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img\n    edges = cv2.Canny(gray, 50, 150, apertureSize=3)\n    lines = cv2.HoughLinesP(edges, 1, np.pi / 180,\n                             threshold=80,\n                             minLineLength=img.shape[0] // 6,\n                             maxLineGap=10)\n    if lines is None:\n        return 0.0\n    lengths, residuals = [], []\n    for x1, y1, x2, y2 in lines[:, 0]:\n        dx, dy = x2 - x1, y2 - y1\n        length = np.hypot(dx, dy)\n        angle = np.abs(np.arctan2(dy, dx))\n        dev = min(angle, np.pi / 2 - angle, np.abs(angle - np.pi / 2))\n        lengths.append(length)\n        residuals.append(dev * length)\n    if not lengths:\n        return 0.0\n    return np.sum(residuals) / np.sum(lengths)\n\n\ndef best_k1(img: np.ndarray,\n             bounds: tuple[float, float] = (-0.6, 0.6)) -> float:\n    scale = 0.25\n    small = cv2.resize(img, None, fx=scale, fy=scale)\n\n    def objective(k1):\n        return line_straightness_score(undistort(small, k1))\n\n    result = minimize_scalar(objective, bounds=bounds, method=\"bounded\",\n                             options={\"xatol\": 5e-3, \"maxiter\": 20})\n    return result.x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T11:43:33.209775Z","iopub.execute_input":"2026-02-21T11:43:33.210753Z","iopub.status.idle":"2026-02-21T11:43:33.222997Z","shell.execute_reply.started":"2026-02-21T11:43:33.210712Z","shell.execute_reply":"2026-02-21T11:43:33.221840Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import random\n\ndef make_val_set(n=200, seed=2026) -> list[Path]:\n    all_orig = sorted(TRAIN_DIR.glob(\"*_original.jpg\"))\n    rng = random.Random(seed)\n    return rng.sample(all_orig, n)\n\nVAL_SET = make_val_set(n=500, seed=2026)\nprint(f\"val set size: {len(VAL_SET)}\")\nprint(VAL_SET[0].name)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T11:46:26.613670Z","iopub.execute_input":"2026-02-21T11:46:26.614053Z","iopub.status.idle":"2026-02-21T11:46:27.045903Z","shell.execute_reply.started":"2026-02-21T11:46:26.614019Z","shell.execute_reply":"2026-02-21T11:46:27.044968Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def validate(orig_paths: list[Path], verbose=False) -> dict:\n    mae_scores, ssim_scores, k1s = [], [], []\n\n    for orig_path in tqdm(orig_paths, desc=\"validating\"):\n        corr_path = Path(str(orig_path).replace(\"_original.jpg\", \"_generated.jpg\"))\n        orig = cv2.imread(str(orig_path))\n        corr = cv2.imread(str(corr_path))\n        if orig is None or corr is None:\n            continue\n\n        k1    = best_k1(orig)\n        fixed = undistort(orig, k1)\n\n        if fixed.shape != corr.shape:\n            corr = cv2.resize(corr, (fixed.shape[1], fixed.shape[0]))\n\n        mae = np.mean(np.abs(fixed.astype(float) - corr.astype(float))) / 255\n        gray_fixed = cv2.cvtColor(fixed, cv2.COLOR_BGR2GRAY).astype(float)\n        gray_corr  = cv2.cvtColor(corr,  cv2.COLOR_BGR2GRAY).astype(float)\n        ssim = float(np.mean(\n            (2 * gray_fixed * gray_corr + 1e-6) /\n            (gray_fixed**2 + gray_corr**2 + 1e-6)\n        ))\n\n        mae_scores.append(mae)\n        ssim_scores.append(ssim)\n        k1s.append(k1)\n\n        if verbose:\n            print(f\"{orig_path.name[:40]}  k1={k1:+.3f}  \"\n                  f\"mae={mae:.4f}  ssim={ssim:.4f}\")\n\n    results = {\n        \"mean_mae\":  float(np.mean(mae_scores)),\n        \"mean_ssim\": float(np.mean(ssim_scores)),\n        \"k1_median\": float(np.median(k1s)),\n        \"k1_std\":    float(np.std(k1s)),\n    }\n    print(f\"\\n--- validation summary (n={len(mae_scores)}) ---\")\n    print(f\"mean MAE  : {results['mean_mae']:.4f}  (lower is better)\")\n    print(f\"mean SSIM : {results['mean_ssim']:.4f}  (higher is better)\")\n    print(f\"k1 median : {results['k1_median']:+.4f}  std={results['k1_std']:.4f}\")\n    return results\n\nval = validate(VAL_SET)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T11:46:28.442223Z","iopub.execute_input":"2026-02-21T11:46:28.443398Z","iopub.status.idle":"2026-02-21T11:50:35.289893Z","shell.execute_reply.started":"2026-02-21T11:46:28.443324Z","shell.execute_reply":"2026-02-21T11:50:35.288871Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_val_samples(n=4):\n    orig_paths = sorted(TRAIN_DIR.glob(\"*_original.jpg\"))[:n]\n    fig, axes = plt.subplots(n, 3, figsize=(15, 5 * n))\n\n    for i, orig_path in enumerate(orig_paths):\n        corr_path = Path(str(orig_path).replace(\"_original.jpg\", \"_generated.jpg\"))\n        orig = cv2.imread(str(orig_path))\n        corr = cv2.imread(str(corr_path))\n        k1   = best_k1(orig)\n        fixed = undistort(orig, k1)\n\n        for ax, im, title in zip(axes[i],\n                                  [orig, fixed, corr],\n                                  [\"distorted\",\n                                   f\"ours (k1={k1:+.3f})\",\n                                   \"ground truth\"]):\n            ax.imshow(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))\n            ax.set_title(title); ax.axis(\"off\")\n\n    plt.tight_layout()\n    plt.show()\n\nplot_val_samples(n=4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T11:50:40.713166Z","iopub.execute_input":"2026-02-21T11:50:40.713510Z","iopub.status.idle":"2026-02-21T11:50:46.438950Z","shell.execute_reply.started":"2026-02-21T11:50:40.713481Z","shell.execute_reply":"2026-02-21T11:50:46.437328Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def process_test():\n    test_imgs = sorted(TEST_DIR.glob(\"*.jpg\"))\n    for path in tqdm(test_imgs, desc=\"correcting\"):\n        img = cv2.imread(str(path))\n        if img is None:\n            continue\n        k1  = best_k1(img)\n        out = undistort(img, k1)\n        cv2.imwrite(str(OUT_DIR / path.name), out)\n\nprocess_test()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T11:51:22.908889Z","iopub.execute_input":"2026-02-21T11:51:22.909226Z","iopub.status.idle":"2026-02-21T11:57:11.490305Z","shell.execute_reply.started":"2026-02-21T11:51:22.909198Z","shell.execute_reply":"2026-02-21T11:57:11.489111Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import zipfile\n\nzip_path = Path(\"/kaggle/working/submission_corrected.zip\")\nwith zipfile.ZipFile(zip_path, \"w\", zipfile.ZIP_DEFLATED) as zf:\n    for f in sorted(OUT_DIR.glob(\"*.jpg\")):\n        zf.write(f, arcname=f.name)\n\nprint(f\"zipped {len(list(OUT_DIR.glob('*.jpg')))} images -> {zip_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T11:57:20.165483Z","iopub.execute_input":"2026-02-21T11:57:20.165830Z","iopub.status.idle":"2026-02-21T11:57:44.802519Z","shell.execute_reply.started":"2026-02-21T11:57:20.165801Z","shell.execute_reply":"2026-02-21T11:57:44.801561Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}