{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":125981,"databundleVersionId":14910697,"sourceType":"competition"}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Imports\nimport json\nimport math\nimport heapq\nfrom collections import deque\nfrom pathlib import Path\nfrom typing import List, Tuple, Dict\n\nfrom pathlib import Path\nimport numpy as np\nfrom PIL import Image\nimport random\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader, ConcatDataset\nfrom torchvision import transforms\nimport torch.nn as nn\nimport torchvision.transforms.functional as TF","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T09:37:24.542585Z","iopub.execute_input":"2025-12-19T09:37:24.542765Z","iopub.status.idle":"2025-12-19T09:37:32.514486Z","shell.execute_reply.started":"2025-12-19T09:37:24.542748Z","shell.execute_reply":"2025-12-19T09:37:32.513906Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"GRID_SIZE = 20\nIMG_SIZE = 320\nNUM_CLASSES = 5  # 0..4\nBATCH_SIZE = 16\nEPOCHS = 40\nLR = 1e-4\n\n\nTRAIN_IMAGES_DIR = Path(\"/kaggle/input/the-blind-flight-synapse-drive-ps-1/SynapseDrive_Dataset/train/images\")\nTRAIN_LABELS_DIR = Path(\"/kaggle/input/the-blind-flight-synapse-drive-ps-1/SynapseDrive_Dataset/train/labels\")\n\nTEST_IMAGES_DIR = Path(\"/kaggle/input/the-blind-flight-synapse-drive-ps-1/SynapseDrive_Dataset/test/images\")\n\nSUBMISSION_PATH = Path(\"/kaggle/working/submission_baseline.csv\")\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Class-weighting to focus on walls/start/goal\n# idx: class_id -> weight\nCLASS_WEIGHTS = torch.tensor(\n    [1.0,  # 0 = walkable\n     2.0,  # 1 = wall\n     5.0,  # 2 = hazard\n     50.0,  # 3 = start\n     50.0], # 4 = goal\n    dtype=torch.float32\n)\n\nCLASS_WALL = 1\nCLASS_START = 3\nCLASS_GOAL = 4","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T09:37:32.515735Z","iopub.execute_input":"2025-12-19T09:37:32.516069Z","iopub.status.idle":"2025-12-19T09:37:32.570815Z","shell.execute_reply.started":"2025-12-19T09:37:32.516049Z","shell.execute_reply":"2025-12-19T09:37:32.570047Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SemanticGridDataset(Dataset):\n    def __init__(self, images_dir: Path, labels_dir: Path, augment: bool = False):\n        self.images_dir = images_dir\n        self.labels_dir = labels_dir\n        self.augment = augment\n        \n        self.image_ids = [\n            p.stem for p in sorted(labels_dir.glob(\"*.json\")) \n            if (images_dir / f\"{p.stem}.png\").exists()\n        ]\n        \n        self.to_tensor = transforms.Compose([\n            transforms.Resize((IMG_SIZE, IMG_SIZE)),\n            transforms.ToTensor(),\n        ])\n\n    def __len__(self):\n        return len(self.image_ids)\n\n    def __getitem__(self, idx):\n        image_id = self.image_ids[idx]\n        img_path = self.images_dir / f\"{image_id}.png\"\n        label_path = self.labels_dir / f\"{image_id}.json\"\n\n        img = Image.open(img_path).convert(\"RGB\")\n        with label_path.open(\"r\") as f:\n            data = json.load(f)\n        grid = np.array(data[\"grid\"], dtype=np.int64)\n        \n        x = self.to_tensor(img)\n        y = torch.from_numpy(grid).long().unsqueeze(0)\n\n        if self.augment:\n            rot_k = random.choice([0, 1, 2, 3])\n            if rot_k > 0:\n                x = torch.rot90(x, k=rot_k, dims=[1, 2])\n                y = torch.rot90(y, k=rot_k, dims=[1, 2])\n            if random.random() > 0.5:\n                x = TF.hflip(x)\n                y = TF.hflip(y)\n            if random.random() > 0.5:\n                x = TF.vflip(x)\n                y = TF.vflip(y)\n                \n        return x, y.squeeze(0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T09:37:32.571588Z","iopub.execute_input":"2025-12-19T09:37:32.572001Z","iopub.status.idle":"2025-12-19T09:37:32.607849Z","shell.execute_reply.started":"2025-12-19T09:37:32.571963Z","shell.execute_reply":"2025-12-19T09:37:32.607114Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n    \"\"\"Helper: (Conv -> BN -> ReLU) * 2\"\"\"\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 UNetGrid(nn.Module):\n    def __init__(self, n_channels=3, n_classes=5):\n        super(UNetGrid, self).__init__()\n        \n        # --- ENCODER ---\n        self.inc = DoubleConv(n_channels, 64)         \n        self.down1 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(64, 128))   \n        self.down2 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(128, 256)) \n        self.down3 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(256, 512))\n\n        # --- DECODER ---\n        self.up1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n        self.conv1 = DoubleConv(512, 256) \n        \n        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n        self.conv2 = DoubleConv(256, 128) \n        \n\n        self.up3 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n        self.conv3 = DoubleConv(128, 64)  \n\n        self.outc = nn.Conv2d(64, n_classes, kernel_size=1)\n        self.grid_pool = nn.AdaptiveAvgPool2d((20, 20))\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\n        x = self.up1(x4)                \n        x = torch.cat([x3, x], dim=1)  \n        x = self.conv1(x)\n\n        x = self.up2(x)                \n        x = torch.cat([x2, x], dim=1)  \n        x = self.conv2(x)\n\n        x = self.up3(x)                \n        x = torch.cat([x1, x], dim=1)\n        x = self.conv3(x)\n        \n        logits_high_res = self.outc(x)\n        logits_grid = self.grid_pool(logits_high_res)\n        \n        return logits_grid","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T09:37:32.608675Z","iopub.execute_input":"2025-12-19T09:37:32.608950Z","iopub.status.idle":"2025-12-19T09:37:32.626827Z","shell.execute_reply.started":"2025-12-19T09:37:32.608931Z","shell.execute_reply":"2025-12-19T09:37:32.626329Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_model(model: nn.Module, loader: DataLoader, epochs: int = EPOCHS):\n    model.to(DEVICE)\n    model.train()\n\n    class_weights = CLASS_WEIGHTS.to(DEVICE)\n    criterion = nn.CrossEntropyLoss(weight=class_weights)\n    \n\n    optimizer = torch.optim.Adam(model.parameters(), lr=LR)\n\n    for epoch in range(1, epochs + 1):\n        running_loss = 0.0\n        for xb, yb in loader:\n            xb = xb.to(DEVICE)               \n            yb = yb.to(DEVICE)               \n\n            optimizer.zero_grad()\n            logits = model(xb)               \n\n            B, C, G, _ = logits.shape\n            logits_flat = logits.view(B, C, G * G)   \n            y_flat = yb.view(B, G * G)              \n\n            loss = criterion(logits_flat, y_flat)\n            if torch.isnan(loss):\n                print(\"Error: Loss turned to NaN! Stopping training.\")\n                return\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n            \n            optimizer.step()\n\n            running_loss += loss.item() * xb.size(0)\n\n        avg_loss = running_loss / len(loader.dataset)\n        print(f\"Epoch {epoch:02d} | loss = {avg_loss:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T09:37:32.628428Z","iopub.execute_input":"2025-12-19T09:37:32.628962Z","iopub.status.idle":"2025-12-19T09:37:32.643432Z","shell.execute_reply.started":"2025-12-19T09:37:32.628944Z","shell.execute_reply":"2025-12-19T09:37:32.642917Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def pick_start_goal_from_logits(logits):\n\n    b, c, h, w = logits.shape\n    flat_logits = logits.view(c, -1)\n    \n    values_s, indices_s = torch.topk(flat_logits[CLASS_START], k=2)\n    s1_idx, s2_idx = indices_s[0].item(), indices_s[1].item()\n    s1_score = values_s[0].item()\n    \n    values_g, indices_g = torch.topk(flat_logits[CLASS_GOAL], k=2)\n    g1_idx, g2_idx = indices_g[0].item(), indices_g[1].item()\n    g1_score = values_g[0].item()\n\n    start_idx = s1_idx\n    goal_idx = g1_idx\n    \n    if start_idx == goal_idx:\n        score_a = s1_score + values_g[1].item()\n        score_b = values_s[1].item() + g1_score\n        \n        if score_a > score_b:\n            goal_idx = g2_idx\n        else:\n            start_idx = s2_idx \n    start_pos = (start_idx // w, start_idx % w)\n    goal_pos = (goal_idx // w, goal_idx % w)\n    \n    return start_pos, goal_pos\n\nimport heapq\n\n# --- NEW COST LOGIC & TERRAIN DETECTION ---\n\n# 1. Define Base Costs from Problem Statement (PS1)\nTERRAIN_COSTS = {\n    \"LAB\":    {0: 1.0, 1: 9999.0, 2: 3.0, 3: 1.0, 4: 2.0},\n    \"FOREST\": {0: 1.5, 1: 9999.0, 2: 2.8, 3: 1.5, 4: 2.5},\n    \"DESERT\": {0: 1.2, 1: 9999.0, 2: 3.7, 3: 1.2, 4: 2.2}\n}\n\ndef detect_terrain(img: Image.Image) -> str:\n    \"\"\"\n    Determines if map is Lab, Forest, or Desert based on average color of the WHOLE image.\n    \"\"\"\n    mean_color = np.array(img).mean(axis=(0, 1)) \n    r, g, b = mean_color[0], mean_color[1], mean_color[2]\n    \n    if g > r and g > b:\n        return \"FOREST\"  \n    elif r > b and g > b and r > 100: \n        return \"DESERT\"  \n    else:\n        return \"LAB\"  \n\ndef build_cost_matrix(grid_classes, terrain_type, boost_matrix):\n    \"\"\"\n    Calculates the actual Step Cost for every cell.\n    Formula: step_cost = base_cost - boost\n    \"\"\"\n    costs = np.zeros((20, 20), dtype=np.float32)\n    base_map = TERRAIN_COSTS[terrain_type]\n    \n    for r in range(20):\n        for c in range(20):\n            class_id = grid_classes[r, c]\n            base = base_map.get(class_id, 1.0) \n            \n            # Apply Boost\n            boost = boost_matrix[r][c]\n            \n            # Final calculation\n            step_cost = base - boost\n            \n            # Safety clamp\n            if step_cost <= 0.01: step_cost = 0.01\n            \n            costs[r, c] = step_cost\n            \n    return costs\n\ndef astar_exact(cost_matrix, start, goal):\n    \"\"\"\n    A* that finds the path with minimum Total Cost.\n    \"\"\"\n    rows, cols = cost_matrix.shape\n    pq = []\n    heapq.heappush(pq, (0, 0, start))\n    came_from = {}\n    g_score = {start: 0}\n    while pq:\n        _, current_g, current = heapq.heappop(pq)\n        if current == goal:\n            break\n        r, c = current\n        for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:\n            nr, nc = r + dr, c + dc\n            if 0 <= nr < rows and 0 <= nc < cols:\n                step_cost = cost_matrix[nr, nc]\n                if step_cost > 1000:\n                    continue \n                new_g = current_g + step_cost  \n                if (nr, nc) not in g_score or new_g < g_score[(nr, nc)]:\n                    g_score[(nr, nc)] = new_g\n                    h = (abs(nr - goal[0]) + abs(nc - goal[1])) * 0.01\n                    f = new_g + h           \n                    heapq.heappush(pq, (f, new_g, (nr, nc)))\n                    came_from[(nr, nc)] = current\n    if goal not in came_from:\n        return [] \n    path = []\n    curr = goal\n    while curr != start:\n        path.append(curr)\n        curr = came_from[curr]\n    path.append(start)\n    return path[::-1]\n\ndef fallback_manhattan_path(start: Tuple[int, int], goal: Tuple[int, int]) -> List[Tuple[int, int]]:\n    \"\"\"\n    Simple L-shaped deterministic path (ignores walls).\n    start->goal by row, then by col.\n    \"\"\"\n    sr, sc = start\n    gr, gc = goal\n    path = []\n    r, c = sr, sc\n    path.append((r, c))\n\n    # move vertically\n    step_r = 1 if gr > r else -1\n    while r != gr:\n        r += step_r\n        path.append((r, c))\n\n    # move horizontally\n    step_c = 1 if gc > c else -1\n    while c != gc:\n        c += step_c\n        path.append((r, c))\n\n    return path\n\n\ndef path_to_lrud(path: List[Tuple[int, int]]) -> str:\n    \"\"\"\n    Convert list of (i,j) positions to lrud sequence.\n    i = row (down), j = col (right).\n    \"\"\"\n    moves = []\n    for (i1, j1), (i2, j2) in zip(path[:-1], path[1:]):\n        di, dj = i2 - i1, j2 - j1\n        if di == 1 and dj == 0:\n            moves.append(\"d\")\n        elif di == -1 and dj == 0:\n            moves.append(\"u\")\n        elif di == 0 and dj == 1:\n            moves.append(\"r\")\n        elif di == 0 and dj == -1:\n            moves.append(\"l\")\n        else:\n            moves.append(\"x\")  # unexpected step\n    return \"\".join(moves)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T09:37:32.644083Z","iopub.execute_input":"2025-12-19T09:37:32.644266Z","iopub.status.idle":"2025-12-19T09:37:32.661364Z","shell.execute_reply.started":"2025-12-19T09:37:32.644251Z","shell.execute_reply":"2025-12-19T09:37:32.660831Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_logits_and_grid(model: nn.Module, img_path: Path) -> Tuple[np.ndarray, torch.Tensor]:\n    \"\"\"\n    Runs inference on a single image using the U-Net.\n    1. Resizes image to 320x320 (same as training).\n    2. Feeds to model.\n    3. Model outputs (1, 5, 20, 20).\n    4. Returns grid prediction (20, 20) and logits.\n    \"\"\"\n    model.eval()\n    \n    img = Image.open(img_path).convert(\"RGB\")\n    transform = transforms.Compose([\n        transforms.Resize((IMG_SIZE, IMG_SIZE)), \n        transforms.ToTensor(),\n    ])\n    x = transform(img).unsqueeze(0).to(DEVICE)  \n    with torch.no_grad():\n        logits = model(x)                  \n        preds = torch.argmax(logits, dim=1)    \n\n    grid_pred = preds.squeeze(0).cpu().numpy().astype(np.int64)\n    logits_cpu = logits.cpu()\n    return grid_pred, logits_cpu\n\ndef run_inference_on_test(model: nn.Module):\n    print(\"Running inference with Physics-Aware A*...\")\n    model.to(DEVICE)\n    model.eval()\n    VELOCITY_DIR = TEST_IMAGES_DIR.parent / \"velocities\"\n    image_paths = sorted(TEST_IMAGES_DIR.glob(\"*.png\"))\n    records = [(\"image_id\", \"path\")]\n    \n    for idx, img_path in enumerate(image_paths):\n        image_id = img_path.stem\n        print(f\"Processing {image_id}...\")\n        # 1. Load Image & Predict Grid\n        img = Image.open(img_path).convert(\"RGB\")\n        grid_pred, logits = predict_logits_and_grid(model, img_path)\n        # 2. Pick Start/Goal\n        start, goal = pick_start_goal_from_logits(logits)\n        if 0 <= start[0] < 20 and 0 <= start[1] < 20:\n            grid_pred[start[0], start[1]] = 0\n        if 0 <= goal[0] < 20 and 0 <= goal[1] < 20:\n            grid_pred[goal[0], goal[1]] = 0\n        # 3. Detect Terrain\n        terrain = detect_terrain(img)\n        # 4. Load Velocity Boost\n        boost_path = VELOCITY_DIR / f\"{image_id}.json\"\n        if boost_path.exists():\n            with open(boost_path, 'r') as f:\n                boost_data = json.load(f)\n            boost_matrix = np.array(boost_data[\"boost\"])\n        else:\n            boost_matrix = np.zeros((20, 20))\n        # 5. Build Cost Matrix & Pathfind\n        cost_matrix = build_cost_matrix(grid_pred, terrain, boost_matrix)\n        path = astar_exact(cost_matrix, start, goal)\n        # 6. Fallback\n        if not path:\n            path = fallback_manhattan_path(start, goal)\n        # 7. Convert to Moves\n        if len(path) < 2:\n             moves = \"r\" \n        else:\n             moves = path_to_lrud(path)\n        records.append((image_id, moves))\n        \n    with SUBMISSION_PATH.open(\"w\", encoding=\"utf-8\") as f:\n        for image_id, path_str in records:\n            f.write(f\"{image_id},{path_str}\\n\")\n\n    print(f\"Submission saved to: {SUBMISSION_PATH}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T09:37:32.662013Z","iopub.execute_input":"2025-12-19T09:37:32.662245Z","iopub.status.idle":"2025-12-19T09:37:32.681520Z","shell.execute_reply.started":"2025-12-19T09:37:32.662224Z","shell.execute_reply":"2025-12-19T09:37:32.680972Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataset = SemanticGridDataset(\n    TRAIN_IMAGES_DIR, \n    TRAIN_LABELS_DIR, \n    augment=True\n)\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n\nmodel = UNetGrid(n_channels=3, n_classes=5).to(DEVICE)\noptimizer = torch.optim.Adam(model.parameters(), lr=LR)\n\n# 3. Train model\nprint(\"Training baseline model (focus: start/goal/walls)...\")\ntrain_model(model, train_loader, epochs=EPOCHS)\n\n# 4. Inference + CSV\nprint(\"Running inference on test set...\")\nrun_inference_on_test(model)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T09:37:32.682316Z","iopub.execute_input":"2025-12-19T09:37:32.682544Z"}},"outputs":[],"execution_count":null}]}