{"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":"none","dataSources":[{"sourceId":125981,"sourceType":"competition"}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Sample Pipeline:\n  - Load small training set (20 maps) from:\n\n        train/images/*.png\n\n        train/labels/*.json     # contains true class grid\n\n        train/velocities/*.json # boost; we ignore it in this baseline\n\n  - Train a tiny NN that predicts a 20x20 class grid from the image.\n\n        * It treats all terrains together (no terrain-specific modeling).\n\n        * It only uses per-cell average colour (super crude).\n\n  - For each test image:\n\n        * Predict class grid.\n\n        * Find start (class 3) and goal (class 4).\n\n        * Compute a path with a dumb cost model (ignoring boosts).\n\n        * Convert path to an 'lrud' sequence and write submission_baseline.csv\n        ","metadata":{}},{"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\nimport numpy as np\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T11:34:36.779466Z","iopub.execute_input":"2025-12-12T11:34:36.779859Z","iopub.status.idle":"2025-12-12T11:34:36.785786Z","shell.execute_reply.started":"2025-12-12T11:34:36.779833Z","shell.execute_reply":"2025-12-12T11:34:36.784704Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"GRID_SIZE = 20\nNUM_CLASSES = 5  # 0..4\nBATCH_SIZE = 4\nEPOCHS = 10\nLR = 1e-3\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    [0.5,  # 0 = walkable\n     3.0,  # 1 = wall\n     0.7,  # 2 = hazard\n     4.0,  # 3 = start\n     4.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-12T11:34:36.787444Z","iopub.execute_input":"2025-12-12T11:34:36.787831Z","iopub.status.idle":"2025-12-12T11:34:36.841156Z","shell.execute_reply.started":"2025-12-12T11:34:36.787797Z","shell.execute_reply":"2025-12-12T11:34:36.839910Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_label_grid(json_path: Path) -> np.ndarray:\n    \"\"\"Load true class grid from label JSON (expects key 'grid').\"\"\"\n    with json_path.open(\"r\", encoding=\"utf-8\") as f:\n        data = json.load(f)\n    if \"grid\" not in data:\n        raise KeyError(\n            f\"'grid' not found in {json_path}. \"\n            \"Make sure your generator writes 'grid': grid.tolist() into labels.\"\n        )\n    grid = np.array(data[\"grid\"], dtype=np.int64)\n    assert grid.shape == (GRID_SIZE, GRID_SIZE), f\"Expected {GRID_SIZE}x{GRID_SIZE}, got {grid.shape}\"\n    return grid\n\n\ndef compute_cell_tensor(img: Image.Image, grid_size: int) -> torch.Tensor:\n    \"\"\"\n    Convert full map image into a (3, grid_size, grid_size) tensor by\n    averaging RGB values in each cell.\n\n    Very crude, but enough for a baseline.\n    \"\"\"\n    img = img.convert(\"RGB\")\n    w, h = img.size\n    img_arr = np.array(img, dtype=np.float32) / 255.0  # (H,W,3) in [0,1]\n\n    cell_w = w // grid_size\n    cell_h = h // grid_size\n\n    cells = np.zeros((grid_size, grid_size, 3), dtype=np.float32)\n\n    for i in range(grid_size):\n        for j in range(grid_size):\n            x0 = j * cell_w\n            x1 = (j + 1) * cell_w if j < grid_size - 1 else w\n            y0 = i * cell_h\n            y1 = (i + 1) * cell_h if i < grid_size - 1 else h\n\n            patch = img_arr[y0:y1, x0:x1, :]\n            if patch.size == 0:\n                continue\n            cells[i, j, :] = patch.mean(axis=(0, 1))\n\n    # (G,G,3) -> (3,G,G)\n    cells = np.transpose(cells, (2, 0, 1))\n    return torch.from_numpy(cells)  # float32\n\n\nclass GridDataset(Dataset):\n    \"\"\"\n    Each sample:\n      X: (3, GRID_SIZE, GRID_SIZE) tensor\n      y: (GRID_SIZE, GRID_SIZE) long tensor with values 0..4\n    \"\"\"\n\n    def __init__(\n        self,\n        images_dir: Path,\n        labels_dir: Path,\n        grid_size: int,\n    ):\n        self.images_dir = images_dir\n        self.labels_dir = labels_dir\n        self.grid_size = grid_size\n\n        self.image_ids: List[str] = []\n        for p in sorted(labels_dir.glob(\"*.json\")):\n            image_id = p.stem  # \"0001\"\n            img_path = images_dir / f\"{image_id}.png\"\n            if img_path.is_file():\n                self.image_ids.append(image_id)\n\n        if not self.image_ids:\n            raise RuntimeError(f\"No training labels/images found in {labels_dir}\")\n\n    def __len__(self) -> int:\n        return len(self.image_ids)\n\n    def __getitem__(self, idx: int):\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)\n        x = compute_cell_tensor(img, self.grid_size).float()  # (3, G, G)\n\n        grid = load_label_grid(label_path)  # (G, G)\n        y = torch.from_numpy(grid).long()\n\n        return x, y","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T11:34:36.842424Z","iopub.execute_input":"2025-12-12T11:34:36.842903Z","iopub.status.idle":"2025-12-12T11:34:36.857268Z","shell.execute_reply.started":"2025-12-12T11:34:36.842867Z","shell.execute_reply":"2025-12-12T11:34:36.856216Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SimpleGridNet(nn.Module):\n    \"\"\"\n    Simple per-grid classifier:\n      - Flatten (3,G,G) -> FC -> ReLU -> FC -> reshape to (C,G,G)\n    This is intentionally basic, but with class-weighted loss we bias\n    it to learn walls/start/goal reasonably.\n    \"\"\"\n\n    def __init__(self, grid_size: int, num_classes: int):\n        super().__init__()\n        self.grid_size = grid_size\n        self.num_classes = num_classes\n        in_features = 3 * grid_size * grid_size\n        hidden = 256\n\n        self.fc1 = nn.Linear(in_features, hidden)\n        self.fc2 = nn.Linear(hidden, num_classes * grid_size * grid_size)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        # x: (B, 3, G, G)\n        b = x.shape[0]\n        x = x.reshape(b, -1)            # (B, 3*G*G)\n        x = torch.relu(self.fc1(x))     # (B, hidden)\n        x = self.fc2(x)                 # (B, num_classes*G*G)\n        x = x.view(b, self.num_classes, self.grid_size, self.grid_size)\n        return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T11:34:36.858252Z","iopub.execute_input":"2025-12-12T11:34:36.858595Z","iopub.status.idle":"2025-12-12T11:34:36.881162Z","shell.execute_reply.started":"2025-12-12T11:34:36.858540Z","shell.execute_reply":"2025-12-12T11:34:36.879869Z"}},"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    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)               # (B, 3, G, G)\n            yb = yb.to(DEVICE)               # (B, G, G)\n\n            optimizer.zero_grad()\n            logits = model(xb)               # (B, C, G, G)\n\n            B, C, G, _ = logits.shape\n            logits_flat = logits.view(B, C, G * G)   # (B, C, G*G)\n            y_flat = yb.view(B, G * G)              # (B, G*G)\n\n            loss = criterion(logits_flat, y_flat)\n            loss.backward()\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-12T11:34:36.882877Z","iopub.execute_input":"2025-12-12T11:34:36.883169Z","iopub.status.idle":"2025-12-12T11:34:36.902196Z","shell.execute_reply.started":"2025-12-12T11:34:36.883146Z","shell.execute_reply":"2025-12-12T11:34:36.901212Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def pick_start_goal_from_logits(logits: torch.Tensor) -> Tuple[Tuple[int, int], Tuple[int, int]]:\n    \"\"\"\n    logits: (1, C, G, G)\n    Returns:\n      start = (row, col) with highest logit for START class\n      goal  = (row, col) with highest logit for GOAL class\n    \"\"\"\n    assert logits.shape[0] == 1\n    _, C, G, _ = logits.shape\n\n    # class index 3 = start, 4 = goal\n    logit_start = logits[0, CLASS_START, :, :]  # (G,G)\n    logit_goal = logits[0, CLASS_GOAL, :, :]    # (G,G)\n\n    # Flatten then argmax -> index -> (i,j)\n    start_idx = torch.argmax(logit_start).item()\n    goal_idx = torch.argmax(logit_goal).item()\n\n    start_row = start_idx // G\n    start_col = start_idx % G\n\n    goal_row = goal_idx // G\n    goal_col = goal_idx % G\n\n    return (int(start_row), int(start_col)), (int(goal_row), int(goal_col))\n\n\ndef bfs_path(grid_pred: np.ndarray, start: Tuple[int, int], goal: Tuple[int, int]) -> List[Tuple[int, int]]:\n    \"\"\"\n    BFS ignoring cost, just avoiding walls (class 1).\n    Returns path as list of (i,j), or [] if no path found.\n    \"\"\"\n    G = grid_pred.shape[0]\n    sr, sc = start\n    gr, gc = goal\n\n    visited = np.zeros((G, G), dtype=bool)\n    prev: Dict[Tuple[int, int], Tuple[int, int]] = {}\n\n    q = deque()\n    q.append((sr, sc))\n    visited[sr, sc] = True\n\n    def neighbours(i, j):\n        for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:\n            ni, nj = i + di, j + dj\n            if 0 <= ni < G and 0 <= nj < G:\n                yield ni, nj\n\n    while q:\n        i, j = q.popleft()\n        if (i, j) == (gr, gc):\n            break\n        for ni, nj in neighbours(i, j):\n            if visited[ni, nj]:\n                continue\n            if grid_pred[ni, nj] == CLASS_WALL:\n                continue  # block walls\n            visited[ni, nj] = True\n            prev[(ni, nj)] = (i, j)\n            q.append((ni, nj))\n\n    if (gr, gc) not in prev and (gr, gc) != (sr, sc):\n        # No path found\n        return []\n\n    # Reconstruct\n    path: List[Tuple[int, int]] = []\n    cur = (gr, gc)\n    path.append(cur)\n    while cur != (sr, sc):\n        cur = prev[cur]\n        path.append(cur)\n    path.reverse()\n    return path\n\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-12T11:34:36.981161Z","iopub.execute_input":"2025-12-12T11:34:36.981517Z","iopub.status.idle":"2025-12-12T11:34:36.998766Z","shell.execute_reply.started":"2025-12-12T11:34:36.981491Z","shell.execute_reply":"2025-12-12T11:34:36.997533Z"}},"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    Returns:\n      grid_pred: (G,G) np.int64 argmax classes\n      logits:    (1,C,G,G) torch.Tensor (on CPU) for start/goal picking\n    \"\"\"\n    img = Image.open(img_path)\n    x = compute_cell_tensor(img, GRID_SIZE).unsqueeze(0).float()  # (1,3,G,G)\n    x = x.to(DEVICE)\n\n    model.eval()\n    with torch.no_grad():\n        logits = model(x)                       # (1,C,G,G)\n        preds = torch.argmax(logits, dim=1)     # (1,G,G)\n\n    grid_pred = preds.squeeze(0).cpu().numpy().astype(np.int64)\n    logits_cpu = logits.cpu()\n    return grid_pred, logits_cpu\n\n\ndef run_inference_on_test(model: nn.Module):\n    \"\"\"\n    For each test image:\n      - Predict grid + logits\n      - Pick start/goal from logits\n      - BFS path avoiding walls, else Manhattan fallback\n      - Write submission CSV\n    \"\"\"\n    model.to(DEVICE)\n    image_paths = sorted(TEST_IMAGES_DIR.glob(\"*.png\"))\n\n    records = [(\"image_id\", \"path\")]\n\n    for img_path in image_paths:\n        image_id = img_path.stem\n        print(f\"Inference on {image_id}...\")\n\n        grid_pred, logits = predict_logits_and_grid(model, img_path)\n\n        # Start/goal via logits (most confident cell for class 3/4)\n        start, goal = pick_start_goal_from_logits(logits)\n\n        # BFS with walls from predicted grid\n        path = bfs_path(grid_pred, start, goal)\n        if not path:\n            # BFS failed -> fallback ignoring walls\n            path = fallback_manhattan_path(start, goal)\n\n        moves = path_to_lrud(path)\n        records.append((image_id, moves))\n\n    # Write CSV\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\"Baseline submission written to: {SUBMISSION_PATH}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T11:34:37.000421Z","iopub.execute_input":"2025-12-12T11:34:37.000739Z","iopub.status.idle":"2025-12-12T11:34:37.021136Z","shell.execute_reply.started":"2025-12-12T11:34:37.000712Z","shell.execute_reply":"2025-12-12T11:34:37.019698Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataset = GridDataset(\n        images_dir=TRAIN_IMAGES_DIR,\n        labels_dir=TRAIN_LABELS_DIR,\n        grid_size=GRID_SIZE,\n    )\n\ntrain_loader = DataLoader(\n        train_dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=0,\n    )\n\n    # 2. Create model\nmodel = SimpleGridNet(grid_size=GRID_SIZE, num_classes=NUM_CLASSES)\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-12T11:34:37.022067Z","iopub.execute_input":"2025-12-12T11:34:37.022340Z","iopub.status.idle":"2025-12-12T11:58:07.966429Z","shell.execute_reply.started":"2025-12-12T11:34:37.022318Z","shell.execute_reply":"2025-12-12T11:58:07.964855Z"}},"outputs":[],"execution_count":null}]}