{"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":"markdown","source":"### Explanation of the Algorithm\n\n1.  **Detection (Hough Transform):**\n    The code assumes that straight lines exist in the real world (edges of buildings, horizons, door frames). It uses `cv2.Canny` and `cv2.HoughLinesP` to detect these lines in the distorted image. In a barrel-distorted image, these lines will appear curved.\n\n2.  **Optimization (The \"Model\"):**\n    Instead of a Neural Network, we use a mathematical optimization model.\n    *   We define a **Loss Function**: We take the detected curved lines and simulate un-distorting them with a variable parameter $k$ (distortion coefficient).\n    *   We measure the \"straightness\" of the resulting lines by calculating the area of triangles formed by the start, middle, and end points of the lines. If a line is perfectly straight, the area is 0.\n    *   We use `scipy.optimize.minimize_scalar` to find the exact value of $k$ that minimizes this error.\n\n3.  **Correction & Cropping:**\n    Once the optimal $k$ is found, we apply `cv2.undistort`. Crucially, we use `cv2.getOptimalNewCameraMatrix` with `alpha=0`. This automatically zooms and crops the image to remove the curved black borders that appear after correction, ensuring high scores on \"Structural Similarity\" and aesthetics.\n\n4.  **Performance:**\n    *   It downscales images (`process_scale=0.5`) during the detection phase to ensure it runs quickly (under 9 hours for 1,000 images).\n    *   It applies the calculated parameters to the full-resolution image for the final output.","metadata":{}},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport shutil\nfrom glob import glob\nfrom tqdm import tqdm\nfrom scipy.optimize import minimize_scalar\n\nclass AutoLensCorrector:\n    def __init__(self):\n        # Configuration for line detection\n        self.canny_threshold1 = 50\n        self.canny_threshold2 = 150\n        self.hough_threshold = 100\n        self.min_line_length = 50\n        self.max_line_gap = 10\n        \n        # Scale factor to speed up optimization (process on small image, apply to large)\n        self.process_scale = 0.5\n\n    def detect_lines(self, img_gray):\n        \"\"\"Detects lines in the image using Canny edges and Hough Transform.\"\"\"\n        edges = cv2.Canny(img_gray, self.canny_threshold1, self.canny_threshold2, apertureSize=3)\n        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, \n                                threshold=self.hough_threshold, \n                                minLineLength=self.min_line_length, \n                                maxLineGap=self.max_line_gap)\n        return lines\n\n    def get_undistorted_points(self, lines, k, width, height, cam_matrix):\n        \"\"\"\n        Applies the distortion model to a set of line points.\n        We optimize k (k1) while assuming k2=p1=p2=0 for stability.\n        \"\"\"\n        if lines is None:\n            return []\n\n        # Extract points from lines\n        points = []\n        for line in lines:\n            x1, y1, x2, y2 = line[0]\n            points.append([[x1, y1]])\n            points.append([[x2, y2]])\n            # Add midpoint to better capture curvature\n            points.append([[(x1+x2)/2, (y1+y2)/2]])\n        \n        points = np.array(points, dtype=np.float32)\n        \n        # Define distortion coefficients (optimizing k1 only)\n        dist_coeffs = np.array([k, 0, 0, 0], dtype=np.float32)\n        \n        # Undistort the points\n        undistorted = cv2.undistortPoints(points, cam_matrix, dist_coeffs, P=cam_matrix)\n        \n        return undistorted\n\n    def error_function(self, k, lines, width, height, cam_matrix):\n        \"\"\"\n        Loss function: Calculates how 'bent' the lines are after correction.\n        Lower error = Straighter lines.\n        \"\"\"\n        undistorted_points = self.get_undistorted_points(lines, k, width, height, cam_matrix)\n        \n        if len(undistorted_points) == 0:\n            return 0\n            \n        error = 0\n        # Process every 3 points (start, end, mid) as a group\n        num_lines = len(undistorted_points) // 3\n        \n        for i in range(num_lines):\n            p1 = undistorted_points[3*i][0]\n            p2 = undistorted_points[3*i+1][0]\n            p3 = undistorted_points[3*i+2][0]\n            \n            # Triangle area formula: 0.5 * |x1(y2 - y3) + x2(y3 - y1) + x3(y1 - y2)|\n            # We want this area to be 0 for perfectly straight lines\n            area = 0.5 * abs(p1[0]*(p2[1] - p3[1]) + p2[0]*(p3[1] - p1[1]) + p3[0]*(p1[1] - p2[1]))\n            error += area\n\n        return error\n\n    def process_image(self, image_path):\n        # 1. Load Image\n        original_img = cv2.imread(image_path)\n        if original_img is None:\n            return None\n        \n        h, w = original_img.shape[:2]\n        \n        # 2. Resize for faster optimization (Speed vs Accuracy trade-off)\n        small_h, small_w = int(h * self.process_scale), int(w * self.process_scale)\n        img_small = cv2.resize(original_img, (small_w, small_h))\n        img_gray = cv2.cvtColor(img_small, cv2.COLOR_BGR2GRAY)\n        \n        # 3. Approximate Camera Matrix (K)\n        cx, cy = small_w / 2.0, small_h / 2.0\n        # Focal length approximation\n        fx = fy = small_w \n        cam_matrix = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)\n        \n        # 4. Detect straight lines (distorted) in the image\n        lines = self.detect_lines(img_gray)\n        \n        # If no lines found, return original (safeguard)\n        if lines is None or len(lines) < 5:\n            return original_img\n            \n        # 5. Optimize k1 to straighten the lines\n        res = minimize_scalar(\n            self.error_function, \n            bounds=(-0.5, 0.5), \n            args=(lines, small_w, small_h, cam_matrix),\n            method='bounded',\n            options={'xatol': 1e-3, 'maxiter': 20}\n        )\n        \n        best_k1 = res.x\n        \n        # 6. Apply Correction to Full Resolution Image\n        # Scale camera matrix to full resolution\n        cam_matrix_full = np.array([[fx/self.process_scale, 0, cx/self.process_scale], \n                                    [0, fy/self.process_scale, cy/self.process_scale], \n                                    [0, 0, 1]], dtype=np.float32)\n        \n        dist_coeffs = np.array([best_k1, 0, 0, 0], dtype=np.float32)\n        \n        # Generate new optimal camera matrix to keep valid pixels (smart cropping)\n        new_cam_matrix, _ = cv2.getOptimalNewCameraMatrix(cam_matrix_full, dist_coeffs, (w, h), alpha=0, newImgSize=(w, h))\n        \n        corrected_img = cv2.undistort(original_img, cam_matrix_full, dist_coeffs, None, new_cam_matrix)\n        \n        return corrected_img\n\ndef main():\n    # Directories\n    input_dir = '/kaggle/input/automatic-lens-correction/test-originals'\n    working_dir = '/kaggle/working/'\n    \n    # Create a temporary directory to store images before zipping\n    temp_output_dir = os.path.join(working_dir, 'temp_images_processing')\n    \n    if not os.path.exists(input_dir):\n        print(f\"Error: Input directory {input_dir} not found.\")\n        return\n        \n    # Ensure clean start for temp dir\n    if os.path.exists(temp_output_dir):\n        shutil.rmtree(temp_output_dir)\n    os.makedirs(temp_output_dir, exist_ok=True)\n    \n    # Get test images\n    image_files = glob(os.path.join(input_dir, '*.jpg'))\n    print(f\"Found {len(image_files)} images to process.\")\n    \n    corrector = AutoLensCorrector()\n    \n    # Process Loop\n    for img_path in tqdm(image_files):\n        try:\n            filename = os.path.basename(img_path)\n            \n            # Predict and Correct\n            result_img = corrector.process_image(img_path)\n            \n            # Save to temporary folder\n            if result_img is not None:\n                save_path = os.path.join(temp_output_dir, filename)\n                cv2.imwrite(save_path, result_img)\n            else:\n                # Fallback: copy original if read failed\n                img = cv2.imread(img_path)\n                cv2.imwrite(os.path.join(temp_output_dir, filename), img)\n                \n        except Exception as e:\n            print(f\"Failed to process {img_path}: {str(e)}\")\n\n    print(\"Processing complete. Zipping files...\")\n    \n    # Create the zip file\n    zip_base_name = os.path.join(working_dir, 'submission')\n    shutil.make_archive(zip_base_name, 'zip', temp_output_dir)\n    \n    print(\"Zip created. Cleaning up temporary files...\")\n    \n    # Delete the temporary directory containing the images so only the zip remains\n    shutil.rmtree(temp_output_dir)\n    \n    print(\"Cleanup complete. Only submission.zip remains.\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T16:09:00.958009Z","iopub.execute_input":"2026-02-22T16:09:00.958886Z","iopub.status.idle":"2026-02-22T16:12:28.737374Z","shell.execute_reply.started":"2026-02-22T16:09:00.958850Z","shell.execute_reply":"2026-02-22T16:12:28.736214Z"}},"outputs":[],"execution_count":null}]}