{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"},{"sourceId":11569755,"sourceType":"datasetVersion","datasetId":7253661},{"sourceId":242815837,"sourceType":"kernelVersion"},{"sourceId":245791096,"sourceType":"kernelVersion"},{"sourceId":436557,"sourceType":"modelInstanceVersion","modelInstanceId":355848,"modelId":377148},{"sourceId":437058,"sourceType":"modelInstanceVersion","modelInstanceId":355848,"modelId":377148},{"sourceId":437311,"sourceType":"modelInstanceVersion","modelInstanceId":355848,"modelId":377148},{"sourceId":438358,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":357627,"modelId":377148},{"sourceId":446471,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":355848,"modelId":377148},{"sourceId":446920,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":355848,"modelId":377148}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n# for dirname, _, filenames in os.walk('/kaggle/input'):\n#     for filename in filenames:\n#         print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-06-15T11:44:58.293958Z","iopub.execute_input":"2025-06-15T11:44:58.294284Z","iopub.status.idle":"2025-06-15T11:45:00.788676Z","shell.execute_reply.started":"2025-06-15T11:44:58.294255Z","shell.execute_reply":"2025-06-15T11:45:00.787631Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile data_helper.py\n# Copying from https://www.kaggle.com/code/tpmeli/exploratory-deep-dive-geo-wfi-data-insights\nimport os\nfrom pathlib import Path\nfrom typing import List, Tuple, Dict\n\nimport numpy as np\nfrom torch.utils.data import Dataset\n\n###############################################################################\n# 1) LIGHT SEISMIC HELPER\n###############################################################################\nclass SeismicDataHelperLight:\n    \"\"\"\n    Minimal version of the Seismic Data Helper:\n      - Scans a root folder for subdirectories (datasets).\n      - Each dataset has matching .npy files: dataXXXX.npy <-> modelXXXX.npy.\n      - Access them through .datasets (list of dataset names) or pairs(dataset).\n    \"\"\"\n\n    def __init__(self, root_dir: str):\n        \"\"\"\n        Args:\n            root_dir (str): Directory path containing subfolders with\n                            seismic/velocity .npy files.\n        \"\"\"\n        self.root_dir = root_dir\n        self._pairs: Dict[str, List[Tuple[str, str]]] = self._scan_pairs()\n\n    def _scan_pairs(self) -> Dict[str, List[Tuple[str, str]]]:\n        \"\"\"\n        Look for subdirectories under root_dir. \n        For each subdirectory, try:\n          - If /data and /model exist, match dataXXXX.npy with modelXXXX.npy\n          - Otherwise, match seisXXXX.npy with velXXXX.npy.\n        Return a dict:  {\"folder_name\": [ (seis_path, vel_path), ... ], ... }\n        \"\"\"\n        pairs: Dict[str, List[Tuple[str, str]]] = {}\n        for folder_name in os.listdir(self.root_dir):\n            ds_dir = os.path.join(self.root_dir, folder_name)\n            if not os.path.isdir(ds_dir):\n                continue\n\n            data_dir = os.path.join(ds_dir, 'data')\n            model_dir = os.path.join(ds_dir, 'model')\n\n            # If /data and /model exist:\n            if os.path.isdir(data_dir) and os.path.isdir(model_dir):\n                matched_list = []\n                for fname in os.listdir(data_dir):\n                    if fname.startswith('data') and fname.endswith('.npy'):\n                        suf = fname[len('data'):-4]  # part after 'data' before '.npy'\n                        seis_file = os.path.join(data_dir, fname)\n                        vel_file  = os.path.join(model_dir, f'model{suf}.npy')\n                        if os.path.exists(vel_file):\n                            matched_list.append((seis_file, vel_file))\n                if matched_list:\n                    pairs[folder_name] = sorted(matched_list)\n                continue\n\n            # Otherwise, try to match seisXXXX.npy with velXXXX.npy\n            all_files = os.listdir(ds_dir)\n            seis_fs = [f for f in all_files if f.startswith('seis')  or f.startswith('data') and f.endswith('.npy')]\n            vel_fs  = [f for f in all_files if f.startswith('vel') or f.startswith('model') and f.endswith('.npy')]\n            if seis_fs and vel_fs:\n                vel_set = set(vel_fs)\n                matched_list = []\n                for sf in seis_fs:\n                    suf = sf[len('seis'):-4]\n                    vf  = f'vel{suf}.npy'\n                    if vf in vel_set:\n                        matched_list.append((os.path.join(ds_dir, sf),\n                                             os.path.join(ds_dir, vf)))\n                if matched_list:\n                    pairs[folder_name] = sorted(matched_list)\n\n        return pairs\n\n    @property\n    def datasets(self) -> List[str]:\n        \"\"\"List all available subfolders that contain matched (seis, vel) pairs.\"\"\"\n        return list(self._pairs.keys())\n\n    def is_there_folder(self, folder_name: str) -> bool:\n        return folder_name in self._pairs\n\n    def pairs(self, folder_name: str) -> List[Tuple[str, str]]:\n        \"\"\"All (seis_path, vel_path) pairs belonging to a named dataset folder.\"\"\"\n        return self._pairs[folder_name]\n\nclass SeismicDataset(Dataset):\n    def __init__(self, inputs_files, output_files, n_examples_per_file=500):\n        assert len(inputs_files) == len(output_files)\n        self.inputs_files = inputs_files\n        self.output_files = output_files\n        self.n_examples_per_file = n_examples_per_file\n\n    def __len__(self):\n        return len(self.inputs_files) * self.n_examples_per_file\n\n    def __getitem__(self, idx):\n        # Calculate file offset and sample offset within file\n        file_idx = idx // self.n_examples_per_file\n        sample_idx = idx % self.n_examples_per_file\n\n        X = np.load(self.inputs_files[file_idx], mmap_mode='r')\n        y = np.load(self.output_files[file_idx], mmap_mode='r')\n\n        try:\n            return X[sample_idx].copy(), y[sample_idx].copy()\n        finally:\n            del X, y\n\n\nclass TestDataset(Dataset):\n    def __init__(self, test_files):\n        self.test_files = test_files\n\n\n    def __len__(self):\n        return len(self.test_files)\n\n\n    def __getitem__(self, i):\n        test_file = self.test_files[i]\n\n        return np.load(test_file), test_file.stem\n\n\n# end of data_helper.py\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-15T11:45:00.790049Z","iopub.execute_input":"2025-06-15T11:45:00.790562Z","iopub.status.idle":"2025-06-15T11:45:00.803011Z","shell.execute_reply.started":"2025-06-15T11:45:00.790532Z","shell.execute_reply":"2025-06-15T11:45:00.802177Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile unet_def.py\n# Just copying from https://www.kaggle.com/code/egortrushin/gwi-unet-with-float16-dataset/notebook\n# Model\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\n\nclass ResidualDoubleConv(nn.Module):\n    \"\"\"(Convolution => [BN] => ReLU) * 2 + Residual Connection\"\"\"\n\n    def __init__(self, in_channels, out_channels, mid_channels=None, use_bn=True):\n        super().__init__()\n        if not mid_channels:\n            mid_channels = out_channels\n\n        # First convolution layer\n        # 3*3畳み込み. padding=1にすることで、出力のサイズは入力と変わらない。チャネル数はin_channelsからmid_channelsに変わる。\n        # テンキーで説明してみる。5,6,2,3の部分が左上の端だとする。でもpadding=1だから1,4,7,8,9の部分が埋められている。ので、5の部分も3*3の畳み込みができる.\n        # (batch_size, in_channels, height, width) -> (batch_size, mid_channels, height, width)\n        self.conv1 = nn.Conv2d(\n            in_channels, mid_channels, kernel_size=3, padding=1, bias=not use_bn\n        )\n        if use_bn:\n            self.bn1 = nn.BatchNorm2d(mid_channels)\n        else:\n            self.bn1 = nn.Identity()\n\n        self.relu = nn.ReLU(inplace=True)\n\n        # Second convolution layer\n        # (batch_size, mid_channels, height, width) -> (batch_size, out_channels, height, width)\n        self.conv2 = nn.Conv2d(\n            mid_channels, out_channels, kernel_size=3, padding=1, bias=not use_bn\n        )\n        if use_bn:\n            self.bn2 = nn.BatchNorm2d(out_channels)\n        else:\n            self.bn2 = nn.Identity()\n\n        # Shortcut connection to handle potential channel mismatch\n        if in_channels == out_channels:\n            self.shortcut = nn.Identity()\n        else:\n            # Projection shortcut: 1x1 conv + BN to match output channels\n            # 1*1畳み込み. padding=0にすることで、出力のサイズは入力と変わらない。チャネル数はin_channelsからout_channelsに変わる。\n            layers = [nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=not use_bn)]\n            if use_bn:\n                layers.append(nn.BatchNorm2d(out_channels))\n\n            self.shortcut = nn.Sequential(*layers)\n\n    def forward(self, x):\n        identity = x  # Store the input for the residual connection\n\n        \"\"\"\n        x-------------------------\\\n        |                         |\n        v                         |\n        conv1                     |\n        |                         |\n        v                         |\n        batch normalization 1     |\n        |                         |\n        v                         |\n        ReLU                      |\n        |                         |\n        v                         |\n        Conv2                     |\n        |                         |\n        v                         |\n        batch normalization 2     |\n        |                         |\n        V                         |\n        Add <---------------------/\n        |\n        v\n        ReLu\n        |\n        v\n        out\n        \"\"\"\n\n        # First conv block\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        # Second conv block (without final ReLU yet)\n        out = self.conv2(out)\n        out = self.bn2(out)\n\n        # Apply shortcut to the identity path\n        identity_mapped = self.shortcut(identity)\n\n        # Add the residual connection\n        out += identity_mapped\n\n        # Apply final ReLU\n        out = self.relu(out)\n        return out\n\n\nclass Up(nn.Module):\n    \"\"\"Upscaling then ResidualDoubleConv\"\"\"\n\n    def __init__(self, in_channels, out_channels, bilinear=True):\n        super().__init__()\n        self.bilinear = bilinear\n\n        if bilinear:\n            self.up = nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False)\n            # Input to ResidualDoubleConv = channels from upsampled layer below + channels from skip connection\n            # Output of ResidualDoubleConv = desired output channels for this decoder stage\n            self.conv = ResidualDoubleConv(\n                in_channels + out_channels, out_channels\n            )  # Use ResidualDoubleConv\n\n        else:  # Using ConvTranspose2d\n            # ConvTranspose halves the channels: in_channels -> in_channels // 2\n            self.up = nn.ConvTranspose2d(\n                in_channels, in_channels // 2, kernel_size=2, stride=2\n            )\n            # Input channels to ResidualDoubleConv\n            conv_in_channels = in_channels // 2  # Channels after ConvTranspose\n            skip_channels = out_channels  # Channels from skip connection\n            total_in_channels = conv_in_channels + skip_channels\n            self.conv = ResidualDoubleConv(\n                total_in_channels, out_channels\n            )  # Use ResidualDoubleConv\n\n    def forward(self, x1, x2):\n        # x1 is the feature map from the layer below (needs upsampling)\n        # x2 is the skip connection from the corresponding encoder layer\n        x1 = self.up(x1)\n\n        # Pad x1 if its dimensions don't match x2 after upsampling\n        # Input is CHW\n        diffY = x2.size(2) - x1.size(2)\n        diffX = x2.size(3) - x1.size(3)\n\n        # Pad format: (padding_left, padding_right, padding_top, padding_bottom)\n        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])\n\n        # Concatenate along the channel dimension\n        x = torch.cat([x2, x1], dim=1)\n        return self.conv(x)\n\n\nclass OutConv(nn.Module):\n    \"\"\"1x1 Convolution for the output layer\"\"\"\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass UNet(nn.Module):\n    \"\"\"U-Net architecture implementation with Residual Blocks\"\"\"\n\n    def __init__(\n        self,\n        n_channels=5,\n        n_classes=1,\n        init_features=32,\n        depth=5,  # number of pooling layers\n        bilinear=True,\n    ):\n        super().__init__()\n        self.n_channels = n_channels\n        self.n_classes = n_classes\n        self.bilinear = bilinear\n        self.depth = depth\n\n        # initial poolingをやめてみる→あんま効果なし.\n        self.initial_pool = nn.AvgPool2d(kernel_size=(14, 1), stride=(14, 1))\n\n        # --- Encoder ---\n        # convは最初のconv以外、チャネルを倍々にしていく\n        self.encoder_convs = nn.ModuleList()  # Store conv blocks (ResidualDoubleConv)\n        self.encoder_pools = nn.ModuleList()  # Store pool layers (MaxPool2d)\n\n        # Initial conv block (no pooling before it)\n        # Use ResidualDoubleConv for the initial convolution block\n        self.inc = ResidualDoubleConv(n_channels, init_features)\n        self.encoder_convs.append(self.inc)\n\n        current_features = init_features\n        for _ in range(depth):\n            # Define convolution block for this stage\n            conv = ResidualDoubleConv(current_features, current_features * 2, use_bn=True)\n            # Define pooling layer for this stage\n            # MaxPool2d をカーネルサイズ2, ストライド2で行う。\n            pool = nn.MaxPool2d(2)\n            self.encoder_convs.append(conv)\n            self.encoder_pools.append(pool)\n            current_features *= 2\n\n        # --- Bottleneck ---\n        # Use ResidualDoubleConv for the bottleneck\n        self.bottleneck = ResidualDoubleConv(current_features, current_features, use_bn=True)\n\n        # --- Decoder ---\n        self.decoder_blocks = nn.ModuleList()\n        # Input features start from bottleneck output features\n        # Output features at each stage are halved\n        for _ in range(depth):\n            # Up block uses ResidualDoubleConv internally and handles channels\n            up_block = Up(current_features, current_features // 2, bilinear)\n            self.decoder_blocks.append(up_block)\n            current_features //= 2  # Halve features for next Up block input\n\n        # --- Output Layer ---\n        # Input features are the output features of the last Up block\n        self.outc = OutConv(current_features, n_classes)\n\n    def _pad_or_crop(self, x, target_h=70, target_w=70):\n        \"\"\"Pads or crops input tensor x to target height and width.\"\"\"\n        _, _, h, w = x.shape\n        # Pad Height if needed\n        if h < target_h:\n            pad_top = (target_h - h) // 2\n            pad_bottom = target_h - h - pad_top\n            x = F.pad(x, (0, 0, pad_top, pad_bottom))  # Pad height only\n            h = target_h\n        # Pad Width if needed\n        if w < target_w:\n            pad_left = (target_w - w) // 2\n            pad_right = target_w - w - pad_left\n            x = F.pad(x, (pad_left, pad_right, 0, 0))  # Pad width only\n            w = target_w\n        # Crop Height if needed\n        if h > target_h:\n            crop_top = (h - target_h) // 2\n            # Use slicing to crop\n            x = x[:, :, crop_top : crop_top + target_h, :]\n            h = target_h\n        # Crop Width if needed\n        if w > target_w:\n            crop_left = (w - target_w) // 2\n            x = x[:, :, :, crop_left : crop_left + target_w]\n            w = target_w\n        return x\n\n    def forward(self, x):\n        # Initial pooling and resizing\n        x_pooled = self.initial_pool(x) # (14,1)のkernelを(14,1)のstrideで avg-poolingする。 サイズが(1/14, 1)倍になる\n        # 画像のサイズが(1000, 70)だったなら、(1000/14, 70) = (71, 70)\n        x_resized = self._pad_or_crop(x_pooled, target_h=70, target_w=70) # cropするので、(70,70)になる\n\n        # --- Encoder Path ---\n        skip_connections = []\n        xi = x_resized\n\n        # Apply initial conv (inc)\n        xi = self.encoder_convs[0](xi)\n        skip_connections.append(xi)  # Store output of inc\n\n        # Apply subsequent encoder convs and pools\n        # self.depth is the number of pooling layers\n        for i in range(self.depth):\n            # Apply conv block for this stage\n\n            # (batch_size, current_features, height, width) -> (batch_size, current_features*2, height, width)になる\n            xi = self.encoder_convs[i + 1](xi)\n            # Store skip connection *before* pooling\n            skip_connections.append(xi)\n            # Apply pooling layer for this stage\n            # MaxPooling2dを行う\n            # (batch_size, current_features*2, height, width) -> (batch_size, current_features*2, height/2, width/2)になる\n            xi = self.encoder_pools[i](xi)\n\n        # Apply bottleneck conv\n        xi = self.bottleneck(xi)\n\n        # --- Decoder Path ---\n        xu = xi  # Start with bottleneck output\n        # Iterate through decoder blocks and corresponding skip connections in reverse\n        for i, block in enumerate(self.decoder_blocks):\n            # Determine the correct skip connection index from the end\n            # Example: depth=5. Skips stored: [inc, enc1, enc2, enc3, enc4] (indices 0-4)\n            # Decoder 0 (Up(1024, 512)) needs skip 4 (enc4)\n            # Decoder 1 (Up(512, 256)) needs skip 3 (enc3) ...\n            # Decoder 4 (Up(64, 32)) needs skip 0 (inc)\n            skip_index = self.depth - 1 - i\n            skip = skip_connections[skip_index]\n            xu = block(xu, skip)  # Up block combines xu (from below) and skip\n\n        # --- Final Output ---\n        logits = self.outc(xu)\n        # Apply scaling and offset specific to the problem's target range\n        # だいたい中央値は3000くらいなので、3000にする。\n        # 標準偏差が800くらいなんだけど速度の分布は正規分布よりももっと裾が広いきがするから、スケールは1000くらい\n        output = logits * 1000.0 + 3000.0\n        return output\n\n\n# end of unet_def.py","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-15T11:46:22.137526Z","iopub.execute_input":"2025-06-15T11:46:22.137832Z","iopub.status.idle":"2025-06-15T11:46:22.150843Z","shell.execute_reply.started":"2025-06-15T11:46:22.137807Z","shell.execute_reply":"2025-06-15T11:46:22.150026Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Add the directory where data_helper.py lives to the system path\nimport sys\nsys.path.append('/kaggle/working')  # or wherever the .py files were saved\n\nfrom data_helper import TestDataset\nfrom unet_def import UNet\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-15T11:46:25.729497Z","iopub.execute_input":"2025-06-15T11:46:25.729770Z","iopub.status.idle":"2025-06-15T11:46:25.739432Z","shell.execute_reply.started":"2025-06-15T11:46:25.729750Z","shell.execute_reply":"2025-06-15T11:46:25.738083Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"credits to https://www.kaggle.com/code/haruiig/train-unet-with-tpu-pytorch-xla","metadata":{}},{"cell_type":"code","source":"# predict_for_submission.py\nimport argparse\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\nfrom collections import OrderedDict\n\nfrom data_helper import TestDataset\nfrom unet_def import UNet\n\nimport torch.serialization\ntorch.serialization.add_safe_globals({'_reconstruct': np.core.multiarray._reconstruct})\n\n\n\ndef load_model(model_path, init_features=32, depth=5, device='cpu'):\n    model = UNet(n_channels=5, n_classes=1, init_features=init_features, depth=depth)\n    # checkpoint = torch.load(model_path, map_location=device)\n    checkpoint = torch.load(model_path, map_location=device, weights_only=False)\n    state_dict = checkpoint[\"model_state_dict\"]\n    new_state = OrderedDict((k.replace(\"module.\", \"\"), v) for k, v in state_dict.items())\n    model.load_state_dict(new_state)\n    model.to(device)\n    model.eval()\n    return model\ndef make_submission(model, test_dir, output_csv, batch_size=32, device='cpu'):\n    from torch.utils.data import DataLoader\n    from data_helper import TestDataset\n    from pathlib import Path\n    import torch\n    import numpy as np\n    import pandas as pd\n    from tqdm import tqdm\n\n    test_files = list(Path(test_dir).glob(\"*.npy\"))\n    dataset = TestDataset(test_files)\n    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n\n    all_rows = []\n    for inputs, oids in tqdm(dataloader, desc=\"Predicting\"):\n        # Ensure all inputs are float32 tensors\n        inputs = torch.tensor(inputs, dtype=torch.float32).to(device)\n\n        with torch.inference_mode():\n            outputs = model(inputs).cpu().numpy()  # Shape: (batch_size, 1, 70, 70)\n\n        for pred, oid in zip(outputs, oids):\n            pred = pred[0]  # Remove channel dimension → shape (70, 70)\n            for y in range(pred.shape[0]):\n                row = [f\"{oid}_y_{y}\"] + pred[y][1::2].tolist()  # odd x_ values\n                all_rows.append(row)\n\n    # Build DataFrame and save\n    columns = [\"oid_ypos\"] + [f\"x_{i}\" for i in range(1, 70, 2)]\n    df = pd.DataFrame(all_rows, columns=columns)\n    df.to_csv(output_csv, index=False)\n    print(f\"✅ Submission saved to {output_csv}\")\n        # Ensure all inputs are float32 tensors\n\n\n\n# def make_submission(model, test_dir, output_csv, batch_size=32, device='cpu'):\n#     test_files = list(Path(test_dir).glob(\"*.npy\"))\n#     dataset = TestDataset(test_files)\n#     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n\n#     all_rows = []\n#     with torch.inference_mode(), torch.autocast(device_type=device):\n#         # for inputs, oids in tqdm(dataloader, desc=\"Predicting\"):\n#         #     inputs = inputs.to(device)\n#         #     outputs = model(inputs).float().cpu().numpy()  # (batch_size, 1, 70, 70)\n#             for inputs, oids in tqdm(dataloader, desc=\"Predicting\"):\n#                 inputs = torch.tensor(inputs).to(dtype=torch.float32, device=device)\n#                 with torch.inference_mode(), torch.autocast(device_type=device):\n#                     outputs = model(inputs).float().cpu().numpy()\n\n\n#             for pred, oid in zip(outputs, oids):\n#                 pred = pred[0]  # shape (70, 70)\n#                 for y_pos in range(70):\n#                     row_id = f\"{oid}_y_{y_pos}\"\n#                     odd_vals = pred[y_pos, 1::2]  # x_1, x_3, ..., x_69\n#                     all_rows.append([row_id] + odd_vals.tolist())\n\n#     columns = [\"oid_ypos\"] + [f\"x_{i}\" for i in range(1, 70, 2)]\n#     df = pd.DataFrame(all_rows, columns=columns)\n#     df.to_csv(output_csv, index=False)\n#     print(f\"Submission saved to {output_csv}\")\n\n# if __name__ == \"__main__\":\n#     parser = argparse.ArgumentParser()\n#     parser.add_argument(\"--model\", type=str, required=True, help=\"Path to model weights\")\n#     parser.add_argument(\"--test-dir\", type=str, required=True, help=\"Directory with .npy input files\")\n#     parser.add_argument(\"--output\", type=str, required=True, help=\"Submission CSV file path\")\n#     parser.add_argument(\"--batch-size\", type=int, default=32)\n#     parser.add_argument(\"--device\", type=str, default=\"cuda\" if torch.cuda.is_available() else \"cpu\")\n#     parser.add_argument(\"--init-features\", type=int, default=32)\n#     parser.add_argument(\"--depth\", type=int, default=5)\n#     args = parser.parse_args()\n\n#     model = load_model(args.model, args.init_features, args.depth, args.device)\n#     make_submission(model, args.test_dir, args.output, args.batch_size, args.device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-15T11:58:33.086582Z","iopub.execute_input":"2025-06-15T11:58:33.086889Z","iopub.status.idle":"2025-06-15T11:58:33.099856Z","shell.execute_reply.started":"2025-06-15T11:58:33.086867Z","shell.execute_reply":"2025-06-15T11:58:33.099036Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Manually defined arguments\nmodel_path = \"/kaggle/input/seismic-wave-inversion/pytorch/default/6/UNet_models/UNetFloat16TPUWithCheckpointMAX/checkpoint_all_UNetFloat16TPUWithCheckpointMAX_bs256_epoch08_valloss82.47.pth.pth\"\ntest_dir = \"/kaggle/input/open-wfi-test/test\"\noutput_csv = \"/kaggle/working/submission.csv\"\nbatch_size = 32\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\ninit_features = 32\ndepth = 5\n# Run prediction\nmodel = load_model(model_path, init_features, depth, device)\nmake_submission(model, test_dir, output_csv, batch_size, device)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-15T11:58:36.412882Z","iopub.execute_input":"2025-06-15T11:58:36.413623Z","execution_failed":"2025-06-15T13:30:10.045Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\nimport random\nfrom pathlib import Path\nfrom data_helper import TestDataset\nfrom torch.utils.data import DataLoader\nimport torch\nfrom unet_def import UNet\nfrom collections import OrderedDict\n\n# Correct the line that caused the error by passing the actual function\nimport numpy as np\ntorch.serialization.add_safe_globals({'_reconstruct': np.core.multiarray._reconstruct})\n\n# # --- Config ---\n# model_path = \"/kaggle/input/seismic-wave-inversion/pytorch/default/4/UNet_models/UNetFloat16TPUWithCheckpointMAX/checkpoint_all_UNetFloat16TPUWithCheckpointMAX_bs256_epoch11_valloss88.56.pth.pth\"\n# test_dir = \"/kaggle/input/open-wfi-test/test\"\n# device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n# init_features = 32\n# depth = 5\n# batch_size = 8\n\n# # --- Load model ---\n# model = UNet(n_channels=5, n_classes=1, init_features=init_features, depth=depth).to(device)\n# state_dict = torch.load(model_path, map_location=device, weights_only=False) # Changed weights_only to False based on previous context \n# state_dict = OrderedDict((k.replace(\"module.\", \"\"), v) for k, v in state_dict.items())\n# model.load_state_dict(state_dict)\n# model.eval()\n\n# --- Load dataset ---\ntest_files = sorted(list(Path(test_dir).glob(\"*.npy\")))\ndataset = TestDataset(test_files)\ndataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n\n# --- Choose 5 random indices ---\nrandom.seed(42)\nindices = random.sample(range(len(dataset)), 5)\n\n# --- Predict and plot ---\nfig, axes = plt.subplots(5, 2, figsize=(10, 15))\nfor i, idx in enumerate(indices):\n    input_data, oid = dataset[idx]\n    # Ensure input_data is a float32 tensor as expected by the model\n    input_tensor = torch.tensor(input_data, dtype=torch.float32).unsqueeze(0).to(device)\n\n    with torch.inference_mode(): # Removed torch.autocast since it might not be needed or could cause issues without proper setup for mixed precision\n        pred = model(input_tensor).squeeze().cpu().numpy()\n\n    axes[i, 0].imshow(input_data[0], cmap=\"gray\")\n    axes[i, 0].set_title(f\"Input (channel 0) - {oid}\")\n    axes[i, 0].axis(\"off\")\n\n    axes[i, 1].imshow(pred, cmap=\"viridis\")\n    axes[i, 1].set_title(f\"Predicted Velocity - {oid}\")\n    axes[i, 1].axis(\"off\")\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-15T11:46:40.656885Z","iopub.status.idle":"2025-06-15T11:46:40.657268Z","shell.execute_reply.started":"2025-06-15T11:46:40.657101Z","shell.execute_reply":"2025-06-15T11:46:40.657115Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}