{"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"}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport glob\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom tqdm import tqdm\nimport random\n\n# --- 1. CONFIGURAÇÃO ---\n# Parâmetros para controlar o treinamento. Modifique-os conforme necessário.\nclass Config:\n    TRAIN_DIR = \"/kaggle/input/waveform-inversion/train_samples/\"\n    TEST_DIR = \"/kaggle/input/waveform-inversion/test/\"\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    BATCH_SIZE = 8\n    EPOCHS = 10 # Comece com poucas épocas e aumente conforme o resultado\n    LEARNING_RATE = 1e-4\n    VALIDATION_SPLIT = 0.1 # 10% dos dados para validação\n    SEED = 42\n\n# Para garantir a reprodutibilidade dos resultados\nrandom.seed(Config.SEED)\nnp.random.seed(Config.SEED)\ntorch.manual_seed(Config.SEED)\nif Config.DEVICE == \"cuda\":\n    torch.cuda.manual_seed(Config.SEED)\n\n# --- 2. PIPELINE DE DADOS (DATASET) ---\n# Esta classe é o coração da pipeline de dados. Ela lida com o carregamento,\n# pré-processamento e transformação dos dados sísmicos em tensores.\n\nclass SeismicDataset(Dataset):\n    \"\"\"\n    Dataset personalizado para carregar os dados sísmicos e os mapas de velocidade.\n    Cada arquivo .npy contém um batch de 500 amostras. Esta classe trata cada uma\n    dessas 500 amostras como um item individual.\n    \"\"\"\n    def __init__(self, file_paths, is_test=False):\n        self.file_paths = file_paths\n        self.is_test = is_test\n        # Cada arquivo contém 500 amostras. O tamanho do dataset é o número de arquivos * 500.\n        self.num_samples_per_file = 500\n        \n        # O diretório de teste tem uma estrutura diferente (cada .npy é uma amostra)\n        if self.is_test:\n            # Nos dados de teste, cada arquivo .npy é uma única amostra 3D (5, 1000, 70), \n            # não um batch de amostras 4D.\n            self.num_samples_per_file = 1\n\n        # Transformação para redimensionar os dados sísmicos de (1000, 70) para (70, 70)\n        self.resizer = transforms.Resize((70, 70), antialias=True)\n        \n    def __len__(self):\n        return len(self.file_paths) * self.num_samples_per_file\n\n    def __getitem__(self, idx):\n        # Mapeia o índice global para um índice de arquivo e um índice de amostra dentro do arquivo\n        file_idx = idx // self.num_samples_per_file\n        sample_idx = idx % self.num_samples_per_file\n        \n        seis_path = self.file_paths[file_idx][0]\n        \n        if self.is_test:\n            # Carrega a amostra de teste única\n            seis_sample_np = np.load(seis_path).astype(np.float32) # Shape: (5, 1000, 70)\n        else:\n            # Carrega o batch de treino e pega a amostra específica\n            seis_data_full = np.load(seis_path).astype(np.float32)\n            seis_sample_np = seis_data_full[sample_idx] # Shape: (5, 1000, 70)\n            \n        seis_sample = torch.from_numpy(seis_sample_np)\n        \n        # --- Pré-processamento dos Dados de Entrada (Sismograma) ---\n        # 1. Redimensionar a altura (time_steps) de 1000 para 70\n        seis_resized = self.resizer(seis_sample) # Shape: (5, 70, 70)\n        \n        # 2. Normalização por amostra\n        max_abs = torch.abs(seis_resized).max()\n        if max_abs > 0:\n            seis_resized /= max_abs\n        \n        if self.is_test:\n            return seis_resized\n        else:\n            vel_path = self.file_paths[file_idx][1]\n            vel_map_full = np.load(vel_path).astype(np.float32)\n            \n            # Pega o mapa de velocidade correspondente (shape: 1, 70, 70)\n            vel_sample = torch.from_numpy(vel_map_full[sample_idx])\n\n            # --- Pré-processamento do Alvo (Mapa de Velocidade) ---\n            # Normalizar para o intervalo [0, 1]\n            vmin, vmax = 1500, 4500\n            vel_sample_normalized = (vel_sample - vmin) / (vmax - vmin)\n            \n            return seis_resized, vel_sample_normalized\n\n\n# --- 3. MODELO (ARQUITETURA U-NET) ---\n# A implementação da U-Net, ideal para tarefas de imagem-para-imagem.\nclass DoubleConv(nn.Module):\n    def __init__(self, in_channels, out_channels, mid_channels=None):\n        super().__init__()\n        if not mid_channels:\n            mid_channels = out_channels\n        self.double_conv = nn.Sequential(\n            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(mid_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n    def forward(self, x):\n        return self.double_conv(x)\n\nclass Down(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.maxpool_conv = nn.Sequential(\n            nn.MaxPool2d(2),\n            DoubleConv(in_channels, out_channels)\n        )\n    def forward(self, x):\n        return self.maxpool_conv(x)\n\nclass Up(nn.Module):\n    def __init__(self, in_channels, out_channels, bilinear=True):\n        super().__init__()\n        if bilinear:\n            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)\n        else:\n            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)\n            self.conv = DoubleConv(in_channels, out_channels)\n\n    def forward(self, x1, x2):\n        x1 = self.up(x1)\n        # Corrige a diferença de tamanho devido ao arredondamento no MaxPool\n        diffY = x2.size()[2] - x1.size()[2]\n        diffX = x2.size()[3] - x1.size()[3]\n\n        # Adiciona padding ao tensor upsampled (x1) para que tenha o mesmo tamanho do tensor da skip connection (x2)\n        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,\n                        diffY // 2, diffY - diffY // 2])\n        \n        x = torch.cat([x2, x1], dim=1)\n        return self.conv(x)\n\nclass UNet(nn.Module):\n    def __init__(self, n_channels, n_classes, bilinear=True):\n        super(UNet, self).__init__()\n        self.n_channels = n_channels\n        self.n_classes = n_classes\n        self.bilinear = bilinear\n        factor = 2 if bilinear else 1\n\n        self.inc = DoubleConv(n_channels, 64)\n        self.down1 = Down(64, 128)\n        self.down2 = Down(128, 256)\n        self.down3 = Down(256, 512)\n        self.down4 = Down(512, 1024 // factor)\n        self.up1 = Up(1024, 512 // factor, bilinear)\n        self.up2 = Up(512, 256 // factor, bilinear)\n        self.up3 = Up(256, 128 // factor, bilinear)\n        self.up4 = Up(128, 64, bilinear)\n        self.outc = nn.Conv2d(64, n_classes, kernel_size=1)\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        x5 = self.down4(x4)\n        x = self.up1(x5, x4)\n        x = self.up2(x, x3)\n        x = self.up3(x, x2)\n        x = self.up4(x, x1)\n        logits = self.outc(x)\n        return logits\n\n\n# --- 4. LÓGICA DE TREINAMENTO ---\n\ndef train_one_epoch(model, dataloader, optimizer, criterion, device):\n    model.train()\n    total_loss = 0\n    progress_bar = tqdm(dataloader, desc=\"Training\", leave=False)\n    for inputs, targets in progress_bar:\n        inputs, targets = inputs.to(device), targets.to(device)\n        \n        optimizer.zero_grad()\n        outputs = model(inputs)\n        loss = criterion(outputs, targets)\n        loss.backward()\n        optimizer.step()\n        \n        total_loss += loss.item()\n        progress_bar.set_postfix(loss=loss.item())\n        \n    return total_loss / len(dataloader)\n\ndef validate_one_epoch(model, dataloader, criterion, device):\n    model.eval()\n    total_loss = 0\n    progress_bar = tqdm(dataloader, desc=\"Validating\", leave=False)\n    with torch.no_grad():\n        for inputs, targets in progress_bar:\n            inputs, targets = inputs.to(device), targets.to(device)\n            outputs = model(inputs)\n            loss = criterion(outputs, targets)\n            total_loss += loss.item()\n            progress_bar.set_postfix(loss=loss.item())\n            \n    return total_loss / len(dataloader)\n\n# --- 5. EXECUÇÃO PRINCIPAL ---\n\nif __name__ == '__main__':\n    print(f\"Usando dispositivo: {Config.DEVICE}\")\n\n    # Coletar todos os pares de arquivos de treino\n    all_files = []\n    for family in os.listdir(Config.TRAIN_DIR):\n        family_path = os.path.join(Config.TRAIN_DIR, family)\n        if os.path.isdir(family_path):\n            data_dir = os.path.join(family_path, 'data')\n            model_dir = os.path.join(family_path, 'model')\n            \n            # Lida com a estrutura de pastas diferente\n            if os.path.exists(data_dir) and os.path.exists(model_dir):\n                # Famílias Vel e Style\n                data_files = sorted(glob.glob(os.path.join(data_dir, '*.npy')))\n                model_files = sorted(glob.glob(os.path.join(model_dir, '*.npy')))\n            else:\n                # Famílias Fault\n                data_files = sorted(glob.glob(os.path.join(family_path, 'seis*.npy')))\n                model_files = sorted(glob.glob(os.path.join(family_path, 'vel*.npy')))\n            \n            for d_file, m_file in zip(data_files, model_files):\n                all_files.append((d_file, m_file))\n\n    # Embaralhar e dividir os arquivos para treino e validação\n    random.shuffle(all_files)\n    split_idx = int(len(all_files) * (1 - Config.VALIDATION_SPLIT))\n    train_files = all_files[:split_idx]\n    val_files = all_files[split_idx:]\n\n    print(f\"Total de arquivos: {len(all_files)}. Treino: {len(train_files)}, Validação: {len(val_files)}\")\n\n    # Criar Datasets e DataLoaders\n    train_dataset = SeismicDataset(train_files)\n    val_dataset = SeismicDataset(val_files)\n    \n    train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True)\n    val_loader = DataLoader(val_dataset, batch_size=Config.BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)\n\n    # Inicializar modelo, otimizador e função de perda\n    # n_channels=5 (5 fontes sísmicas), n_classes=1 (mapa de velocidade monocromático)\n    model = UNet(n_channels=5, n_classes=1).to(Config.DEVICE)\n    optimizer = torch.optim.AdamW(model.parameters(), lr=Config.LEARNING_RATE)\n    criterion = nn.L1Loss() # MAE Loss, conforme a métrica da competição\n\n    # Loop de treinamento\n    best_val_loss = float('inf')\n    for epoch in range(Config.EPOCHS):\n        print(f\"\\n--- Epoch {epoch+1}/{Config.EPOCHS} ---\")\n        train_loss = train_one_epoch(model, train_loader, optimizer, criterion, Config.DEVICE)\n        val_loss = validate_one_epoch(model, val_loader, criterion, Config.DEVICE)\n        \n        print(f\"Epoch {epoch+1}: Train Loss = {train_loss:.6f}, Val Loss = {val_loss:.6f}\")\n        \n        if val_loss < best_val_loss:\n            best_val_loss = val_loss\n            torch.save(model.state_dict(), 'best_model.pth')\n            print(\"✨ Modelo salvo! Melhor perda de validação até agora.\")\n\n    print(\"\\nTreinamento concluído!\")\n\n# --- 6. GERAÇÃO DA SUBMISSÃO (Exemplo) ---\n# Esta parte deve ser executada após o treinamento, num script separado ou célula.\n\ndef generate_submission():\n    print(\"Gerando arquivo de submissão...\")\n    model = UNet(n_channels=5, n_classes=1).to(Config.DEVICE)\n    model.load_state_dict(torch.load('best_model.pth'))\n    model.eval()\n\n    test_files = [(f, '') for f in glob.glob(os.path.join(Config.TEST_DIR, '*.npy'))]\n    test_dataset = SeismicDataset(test_files, is_test=True)\n    test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n\n    predictions = []\n    file_ids = [os.path.basename(f[0]).replace('.npy', '') for f in test_files]\n\n    with torch.no_grad():\n        for i, inputs in enumerate(tqdm(test_loader, desc=\"Predicting\")):\n            inputs = inputs.to(Config.DEVICE)\n            output = model(inputs) # Shape: (1, 1, 70, 70)\n            \n            # --- Pós-processamento ---\n            # 1. Desnormalizar a predição\n            vmin, vmax = 1500, 4500\n            unnormalized_prediction = output * (vmax - vmin) + vmin\n            \n            # 2. Garantir que os valores estão no intervalo físico\n            final_prediction = torch.clamp(unnormalized_prediction, vmin, vmax)\n            \n            pred_np = final_prediction.cpu().numpy().squeeze() # Shape: (70, 70)\n            \n            # Formatar para o arquivo de submissão\n            oid = file_ids[i]\n            for y_pos in range(pred_np.shape[0]):\n                row_id = f\"{oid}_y_{y_pos}\"\n                row_data = {'oid_ypos': row_id}\n                # A submissão exige apenas as colunas ímpares\n                for x_pos in range(1, pred_np.shape[1], 2):\n                    row_data[f'x_{x_pos}'] = pred_np[y_pos, x_pos]\n                predictions.append(row_data)\n\n    submission_df = pd.DataFrame(predictions)\n    submission_df.to_csv('submission.csv', index=False)\n    print(\"Arquivo submission.csv gerado com sucesso!\")\n\n\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-06-22T20:52:17.924851Z","iopub.execute_input":"2025-06-22T20:52:17.925966Z","iopub.status.idle":"2025-06-23T07:09:07.408195Z","shell.execute_reply.started":"2025-06-22T20:52:17.925921Z","shell.execute_reply":"2025-06-23T07:09:07.405989Z"}},"outputs":[],"execution_count":null}]}