{"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":11367935,"sourceType":"datasetVersion","datasetId":7116013},{"sourceId":11368083,"sourceType":"datasetVersion","datasetId":7116134},{"sourceId":11368169,"sourceType":"datasetVersion","datasetId":7116196},{"sourceId":11368268,"sourceType":"datasetVersion","datasetId":7116272},{"sourceId":11368499,"sourceType":"datasetVersion","datasetId":7116445},{"sourceId":11368545,"sourceType":"datasetVersion","datasetId":7116479},{"sourceId":11368547,"sourceType":"datasetVersion","datasetId":7116481},{"sourceId":11376433,"sourceType":"datasetVersion","datasetId":7122462},{"sourceId":11376448,"sourceType":"datasetVersion","datasetId":7122476},{"sourceId":11376464,"sourceType":"datasetVersion","datasetId":7122489},{"sourceId":11376742,"sourceType":"datasetVersion","datasetId":7122712},{"sourceId":11376868,"sourceType":"datasetVersion","datasetId":7122812},{"sourceId":11376871,"sourceType":"datasetVersion","datasetId":7122814},{"sourceId":11376872,"sourceType":"datasetVersion","datasetId":7122815},{"sourceId":11376935,"sourceType":"datasetVersion","datasetId":7122866},{"sourceId":11377083,"sourceType":"datasetVersion","datasetId":7122981},{"sourceId":11377231,"sourceType":"datasetVersion","datasetId":7123090},{"sourceId":11377291,"sourceType":"datasetVersion","datasetId":7123138},{"sourceId":11377325,"sourceType":"datasetVersion","datasetId":7123163},{"sourceId":11377334,"sourceType":"datasetVersion","datasetId":7123172},{"sourceId":11377594,"sourceType":"datasetVersion","datasetId":7123380},{"sourceId":11377614,"sourceType":"datasetVersion","datasetId":7123394},{"sourceId":11377741,"sourceType":"datasetVersion","datasetId":7123490},{"sourceId":11377752,"sourceType":"datasetVersion","datasetId":7123499},{"sourceId":11377756,"sourceType":"datasetVersion","datasetId":7123503},{"sourceId":11377935,"sourceType":"datasetVersion","datasetId":7123649},{"sourceId":11377970,"sourceType":"datasetVersion","datasetId":7123675},{"sourceId":11378141,"sourceType":"datasetVersion","datasetId":7123813},{"sourceId":11378162,"sourceType":"datasetVersion","datasetId":7123831},{"sourceId":11378178,"sourceType":"datasetVersion","datasetId":7123842}],"dockerImageVersionId":31012,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Initialize","metadata":{}},{"cell_type":"code","source":"!pip install transformers torchvision faiss-cpu\n\nimport numpy as np\nimport pandas as pd\nimport os\nfrom pathlib import Path\nimport torch\nimport faiss\nimport matplotlib.pyplot as plt\nimport matplotlib.colors as mcolors\nimport gc\nfrom transformers import CLIPModel, CLIPProcessor\nfrom torchvision.transforms import Compose, Resize, ToTensor, Normalize, ToPILImage\nfrom sklearn.model_selection import train_test_split\nfrom tqdm import tqdm\n\n# Verify GPU availability\nprint(f\"PyTorch Version: {torch.__version__}\")\nprint(f\"CUDA Available: {torch.cuda.is_available()}\")\n\nif torch.cuda.is_available(): \n    device = torch.device(\"cuda:0\") \n    print(f\"GPU Name: {torch.cuda.get_device_name(0)}\") \n    print(f\"CUDA Version: {torch.version.cuda}\")\n    print(f\"Number of GPUs: {torch.cuda.device_count()}\")\nelse: \n    device = torch.device(\"cpu\") \n    print(f\"Using device: {device}\")\n    \n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T16:02:45.156121Z","iopub.execute_input":"2025-05-08T16:02:45.157415Z","iopub.status.idle":"2025-05-08T16:02:52.938342Z","shell.execute_reply.started":"2025-05-08T16:02:45.157383Z","shell.execute_reply":"2025-05-08T16:02:52.936922Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Find all sample directories\n### Use Kaggle Original and Open FWI database","metadata":{}},{"cell_type":"code","source":"# New updated \n# Collect Original Kaggle input files\ninput_dir = Path('/kaggle/input/waveform-inversion/train_samples')\n\ndata_files = []\nmodel_files = []\n\n# Add files from Vel and Style families\n\ndata_files = sorted(input_dir.rglob('data*.npy'))\nmodel_files = sorted(input_dir.rglob('model*.npy'))\n\n# Add Fault family\nfault_data_files = sorted(input_dir.rglob('seis*.npy'))\nfault_model_files = sorted(input_dir.rglob('vel*.npy'))\ndata_files += fault_data_files\nmodel_files += fault_model_files\n\n#LOAD the OPEN FWI training files\nroot_dir = Path(\"/kaggle/input\")\n\nall_data_files = []\nall_model_files = []\nseis_files = []\nvel_files = []\nseis2_files = []\nvel2_files = []\n\n# Loop over waveform-inversion-1 to waveform-inversion-30\nfor i in range(1,31):\n    subroot = root_dir / f\"waveform-inversion-{i}\"\n    seis_files = sorted(subroot.rglob(\"seis*.npy\"))\n    vel_files = sorted(subroot.rglob(\"vel*.npy\"))\n    all_data_files.extend(seis_files)\n    all_model_files.extend(vel_files)\n\n    seis2_files = sorted(subroot.rglob(\"data*.npy\"))\n    vel2_files = sorted(subroot.rglob(\"model*.npy\"))\n    all_data_files.extend(seis2_files)\n    all_model_files.extend(vel2_files)\n    seis_files = []\n    vel_files = [] \n    seis2_files = []\n    vel2_files = []\n  \n## File check\n#for i in range(1, 20):\n#    print(f\"index: {i}\" )\n#    print(\"data\", all_data_files[i])\n#    print(\"model\", all_model_files[i])   \n\n# Randomly select some Open FWI consistent pairs\n\nprint(\"Number of Open FWI seismic training files found:\", len(all_data_files))\n\nimport random\n# Set a fixed seed for reproducibility\nrandom.seed(42)\ndata_FWIfiles = []\nmodel_FWIfiles = []\n\n#Max number tested is 80, larger may crash memory\nselect_FWI  = 80\nprint(\"Number of Open FWI to add to original set:\", select_FWI)\nindices = random.sample(range(len(all_data_files)), select_FWI)\ndata_FWIfiles = [all_data_files[i] for i in indices]\nmodel_FWIfiles = [all_model_files[i] for i in indices]\n\n# Combine Original Kaggle and select Open FWI files\ndata_files += data_FWIfiles\nmodel_files += model_FWIfiles\n\nprint(\"Number of combined origianl/Open FWI training files found:\", len(data_files))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T16:03:08.358372Z","iopub.execute_input":"2025-05-08T16:03:08.35874Z","iopub.status.idle":"2025-05-08T16:03:14.360585Z","shell.execute_reply.started":"2025-05-08T16:03:08.358714Z","shell.execute_reply":"2025-05-08T16:03:14.359344Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load and preprocess sample files ","metadata":{}},{"cell_type":"code","source":"\nX_list = []\ny_list = []\n\nfor count, (data_path, model_path) in enumerate(zip(data_files, model_files), start=1):\n    try:\n        #print(f\"Loading file {count} of {len(data_files)}\")\n        #print(\"File\", data_path)\n        \n        # Memory-mapped loading\n        X = np.load(data_path, mmap_mode='r')  # (500, 5, 980, 70)\n        y = np.load(model_path, mmap_mode='r') # (500, 1, 70, 70)\n\n        # Preprocessing\n        X = X[:, :, 20:1000, :]        # -> (500, 5, 980, 70)\n        X = X[:, :, ::14, :]           # -> (500, 5, 70, 70)\n        X = X.transpose(0, 2, 1, 3)    # -> (500, 70, 5, 70)\n        X = X.reshape(500, 70, 350)    # -> (500, 70, 350)\n        X = X[:, :, ::5]               # -> (500, 70, 70)\n\n        if y.ndim == 4:\n            y = np.squeeze(y, axis=1)  # -> (500, 70, 70)\n\n        y = (y - 1500.0) / (4500.0 - 1500.0)\n\n        X_list.append(X)\n        y_list.append(y)\n\n    except Exception as e:\n        print(f\"Error loading {data_path} or {model_path}: {str(e)}\")\n        continue\n\n# Final conversion (only once, fast)\nX_data = np.concatenate(X_list, axis=0)\ny_data = np.concatenate(y_list, axis=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T16:03:42.465818Z","iopub.execute_input":"2025-05-08T16:03:42.466251Z","iopub.status.idle":"2025-05-08T16:20:48.771028Z","shell.execute_reply.started":"2025-05-08T16:03:42.466225Z","shell.execute_reply":"2025-05-08T16:20:48.769685Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\n# Train/Val Split (80/20)\nX_train, X_val, y_train, y_val = train_test_split(\n    X_data, y_data, test_size=0.2, random_state=42\n)\n\nprint(\"X_train shape:\", X_train.shape)\nprint(\"X_val shape:\", X_val.shape)\nprint(\"y_train shape:\", y_train.shape)\nprint(\"y_val shape:\", y_val.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T16:23:22.939333Z","iopub.execute_input":"2025-05-08T16:23:22.940291Z","iopub.status.idle":"2025-05-08T16:23:26.800447Z","shell.execute_reply.started":"2025-05-08T16:23:22.940256Z","shell.execute_reply":"2025-05-08T16:23:26.799258Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualize training samples","metadata":{}},{"cell_type":"code","source":"# Loop through all samples\nfor sample_idx in range(0, len(X_train), 5000):\n    seismic = X_train[sample_idx]      # shape (70, 70)\n    velocity = y_train[sample_idx]     # shape (70, 70)\n\n    # Plot them\n    plt.figure(figsize=(12, 5))\n\n    # Seismic Input\n    plt.subplot(1, 2, 1)\n    plt.imshow(seismic, aspect='auto', cmap='seismic', origin='lower')\n    plt.title(f\"Seismic Input (Sample {sample_idx})\")\n    plt.xlabel(\"Fused Shots Offset\")\n    plt.ylabel(\"Time\")\n    plt.colorbar(label=\"Amplitude\")\n\n    # Velocity Model (Target)\n    plt.subplot(1, 2, 2)\n    plt.imshow(velocity, aspect='equal', cmap='viridis', origin='lower', vmin=0, vmax=1)\n    plt.title(\"Velocity Model (Target)\")\n    plt.xlabel(\"X Position\")\n    plt.ylabel(\"Depth\")\n    plt.colorbar(label=\"Velocity\")\n\n    plt.tight_layout()\n    plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T16:23:30.850748Z","iopub.execute_input":"2025-05-08T16:23:30.851141Z","iopub.status.idle":"2025-05-08T16:23:35.785307Z","shell.execute_reply.started":"2025-05-08T16:23:30.851112Z","shell.execute_reply":"2025-05-08T16:23:35.784161Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## CLIP","metadata":{}},{"cell_type":"code","source":"from torchvision.transforms import Compose, Resize, ToTensor, Normalize, ToPILImage\n\n# Load CLIP\nclip_model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\").to(device)\nclip_processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n\n# Simulated example data — replace with your real arrays\n# X_train: (8000, 70, 70), y_train: (8000, 70, 70)\n# X_val: (2000, 70, 70), y_val: (2000, 70, 70)\n\n# Transform: grayscale 70×70 ➜ RGB 224×224\ntransform = Compose([\n    ToPILImage(),\n    Resize((224, 224)),\n    ToTensor(),\n    Normalize((0.5,), (0.5,))\n])\n\ndef preprocess_images(imgs):\n    \"\"\"Convert (N, 70, 70) grayscale to (N, 3, 224, 224) RGB tensors.\"\"\"\n    tensors = []\n    for img in imgs:\n        img_rgb = np.repeat(img[:, :, np.newaxis], 3, axis=2)  # (70, 70, 3)\n        img_tensor = transform(img_rgb)\n        tensors.append(img_tensor)\n    return torch.stack(tensors)\n\ndef extract_clip_features(X):\n    \"\"\"Convert images to CLIP features.\"\"\"\n    clip_features = []\n    for i in tqdm(range(0, len(X), 64)):\n        batch = preprocess_images(X[i:i+64]).to(device)\n        with torch.no_grad():\n            features = clip_model.get_image_features(batch)\n            features = features / features.norm(dim=1, keepdim=True)\n            clip_features.append(features.cpu().numpy())\n    return np.vstack(clip_features)\n\n# Extract features\nclip_train_features = extract_clip_features(X_train)\nclip_val_features = extract_clip_features(X_val)\n\n# FAISS: cosine similarity via normalized dot product\nfaiss_index = faiss.IndexFlatIP(clip_train_features.shape[1])\nfaiss_index.add(clip_train_features)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T16:23:40.994976Z","iopub.execute_input":"2025-05-08T16:23:40.995461Z","iopub.status.idle":"2025-05-08T17:46:47.505236Z","shell.execute_reply.started":"2025-05-08T16:23:40.995427Z","shell.execute_reply":"2025-05-08T17:46:47.501168Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Prediction vs Validation Plots\n","metadata":{}},{"cell_type":"code","source":"# Predict from nearest neighbor\nval_preds = []\nnn_indices = []\n\nfor vec in clip_val_features:\n    D, I = faiss_index.search(vec[None, :], 1)\n    val_preds.append(y_train[I[0][0]])\n    nn_indices.append(I[0][0])\n\nval_preds = np.array(val_preds)\n\n\n\n# Visualize\nn = 5\nsamples = np.random.choice(len(X_val), n, replace=False)\n\nplt.figure(figsize=(15, 4 * n))\nfor i, idx in enumerate(samples):\n    gt = y_val[idx] * (4500 - 1500) + 1500\n    pred = val_preds[idx] * (4500 - 1500) + 1500\n    diff = pred - gt\n\n    plt.subplot(n, 3, 3*i + 1)\n    plt.imshow(gt, cmap='viridis', vmin=1500, vmax=4500)\n    plt.title(f\"Ground Truth #{idx}\")\n    plt.colorbar()\n\n    plt.subplot(n, 3, 3*i + 2)\n    plt.imshow(pred, cmap='viridis', vmin=1500, vmax=4500)\n    plt.title(f\"CLIP NN Prediction #{idx}\")\n    plt.colorbar()\n\n    plt.subplot(n, 3, 3*i + 3)\n    plt.imshow(diff, cmap='bwr')\n    plt.title(\"Difference\")\n    plt.colorbar()\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T17:47:44.331326Z","iopub.execute_input":"2025-05-08T17:47:44.331993Z","iopub.status.idle":"2025-05-08T17:49:36.432649Z","shell.execute_reply.started":"2025-05-08T17:47:44.331958Z","shell.execute_reply":"2025-05-08T17:49:36.431391Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Prepare Submission File","metadata":{}},{"cell_type":"code","source":"# CLIP transform: grayscale (70x70) ➜ RGB (224x224)\nclip_transform = Compose([\n    ToPILImage(),\n    Resize((224, 224)),\n    ToTensor(),\n    Normalize((0.5,), (0.5,))\n])\n\n\ndef preprocess_clip_input(X):\n    \"\"\"Preprocess raw test data into 70x70 grayscale, then into 3x224x224 tensor.\"\"\"\n    X = X[:, 20:1000, :]\n    X = X[:, ::14, :]\n    X = X.transpose(1, 0, 2).reshape(70, -1)\n    X = X[:, ::5]  # now (70, 70)\n    X_rgb = np.repeat(X[:, :, np.newaxis], 3, axis=2)  # (70, 70, 3)\n    X_tensor = clip_transform(X_rgb).unsqueeze(0).to(device)  # (1, 3, 224, 224)\n    return X_tensor, X  # return both CLIP input and the original grayscale\n\n# Load test files\ntest_dir = Path('/kaggle/input/waveform-inversion/test')\nsubmission_rows = []\ncount = 0\ntest_examples = []\n\nfor file_path in tqdm(list(test_dir.rglob('*.npy'))):\n    # Load and preprocess\n    X_raw = np.load(file_path)\n    \n    count = count + 1\n    if count % 5000 == 0:\n        print(\"Loading file number:\", count)\n        test_examples.append(file_path)\n\n    \n    clip_tensor, _ = preprocess_clip_input(X_raw)   \n    # Get CLIP features\n    with torch.no_grad():\n        clip_feat = clip_model.get_image_features(clip_tensor).cpu().numpy()\n        clip_feat /= np.linalg.norm(clip_feat, axis=1, keepdims=True)\n\n    # Nearest neighbor search\n    D, I = faiss_index.search(clip_feat, 1)\n    nn_pred = y_train[I[0][0]]  # (70, 70)\n\n    # Rescale velocity prediction\n    pred = nn_pred * (4500.0 - 1500.0) + 1500.0\n\n    # Format for submission\n    file_stem = file_path.stem\n    for i in range(pred.shape[0]):\n        oid_ypos = f\"{file_stem}_y_{i}\"\n        x_values = pred[i, 1::2]  # x_1, x_3, ..., x_69\n        row = [oid_ypos] + x_values.tolist()\n        submission_rows.append(row)\n\n# Build submission DataFrame\ncolumns = ['oid_ypos'] + [f'x_{i}' for i in range(1, 70, 2)]\nsubmission_df = pd.DataFrame(submission_rows, columns=columns)\nsubmission_df.to_csv(\"submission.csv\", index=False)\nprint(\"✅ CLIP-based submission file saved!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T14:29:51.741171Z","iopub.execute_input":"2025-05-08T14:29:51.741528Z","iopub.status.idle":"2025-05-08T16:02:21.133586Z","shell.execute_reply.started":"2025-05-08T14:29:51.741502Z","shell.execute_reply":"2025-05-08T16:02:21.127989Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualize a test sample","metadata":{}},{"cell_type":"code","source":"\n# Put model in evaluation mode\nmodel.eval()\n\nfor i in range (len(test_examples)):\n    X = np.load(test_examples[i])  # shape (5, 1000, 70)\n    print(test_examples[i])\n    X = X[:, 20:1000, :]\n    X = X[:, ::14, :]\n    X = X.transpose(1, 0, 2).reshape(70, -1)  # (70, 350)\n    X = X[:, ::5]  # (70, 70)\n\n    X_tensor = torch.tensor(X, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)  # (1, 1, 70, 70)\n\n    # Extract embedding and search\n    with torch.no_grad():\n        embedding = feature_extractor(X_tensor).cpu().numpy()\n    _, indices = faiss_index.search(embedding, 1)\n    pred = y_train[indices[0][0]] * (4500 - 1500) + 1500\n\n    # Plot Prediction\n    plt.imshow(pred, cmap='viridis', aspect='equal', origin='lower', vmin=1500, vmax=4500)\n    plt.colorbar()\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T01:51:46.828526Z","iopub.execute_input":"2025-05-08T01:51:46.828958Z","iopub.status.idle":"2025-05-08T01:51:51.067602Z","shell.execute_reply.started":"2025-05-08T01:51:46.828931Z","shell.execute_reply":"2025-05-08T01:51:51.066447Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}