{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13704713,"sourceType":"datasetVersion","datasetId":8718010},{"sourceId":13828653,"sourceType":"datasetVersion","datasetId":8807057},{"sourceId":13828654,"sourceType":"datasetVersion","datasetId":8807058},{"sourceId":13828677,"sourceType":"datasetVersion","datasetId":8807076},{"sourceId":14107165,"sourceType":"datasetVersion","datasetId":8985434},{"sourceId":14224921,"sourceType":"datasetVersion","datasetId":9074511},{"sourceId":14276753,"sourceType":"datasetVersion","datasetId":9111572},{"sourceId":14277062,"sourceType":"datasetVersion","datasetId":8720307},{"sourceId":14633554,"sourceType":"datasetVersion","datasetId":9347648},{"sourceId":14652719,"sourceType":"datasetVersion","datasetId":9009687},{"sourceId":4537,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":3329,"modelId":986},{"sourceId":636705,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":480070,"modelId":495773}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Install Dependencies","metadata":{}},{"cell_type":"code","source":"!pip install -U /kaggle/input/transformers-4-57-1/huggingface_hub-0.36.0-py3-none-any.whl\n!pip install -U /kaggle/input/transformers-4-57-1/tokenizers-0.22.1-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n!pip install -U /kaggle/input/transformers-4-57-1/transformers-4.57.1-py3-none-any.whl\n!pip install -U /kaggle/input/transformers-4-57-1/hydra_core-1.3.2-py3-none-any.whl\n\n# !cp /kaggle/input/transformers-4-57-1/flash_attn-2.8.3cu124torch2.6-cp311-cp311-linux_x86_64.whl flash_attn-2.8.3+cu124torch2.6-cp311-cp311-linux_x86_64.whl\n# !pip install -U --no-deps flash_attn-2.8.3+cu124torch2.6-cp311-cp311-linux_x86_64.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-26T04:35:07.262560Z","iopub.execute_input":"2025-12-26T04:35:07.262885Z","iopub.status.idle":"2025-12-26T04:35:40.208753Z","shell.execute_reply.started":"2025-12-26T04:35:07.262860Z","shell.execute_reply":"2025-12-26T04:35:40.207608Z"},"_kg_hide-output":true,"collapsed":true,"jupyter":{"outputs_hidden":true,"source_hidden":true},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import tokenizers\nprint(f\"{tokenizers.__version__=}\")\n\nimport huggingface_hub\nprint(f\"{huggingface_hub.__version__=}\")\n\nimport transformers\nprint(f\"{transformers.__version__=}\")\n\nimport hydra\nprint(f\"{hydra.__version__=}\")\n\nimport torch\nprint(f\"{torch.__version__=}\")\n\n# import flash_attn\n# print(f\"{flash_attn.__version__=}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-26T04:35:40.210797Z","iopub.execute_input":"2025-12-26T04:35:40.211123Z","iopub.status.idle":"2025-12-26T04:35:47.445239Z","shell.execute_reply.started":"2025-12-26T04:35:40.211094Z","shell.execute_reply":"2025-12-26T04:35:47.444242Z"},"jupyter":{"source_hidden":true},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Write Config Files and Scripts","metadata":{}},{"cell_type":"code","source":"%%writefile siglip2-giant.yml\n\ncompute_features: true\n\nmodels:\n  - name: siglip2-giant\n    pretrained_model_name_or_path: google/siglip2-giant-opt-patch16-384\n    # pretrained_model_name_or_path: /kaggle/input/google-siglip2-giant-opt-patch16-384\n\nregressors:\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n      - - model\n        - _target_: sklearn.linear_model.Ridge\n          alpha: 630.0\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n      - - model\n        - _target_: sklearn.linear_model.MultiTaskLasso\n          alpha: 1.0 # 0.2\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n      - - model\n        - _target_: sklearn.linear_model.MultiTaskElasticNet\n          alpha: 1.7\n          l1_ratio: 0.1\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - - model\n        - _target_: sklearn.svm.SVR\n          C: 190.0\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: sklearn.ensemble.GradientBoostingRegressor\n  #         learning_rate: 0.1\n  #         n_estimators: 100\n  #         subsample: 0.9\n  #         max_depth: 3\n  #         random_state: 42\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: catboost.CatBoostRegressor\n  #         iterations: 100\n  #         l2_leaf_reg:\n  #         subsample: 0.9\n  #         random_state: 42\n  #         verbose: false\n# (\"pca\", PCA(n_components=256, svd_solver=\"randomized\"))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-01T15:43:53.518061Z","iopub.execute_input":"2025-12-01T15:43:53.518293Z","iopub.status.idle":"2025-12-01T15:43:53.536888Z","shell.execute_reply.started":"2025-12-01T15:43:53.518277Z","shell.execute_reply":"2025-12-01T15:43:53.536287Z"},"jupyter":{"source_hidden":true},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile pe_spatial_gigantic.yml\n\ncompute_features: false\n\nmodels:\n  - name: pe_spatial_gigantic\n    # pretrained_model_name_or_path: timm/vit_pe_spatial_gigantic_patch14_448.fb\n    pretrained_model_name_or_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n    crops:\n      # - patch_size: 1000\n      #   stride: 1000\n      #   flip: false\n      # - patch_size: 500\n      #   stride: 500\n      #   flip: false\n      - patch_size: 250\n        stride: 250\n        flip: false\n\nmlp:\n  hidden_dim: 1024\n  dropout1: 0.3\n  dropout2: 0.1\n  mixup_alpha: 0.45\n  base_lr: 1e-2\n  weight_decay: 1e-6\n  batch_size: 12\n  max_epochs: 80\n  patience: 10 # Early stopping patience\n  clip_max: 150.0\n  beta: 5.0\n\nregressors:\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n      # - - pca\n      #   - _target_: sklearn.decomposition.PCA\n      #     n_components: 128\n      #     svd_solver: randomized\n      - - model\n        - _target_: sklearn.linear_model.Ridge\n          alpha: 1200.0\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - - model\n        - _target_: sklearn.linear_model.MultiTaskLasso\n          alpha: 0.006\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n      - - model\n        - _target_: sklearn.linear_model.MultiTaskElasticNet\n          alpha: 3.0\n          l1_ratio: 0.1\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - - model\n        - _target_: sklearn.svm.SVR\n          C: 200.0\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: lightgbm.LGBMRegressor\n  #         n_estimators: 400\n  #         learning_rate: 0.05\n  #         num_leaves: 31\n  #         max_depth: 6\n  #         subsample: 0.8\n  #         colsample_bytree: 0.7\n  #         reg_alpha: 2.0\n  #         reg_lambda: 10.0\n  #         min_child_samples: 50\n  #         min_child_weight: 0.001\n  #         random_state: 42\n  #         verbosity: -1\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: sklearn.ensemble.GradientBoostingRegressor\n  #         learning_rate: 0.1\n  #         n_estimators: 100\n  #         subsample: 0.9\n  #         max_depth: 3\n  #         random_state: 42\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: catboost.CatBoostRegressor\n  #         iterations: 100\n  #         l2_leaf_reg:\n  #         subsample: 0.9\n  #         random_state: 42\n  #         verbose: False\n","metadata":{"trusted":true,"_kg_hide-input":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile dinov3-vit7b16.yml\n\ncompute_features: false\n\nmodels:\n  - name: dinov3-vit7b16\n    # pretrained_model_name_or_path: facebook/dinov3-vit7b16-pretrain-lvd1689m\n    pretrained_model_name_or_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n    crops:\n      - patch_size: -1\n        stride: -1\n        flip: false\n      - patch_size: 1000\n        stride: 1000\n        flip: false\n      - patch_size: 500\n        stride: 500\n        flip: false\n      - patch_size: 250\n        stride: 250\n        flip: false\n      # - patch_size: -1\n      #   stride: -1\n      #   flip: true\n      # - patch_size: 1000\n      #   stride: 1000\n      #   flip: true\n      # - patch_size: 500\n      #   stride: 500\n      #   flip: true\n      # - patch_size: 250\n      #   stride: 250\n      #   flip: true\n\nmlp:\n  hidden_dim: 768\n  dropout1: 0.3\n  dropout2: 0.1\n  mixup_alpha: 0.45\n  base_lr: 8e-3\n  weight_decay: 1e-3\n  batch_size: 12\n  max_epochs: 80\n  patience: 10 # Early stopping patience\n  clip_max: 130.0\n  beta: 5.0\n\nregressors:\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n      # - - pca\n      #   - _target_: sklearn.decomposition.PCA\n      #     n_components: 128\n      #     svd_solver: randomized\n      - - model\n        - _target_: sklearn.linear_model.Ridge\n          alpha: 2000.0\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - - model\n        - _target_: sklearn.linear_model.MultiTaskLasso\n          alpha: 0.026\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n      - - model\n        - _target_: sklearn.linear_model.MultiTaskElasticNet\n          alpha: 3.0\n          l1_ratio: 0.1\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - - model\n        - _target_: sklearn.svm.SVR\n          C: 250.0\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: lightgbm.LGBMRegressor\n  #         n_estimators: 400\n  #         learning_rate: 0.05\n  #         num_leaves: 31\n  #         max_depth: 6\n  #         subsample: 0.8\n  #         colsample_bytree: 0.7\n  #         reg_alpha: 2.0\n  #         reg_lambda: 10.0\n  #         min_child_samples: 50\n  #         min_child_weight: 0.001\n  #         random_state: 42\n  #         verbosity: -1\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: sklearn.ensemble.GradientBoostingRegressor\n  #         learning_rate: 0.1\n  #         n_estimators: 100\n  #         subsample: 0.9\n  #         max_depth: 3\n  #         random_state: 42\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: catboost.CatBoostRegressor\n  #         iterations: 100\n  #         l2_leaf_reg:\n  #         subsample: 0.9\n  #         random_state: 42\n  #         verbose: False\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-01T15:43:53.559922Z","iopub.execute_input":"2025-12-01T15:43:53.560207Z","iopub.status.idle":"2025-12-01T15:43:53.582246Z","shell.execute_reply.started":"2025-12-01T15:43:53.560189Z","shell.execute_reply":"2025-12-01T15:43:53.581481Z"},"_kg_hide-input":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile multi-model.yml\n\n# -----------------------------------------------------------------------------\n# Global Settings\n# -----------------------------------------------------------------------------\nuse_grass_clover: false\nuse_irish: false\n\n# -----------------------------------------------------------------------------\n# Feature Extraction\n# -----------------------------------------------------------------------------\ncompute_features: false\n\nmodels:\n  # --- dinov3-vit7b16 ---\n  # - key: dinov3-vit7b16\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: dinov3-vit7b16-whole-flip\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: true\n  # - key: dinov3-vit7b16-1000\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  # - key: dinov3-vit7b16-500\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: dinov3-vit7b16-250\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   avg_pool_pct: 90\n  #   max_pool_pct: 20\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  - key: dinov3-vit7b16-200\n    model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n    kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n    avg_pool_pct: 100\n    max_pool_pct: 10\n    min_pool_pct: 10\n    std_pool_pct: 0\n    # total_pct: 40\n    crops:\n      - patch_size: 200\n        stride: 200\n        flip: false\n\n  # # --- siglip2-giant ---\n  # - key: siglip2-giant-text\n  #   model_name: google/siglip2-giant-opt-patch16-384\n  #   kaggle_path: /kaggle/input/google-siglip2-giant-opt-patch16-384\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 500\n  #       flip: false\n  #     # - patch_size: 500\n  #     #   stride: 500\n  #     #   flip: false\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: true\n  #     - patch_size: 1000\n  #       stride: 500\n  #       flip: true\n  #     # - patch_size: 500\n  #     #   stride: 500\n  #     #   flip: true\n  # - key: siglip2-giant-whole\n  #   model_name: google/siglip2-giant-opt-patch16-384\n  #   kaggle_path: /kaggle/input/google-siglip2-giant-opt-patch16-384\n  #   # avg_pool_pct: 90\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: true\n  # - key: siglip2-giant-1000\n  #   model_name: google/siglip2-giant-opt-patch16-384\n  #   kaggle_path: /kaggle/input/google-siglip2-giant-opt-patch16-384\n  #   avg_pool_pct: 90\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 500\n  #       flip: true\n\n  # # --- pe_spatial_gigantic ---\n  # - key: pe_spatial_gigantic\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: pe_spatial_gigantic-whole\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  # - key: pe_spatial_gigantic-1000-flip\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   max_pool_pct: 5\n  #   min_pool_pct: 5\n  #   std_pool_pct: 0\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: true\n  # - key: pe_spatial_gigantic-500\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   # avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: pe_spatial_gigantic-333\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 333\n  #       stride: 333\n  #       flip: false\n  - key: pe_spatial_gigantic-250\n    model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n    kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n    avg_pool_pct: 100\n    max_pool_pct: 10\n    min_pool_pct: 10\n    std_pool_pct: 0\n    crops:\n      - patch_size: 250\n        stride: 250\n        flip: false\n  # - key: pe_spatial_gigantic-200\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   # avg_pool_pct: 90\n  #   crops:\n  #     - patch_size: 200\n  #       stride: 200\n  #       flip: false\n\n  # # --- pe_core_gigantic ---\n  # - key: pe_core_gigantic\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: pe_core_gigantic-whole\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  # - key: pe_core_gigantic-1000\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  # - key: pe_core_gigantic-500\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: pe_core_gigantic-333\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 333\n  #       stride: 333\n  #       flip: false\n  # - key: pe_core_gigantic-250-flip\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   # avg_pool_pct: 10\n  #   # max_pool_pct: 10\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: true\n  # - key: pe_core_gigantic-200\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 200\n  #       stride: 200\n  #       flip: false\n\n  # --- c-radiov3-g ---\n  # - key: c-radiov3-g\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: c-radiov3-g-whole\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   # avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  # - key: c-radiov3-g-1000\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   # avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  # - key: c-radiov3-g-500\n  #   model_name: nvidia/C-RADIOv3-g\n  #   # model_name: data/models/nvidia__C-RADIOv3-g/\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   # avg_pool_pct: 50\n  #   # max_pool_pct: 10\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: c-radiov3-g-250\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   # avg_pool_pct: 0.1\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: c-radiov3-g-200\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   avg_pool_pct: 0.1\n  #   crops:\n  #     - patch_size: 200\n  #       stride: 200\n  #       flip: false\n\n  # --- internvit-6b ---\n  # - key: internvit-6b-500\n  #   model_name: OpenGVLab/InternViT-6B-448px-V2_5\n  #   kaggle_path: /kaggle/input/opengvlab-internvit-6b-448px-v2-5\n  #   avg_pool_pct: 0.9\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: internvit-6b-250\n  #   model_name: OpenGVLab/InternViT-6B-448px-V2_5\n  #   kaggle_path: /kaggle/input/opengvlab-internvit-6b-448px-v2-5\n  #   avg_pool_pct: 0.9\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n\n  # # --- dinov3-vith16plus ---\n  # - key: dinov3-vith16plus-250\n  #   model_name: facebook/dinov3-vith16plus-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vith16plus-pretrain-lvd1689m\n  #   # avg_pool_pct: 0.1\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n\n  # # --- eva02-large ---\n  # - key: eva02-large-250\n  #   model_name: timm/eva02_large_patch14_448.mim_m38m_ft_in22k_in1k\n  #   kaggle_path: /kaggle/input/timm-eva02-large-patch14-448\n  #   avg_pool_pct: 10\n  #   max_pool_pct: 0\n  #   min_pool_pct: 0\n  #   std_pool_pct: 0\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n\n# -----------------------------------------------------------------------------\n# MLP\n# -----------------------------------------------------------------------------\n# Training\nmlp:\n  spatial: false\n  hidden_dim: 1280\n  dropout1: 0.05\n  dropout2: 0.05\n  act_layer:\n    _target_: torch.nn.LeakyReLU\n  norm_layer:\n    _target_: torch.nn.BatchNorm1d\n    affine: true\n  num_layers: 2\n  drop_path_rate: 0.3\n  out_act:\n    _target_: torch.nn.Softplus\n  mixup_alpha: 0.3\n  clip_max: 160.0\n  beta: 2.0\n  base_lr: 1.5e-2\n  weight_decay: 1e-5\n  disasble_bias_wd: true\n  optimizer:\n    _target_: torch.optim.AdamW\n    lr: ${mlp.base_lr}\n    weight_decay: ${mlp.weight_decay}\n    betas: [0.9, 0.998]\n    eps: 4e-6\n  # optimizer:\n  #   _target_: timm.optim.lion.Lion\n  #   lr: ${mlp.base_lr}\n  #   weight_decay: ${mlp.weight_decay}\n  #   betas: [0.85, 0.94]\n  scheduler_params:\n    pct_start: 0.1\n    div_factor: 1e7\n    final_div_factor: ${eval:'5.0 / 1e7'}\n  batch_size: 16\n  max_epochs: 80\n  ema_decay: null\n  patience: 20 # Early stopping patience\n\nlog_dir: logs/mlp/multi-model\n\n# Hyperparameter Optimization\nopt: false\nstorage: sqlite:///${log_dir}/optuna_study.db\nstudy_name: null\nsampler:\n  _target_: optuna.samplers.TPESampler\n  n_startup_trials: 20\n  n_ei_candidates: 24\n  # gamma:\n  #   _target_: tools.utils.custom_gamma\n  #   _partial_: true\nload_if_exists: false\nn_trials: 500\nval_fold: null\n\n# Inference\ninfer: false\nckpt_dir: ${log_dir} # data/models/mlp/xx\nckpt_dir_kaggle: /kaggle/input/csiro-mlp-checkpoints\n\n# -----------------------------------------------------------------------------\n# Other Regressors\n# -----------------------------------------------------------------------------\nregressors:\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n      # - - pca\n      #   - _target_: sklearn.decomposition.PCA\n      #     n_components: 128\n      #     svd_solver: randomized\n      - - model\n        - _target_: sklearn.linear_model.Ridge\n          alpha: 2000.0\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - - model\n        - _target_: sklearn.linear_model.MultiTaskLasso\n          alpha: 0.026\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n  #     - - model\n  #       - _target_: sklearn.linear_model.MultiTaskElasticNet\n  #         alpha: 3.0\n  #         l1_ratio: 0.1\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: sklearn.svm.SVR\n  #         C: 250.0\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: lightgbm.LGBMRegressor\n  #         n_estimators: 400\n  #         learning_rate: 0.05\n  #         num_leaves: 31\n  #         max_depth: 6\n  #         subsample: 0.8\n  #         colsample_bytree: 0.7\n  #         reg_alpha: 2.0\n  #         reg_lambda: 10.0\n  #         min_child_samples: 50\n  #         min_child_weight: 0.001\n  #         random_state: 42\n  #         verbosity: -1\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: sklearn.ensemble.GradientBoostingRegressor\n  #         learning_rate: 0.1\n  #         n_estimators: 100\n  #         subsample: 0.9\n  #         max_depth: 3\n  #         random_state: 42\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: catboost.CatBoostRegressor\n  #         iterations: 100\n  #         l2_leaf_reg:\n  #         subsample: 0.9\n  #         random_state: 42\n  #         verbose: False\n","metadata":{"trusted":true,"_kg_hide-input":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile spatial.yml\n\n# -----------------------------------------------------------------------------\n# Global Settings\n# -----------------------------------------------------------------------------\nuse_grass_clover: false\nuse_irish: false\n\n# -----------------------------------------------------------------------------\n# Feature Extraction\n# -----------------------------------------------------------------------------\ncompute_features: false\n\nmodels:\n  # --- dinov3-vit7b16 ---\n  # - key: dinov3-vit7b16\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: dinov3-vit7b16-whole-flip\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: true\n  # - key: dinov3-vit7b16-1000\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  # - key: dinov3-vit7b16-500\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: dinov3-vit7b16-250\n  #   model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n  #   avg_pool_pct: 90\n  #   max_pool_pct: 20\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  - key: dinov3-vit7b16-200\n    model_name: facebook/dinov3-vit7b16-pretrain-lvd1689m\n    kaggle_path: /kaggle/input/facebook-dinov3-vit7b16-pretrain-lvd1689m\n    avg_pool_pct: 100\n    # max_pool_pct: 10\n    # min_pool_pct: 10\n    # std_pool_pct: 0\n    # total_pct: 40\n    crops:\n      - patch_size: 200\n        stride: 200\n        flip: false\n\n  # # --- siglip2-giant ---\n  # - key: siglip2-giant-text\n  #   model_name: google/siglip2-giant-opt-patch16-384\n  #   kaggle_path: /kaggle/input/google-siglip2-giant-opt-patch16-384\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 500\n  #       flip: false\n  #     # - patch_size: 500\n  #     #   stride: 500\n  #     #   flip: false\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: true\n  #     - patch_size: 1000\n  #       stride: 500\n  #       flip: true\n  #     # - patch_size: 500\n  #     #   stride: 500\n  #     #   flip: true\n  # - key: siglip2-giant-whole\n  #   model_name: google/siglip2-giant-opt-patch16-384\n  #   kaggle_path: /kaggle/input/google-siglip2-giant-opt-patch16-384\n  #   # avg_pool_pct: 90\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: true\n  # - key: siglip2-giant-1000\n  #   model_name: google/siglip2-giant-opt-patch16-384\n  #   kaggle_path: /kaggle/input/google-siglip2-giant-opt-patch16-384\n  #   avg_pool_pct: 90\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 500\n  #       flip: true\n\n  # # --- pe_spatial_gigantic ---\n  # - key: pe_spatial_gigantic\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: pe_spatial_gigantic-whole\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  # - key: pe_spatial_gigantic-1000-flip\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   max_pool_pct: 5\n  #   min_pool_pct: 5\n  #   std_pool_pct: 0\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: true\n  # - key: pe_spatial_gigantic-500\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   # avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: pe_spatial_gigantic-333\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 333\n  #       stride: 333\n  #       flip: false\n  - key: pe_spatial_gigantic-250\n    model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n    kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n    avg_pool_pct: 100\n    # max_pool_pct: 10\n    # min_pool_pct: 10\n    # std_pool_pct: 0\n    crops:\n      - patch_size: 250\n        stride: 250\n        flip: false\n  # - key: pe_spatial_gigantic-200\n  #   model_name: timm/vit_pe_spatial_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-spatial-gigantic-patch14-448\n  #   # avg_pool_pct: 90\n  #   crops:\n  #     - patch_size: 200\n  #       stride: 200\n  #       flip: false\n\n  # # --- pe_core_gigantic ---\n  # - key: pe_core_gigantic\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: pe_core_gigantic-whole\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  # - key: pe_core_gigantic-1000\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  # - key: pe_core_gigantic-500\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: pe_core_gigantic-333\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 333\n  #       stride: 333\n  #       flip: false\n  # - key: pe_core_gigantic-250-flip\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   # avg_pool_pct: 10\n  #   # max_pool_pct: 10\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: true\n  # - key: pe_core_gigantic-200\n  #   model_name: timm/vit_pe_core_gigantic_patch14_448.fb\n  #   kaggle_path: /kaggle/input/timm-vit-pe-core-gigantic-patch14-448\n  #   avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 200\n  #       stride: 200\n  #       flip: false\n\n  # --- c-radiov3-g ---\n  # - key: c-radiov3-g\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: c-radiov3-g-whole\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   # avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: -1\n  #       stride: -1\n  #       flip: false\n  # - key: c-radiov3-g-1000\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   # avg_pool_pct: 10\n  #   crops:\n  #     - patch_size: 1000\n  #       stride: 1000\n  #       flip: false\n  # - key: c-radiov3-g-500\n  #   model_name: nvidia/C-RADIOv3-g\n  #   # model_name: data/models/nvidia__C-RADIOv3-g/\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   # avg_pool_pct: 50\n  #   # max_pool_pct: 10\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: c-radiov3-g-250\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   # avg_pool_pct: 0.1\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n  # - key: c-radiov3-g-200\n  #   model_name: nvidia/C-RADIOv3-g\n  #   kaggle_path: /kaggle/input/nvidia-c-radiov3-g/\n  #   avg_pool_pct: 0.1\n  #   crops:\n  #     - patch_size: 200\n  #       stride: 200\n  #       flip: false\n\n  # --- internvit-6b ---\n  # - key: internvit-6b-500\n  #   model_name: OpenGVLab/InternViT-6B-448px-V2_5\n  #   kaggle_path: /kaggle/input/opengvlab-internvit-6b-448px-v2-5\n  #   avg_pool_pct: 0.9\n  #   crops:\n  #     - patch_size: 500\n  #       stride: 500\n  #       flip: false\n  # - key: internvit-6b-250\n  #   model_name: OpenGVLab/InternViT-6B-448px-V2_5\n  #   kaggle_path: /kaggle/input/opengvlab-internvit-6b-448px-v2-5\n  #   avg_pool_pct: 0.9\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n\n  # # --- dinov3-vith16plus ---\n  # - key: dinov3-vith16plus-250\n  #   model_name: facebook/dinov3-vith16plus-pretrain-lvd1689m\n  #   kaggle_path: /kaggle/input/facebook-dinov3-vith16plus-pretrain-lvd1689m\n  #   # avg_pool_pct: 0.1\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n\n  # # --- eva02-large ---\n  # - key: eva02-large-250\n  #   model_name: timm/eva02_large_patch14_448.mim_m38m_ft_in22k_in1k\n  #   kaggle_path: /kaggle/input/timm-eva02-large-patch14-448\n  #   avg_pool_pct: 10\n  #   max_pool_pct: 0\n  #   min_pool_pct: 0\n  #   std_pool_pct: 0\n  #   crops:\n  #     - patch_size: 250\n  #       stride: 250\n  #       flip: false\n\n# -----------------------------------------------------------------------------\n# MLP\n# -----------------------------------------------------------------------------\n# Training\nmlp:\n  spatial: true\n  hidden_dim: 1280\n  dropout1: 0.05\n  dropout2: 0.05\n  act_layer:\n    _target_: torch.nn.LeakyReLU\n  norm_layer:\n    _target_: torch.nn.BatchNorm1d\n    affine: true\n  num_layers: 2\n  drop_path_rate: 0.3\n  out_act:\n    _target_: torch.nn.Softplus\n  mixup_alpha: 0.3\n  clip_max: 160.0\n  beta: 2.0\n  base_lr: 1.5e-2\n  weight_decay: 1e-5\n  disasble_bias_wd: true\n  optimizer:\n    _target_: torch.optim.AdamW\n    lr: ${mlp.base_lr}\n    weight_decay: ${mlp.weight_decay}\n    betas: [0.9, 0.998]\n    eps: 4e-6\n  # optimizer:\n  #   _target_: timm.optim.lion.Lion\n  #   lr: ${mlp.base_lr}\n  #   weight_decay: ${mlp.weight_decay}\n  #   betas: [0.85, 0.94]\n  scheduler_params:\n    pct_start: 0.1\n    div_factor: 1e7\n    final_div_factor: ${eval:'5.0 / 1e7'}\n  batch_size: 16\n  max_epochs: 80\n  ema_decay: null\n  patience: 20 # Early stopping patience\n\nlog_dir: logs/mlp/multi-model\n\n# Hyperparameter Optimization\nopt: false\nstorage: sqlite:///${log_dir}/optuna_study.db\nstudy_name: null\nsampler:\n  _target_: optuna.samplers.TPESampler\n  n_startup_trials: 20\n  n_ei_candidates: 24\n  # gamma:\n  #   _target_: tools.utils.custom_gamma\n  #   _partial_: true\nload_if_exists: false\nn_trials: 500\nval_fold: null\n\n# Inference\ninfer: false\nckpt_dir: ${log_dir} # data/models/mlp/xx\nckpt_dir_kaggle: /kaggle/input/csiro-mlp-checkpoints2\n\n# -----------------------------------------------------------------------------\n# Other Regressors\n# -----------------------------------------------------------------------------\nregressors:\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n      # - - pca\n      #   - _target_: sklearn.decomposition.PCA\n      #     n_components: 128\n      #     svd_solver: randomized\n      - - model\n        - _target_: sklearn.linear_model.Ridge\n          alpha: 2000.0\n  - _target_: sklearn.pipeline.Pipeline\n    steps:\n      - - model\n        - _target_: sklearn.linear_model.MultiTaskLasso\n          alpha: 0.026\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - [scaler, _target_: sklearn.preprocessing.StandardScaler]\n  #     - - model\n  #       - _target_: sklearn.linear_model.MultiTaskElasticNet\n  #         alpha: 3.0\n  #         l1_ratio: 0.1\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: sklearn.svm.SVR\n  #         C: 250.0\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: lightgbm.LGBMRegressor\n  #         n_estimators: 400\n  #         learning_rate: 0.05\n  #         num_leaves: 31\n  #         max_depth: 6\n  #         subsample: 0.8\n  #         colsample_bytree: 0.7\n  #         reg_alpha: 2.0\n  #         reg_lambda: 10.0\n  #         min_child_samples: 50\n  #         min_child_weight: 0.001\n  #         random_state: 42\n  #         verbosity: -1\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: sklearn.ensemble.GradientBoostingRegressor\n  #         learning_rate: 0.1\n  #         n_estimators: 100\n  #         subsample: 0.9\n  #         max_depth: 3\n  #         random_state: 42\n  # - _target_: sklearn.pipeline.Pipeline\n  #   steps:\n  #     - - model\n  #       - _target_: catboost.CatBoostRegressor\n  #         iterations: 100\n  #         l2_leaf_reg:\n  #         subsample: 0.9\n  #         random_state: 42\n  #         verbose: False\n","metadata":{"trusted":true,"_kg_hide-input":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile utils.py\n\nimport math\nimport os\nfrom dataclasses import dataclass\nfrom typing import Sequence\n\nimport hydra\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchmetrics\nfrom omegaconf import DictConfig, ListConfig, OmegaConf\nfrom scipy.optimize import minimize\nfrom torchmetrics.utilities import dim_zero_cat\n\n\ndef is_kaggle():\n    # Kaggle 環境では KAGGLE_URL_BASE が定義されている\n    return bool(os.environ.get(\"KAGGLE_URL_BASE\", \"\"))\n\n\ndef custom_gamma(n_finished_trials: int) -> int:\n    \"\"\"Custom gamma function for optuna.samplers.TPESampler.\"\"\"\n    return min(math.ceil(0.2 * n_finished_trials), 50)\n\n\ndef load_config(config_path: str, dotlist: list[str] | None = None) -> DictConfig:\n    def _load_config_recursive(config_path: str) -> DictConfig | ListConfig:\n        \"\"\"_base_ フィールドに指定された config を継承する形で config を再帰的に読み込む.\"\"\"\n        config = OmegaConf.load(config_path)\n        if \"_base_\" in config and config._base_ is not None:\n            base_config_path = os.path.join(os.path.dirname(config_path), config._base_)\n            base_config = _load_config_recursive(base_config_path)\n            config = OmegaConf.merge(base_config, config)\n            del config._base_\n        return config\n\n    # Enable arithmetic in config by registering eval as a resolver\n    OmegaConf.register_new_resolver(\"eval\", eval, replace=True)\n\n    config = _load_config_recursive(config_path)\n\n    # Prevent the creation of fields that do not exist\n    OmegaConf.set_struct(config, True)\n\n    # Override configs with dotlist\n    if dotlist is not None:\n        config.merge_with_dotlist(dotlist)\n\n    OmegaConf.resolve(config)\n    OmegaConf.set_readonly(config, True)\n    return config\n\n\nweights = {\n    \"Dry_Green_g\": 0.1,\n    \"Dry_Dead_g\": 0.1,\n    \"Dry_Clover_g\": 0.1,\n    \"GDM_g\": 0.2,\n    \"Dry_Total_g\": 0.5,\n}\n\n\ndef competition_metric(\n    target,\n    pred,\n    labels=[\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"],\n):\n    \"\"\"Function to calculate the competition's official evaluation metric (weighted R2 score).\"\"\"\n    weights_array = np.array([weights[x] for x in labels])\n\n    # Align with this calculation method\n    y_weighted_mean = np.average(target, weights=weights_array, axis=1).mean()\n\n    # For ss_res and ss_tot, also take the weighted average on axis=1, then the mean of the result\n    ss_res = np.average((target - pred) ** 2, weights=weights_array, axis=1).mean()\n    ss_tot = np.average(\n        (target - y_weighted_mean) ** 2, weights=weights_array, axis=1\n    ).mean()\n    return 1 - ss_res / ss_tot\n\n\nclass WeightedR2(torchmetrics.Metric):\n    def __init__(self):\n        super().__init__()\n        self.weights = [0.1, 0.1, 0.1, 0.2, 0.5]\n\n        self.add_state(\"preds\", default=[], dist_reduce_fx=\"cat\")\n        self.add_state(\"labels\", default=[], dist_reduce_fx=\"cat\")\n\n    def update(self, preds: torch.Tensor, labels: torch.Tensor) -> None:\n        assert preds.shape == labels.shape\n\n        self.preds.append(preds.detach().cpu())\n        self.labels.append(labels.detach().cpu())\n\n    def compute(self) -> torch.Tensor:\n        preds = dim_zero_cat(self.preds)\n        labels = dim_zero_cat(self.labels)\n\n        weights = torch.tensor(self.weights, device=labels.device).view(1, -1)\n        weights = weights.expand_as(labels)\n\n        y_f = labels.reshape(-1)\n        p_f = preds.reshape(-1)\n        w_f = weights.reshape(-1)\n\n        y_bar = torch.sum(w_f * y_f) / torch.sum(w_f)\n        sse = torch.sum(w_f * (y_f - p_f) ** 2)\n        sst = torch.sum(w_f * (y_f - y_bar) ** 2)\n        return 1.0 - (sse / (sst + 1e-8))\n\n\nclass WeightedSmoothL1Loss(nn.Module):\n    def __init__(self, beta=1.0, weights=[0.2, 0.2, 0.2, 0.4, 1.0]) -> None:\n        super().__init__()\n        self.beta = beta\n        self.weights = weights\n\n    def forward(self, preds: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n        loss_elem = F.smooth_l1_loss(preds, labels, reduction=\"none\", beta=self.beta)\n        weights = torch.tensor(self.weights, device=preds.device)\n        weighted_loss = loss_elem * weights.view(1, -1)\n        return weighted_loss.mean()\n\n\nclass WeightedMSELoss(nn.Module):\n    def __init__(self, weights=[0.2, 0.2, 0.2, 0.4, 1.0]) -> None:\n        super().__init__()\n        self.weights = weights\n\n    def forward(self, preds: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n        loss_elem = F.mse_loss(preds, labels, reduction=\"none\")\n        weights = torch.tensor(self.weights, device=preds.device)\n        weighted_loss = loss_elem * weights.view(1, -1)\n        return weighted_loss.mean()\n\n\nclass DropPath(nn.Module):\n    \"\"\"Drop paths (Stochastic Depth) per sample\"\"\"\n\n    def __init__(self, drop_prob=0.0):\n        super().__init__()\n        self.drop_prob = drop_prob\n\n    def forward(self, x):\n        if self.drop_prob == 0.0 or not self.training:\n            return x\n        keep_prob = 1 - self.drop_prob\n        shape = (x.shape[0],) + (1,) * (x.ndim - 1)\n        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)\n        random_tensor.floor_()\n        return x.div(keep_prob) * random_tensor\n\n\nclass MLPBlock(nn.Module):\n    \"\"\"MLP block with optional residual connections and drop path\"\"\"\n\n    def __init__(\n        self,\n        in_features: int,\n        out_features: int,\n        dropout: float,\n        act_layer: DictConfig,\n        norm_layer: DictConfig,\n        layer_scale_init: float | None = None,\n        use_residual: bool = True,\n        drop_path: float | None = None,\n    ):\n        super().__init__()\n        self.use_residual = use_residual\n        if self.use_residual:\n            assert in_features == out_features, (\n                \"Residual connection requires same input and output dimensions\"\n            )\n\n        # Main path\n        self.drop = nn.Dropout(dropout)\n        self.linear = nn.Linear(in_features, out_features, bias=False)\n        self.norm = hydra.utils.instantiate(norm_layer, num_features=out_features)\n        self.act = hydra.utils.instantiate(act_layer)\n\n        # Optional Layer Scale\n        self.layer_scale = None\n        if layer_scale_init is not None and layer_scale_init > 0:\n            self.layer_scale = nn.Parameter(layer_scale_init * torch.ones(out_features))\n\n        # Optional Drop Path\n        self.drop_path = (\n            DropPath(drop_path) if drop_path is not None and self.use_residual else None\n        )\n\n    def forward(self, x):\n        shortcut = x\n\n        x = self.drop(x)\n        x = self.linear(x)\n        x = self.norm(x)\n        x = self.act(x)\n\n        if self.layer_scale is not None:\n            x = x * self.layer_scale\n\n        if self.use_residual:\n            if self.drop_path is not None:\n                x = self.drop_path(x)\n            return x + shortcut\n        return x\n\n\nclass GroupedLinear(nn.Module):\n    def __init__(self, in_features, out_features, groups, bias=True):\n        super().__init__()\n        assert in_features % groups == 0, \"in_features must be divisible by groups\"\n        assert out_features % groups == 0, \"out_features must be divisible by groups\"\n\n        self.in_features = in_features\n        self.out_features = out_features\n        self.groups = groups\n\n        self.in_per_group = in_features // groups\n        self.out_per_group = out_features // groups\n\n        # shape: (groups, out_per_group, in_per_group)\n        self.weight = nn.Parameter(\n            torch.empty(groups, self.out_per_group, self.in_per_group)\n        )\n        if bias:\n            # shape: (groups, out_per_group)\n            self.bias = nn.Parameter(torch.empty(groups, self.out_per_group))\n        else:\n            self.bias = None\n\n        self.reset_parameters()\n\n    def reset_parameters(self):\n        # Linear の初期化と同じような感じ\n        for g in range(self.groups):\n            nn.init.kaiming_uniform_(self.weight[g], a=math.sqrt(5))\n        if self.bias is not None:\n            fan_in = self.in_per_group\n            bound = 1 / math.sqrt(fan_in)\n            nn.init.uniform_(self.bias, -bound, bound)\n\n    def forward(self, x):\n        # x: (B, in_features)\n        B = x.size(0)\n\n        # (B, groups, in_per_group) に reshape\n        x = x.view(B, self.groups, self.in_per_group)\n\n        # einsum で各グループごとの matmul: (B, G, I) × (G, O, I)^T -> (B, G, O)\n        # \"bgi,goi->bgo\" で i を内積\n        out = torch.einsum(\"bgi,goi->bgo\", x, self.weight)\n        if self.bias is not None:\n            out = out + self.bias  # (G, O) が broadcast される\n\n        # (B, out_features) に戻す\n        out = out.reshape(B, self.out_features)\n        return out\n\n\n@dataclass(frozen=True)\nclass BlockSpec:\n    \"\"\"Concatenated embedding vector is split into blocks in this order.\"\"\"\n\n    name: str\n    dim: int\n\n\nclass SeparateLinear(nn.Module):\n    \"\"\"Blockごとに別々にLinearを適用し, 出力をconcatする.\"\"\"\n\n    def __init__(\n        self,\n        blocks: Sequence[BlockSpec],\n        out_features: list[int],\n        bias: bool = True,\n    ):\n        super().__init__()\n        self.blocks = list(blocks)\n        assert len(self.blocks) == len(out_features), (\n            \"out_features must match number of blocks\"\n        )\n\n        # Create linear layers for each block\n        self.linears = nn.ModuleList(\n            [\n                # nn.Linear(block.dim, out_feat, bias=bias)\n                MLPBlock(\n                    in_features=block.dim,\n                    out_features=out_feat,\n                    dropout=0.05,\n                    act_layer=OmegaConf.create({\"_target_\": \"torch.nn.LeakyReLU\"}),\n                    norm_layer=OmegaConf.create({\"_target_\": \"torch.nn.BatchNorm1d\"}),\n                    use_residual=False,\n                )\n                for block, out_feat in zip(self.blocks, out_features, strict=True)\n            ]\n        )\n\n        # Precompute slices for fast split\n        start = 0\n        self.slices: list[slice] = []\n        for block in self.blocks:\n            self.slices.append(slice(start, start + block.dim))\n            start += block.dim\n        self.total_in_dim = start\n        self.total_out_dim = sum(out_features)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"Forward.\n\n        Args:\n            x: (B, total_in_dim) where total_in_dim = sum(block.dim)\n\n        Returns:\n            (B, total_out_dim) where total_out_dim = sum(out_features)\n        \"\"\"\n        B, D = x.shape\n        assert D == self.total_in_dim, f\"Expected D={self.total_in_dim}, got {D}\"\n\n        outputs = []\n        for _, (sl, linear) in enumerate(zip(self.slices, self.linears, strict=True)):\n            block_input = x[:, sl]  # (B, block.dim)\n            block_output = linear(block_input)  # (B, out_features[i])\n            outputs.append(block_output)\n\n        return torch.cat(outputs, dim=1)  # (B, total_out_dim)\n\n\nclass BlockDropout(nn.Module):\n    \"\"\"Drop embedding blocks (model-wise) instead of individual features.\n\n    - Input:  (B, D_total) where D_total = sum(block.dim)\n    - Output: (B, D_total), some blocks are zeroed.\n    - Typical: per-sample block mask (each sample sees different dropped blocks)\n    \"\"\"\n\n    def __init__(\n        self,\n        blocks: Sequence[BlockSpec],\n        p: float | Sequence[float] = 0.2,\n        mode: str = \"per_sample\",  # \"per_sample\" or \"per_batch\"\n        rescale: bool = True,  # keep expected activation similar to Dropout\n        always_keep_one: bool = True,  # avoid all blocks dropped for a sample\n    ):\n        super().__init__()\n        assert mode in (\"per_sample\", \"per_batch\")\n        self.blocks = list(blocks)\n        self.mode = mode\n        self.rescale = rescale\n        self.always_keep_one = always_keep_one\n\n        if isinstance(p, (float, int)):\n            ps = [float(p)] * len(self.blocks)\n        else:\n            ps = list(map(float, p))\n            assert len(ps) == len(self.blocks), \"p must match number of blocks\"\n        for pi in ps:\n            assert 0.0 <= pi < 1.0, \"p must be in [0, 1)\"\n        self.register_buffer(\n            \"p_tensor\", torch.tensor(ps, dtype=torch.float32), persistent=False\n        )\n\n        # Precompute slices for fast split\n        start = 0\n        self.slices: list[slice] = []\n        for b in self.blocks:\n            self.slices.append(slice(start, start + b.dim))\n            start += b.dim\n        self.total_dim = start\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"x: (B, D_total)\"\"\"\n        if (not self.training) or self.p_tensor.max().item() <= 0.0:\n            return x\n\n        B, D = x.shape\n        assert D == self.total_dim, f\"Expected D={self.total_dim}, got {D}\"\n\n        device = x.device\n        p = self.p_tensor.to(device=device)\n        keep_prob = (1.0 - p).clamp_min(1e-6)  # (num_blocks,)\n\n        num_blocks = len(self.blocks)\n\n        if self.mode == \"per_batch\":\n            # mask shape: (num_blocks,)\n            keep = (torch.rand(num_blocks, device=device) < keep_prob).float()\n            keep = keep.unsqueeze(0).expand(B, -1)  # (B, num_blocks)\n        else:\n            # per_sample: (B, num_blocks)\n            keep = (torch.rand(B, num_blocks, device=device) < keep_prob).float()\n\n        if self.always_keep_one:\n            # If a sample drops all blocks, force-keep one random block\n            dropped_all = keep.sum(dim=1) == 0  # (B,)\n            if dropped_all.any():\n                idx = torch.randint(\n                    0, num_blocks, (int(dropped_all.sum().item()),), device=device\n                )\n                keep[dropped_all] = 0.0\n                keep[dropped_all, idx] = 1.0\n\n        # Build a feature mask of shape (B, D_total)\n        mask = x.new_ones((B, D))\n        for j, sl in enumerate(self.slices):\n            mj = keep[:, j].unsqueeze(1)  # (B, 1)\n            if self.rescale:\n                # Standard dropout-style scaling per block\n                mj = mj / keep_prob[j]\n            mask[:, sl] = mj\n\n        return x * mask\n\n\ndef mixup(x, batch, alpha=1.0):\n    \"\"\"Mixup data augmentation for features.\n\n    Args:\n        x (torch.Tensor): Input features with shape (N, C).\n        batch (dict): Batch dictionary containing labels and auxiliary targets.\n        alpha (float): Mixup interpolation coefficient.\n\n    Returns:\n        mixed_x (torch.Tensor): Mixed features.\n        batch_a (dict): First batch.\n        batch_b (dict): Second batch (permuted).\n        lam (float): Mixup lambda.\n    \"\"\"\n    if alpha > 0:\n        lam = np.random.beta(alpha, alpha)\n    else:\n        lam = 1\n\n    batch_size = x.size(0)\n    index = torch.randperm(batch_size).to(x.device)\n\n    mixed_x = lam * x + (1 - lam) * x[index, :]\n\n    # Create permuted batch\n    batch_b = {}\n    for key, value in batch.items():\n        if key != \"input\" and isinstance(value, torch.Tensor):\n            batch_b[key] = value[index]\n\n    return mixed_x, batch, batch_b, lam\n\n\ndef make_projection():\n    \"\"\"後処理用の射影行列を作成.\n\n    C y = 0 を満たす最近接点へ直交射影.\n    ターゲット順は [v, w, x, y, z].\n    \"\"\"\n    C = np.array(\n        [\n            [1, 0, 1, -1, 0],  # v + x - y = 0  <=>  v = y - x\n            [0, 1, 0, 1, -1],  # w + y - z = 0  <=>  w = z - y\n        ],\n        dtype=float,\n    )\n    Ct = C.T\n    M = np.linalg.inv(C @ Ct)\n    P = np.eye(5) - Ct @ M @ C\n    return P\n\n\ndef project_nonneg_equalities(Y_hat):\n    \"\"\"バッチ適用\"\"\"\n    # 等式制約 Ay = 0 （順序: v,w,x,y,z）\n    A = np.array(\n        [\n            [1, 0, 1, -1, 0],  # v - y + x = 0\n            [0, 1, 0, 1, -1],  # w - z + y = 0\n        ],\n        dtype=float,\n    )\n\n    # 目的: 0.5 * ||y - yhat||^2（yhatに最も近い可行点）\n    def _quad_obj(y, yhat):\n        r = y - yhat\n        return 0.5 * np.dot(r, r)\n\n    def _quad_jac(y, yhat):\n        return y - yhat\n\n    # 1サンプル分を投影（等式＋非負）\n    def _project_one(yhat, A, maxiter=200, tol=1e-9):\n        cons = [{\"type\": \"eq\", \"fun\": lambda y, A=A: A @ y, \"jac\": lambda y, A=A: A}]\n        bnds = [(0.0, None)] * 5  # 非負制約\n        res = minimize(\n            _quad_obj,\n            x0=np.maximum(yhat, 0.0),\n            args=(yhat,),\n            method=\"SLSQP\",\n            jac=_quad_jac,\n            constraints=cons,\n            bounds=bnds,\n            options={\"maxiter\": maxiter, \"ftol\": tol, \"disp\": False},\n        )\n        if not res.success:\n            # フォールバック: 非負にクリップ → 等式のみ直交射影\n            y0 = np.maximum(yhat, 0.0)\n            # 等式のみの直交射影 P = I - A^T (A A^T)^{-1} A\n            P = np.eye(5) - A.T @ np.linalg.solve(A @ A.T, A)\n            return P @ y0\n        return res.x\n\n    # Y_hat: (n, 5)  [v,w,x,y,z]\n    Y_proj = np.empty_like(Y_hat, dtype=float)\n    for i in range(Y_hat.shape[0]):\n        Y_proj[i] = _project_one(Y_hat[i], A)\n    return Y_proj\n","metadata":{"trusted":true,"_kg_hide-input":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile model_3_embed.py\n\n# ==============================================================================\n# Library Imports and Initial Setup\n# ==============================================================================\nimport argparse\nimport gc\nimport json\nimport logging\nimport math\nimport multiprocessing as mp\nimport os\nimport pickle\nimport shutil\nimport warnings\nfrom collections import defaultdict\nfrom contextlib import contextmanager\nfrom copy import deepcopy\nfrom multiprocessing import Process, Queue\nfrom pathlib import Path\nfrom typing import Literal\n\nos.environ[\"HF_MODULES_CACHE\"] = \"/kaggle/working/hf/modules\"\n\nimport albumentations as A\nimport catboost\nimport cv2\nimport hydra\nimport numpy as np\nimport optuna\nimport polars as pl\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport transformers\nfrom lightgbm import LGBMRegressor\nfrom omegaconf import DictConfig, OmegaConf\nfrom PIL import Image, ImageOps\nfrom pytorch_lightning import (\n    LightningDataModule,\n    LightningModule,\n    Trainer,\n    seed_everything,\n)\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint\nfrom pytorch_lightning.utilities.combined_loader import CombinedLoader\nfrom scipy.optimize import minimize, nnls\nfrom sklearn.decomposition import PCA\nfrom sklearn.dummy import DummyRegressor\nfrom sklearn.ensemble import GradientBoostingRegressor\nfrom sklearn.feature_selection import SelectPercentile, f_regression\nfrom sklearn.linear_model import Lasso, MultiTaskElasticNet, MultiTaskLasso, Ridge\nfrom sklearn.model_selection import GroupKFold, KFold, StratifiedGroupKFold\nfrom sklearn.pipeline import Pipeline\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.svm import SVR\nfrom torch.utils.data import DataLoader, Dataset, TensorDataset\nfrom tqdm.auto import tqdm\nfrom transformers import (\n    AutoBackbone,\n    AutoImageProcessor,\n    AutoModel,\n    AutoProcessor,\n    CLIPImageProcessor,\n)\nfrom utils import (\n    BlockDropout,\n    BlockSpec,\n    GroupedLinear,\n    MLPBlock,\n    SeparateLinear,\n    WeightedMSELoss,\n    WeightedR2,\n    WeightedSmoothL1Loss,\n    competition_metric,\n    is_kaggle,\n    load_config,\n    make_projection,\n    mixup,\n    project_nonneg_equalities,\n)\n\n# from biomass.dsbn import convert_to_dsbn, set_model_domain\n# from biomass.ema import EMACallback\n# from biomass.moe import cross_validate_moe, cross_validate_moe2\n\nwarnings.filterwarnings(\"ignore\")\n\nlogging.getLogger(\"pytorch_lightning\").setLevel(logging.ERROR)\n\n# --- Initial Setup ---\n\n# Dataset root\nDATA_DIR = (\n    Path(\"/kaggle/input/csiro-biomass\") if is_kaggle() else Path(\"data/csiro-biomass\")\n)\n\n# Directory to store features\nFEATS_DIR = Path(\"data/features\")\nos.makedirs(FEATS_DIR, exist_ok=True)\n\n# Target variables to predict\nLABELS = [\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"]\n\nassert torch.cuda.is_available()\nDEVICE = \"cuda\"\n\n\ndef parse_args(args: list[str] | None = None) -> argparse.Namespace:\n    parser = argparse.ArgumentParser(description=\"NN embeddings + regressors.\")\n    parser.add_argument(\"config\", type=str, help=\"config file\")\n    parser.add_argument(\n        \"options\",\n        nargs=argparse.REMAINDER,\n        help=\"override settings in used config with key-value pairs in xxx=yyy format \"\n        \"(e.g. run_name=your_run_name)\",\n    )\n    return parser.parse_args(args)\n\n\n# Load config\nargs = parse_args()\ncfg = load_config(args.config, dotlist=args.options)\n\ncfg_name = Path(args.config).stem\n\nos.makedirs(cfg.log_dir, exist_ok=True)\n\nuse_grass_clover = cfg.get(\"use_grass_clover\", False)\nuse_irish = cfg.get(\"use_irish\", False)\n\n# ==============================================================================\n# Load Dataset\n# ==============================================================================\n\n\ndef add_folds(df: pl.DataFrame, n_folds: int, random_state: int = 42) -> pl.DataFrame:\n    df = deepcopy(df)\n\n    # Split using KFold (more robust than random split)\n    kf = KFold(n_splits=n_folds, shuffle=True, random_state=random_state)\n    x_dummy = df.select(\"image_path\").to_numpy()\n    fold_assignments = np.full(len(df), -1)\n    for fold, (_, val_idxs) in enumerate(kf.split(x_dummy)):\n        fold_assignments[val_idxs] = fold\n    df = df.with_columns(pl.Series(\"fold\", fold_assignments))\n    assert not df.select(pl.col(\"fold\") == -1).to_series().any()\n    return df\n\n\ndef add_group_folds(df: pl.DataFrame, n_folds: int) -> pl.DataFrame:\n    \"\"\"Create fold column using GroupKFold.\"\"\"\n    df = deepcopy(df)\n\n    groups = df.select(\"group\").to_numpy().flatten()\n\n    kf = StratifiedGroupKFold(n_splits=n_folds)\n\n    x_dummy = df.select(\"image_path\").to_numpy()\n\n    fold_assignments = np.full(len(df), -1)\n    for fold, (_, val_idxs) in enumerate(kf.split(x_dummy, df[\"State\"], groups=groups)):\n        fold_assignments[val_idxs] = fold\n    df = df.with_columns(pl.Series(\"fold\", fold_assignments))\n    assert not df.select(pl.col(\"fold\") == -1).to_series().any()\n    return df\n\n\n# --- Load and Preprocess Training Data ---\nprint(\"Processing training data...\")\ntrain = pl.read_csv(DATA_DIR / \"train.csv\")\ndf = train.pivot(\n    on=\"target_name\",\n    index=[\n        \"image_path\",\n        \"Sampling_Date\",\n        \"State\",\n        \"Species\",\n        \"Pre_GSHH_NDVI\",\n        \"Height_Ave_cm\",\n    ],\n    values=\"target\",\n    aggregate_function=\"mean\",\n)\ndf = df.with_columns(\n    (pl.lit(str(DATA_DIR)) + \"/\" + pl.col(\"image_path\")).alias(\"rel_image_path\")\n)\n\n# # Random fold assignment\n# n_folds = 20\n# df = add_folds(df, n_folds=n_folds, random_state=0)\n# # Group on specified column\n# df = df.with_columns(\n#     pl.concat_str([\"Sampling_Date\", \"State\"], separator=\" \").alias(\"group\")\n#     # pl.col(\"Sampling_Date\").alias(\"group\")\n#     # pl.col(\"State\").alias(\"group\")\n# )\n# Group on binned Sampling_Date (for local exp)\nn_folds = 19  # if cfg.opt or cfg.infer or is_kaggle() else 5\ndf = df.with_columns(\n    (pl.col(\"Sampling_Date\").str.to_datetime(\"%Y/%m/%d\").dt.ordinal_day() // 7).alias(\n        \"group\"\n    )\n)\n# # Group on Sampling_Date using quantiles (for submission)\n# n_folds = 12\n# df = df.with_columns(\n#     pl.col(\"Sampling_Date\")\n#     .str.to_datetime(\"%Y/%m/%d\")\n#     .dt.ordinal_day()\n#     .qcut(\n#         n_folds,\n#         labels=[str(x) for x in list(range(n_folds))],\n#         left_closed=False,\n#         allow_duplicates=False,\n#     )\n#     .alias(\"group\")\n# )\n\ndf = add_group_folds(df, n_folds=n_folds)\n\nif is_kaggle:\n    print(df[\"fold\"].to_pandas().value_counts().sort_index())\n\n# df = df[:4]  # For testing\n\n# --- Load and Preprocess Test Data ---\nprint(\"Processing test data...\")\ntest = pl.read_csv(DATA_DIR / \"test.csv\")\ndf_test = (\n    test.group_by(\"image_path\")\n    .len()  # Get the number of targets for each image (5 in this case)\n    .sort(\"image_path\")\n)\ndf_test = df_test.with_columns(\n    (pl.lit(str(DATA_DIR)) + \"/\" + pl.col(\"image_path\")).alias(\"rel_image_path\")\n)\n\n# --- Load and Preprocess Grass-Clover Data ---\nprint(f\"{use_grass_clover=}\")\ngrass_clover_root = Path(\"data/grassclover-dataset/biomass_data/\")\ndf_grass_clover = None\nif use_grass_clover:\n    df_grass_clover = pl.read_csv(\n        grass_clover_root / \"train/biomass_train_data.csv\", separator=\";\"\n    )\n    df_grass_clover = df_grass_clover.with_columns(\n        (\n            pl.lit(str(grass_clover_root))\n            + \"/train/images/\"\n            + pl.col(\"image_file_name\")\n        ).alias(\"rel_image_path\")\n    )\n\n# --- Load and Preprocess Irish Grass Clover Data ---\nprint(f\"{use_irish=}\")\nirish_root = Path(\"data/irish-grass-clover/\")\ndf_irish = None\nif use_irish:\n    df_irish_train = pl.read_csv(irish_root / \"camera/train.csv\")\n    df_irish_val = pl.read_csv(irish_root / \"camera/val.csv\")\n    df_irish_train = df_irish_train.select(\n        [\n            \"Image Name\",\n            \"Herbage Mass (kg DM/ha)\",\n            \"Grass Dried\",\n            \"Clover Dried\",\n            \"Weeds Dried\",\n            \"Quadrat\",\n            \"Plot\",\n        ]\n    )\n    df_irish_val = df_irish_val.select(\n        [\n            \"Image Name\",\n            \"Herbage Mass (kg DM/ha)\",\n            \"Grass Dried\",\n            \"Clover Dried\",\n            \"Weeds Dried\",\n            \"Quadrat\",\n            \"Plot\",\n        ]\n    )\n    df_irish = pl.concat([df_irish_train, df_irish_val])\n    df_irish = df_irish.with_columns(\n        (pl.lit(str(irish_root)) + \"/camera/images/\" + pl.col(\"Image Name\")).alias(\n            \"rel_image_path\"\n        )\n    )\n\n    # Preprocess grams\n    # row[['Grass Dried', 'Clover Dried', 'Weeds Dried']] * 0.01 * row['Herbage Mass (kg DM/ha)'] * 0.025\n    df_irish = df_irish.with_columns(\n        (\n            df_irish[\"Grass Dried\"] * df_irish[\"Herbage Mass (kg DM/ha)\"] * 0.01 * 0.02\n        ).alias(\"Grass Dried g\"),\n        (\n            df_irish[\"Clover Dried\"] * df_irish[\"Herbage Mass (kg DM/ha)\"] * 0.01 * 0.02\n        ).alias(\"Clover Dried g\"),\n        (\n            df_irish[\"Weeds Dried\"] * df_irish[\"Herbage Mass (kg DM/ha)\"] * 0.01 * 0.02\n        ).alias(\"Weeds Dried g\"),\n    )\n\n# ==============================================================================\n# Feature Preparation\n# ==============================================================================\n\n\ndef get_crops(\n    image: np.ndarray, patch_size: int, stride: int, flip: bool = False\n) -> np.ndarray:\n    \"\"\"Extract overlapping patches from an image using sliding window.\n\n    Returns:\n        Array of patches with shape (n_patches, patch_size, patch_size, C)\n    \"\"\"\n    crops = []\n\n    if patch_size == -1:\n        # If patch_size is -1, return the whole image as a single patch\n        crops.append(image)\n    else:\n        assert patch_size > 0\n        assert stride > 0\n\n        H, W, C = image.shape\n        for y in range(0, H - patch_size + 1, stride):\n            for x in range(0, W - patch_size + 1, stride):\n                crop = image[y : y + patch_size, x : x + patch_size, :]\n                # pad = 2\n                # y1 = max(0, y - pad)\n                # y2 = min(H, y + patch_size + pad)\n                # x1 = max(0, x - pad)\n                # x2 = min(W, x + patch_size + pad)\n                # crop = image[y1:y2, x1:x2, :]\n                crops.append(crop)\n\n    if flip:\n        return [np.fliplr(x).copy() for x in crops]\n    return crops\n\n\nclass ImageCropDataset(Dataset):\n    \"\"\"Dataset for loading images and extracting crops in parallel.\"\"\"\n\n    def __init__(\n        self,\n        image_paths: list,\n        crop_config: list[dict],\n        model: nn.Module,\n        image_width: int,\n        image_height: int,\n    ):\n        self.image_paths = image_paths\n        self.crop_config = crop_config\n        self.model = model\n        self.image_width = image_width\n        self.image_height = image_height\n\n    def __len__(self):\n        return len(self.image_paths)\n\n    def __getitem__(self, idx):\n        \"\"\"Load image, extract crops, and preprocess.\"\"\"\n        image_path = self.image_paths[idx]\n        image = cv2.imread(image_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n        # if image.shape != (1000, 2000, 3):\n        #     # grass-clover dataset\n        #     image = cv2.resize(image, (2000, 2000))\n\n        if image.shape != (self.image_height, self.image_width, 3):\n            image = cv2.resize(image, (self.image_width, self.image_height))\n\n        transforms = {\n            \"identity\": A.Compose([]),\n            # \"hflip\": A.Compose([A.HorizontalFlip(p=1.0)]),\n            # \"vflip\": A.Compose([A.VerticalFlip(p=1.0)]),\n            # \"hvflip\": A.Compose([A.HorizontalFlip(p=1.0), A.VerticalFlip(p=1.0)]),\n            # \"rot90\": A.Compose(\n            #     [\n            #         A.Lambda(\n            #             image=lambda img, **k: cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)\n            #         )\n            #     ]\n            # ),\n            # \"hflip-rot90\": A.Compose(\n            #     [\n            #         A.HorizontalFlip(p=1.0),\n            #         A.Lambda(\n            #             image=lambda img, **k: cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)\n            #         ),\n            #     ]\n            # ),\n            # \"vflip-rot90\": A.Compose(\n            #     [\n            #         A.VerticalFlip(p=1.0),\n            #         A.Lambda(\n            #             image=lambda img, **k: cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)\n            #         ),\n            #     ]\n            # ),\n            # \"hvflip-rot90\": A.Compose(\n            #     [\n            #         A.HorizontalFlip(p=1.0),\n            #         A.VerticalFlip(p=1.0),\n            #         A.Lambda(\n            #             image=lambda img, **k: cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)\n            #         ),\n            #     ]\n            # ),\n        }\n\n        aug_crops = {}\n        for key, transform in transforms.items():\n            all_crops = []\n            n_crops = 0\n            for cfg_ in self.crop_config:\n                crops = get_crops(\n                    transform(image=image)[\"image\"],\n                    patch_size=cfg_.patch_size,\n                    stride=cfg_.stride,\n                    flip=cfg_.flip,\n                )\n                all_crops.extend(crops)\n                n_crops += len(crops)\n\n            # Preprocess crops\n            processed = self.model.preprocess(all_crops)\n            aug_crops[key] = processed\n\n        return {\n            \"aug_crops\": aug_crops,\n            \"n_crops\": n_crops,\n            \"image_path\": image_path,\n        }\n\n\ndef collate_crops(batch: dict) -> dict:\n    \"\"\"Custom collate function for ImageCropDataset.\n\n    Args:\n        batch: Dict of tuples (preprocessed_inputs, n_crops) from __getitem__\n\n    Returns:\n        Dict of (combined preprocessed inputs, list of n_crops)\n    \"\"\"\n    aug_crops = [item[\"aug_crops\"] for item in batch]\n    n_crops = [item[\"n_crops\"] for item in batch]\n\n    collated = {}\n    for aug_key in aug_crops[0].keys():\n        crops_list = [x[aug_key] for x in aug_crops]\n\n        # Combine preprocessed inputs\n        if isinstance(crops_list[0], torch.Tensor):\n            # OpenCLIPModel case: single tensor\n            combined = torch.cat(crops_list, dim=0)\n        else:\n            # TransformersModel case: dict-like object (dict or BatchFeature) of tensors\n            combined = {}\n            for key in crops_list[0].keys():\n                combined[key] = torch.cat([p[key] for p in crops_list], dim=0)\n\n        collated[aug_key] = combined\n\n    return {\n        \"aug_crops\": collated,\n        \"n_crops\": n_crops,\n        \"image_path\": [item[\"image_path\"] for item in batch],\n    }\n\n\ndef pool_feature_maps(\n    feature_maps: torch.Tensor,\n    crop_config: list[dict],\n    batch_size: int = 1,\n    image_height: int = 1000,\n    image_width: int = 2000,\n) -> torch.Tensor:\n    \"\"\"Pool feature maps from different crop configurations.\n\n    Args:\n        feature_maps: Feature maps of shape (total_crops, C, H, W), where total_crops = batch_size * n_crops_per_image\n            Crops are ordered as: [img1_crop1, img1_crop2, ..., img2_crop1, img2_crop2, ...]\n        crop_config: List of crop configurations\n        batch_size: Batch size (number of images)\n        image_height: Image height\n        image_width: Image width\n\n    Returns:\n        Pooled tensor of shape (batch_size, C, target_H, target_W)\n    \"\"\"\n    if feature_maps.ndim != 4:\n        raise ValueError(f\"Expected 4D tensor, got {feature_maps.dim()}D\")\n\n    total_crops, C, H, W = feature_maps.shape\n\n    # Calculate number of crops per image\n    n_crops_per_image = 0\n    for cfg_ in crop_config:\n        if cfg_.patch_size == -1:\n            n_crops_per_image += 1\n        else:\n            n_rows = image_height // cfg_.patch_size\n            n_cols = image_width // cfg_.patch_size\n            n_crops_per_image += n_rows * n_cols\n\n    # Verify that total_crops matches batch_size * n_crops_per_image\n    if total_crops != batch_size * n_crops_per_image:\n        raise ValueError(\n            f\"Expected {batch_size * n_crops_per_image} crops \"\n            f\"({batch_size} images * {n_crops_per_image} crops/image), \"\n            f\"but got {total_crops}\"\n        )\n\n    # 1. Determine target spatial size\n    # Use the size of the first feature map as target\n    target_h, target_w = H, W\n\n    # 2. Process each image in the batch\n    batch_results = []\n    for batch_idx in range(batch_size):\n        # Extract crops for this image\n        start_idx = batch_idx * n_crops_per_image\n        end_idx = start_idx + n_crops_per_image\n        image_crops = feature_maps[start_idx:end_idx]  # (n_crops_per_image, C, H, W)\n\n        # Interpolate all feature maps to the same spatial size\n        interpolated_features = []\n        counter = 0\n        for cfg_ in crop_config:\n            if cfg_.patch_size == -1:\n                n_rows = n_cols = 1\n            else:\n                n_rows = image_height // cfg_.patch_size\n                n_cols = image_width // cfg_.patch_size\n\n            feats = image_crops[counter : counter + n_rows * n_cols]  # (L, C, H, W)\n            if cfg_.flip:\n                feats = feats.flip(-1)\n            # Reshape feats to (1, C, H*n_rows, W*n_cols)\n            feats = feats.reshape(n_rows, n_cols, C, H, W)\n            feats = feats.permute(2, 0, 3, 1, 4)  # (C, n_rows, H, n_cols, W)\n            feat = feats.reshape(\n                1, C, H * n_rows, W * n_cols\n            )  # (1, C, H*n_rows, W*n_cols)\n            if feat.shape[2] != target_h or feat.shape[3] != target_w:\n                feat = F.interpolate(\n                    feat,\n                    size=(target_h, target_w),\n                    mode=\"bilinear\",\n                    align_corners=False,\n                )\n            interpolated_features.append(feat)\n            counter += n_rows * n_cols\n\n        # Stack all interpolated features for this image\n        stacked = torch.cat(interpolated_features, dim=0)  # (n_crops, C, H, W)\n\n        # 3. Pool across all crops (average pooling)\n        pooled = stacked.mean(dim=0, keepdim=True)  # (1, C, H, W)\n        batch_results.append(pooled)\n\n    # Concatenate results for all images in batch\n    return torch.cat(batch_results, dim=0)  # (batch_size, C, H, W)\n\n\nclass TransformersModel:\n    texts = [\n        # \"very low above-ground biomass\",\n        # \"low above-ground biomass\",\n        # \"moderate above-ground biomass\",\n        # \"high above-ground biomass\",\n        # \"very high above-ground biomass\",\n        # \"very low dry clover biomass\",\n        # \"low dry clover biomass\",\n        # \"moderate dry clover biomass\",\n        # \"high dry clover biomass\",\n        # \"very high dry clover biomass\",\n        # # Dry_Green_g\n        # \"Top-down photo of a pasture with almost no green grass other than clover, mostly bare soil and a few tiny green shoots.\",\n        # \"Top-down photo of a pasture with sparse patches of green grass (non-clover) and many gaps of exposed soil and dead vegetation.\",\n        # \"Top-down photo of a pasture with dense green grass (non-clover) covering most of the area, with only small patches of clover or dead material.\",\n        # \"Top-down photo of a pasture completely filled with tall, thick green grass (non-clover), forming a dense canopy and almost hiding the soil and other components.\",\n        # # Dry_Dead_g\n        # \"Top-down photo of a pasture with almost no brown or dead vegetation, dominated by fresh green grass and clover.\",\n        # \"Top-down photo of a pasture with a mix of green leaves and noticeable patches of brown, senescent stems and leaves scattered across the ground.\",\n        # \"Top-down photo of a pasture where brown, senescent vegetation and straw-like material cover most of the area, with only some green plants remaining.\",\n        # \"Top-down photo of a pasture almost completely covered by thick layers of dry, brown, straw-like dead vegetation, with very little visible green material.\",\n        # # Dry_Clover_g\n        # \"Top-down photo of a pasture with almost no clover, mostly narrow-leaf grasses and some bare soil.\",\n        # \"Top-down photo of a pasture with scattered patches of clover plants, visible trifoliate leaves mixed with grasses and some dead material.\",\n        # \"Top-down photo of a pasture where clover plants with many trifoliate leaves cover a large part of the area, with grasses and dead material in between.\",\n        # \"Top-down photo of a pasture almost completely dominated by dense clover, a solid carpet of trifoliate leaves and clover stems with very little bare soil or non-clover grass visible.\",\n        # # GDM_g\n        # \"Top-down photo of a pasture with very little green vegetation, mostly bare soil and brown, dead material and only a few small green plants.\",\n        # \"Top-down photo of a pasture with patchy green cover, a mixture of green grass and clover interspersed with exposed soil and dead material.\",\n        # \"Top-down photo of a pasture where lush green grass and clover provide nearly continuous green canopy with only small patches of soil or dead vegetation showing through.\",\n        # \"Top-down photo of a pasture with extremely dense, thick green canopy from both grass and clover, leaves overlapping and completely hiding the soil and most dead material.\",\n        # # Dry_Total_g\n        # \"Top-down photo of a pasture with very low total biomass, mostly bare soil with only a few short plants and thin scattered stems.\",\n        # \"Top-down photo of a pasture with moderate total biomass, patches of green grass, clover and some dead material forming a discontinuous but noticeable plant cover.\",\n        # \"Top-down photo of a pasture with high total biomass, a thick layer of vegetation that looks tall and dense, combining green grass, clover and dead material with little soil visible.\",\n        # \"Top-down photo of a pasture with extremely high total biomass, a very thick, multi-layered canopy of tall plants and accumulated dead material, giving a deep, bulky sward with no visible soil.\",\n        # # Proportion\n        # # Clover-dominated\n        # \"Top-down photo of a pasture dominated by clover, many trifoliate leaves forming most of the canopy with only a small amount of grass and dead material.\",\n        # # Grass-dominated, little clover\n        # \"Top-down photo of a pasture dominated by non-clover grasses, narrow green leaves covering the ground with almost no clover visible.\",\n        # # Dead-dominated\n        # \"Top-down photo of a pasture dominated by dead, brown vegetation, straw-like stems and litter on the ground with very few green plants.\",\n        # # Balanced mixture\n        # \"Top-down photo of a pasture with a balanced mixture of green grass, clover and dead material, none of the components clearly dominating.\",\n        # # Sparse, low-biomass sward\n        # \"Top-down photo of a pasture with a very open, sparse sward, many gaps of bare soil and only small, scattered patches of vegetation.\",\n        # Other\n        \"bare soil\",\n        \"dirt ground\",\n        \"sparse vegetation\",\n        \"exposed earth\",\n        \"low density pasture\",\n        \"thin grass\",\n        \"short clipped grass\",\n        \"average pasture cover\",\n        \"medium height grass\",\n        \"grazed pasture\",\n        \"dense tall pasture\",\n        \"thick grassy volume\",\n        \"high biomass\",\n        \"overgrown vegetation\",\n        \"lush green vibrant pasture\",\n        \"photosynthesizing leaves\",\n        \"fresh growth\",\n        \"dry brown dead grass\",\n        \"yellow straw\",\n        \"senesced material\",\n        \"standing hay\",\n        \"white clover\",\n        \"trifolium repens\",\n        \"broadleaf legume\",\n        \"clover flowers\",\n        \"ryegrass\",\n        \"blade-like leaves\",\n        \"fescue\",\n        \"grassy sward\",\n        \"broadleaf weeds\",\n        \"thistles\",\n        \"non-pasture vegetation\",\n    ]\n\n    def __init__(self, model_name: str, device: str = \"cuda\", compile: bool = False):\n        self.model_name = model_name\n\n        if \"radio\" in model_name.lower():\n            self.processor = CLIPImageProcessor.from_pretrained(model_name)\n            self.processor.size = {\n                \"height\": self.processor.crop_size[\"height\"],\n                \"width\": self.processor.crop_size[\"width\"],\n            }\n            self.processor.do_resize = True\n        elif \"internvit\" in model_name.lower():\n            self.processor = CLIPImageProcessor.from_pretrained(model_name)\n        else:\n            # self.processor = AutoProcessor.from_pretrained(model_name, use_fast=True)\n            self.processor = AutoImageProcessor.from_pretrained(\n                model_name, use_fast=True\n            )\n\n        if \"radio\" in model_name and model_name.startswith(\"/kaggle/input/\"):\n            # Copy all files under model_name to cache_dir\n            cache_dir = Path(os.environ[\"HF_MODULES_CACHE\"], \"transformers_modules\")\n            cache_dir.mkdir(parents=True, exist_ok=True)\n\n            for file_path in Path(model_name).iterdir():\n                if file_path.is_file():\n                    target_path = cache_dir / file_path.name\n                    if not target_path.exists():\n                        shutil.copy2(file_path, target_path)\n\n        self.model = AutoModel.from_pretrained(\n            model_name,\n            dtype=torch.float16,\n            low_cpu_mem_usage=True,\n            trust_remote_code=True,\n            local_files_only=True,\n        )\n        self.model = self.model.to(memory_format=torch.channels_last)\n        self.model = self.model.to(device)\n        self.model.eval()\n        if compile:\n            self.model = torch.compile(self.model)\n\n    def preprocess(self, images):\n        # self.processor.size = {\"height\": 1000, \"width\": 1000}\n\n        kwargs = {}\n        # if \"siglip2\" in self.model_name:\n        #     kwargs.update(dict(text=self.texts, padding=\"max_length\", max_length=64))\n        return self.processor(images=images, return_tensors=\"pt\", **kwargs)\n\n    @torch.inference_mode()\n    def __call__(\n        self,\n        inputs,\n        batch_size,\n        n_crops,\n        crop_config=None,\n        image_height=1000,\n        image_width=2000,\n    ):\n        B = len(inputs)\n\n        # Move to device\n        if isinstance(inputs, dict):\n            inputs = {\n                k: v.to(self.model.device, non_blocking=True) for k, v in inputs.items()\n            }\n        else:\n            inputs = inputs.to(self.model.device, non_blocking=True)\n\n        feature_maps = None\n        # if \"siglip2\" in self.model_name:\n        #     with torch.cuda.amp.autocast(dtype=torch.float16):\n        #         outputs = self.model(**inputs)\n        #     features = outputs.vision_model_output.pooler_output\n        #     text_similarities = outputs.logits_per_image\n        #     # features = torch.cat([features, text_similarities], dim=1)\n        #     features = text_similarities\n        if hasattr(self.model, \"get_image_features\"):\n            with torch.cuda.amp.autocast(dtype=torch.float16):\n                features = self.model.get_image_features(**inputs)\n        elif \"radio\" in self.model_name.lower():\n            with torch.cuda.amp.autocast(dtype=torch.float16):\n                features, spatial_features = self.model(inputs[\"pixel_values\"])\n            # feature_maps = rearrange(\n            #     spatial_features,\n            #     \"b (h w) d -> b d h w\",\n            #     h=432 // 16,\n            #     w=864 // 16,\n            # )\n            # features = feature_maps.mean(dim=[2, 3])\n        else:\n            with torch.cuda.amp.autocast(dtype=torch.float16):\n                outputs = self.model(**inputs, output_hidden_states=False)\n                features = outputs.pooler_output\n\n            # if \"dinov3-vit\" in self.model_name.lower():\n            #     # CLS token + mean patch tokens\n            #     S = 1 + self.model.config.num_register_tokens\n            #     patch_tokens = outputs.last_hidden_state[:, S:, :]\n            #     features = torch.cat([features, patch_tokens.mean(dim=1)], dim=1)\n\n            if (\n                \"dinov3-vit\" in self.model_name.lower()\n                or \"pe_spatial\" in self.model_name.lower()\n                or \"pe-spatial\" in self.model_name.lower()\n                or \"internvit\" in self.model_name.lower()\n            ):\n                # Feature maps (1, C, H, W)\n                B, _, C = outputs.last_hidden_state.shape\n                S = 0\n                if \"dinov3-vit\" in self.model_name.lower():\n                    S = 1 + self.model.config.num_register_tokens\n                elif \"internvit\" in self.model_name.lower():\n                    S = 1\n                patch_tokens = outputs.last_hidden_state[:, S:, :]\n                H = W = int(patch_tokens.size(1) ** 0.5)\n                # patch_size = self.model.config.patch_size\n                # (n_crops, L, C) -> (n_crops, C, H, W)\n                feature_maps = patch_tokens.reshape(B, H, W, C).permute(0, 3, 1, 2)\n                feature_maps = pool_feature_maps(\n                    feature_maps,\n                    crop_config,\n                    batch_size,\n                    image_height=image_height,\n                    image_width=image_width,\n                )\n\n        # def gem(x: torch.Tensor, p: float, eps: float = 1e-6) -> torch.Tensor:\n        #     return x.clamp(min=eps).pow(p).mean(dim=1).pow(1.0 / p)\n\n        # Pool multi-crop features\n        features = features.view(batch_size, n_crops, -1)\n        features = torch.cat(\n            [\n                features.mean(dim=1),\n                features.amax(dim=1),\n                features.amin(dim=1),\n                features.std(dim=1),\n            ],\n            dim=1,\n        )\n\n        # if \"dinov3-vit\" in self.model_name.lower():\n        #     patch_tokens = patch_tokens.view(batch_size, n_crops, 196, -1)\n        #     f_patch = torch.cat(\n        #         [\n        #             patch_tokens.mean(dim=(1, 2)),\n        #             patch_tokens.amax(dim=(1, 2)),\n        #             patch_tokens.amin(dim=(1, 2)),\n        #             patch_tokens.std(dim=(1, 2)),\n        #         ],\n        #         dim=1,\n        #     )\n        #     features = torch.dstack(\n        #         [features.view(B, 4, -1), f_patch.view(B, 4, -1)]\n        #     ).view(B, -1)\n\n        return features, feature_maps\n\n\nclass OpenCLIPModel(TransformersModel):\n    def __init__(self, model_name, device: str = \"cuda\", compile: bool = False):\n        import open_clip\n\n        self.model, _, self.transform = open_clip.create_model_and_transforms(\n            \"TULIP-so400m-14-384\",\n            pretrained=model_name,\n        )\n        self.model = self.model.to(memory_format=torch.channels_last)\n        self.model = self.model.to(device)\n        self.model.eval()\n        if compile:\n            self.model = torch.compile(self.model)\n\n    def preprocess(self, images):\n        return torch.stack([self.transform(x) for x in images], dim=0)\n\n    @torch.inference_mode()\n    def __call__(self, inputs, batch_size, n_crops, crop_config=None):\n        features = self.model.encode_image(inputs)\n        features /= features.norm(dim=-1, keepdim=True)\n\n        # Average multi-crop features\n        features = features.view(batch_size, n_crops, -1).mean(dim=1)\n        return features, None\n\n\n@torch.inference_mode()\ndef compute_features(\n    model,\n    image_paths: list,\n    save_path: Path,\n    crop_config: list[dict],\n    batch_size: int = 1,\n    num_workers: int = 4,\n    save_feature_maps: bool = True,\n    image_height: int = 1000,\n    image_width: int = 2000,\n):\n    \"\"\"Compute image features and save them to a npy file.\"\"\"\n    # Create dataset and dataloader for parallel image loading and preprocessing\n    dataset = ImageCropDataset(\n        image_paths,\n        crop_config,\n        model,\n        image_width=image_width,\n        image_height=image_height,\n    )\n    dataloader = DataLoader(\n        dataset,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=num_workers,\n        pin_memory=True if torch.cuda.is_available() else False,\n        collate_fn=collate_crops,\n    )\n\n    all_features = defaultdict(list)\n    all_feature_maps = defaultdict(list) if save_feature_maps else None\n    for batch in tqdm(dataloader, desc=f\"Extracting {save_path.name}\"):\n        aug_crops = batch[\"aug_crops\"]\n        n_crops_list = batch[\"n_crops\"]\n\n        n_crops = n_crops_list[0]  # Same crop number for all images in batch\n        B = len(n_crops_list)  # Actual batch size (may differ for last batch)\n\n        for key, inputs in aug_crops.items():\n            # Extract features\n            features, feature_maps = model(\n                inputs,\n                B,\n                n_crops,\n                crop_config,\n                image_height=image_height,\n                image_width=image_width,\n            )\n\n            # Collect features\n            features = features.cpu().numpy()\n            all_features[key].append(features)\n            if save_feature_maps and feature_maps is not None and key == \"identity\":\n                feature_maps = feature_maps.cpu().numpy()\n                # all_feature_maps[key].append(feature_maps)\n                for i in range(B):\n                    image_id = Path(batch[\"image_path\"][i]).stem\n                    path = Path(str(save_path).replace(\"features_\", \"feature_maps_\"))\n                    path = path.parent / (path.stem + f\"_{image_id}.npy\")\n                    np.save(path, feature_maps[i])\n\n    # Save all features\n    for key in all_features.keys():\n        all_features[key] = np.concatenate(all_features[key], axis=0)\n    # np.save(save_path, all_features[key])\n    with open(save_path, \"wb\") as f:\n        pickle.dump(all_features, f, protocol=pickle.HIGHEST_PROTOCOL)\n    print(f\"Saved features to {save_path} with shape {all_features['identity'].shape}\")\n    # print(\n    #     f\"  Mean: {all_features.mean():.6f}, Std: {all_features.std():.6f}, \"\n    #     f\"Min: {all_features.min():.6f},  Max: {all_features.max():.6f}\"\n    # )\n\n    # # Save feature maps if requested\n    # if save_feature_maps and all_feature_maps:\n    #     path = str(save_path).replace(\"features_\", \"feature_maps_\")\n    #     all_feature_maps = np.concatenate(all_feature_maps, axis=0)\n    #     np.save(path, all_feature_maps)\n    #     print(f\"Saved feature maps to {path} with shape {all_feature_maps.shape}\")\n\n\ndef _compute_features_on_device(\n    device_id: int,\n    model_name: str,\n    key: str,\n    image_paths: list,\n    save_path: Path,\n    crop_config: list[dict],\n    batch_size: int,\n    compile: bool,\n    image_height: int = 1000,\n    image_width: int = 2000,\n):\n    \"\"\"Helper function to compute features on a specific GPU device.\"\"\"\n    # Load model on specified device\n    print(f\"GPU {device_id}: Loading model {model_name}\")\n    model = TransformersModel(model_name, device=f\"cuda:{device_id}\", compile=compile)\n\n    # Compute features\n    compute_features(\n        model=model,\n        image_paths=image_paths,\n        save_path=save_path,\n        crop_config=crop_config,\n        batch_size=batch_size,\n        image_height=image_height,\n        image_width=image_width,\n    )\n\n    del model\n    torch.cuda.empty_cache()\n\n\ndef compute_features_parallel(\n    model_name: str,\n    key: str,\n    image_paths: list,\n    save_path: Path,\n    crop_config: list[dict],\n    batch_size: int,\n    compile: bool,\n    image_height: int = 1000,\n    image_width: int = 2000,\n    num_gpus: int = 2,\n):\n    \"\"\"Compute features in parallel across multiple GPUs.\"\"\"\n    if (\n        num_gpus <= 1\n        or not torch.cuda.is_available()\n        or torch.cuda.device_count() < num_gpus\n    ):\n        raise RuntimeError(\n            f\"Parallel GPU extraction not available. \"\n            f\"num_gpus={num_gpus}, \"\n            f\"torch.cuda.is_available()={torch.cuda.is_available()}, \"\n            f\"torch.cuda.device_count()={torch.cuda.device_count() if torch.cuda.is_available() else 0}\"\n        )\n\n    # Compute features in parallel\n    n_images = len(image_paths)\n    split_size = n_images // num_gpus\n    processes = []\n    temp_paths = []\n    for gpu_id in range(num_gpus):\n        start_idx = gpu_id * split_size\n        end_idx = start_idx + split_size if gpu_id < num_gpus - 1 else n_images\n\n        split_image_paths = image_paths[start_idx:end_idx]\n        if len(split_image_paths) == 0:\n            continue\n        temp_save_path = save_path.parent / f\"{save_path.stem}_gpu{gpu_id}.npy\"\n        temp_paths.append(temp_save_path)\n\n        print(\n            f\"GPU {gpu_id}: Processing {len(split_image_paths)} images ({start_idx} to {end_idx})\"\n        )\n\n        p = Process(\n            target=_compute_features_on_device,\n            args=(\n                gpu_id,\n                model_name,\n                key,\n                split_image_paths,\n                temp_save_path,\n                crop_config,\n                batch_size,\n                compile,\n                image_height,\n                image_width,\n            ),\n        )\n        p.start()\n        processes.append(p)\n\n    # Wait for all processes to complete\n    for p in processes:\n        p.join()\n\n    # Merge results\n    print(f\"Merging results from {num_gpus} GPUs...\")\n    all_features = []\n    all_feature_maps = []\n    has_feature_maps = False\n\n    for temp_path in temp_paths:\n        features = np.load(temp_path)\n        all_features.append(features)\n\n        # Check for feature maps\n        feature_map_path = Path(str(temp_path).replace(\"features_\", \"feature_maps_\"))\n        if feature_map_path.exists():\n            has_feature_maps = True\n            feature_maps = np.load(feature_map_path)\n            all_feature_maps.append(feature_maps)\n\n    # Save merged features\n    merged_features = np.concatenate(all_features, axis=0)\n    np.save(save_path, merged_features)\n    print(f\"Saved merged features to {save_path} with shape {merged_features.shape}\")\n\n    # Save merged feature maps if they exist\n    if has_feature_maps and all_feature_maps:\n        feature_map_save_path = Path(\n            str(save_path).replace(\"features_\", \"feature_maps_\")\n        )\n        merged_feature_maps = np.concatenate(all_feature_maps, axis=0)\n        np.save(feature_map_save_path, merged_feature_maps)\n        print(\n            f\"Saved merged feature maps to {feature_map_save_path} with shape {merged_feature_maps.shape}\"\n        )\n\n    # Clean up temporary files\n    for temp_path in temp_paths:\n        if temp_path.exists():\n            temp_path.unlink()\n        feature_map_path = Path(str(temp_path).replace(\"features_\", \"feature_maps_\"))\n        if feature_map_path.exists():\n            feature_map_path.unlink()\n\n    print(\"Parallel feature extraction complete.\")\n\n\n# --- Load Models and Extract Image Features ---\nif cfg.compute_features:\n    compile = is_kaggle() or len(cfg.models) >= 2\n    # compile = True\n\n    # Enable parallel GPU if 2+ GPUs available\n    # use_parallel_gpu = torch.cuda.device_count() >= 2\n    use_parallel_gpu = False\n\n    prev_model_name = None\n    model = None\n    for model_cfg in cfg.models:\n        key = model_cfg.key\n        model_name = model_cfg.kaggle_path if is_kaggle() else model_cfg.model_name\n\n        if model_name != prev_model_name and not use_parallel_gpu:\n            del model\n            gc.collect()\n            torch.cuda.empty_cache()\n\n            print(f\"Loading {model_name}\")\n            model = TransformersModel(model_name, device=DEVICE, compile=compile)\n\n        print(f\"Processing model: {model_name}\")\n\n        # Compute features for train data\n        if use_parallel_gpu:\n            compute_features_parallel(\n                model_name=model_name,\n                key=key,\n                image_paths=df[\"rel_image_path\"].to_list(),\n                save_path=FEATS_DIR / f\"features_{key}.pkl\",\n                crop_config=model_cfg.crops,\n                batch_size=model_cfg.get(\"batch_size\", 1),\n                compile=compile,\n                num_gpus=torch.cuda.device_count(),\n            )\n        else:\n            compute_features(\n                model=model,\n                image_paths=df[\"rel_image_path\"],\n                save_path=FEATS_DIR / f\"features_{key}.pkl\",\n                crop_config=model_cfg.crops,\n                batch_size=model_cfg.get(\"batch_size\", 1),\n            )\n        gc.collect()\n        torch.cuda.empty_cache()\n\n        # Compute features for test data\n        if use_parallel_gpu:\n            compute_features_parallel(\n                model_name=model_name,\n                key=key,\n                image_paths=df_test[\"rel_image_path\"].to_list(),\n                save_path=FEATS_DIR / f\"features_test_{key}.pkl\",\n                crop_config=model_cfg.crops,\n                batch_size=model_cfg.get(\"batch_size\", 1),\n                compile=compile,\n                num_gpus=torch.cuda.device_count(),\n            )\n        else:\n            compute_features(\n                model=model,\n                image_paths=df_test[\"rel_image_path\"],\n                save_path=FEATS_DIR / f\"features_test_{key}.pkl\",\n                crop_config=model_cfg.crops,\n                batch_size=model_cfg.get(\"batch_size\", 1),\n            )\n        gc.collect()\n        torch.cuda.empty_cache()\n\n        if use_grass_clover:\n            # Compute features for grass-clover data\n            if use_parallel_gpu:\n                compute_features_parallel(\n                    model_name=model_name,\n                    key=key,\n                    image_paths=df_grass_clover[\"rel_image_path\"].to_list(),\n                    save_path=FEATS_DIR / f\"features_grass_clover_{key}.pkl\",\n                    crop_config=model_cfg.crops,\n                    batch_size=model_cfg.get(\"batch_size\", 1),\n                    compile=compile,\n                    image_height=2000,\n                    image_width=2000,\n                    num_gpus=torch.cuda.device_count(),\n                )\n            else:\n                compute_features(\n                    model=model,\n                    image_paths=df_grass_clover[\"rel_image_path\"],\n                    save_path=FEATS_DIR / f\"features_grass_clover_{key}.pkl\",\n                    crop_config=model_cfg.crops,\n                    batch_size=model_cfg.get(\"batch_size\", 1),\n                    image_height=2000,\n                    image_width=2000,\n                )\n            gc.collect()\n            torch.cuda.empty_cache()\n\n        if use_irish:\n            # Compute features for grass-clover data\n            if use_parallel_gpu:\n                compute_features_parallel(\n                    model_name=model_name,\n                    key=key,\n                    image_paths=df_irish[\"rel_image_path\"].to_list(),\n                    save_path=FEATS_DIR / f\"features_irish_{key}.pkl\",\n                    crop_config=model_cfg.crops,\n                    batch_size=model_cfg.get(\"batch_size\", 1),\n                    compile=compile,\n                    image_height=2000,\n                    image_width=2000,\n                    num_gpus=torch.cuda.device_count(),\n                )\n            else:\n                compute_features(\n                    model=model,\n                    image_paths=df_irish[\"rel_image_path\"],\n                    save_path=FEATS_DIR / f\"features_irish_{key}.pkl\",\n                    crop_config=model_cfg.crops,\n                    batch_size=model_cfg.get(\"batch_size\", 1),\n                    image_height=2000,\n                    image_width=2000,\n                )\n            gc.collect()\n            torch.cuda.empty_cache()\n\n        prev_model_name = model_name\n\n    del model\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    print(\"Feature extraction complete.\")\n\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n\n\n# --- Load Features ---\ntarget = df[\"Dry_Total_g\"].to_numpy()\n\nHW = 32\n\nfeatures = defaultdict(list)\nfeatures_test = defaultdict(list)\nfeatures_grass_clover = defaultdict(list) if use_grass_clover else None\nfeatures_irish = defaultdict(list) if use_irish else None\nfor model_cfg in cfg.models:\n    key = model_cfg.key\n\n    print(f\"Loading {key} features\")\n    # feats = np.load(f\"{FEATS_DIR}/features_{key}.npy\")\n    # feats_test = np.load(f\"{FEATS_DIR}/features_test_{key}.npy\")\n    with open(f\"{FEATS_DIR}/features_{key}.pkl\", \"rb\") as f:\n        aug_feats = pickle.load(f)\n    with open(f\"{FEATS_DIR}/features_test_{key}.pkl\", \"rb\") as f:\n        aug_feats_test = pickle.load(f)\n\n    # for aug_key in [\n    #     \"identity\",\n    #     # \"hflip\",\n    #     # \"vflip\",\n    # ]:\n    for aug_key in aug_feats.keys():\n        feats = aug_feats[aug_key]\n        feats_test = aug_feats_test[aug_key]\n\n        # if \"siglip2\" in key:\n        #     n_texts = len(TransformersModel.texts)\n        #     feats = feats[:, -n_texts:]\n        #     feats_test = feats_test[:, -n_texts:]\n\n        feats, max_feats, min_feats, std_feats = np.split(feats, 4, axis=1)\n        # feats = F.normalize(torch.tensor(feats), dim=1).numpy()\n        # max_feats = F.normalize(torch.tensor(max_feats), dim=1).numpy()\n        # min_feats = F.normalize(torch.tensor(min_feats), dim=1).numpy()\n        # std_feats = F.normalize(torch.tensor(std_feats), dim=1).numpy()\n        print(\n            f\"  Mean: {feats.mean():.6f}, Std: {feats.std():.6f}, \"\n            f\"Min: {feats.min():.6f}, Max: {feats.max():.6f}\"\n        )\n        feats_test, max_feats_test, min_feats_test, std_feats_test = np.split(\n            feats_test, 4, axis=1\n        )\n        # feats_test = F.normalize(torch.tensor(feats_test), dim=1).numpy()\n        # max_feats_test = F.normalize(torch.tensor(max_feats_test), dim=1).numpy()\n        # min_feats_test = F.normalize(torch.tensor(min_feats_test), dim=1).numpy()\n        # std_feats_test = F.normalize(torch.tensor(std_feats_test), dim=1).numpy()\n        print(\n            f\"  Mean: {feats_test.mean():.6f}, Std: {feats_test.std():.6f}, \"\n            f\"Min: {feats_test.min():.6f}, Max: {feats_test.max():.6f}\"\n        )\n\n        # Filter features by correlation with target\n        if hasattr(model_cfg, \"avg_pool_pct\"):\n            sel = SelectPercentile(f_regression, percentile=model_cfg.avg_pool_pct)\n            feats = sel.fit_transform(feats, target)\n            feats_test = sel.transform(feats_test)\n            mask0 = sel.get_support()\n            print(\n                f\"  Filtered avg pool features from {len(mask0)} to {mask0.sum()} \"\n                f\"based on correlation (proportion={model_cfg.avg_pool_pct})\"\n            )\n\n        if hasattr(model_cfg, \"max_pool_pct\"):\n            sel = SelectPercentile(f_regression, percentile=model_cfg.max_pool_pct)\n            max_feats = sel.fit_transform(max_feats, target)\n            max_feats_test = sel.transform(max_feats_test)\n            mask1 = sel.get_support()\n            print(\n                f\"  Filtered max pool features from {len(mask1)} to {mask1.sum()} \"\n                f\"based on correlation (proportion={model_cfg.max_pool_pct})\"\n            )\n            feats = np.concatenate([feats, max_feats], axis=1)\n            feats_test = np.concatenate([feats_test, max_feats_test], axis=1)\n        if hasattr(model_cfg, \"min_pool_pct\"):\n            sel = SelectPercentile(f_regression, percentile=model_cfg.min_pool_pct)\n            min_feats = sel.fit_transform(min_feats, target)\n            min_feats_test = sel.transform(min_feats_test)\n            mask2 = sel.get_support()\n            print(\n                f\"  Filtered min pool features from {len(mask2)} to {mask2.sum()} \"\n                f\"based on correlation (proportion={model_cfg.min_pool_pct})\"\n            )\n            feats = np.concatenate([feats, min_feats], axis=1)\n            feats_test = np.concatenate([feats_test, min_feats_test], axis=1)\n        if hasattr(model_cfg, \"std_pool_pct\"):\n            sel = SelectPercentile(f_regression, percentile=model_cfg.std_pool_pct)\n            std_feats = sel.fit_transform(std_feats, target)\n            std_feats_test = sel.transform(std_feats_test)\n            mask3 = sel.get_support()\n            print(\n                f\"  Filtered std pool features from {len(mask3)} to {mask3.sum()} \"\n                f\"based on correlation (proportion={model_cfg.std_pool_pct})\"\n            )\n            feats = np.concatenate([feats, std_feats], axis=1)\n            feats_test = np.concatenate([feats_test, std_feats_test], axis=1)\n\n        if hasattr(model_cfg, \"total_pct\"):\n            sel = SelectPercentile(f_regression, percentile=model_cfg.total_pct)\n            feats = sel.fit_transform(feats, target)\n            feats_test = sel.transform(feats_test)\n            mask = sel.get_support()\n            print(\n                f\"  Filtered overall features from {len(mask)} to {mask.sum()} \"\n                f\"based on correlation (proportion={model_cfg.total_pct})\"\n            )\n\n        features[aug_key].append(feats)\n        features_test[aug_key].append(feats_test)\n\n        del (\n            feats,\n            max_feats,\n            min_feats,\n            std_feats,\n            feats_test,\n            max_feats_test,\n            min_feats_test,\n            std_feats_test,\n        )\n        gc.collect()\n\n        if use_grass_clover:\n            # feats = np.load(f\"{FEATS_DIR}/features_grass_clover_{key}.npy\")\n            with open(f\"{FEATS_DIR}/features_grass_clover_{key}.pkl\", \"rb\") as f:\n                aug_feats_gs = pickle.load(f)\n\n            if aug_key in aug_feats_gs:\n                feats = aug_feats_gs[aug_key]\n\n                print(\n                    f\"  Mean: {feats.mean():.6f}, Std: {feats.std():.6f}, \"\n                    f\"Min: {feats.min():.6f}, Max: {feats.max():.6f}\"\n                )\n\n                feats, max_feats, min_feats, std_feats = np.split(feats, 4, axis=1)\n                # feats = F.normalize(torch.tensor(feats), dim=1).numpy()\n                # max_feats = F.normalize(torch.tensor(max_feats), dim=1).numpy()\n                # min_feats = F.normalize(torch.tensor(min_feats), dim=1).numpy()\n                # std_feats = F.normalize(torch.tensor(std_feats), dim=1).numpy()\n\n                # Filter features by correlation with target\n                if hasattr(model_cfg, \"avg_pool_pct\"):\n                    feats = feats[:, mask0]\n\n                    print(\n                        f\"  Filtered avg pool features from {len(mask0)} to {mask0.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.avg_pool_pct})\"\n                    )\n\n                if hasattr(model_cfg, \"max_pool_pct\"):\n                    max_feats = max_feats[:, mask1]\n                    print(\n                        f\"  Filtered max pooled features from {len(mask1)} to {mask1.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.max_pool_pct})\"\n                    )\n                    feats = np.concatenate([feats, max_feats], axis=1)\n                if hasattr(model_cfg, \"min_pool_pct\"):\n                    min_feats = min_feats[:, mask2]\n                    print(\n                        f\"  Filtered min pooled features from {len(mask2)} to {mask2.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.min_pool_pct})\"\n                    )\n                    feats = np.concatenate([feats, min_feats], axis=1)\n                if hasattr(model_cfg, \"std_pool_pct\"):\n                    std_feats = std_feats[:, mask3]\n                    print(\n                        f\"  Filtered std pooled features from {len(mask3)} to {mask3.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.std_pool_pct})\"\n                    )\n                    feats = np.concatenate([feats, std_feats], axis=1)\n                if hasattr(model_cfg, \"total_pct\"):\n                    std_feats = std_feats[:, mask]\n                    print(\n                        f\"  Filtered overall features from {len(mask)} to {mask.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.total_pct})\"\n                    )\n                    feats = np.concatenate([feats, std_feats], axis=1)\n\n                features_grass_clover[aug_key].append(feats)\n\n                del feats, max_feats, min_feats, std_feats\n            del aug_feats_gs\n            gc.collect()\n\n        if use_irish:\n            # feats = np.load(f\"{FEATS_DIR}/features_irish_{key}.npy\")\n            with open(f\"{FEATS_DIR}/features_irish_{key}.pkl\", \"rb\") as f:\n                aug_feats_irish = pickle.load(f)\n\n            if aug_key in aug_feats_irish:\n                feats = aug_feats_irish[aug_key]\n\n                print(\n                    f\"  Mean: {feats.mean():.6f}, Std: {feats.std():.6f}, \"\n                    f\"Min: {feats.min():.6f}, Max: {feats.max():.6f}\"\n                )\n\n                feats, max_feats, min_feats, std_feats = np.split(feats, 4, axis=1)\n                # feats = F.normalize(torch.tensor(feats), dim=1).numpy()\n                # max_feats = F.normalize(torch.tensor(max_feats), dim=1).numpy()\n                # min_feats = F.normalize(torch.tensor(min_feats), dim=1).numpy()\n                # std_feats = F.normalize(torch.tensor(std_feats), dim=1).numpy()\n\n                # Filter features by correlation with target\n                if hasattr(model_cfg, \"avg_pool_pct\"):\n                    feats = feats[:, mask0]\n\n                    print(\n                        f\"  Filtered avg pool features from {len(mask0)} to {mask0.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.avg_pool_pct})\"\n                    )\n\n                if hasattr(model_cfg, \"max_pool_pct\"):\n                    max_feats = max_feats[:, mask1]\n                    print(\n                        f\"  Filtered max pooled features from {len(mask1)} to {mask1.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.max_pool_pct})\"\n                    )\n                    feats = np.concatenate([feats, max_feats], axis=1)\n                if hasattr(model_cfg, \"min_pool_pct\"):\n                    min_feats = min_feats[:, mask2]\n                    print(\n                        f\"  Filtered min pooled features from {len(mask2)} to {mask2.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.min_pool_pct})\"\n                    )\n                    feats = np.concatenate([feats, min_feats], axis=1)\n                if hasattr(model_cfg, \"std_pool_pct\"):\n                    std_feats = std_feats[:, mask3]\n                    print(\n                        f\"  Filtered std pooled features from {len(mask3)} to {mask3.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.std_pool_pct})\"\n                    )\n                    feats = np.concatenate([feats, std_feats], axis=1)\n                if hasattr(model_cfg, \"total_pct\"):\n                    std_feats = std_feats[:, mask]\n                    print(\n                        f\"  Filtered overall features from {len(mask)} to {mask.sum()} \"\n                        f\"based on correlation (proportion={model_cfg.total_pct})\"\n                    )\n                    feats = np.concatenate([feats, std_feats], axis=1)\n\n                features_irish[aug_key].append(feats)\n\n                del feats, max_feats, min_feats, std_feats\n            del aug_feats_irish\n            gc.collect()\n\n# Concatenate all features from different models horizontally\nfor aug_key in features.keys():\n    features[aug_key] = np.concatenate(features[aug_key], axis=1)\n    print(f\"\\n{features[aug_key].shape=}\")\n    print(\n        f\"  Mean: {features[aug_key].mean():.6f}, Std: {features[aug_key].std():.6f}, \"\n        f\"Min: {features[aug_key].min():.6f}, Max: {features[aug_key].max():.6f}\"\n    )\n    features_test[aug_key] = np.concatenate(features_test[aug_key], axis=1)\n    print(f\"{features_test[aug_key].shape=}\")\n    print(\n        f\"  Mean: {features_test[aug_key].mean():.6f}, Std: {features_test[aug_key].std():.6f}, \"\n        f\"Min: {features_test[aug_key].min():.6f}, Max: {features_test[aug_key].max():.6f}\"\n    )\nif use_grass_clover:\n    for aug_key in features_grass_clover.keys():\n        features_grass_clover[aug_key] = np.concatenate(\n            features_grass_clover[aug_key], axis=1\n        )\n        print(f\"{features_grass_clover[aug_key].shape=}\")\n        print(\n            f\"  Mean: {features_grass_clover[aug_key].mean():.6f}, Std: {features_grass_clover[aug_key].std():.6f}, \"\n            f\"Min: {features_grass_clover[aug_key].min():.6f}, Max: {features_grass_clover[aug_key].max():.6f}\"\n        )\nif use_irish:\n    for aug_key in features_irish.keys():\n        features_irish[aug_key] = np.concatenate(features_irish[aug_key], axis=1)\n        print(f\"{features_irish[aug_key].shape=}\")\n        print(\n            f\"  Mean: {features_irish[aug_key].mean():.6f}, Std: {features_irish[aug_key].std():.6f}, \"\n            f\"Min: {features_irish[aug_key].min():.6f}, Max: {features_irish[aug_key].max():.6f}\"\n        )\n\n# # Interleave features for grouped processing\n# group_size = len(cfg.models)\n# N, D = features.shape\n# features = np.stack(np.split(features, group_size, axis=1), axis=2).reshape(N, D)\n# N, D = features_test.shape\n# features_test = np.stack(np.split(features_test, group_size, axis=1), axis=2).reshape(\n#     N, D\n# )\n\n# ==============================================================================\n# Validation Setup\n# ==============================================================================\n\n\n@contextmanager\ndef bn_eval_only(model: nn.Module):\n    bns = []\n    for m in model.modules():\n        if isinstance(m, nn.modules.batchnorm._BatchNorm):\n            bns.append((m, m.training))\n            m.eval()  # running statsを更新しない\n    try:\n        yield\n    finally:\n        for m, was_training in bns:\n            m.train(was_training)\n\n\nclass MLPRegressor(LightningModule):\n    def __init__(\n        self,\n        input_dim: int,\n        output_dim: int,\n        hidden_dim: int = 256,\n        dropout1: float = 0.3,\n        dropout2: float = 0.3,\n        act_layer: DictConfig = cfg.mlp.act_layer,\n        norm_layer: DictConfig = cfg.mlp.norm_layer,\n        num_layers: int = 2,\n        drop_path_rate: float | None = 0.4,\n        out_act: DictConfig = cfg.mlp.out_act,\n        mixup_alpha: float = 0.3,\n        clip_max: float = 126.5,\n        beta: float = 1.0,\n        base_lr: float = 1e-3,\n        weight_decay: float = 1e-4,\n        disasble_bias_wd: bool = False,\n        optim_cfg: DictConfig = cfg.mlp.optimizer,\n        scheduler_params: DictConfig = cfg.mlp.scheduler_params,\n        spatial: bool = False,\n    ):\n        super().__init__()\n        self.save_hyperparameters()\n        self.mixup_alpha = mixup_alpha\n        self.clip_max = clip_max\n        self.beta = beta\n        self.base_lr = base_lr\n        self.weight_decay = weight_decay\n        self.disasble_bias_wd = disasble_bias_wd\n        self.optim_cfg = optim_cfg\n        self.scheduler_params = scheduler_params\n\n        self.use_mixup = self.mixup_alpha > 0\n\n        if spatial:\n            self.build_spatial_mlp(\n                input_dim=input_dim,\n                output_dim=output_dim,\n                hidden_dim=hidden_dim,\n                dropout1=dropout1,\n                dropout2=dropout2,\n                act_layer=act_layer,\n                norm_layer=norm_layer,\n                num_layers=num_layers,\n                out_act=out_act,\n                drop_path_rate=drop_path_rate,\n            )\n        else:\n            self.build_mlp(\n                input_dim=input_dim,\n                output_dim=output_dim,\n                hidden_dim=hidden_dim,\n                dropout1=dropout1,\n                dropout2=dropout2,\n                act_layer=act_layer,\n                norm_layer=norm_layer,\n                num_layers=num_layers,\n                out_act=out_act,\n                drop_path_rate=drop_path_rate,\n            )\n\n        # convert_to_dsbn(self)\n\n        self.train_r2 = WeightedR2()\n        self.val_r2 = WeightedR2()\n\n        self.best_score = float(\"-inf\")\n\n        self.training_step_outputs = []\n        self.validation_step_outputs = []\n\n    def build_mlp(\n        self,\n        input_dim: int,\n        output_dim: int,\n        hidden_dim: int,\n        dropout1: float,\n        dropout2: float,\n        act_layer: DictConfig,\n        norm_layer: DictConfig,\n        num_layers: int,\n        out_act: DictConfig,\n        layer_scale_init: float | None = None,\n        use_residual: bool = True,\n        drop_path_rate: float | None = 0.4,\n    ):\n        \"\"\"Build MLP with optional enhancements.\"\"\"\n        # # BlockDropout\n        # blocks = [\n        #     # BlockSpec(\"dinov3-vit7b16-whole-flip\", dim=410),\n        #     # BlockSpec(\"dinov3-vit7b16-200\", dim=5325),\n        #     # BlockSpec(\"pe_spatial_gigantic-1000-flip\", dim=385),\n        #     # BlockSpec(\"pe_spatial_gigantic-250\", dim=1998),\n        #     BlockSpec(\"dinov3-vit7b16\", dim=5735),\n        #     BlockSpec(\"pe_spatial_gigantic\", dim=2383),\n        # ]\n        # assert sum(b.dim for b in blocks) == input_dim\n        self.block_dropout = None\n        # self.block_dropout = BlockDropout(\n        #     blocks=blocks,\n        #     p=0.1,\n        #     mode=\"per_sample\",\n        #     rescale=True,\n        #     always_keep_one=True,\n        # )\n\n        # Total: 7478\n        # - 410\n        # - 4916 (4096 + 410 + 410)\n        # - 308 (154 + 77 + 77)\n        # - 1844 (1536 + 154 + 154)\n        self.sep_linear = None\n        # blocks = [\n        #     # BlockSpec(\"dinov3-vit7b16-whole-flip\", dim=410),\n        #     BlockSpec(\"dinov3-vit7b16-200\", dim=4916),\n        #     # BlockSpec(\"pe_spatial_gigantic-1000-flip\", dim=308),\n        #     BlockSpec(\"pe_spatial_gigantic-250\", dim=1844),\n        # ]\n        # assert sum(b.dim for b in blocks) == input_dim\n        # self.sep_linear = SeparateLinear(blocks, out_features=[640, 640], bias=True)\n\n        # Build layers\n        layers = nn.ModuleList()\n\n        # Stochastic depth: linearly increase drop path rate\n        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]\n\n        for i in range(num_layers):\n            in_features = input_dim if i == 0 else hidden_dim\n            dropout = dropout1 if i == 0 else dropout2\n\n            if i == 0 and isinstance(self.sep_linear, SeparateLinear):\n                in_features = 640 + 640\n\n            layers.append(\n                MLPBlock(\n                    in_features=in_features,\n                    out_features=hidden_dim,\n                    dropout=dropout,\n                    act_layer=act_layer,\n                    norm_layer=norm_layer,\n                    layer_scale_init=layer_scale_init,\n                    use_residual=use_residual and (i != 0),\n                    drop_path=dpr[i],\n                )\n            )\n\n        # Head layer\n        head = nn.Sequential(\n            nn.Dropout(dropout2),\n            nn.Linear(hidden_dim, output_dim, bias=True),\n            hydra.utils.instantiate(out_act),\n        )\n\n        # Combine into sequential model\n        self.model = nn.Sequential(*layers, head)\n\n        # Initialize weights\n        for m in self.model.modules():\n            if isinstance(m, nn.Linear):\n                if m.out_features == output_dim:\n                    # Head layer\n                    nn.init.zeros_(m.weight)\n                    if m.bias is not None:\n                        nn.init.zeros_(m.bias)\n                else:\n                    # Hidden layers\n                    nn.init.kaiming_normal_(m.weight)\n                    if m.bias is not None:\n                        nn.init.zeros_(m.bias)\n\n        # self.linear = nn.Linear(input_dim, output_dim, bias=True)\n        # nn.init.kaiming_normal_(self.linear.weight)\n        # nn.init.zeros_(self.linear.bias)\n\n        # self.ratio = nn.Sequential(\n        #     nn.Dropout(dropout1),\n        #     nn.Linear(input_dim, hidden_dim, bias=False),\n        #     hydra.utils.instantiate(norm_layer, num_features=hidden_dim),\n        #     hydra.utils.instantiate(act_layer),\n        #     nn.Dropout(dropout2),\n        #     nn.Linear(hidden_dim, 3, bias=True),\n        # )\n\n    def build_spatial_mlp(\n        self,\n        input_dim: int,\n        output_dim: int,\n        hidden_dim: int,\n        dropout1: float,\n        dropout2: float,\n        act_layer: DictConfig,\n        norm_layer: DictConfig,\n        num_layers: int,\n        out_act: DictConfig,\n        drop_path_rate: float | None = None,\n    ):\n        # Build layers\n        layers = []\n        for i in range(num_layers):\n            if i == 0:\n                layers.extend(\n                    [\n                        nn.Dropout2d(dropout1),\n                        nn.Conv2d(input_dim, hidden_dim, kernel_size=1, bias=False),\n                        nn.Conv2d(\n                            hidden_dim,\n                            hidden_dim,\n                            3,\n                            1,\n                            1,\n                            bias=False,\n                            groups=hidden_dim,\n                        ),\n                        nn.BatchNorm2d(hidden_dim, affine=True),\n                        hydra.utils.instantiate(act_layer),\n                    ]\n                )\n            else:\n                layers.extend(\n                    [\n                        nn.Dropout2d(dropout2),\n                        nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1, bias=False),\n                        nn.Conv2d(\n                            hidden_dim,\n                            hidden_dim,\n                            3,\n                            1,\n                            1,\n                            bias=False,\n                            groups=hidden_dim,\n                        ),\n                        nn.BatchNorm2d(hidden_dim, affine=True),\n                        hydra.utils.instantiate(act_layer),\n                    ]\n                )\n\n        # Add head layer\n        layers.extend(\n            [\n                nn.Dropout2d(dropout2),\n                nn.Conv2d(hidden_dim, output_dim, kernel_size=1, bias=True),\n                # nn.AdaptiveAvgPool2d(1),\n                # nn.Flatten(),\n                # nn.Linear(hidden_dim, output_dim, bias=True),\n                hydra.utils.instantiate(out_act),\n            ]\n        )\n\n        self.model = nn.Sequential(*layers)\n\n        # Initialize weights\n        for m in self.model.modules():\n            if isinstance(m, nn.Conv2d):\n                if m.out_channels == output_dim:\n                    # Head layer\n                    # nn.init.normal_(m.weight, 0, 1e-2)\n                    nn.init.zeros_(m.weight)\n                    if m.bias is not None:\n                        nn.init.zeros_(m.bias)\n                else:\n                    # Hidden layers\n                    nn.init.kaiming_normal_(m.weight)\n                    if m.bias is not None:\n                        nn.init.zeros_(m.bias)\n\n    # def build_spatial_mlp(\n    #     self,\n    #     input_dim: int,\n    #     output_dim: int,\n    #     hidden_dim: int,\n    #     dropout1: float,\n    #     dropout2: float,\n    #     act_layer: DictConfig,\n    #     norm_layer: DictConfig,\n    #     num_layers: int,\n    #     out_act: DictConfig,\n    #     use_se: bool = False,  # Squeeze-and-Excitation\n    #     se_ratio: float = 0.25,  # SE reduction ratio\n    #     use_residual: bool = True,  # Residual connections\n    #     layer_scale_init: float | None = None,  # Layer Scale initialization\n    #     drop_path_rate: float | None = 0.4,  # Stochastic Depth rate\n    # ):\n    #     \"\"\"Build spatial MLP with depthwise convolutions and optional enhancements.\"\"\"\n\n    #     class SEBlock(nn.Module):\n    #         \"\"\"Squeeze-and-Excitation block for 2D spatial features\"\"\"\n\n    #         def __init__(self, channels, reduction=4):\n    #             super().__init__()\n    #             self.squeeze = nn.AdaptiveAvgPool2d(1)\n    #             self.excitation = nn.Sequential(\n    #                 nn.Conv2d(channels, channels // reduction, 1),\n    #                 nn.ReLU(inplace=True),\n    #                 nn.Conv2d(channels // reduction, channels, 1),\n    #                 nn.Sigmoid(),\n    #             )\n\n    #         def forward(self, x):\n    #             scale = self.squeeze(x)\n    #             scale = self.excitation(scale)\n    #             return x * scale\n\n    #     class ConvBlock(nn.Module):\n    #         \"\"\"Enhanced convolution block with optional SE and residual\"\"\"\n\n    #         def __init__(\n    #             self,\n    #             in_channels,\n    #             out_channels,\n    #             dropout,\n    #             use_se,\n    #             se_ratio,\n    #             layer_scale_init,\n    #             drop_path,\n    #             is_first=False,\n    #         ):\n    #             super().__init__()\n    #             self.is_first = is_first\n    #             self.use_residual = (\n    #                 use_residual and (in_channels == out_channels) and not is_first\n    #             )\n\n    #             # Main path\n    #             self.drop = nn.Dropout2d(dropout)\n    #             self.conv1 = nn.Conv2d(\n    #                 in_channels, out_channels, kernel_size=1, bias=False\n    #             )\n    #             self.dwconv = nn.Conv2d(\n    #                 out_channels, out_channels, 3, 1, 1, bias=False, groups=out_channels\n    #             )\n    #             self.norm = nn.BatchNorm2d(out_channels, affine=True)\n    #             self.act = hydra.utils.instantiate(act_layer)\n\n    #             # Optional SE block\n    #             self.se = SEBlock(out_channels, int(1 / se_ratio)) if use_se else None\n\n    #             # Optional Layer Scale\n    #             self.layer_scale = None\n    #             if layer_scale_init is not None and layer_scale_init > 0:\n    #                 self.layer_scale = nn.Parameter(\n    #                     layer_scale_init * torch.ones(out_channels, 1, 1)\n    #                 )\n\n    #             # Optional Drop Path\n    #             self.drop_path = (\n    #                 DropPath(drop_path) if drop_path > 0 and self.use_residual else None\n    #             )\n\n    #         def forward(self, x):\n    #             shortcut = x\n\n    #             x = self.drop(x)\n    #             x = self.conv1(x)\n    #             x = self.dwconv(x)\n    #             x = self.norm(x)\n    #             x = self.act(x)\n\n    #             if self.se is not None:\n    #                 x = self.se(x)\n\n    #             if self.layer_scale is not None:\n    #                 x = x * self.layer_scale\n\n    #             if self.use_residual:\n    #                 if self.drop_path is not None:\n    #                     x = self.drop_path(x)\n    #                 x = x + shortcut\n\n    #             return x\n\n    #     # Build blocks\n    #     blocks = nn.ModuleList()\n\n    #     # Stochastic depth: linearly increase drop path rate\n    #     dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]\n\n    #     for i in range(num_layers):\n    #         in_ch = input_dim if i == 0 else hidden_dim\n    #         dropout = dropout1 if i == 0 else dropout2\n\n    #         blocks.append(\n    #             ConvBlock(\n    #                 in_ch,\n    #                 hidden_dim,\n    #                 dropout,\n    #                 use_se,\n    #                 se_ratio,\n    #                 layer_scale_init,\n    #                 dpr[i],\n    #                 is_first=(i == 0),\n    #             )\n    #         )\n\n    #     # Head layer\n    #     head = nn.Sequential(\n    #         nn.Dropout2d(dropout2),\n    #         nn.Conv2d(hidden_dim, output_dim, kernel_size=1, bias=True),\n    #         hydra.utils.instantiate(out_act),\n    #     )\n\n    #     # Combine into sequential model\n    #     self.model = nn.Sequential(*blocks, head)\n\n    #     # Initialize weights\n    #     for m in self.model.modules():\n    #         if isinstance(m, nn.Conv2d):\n    #             if m.out_channels == output_dim:\n    #                 # Head layer\n    #                 nn.init.zeros_(m.weight)\n    #                 if m.bias is not None:\n    #                     nn.init.zeros_(m.bias)\n    #             elif m.groups == m.in_channels:\n    #                 # Depthwise conv: use smaller std for small datasets\n    #                 nn.init.kaiming_normal_(\n    #                     m.weight,\n    #                     # mode=\"fan_out\",\n    #                     # nonlinearity=\"relu\",\n    #                 )\n    #                 # m.weight.data *= 0.5  # Scale down for stability\n    #             else:\n    #                 # Pointwise conv\n    #                 nn.init.kaiming_normal_(m.weight)\n    #                 if m.bias is not None:\n    #                     nn.init.zeros_(m.bias)\n\n    def forward(self, x):\n        if hasattr(self, \"block_dropout\") and self.block_dropout is not None:\n            x = self.block_dropout(x)\n\n        if hasattr(self, \"sep_linear\") and self.sep_linear is not None:\n            x = self.sep_linear(x)\n\n        # x = self.linear(F.dropout(x, p=0.0))\n        # x = self.grouped(x)\n\n        x = self.model(x)\n\n        if x.ndim == 4:\n            return x.mean(dim=[2, 3])  # Average pooling\n        return x\n        # return x, self.ratio(x)\n\n    def training_step(self, batch, batch_idx):\n        input, target = (\n            batch[\"csiro-biomass\"][\"input\"],\n            batch[\"csiro-biomass\"][\"target\"],\n        )\n\n        if self.use_mixup and self.training:\n            batch_dict = {\"input\": input, \"target\": target}\n            mixed_x, batch_a, batch_b, lam = mixup(input, batch_dict, self.mixup_alpha)\n\n            # set_model_domain(self, \"csiro\")\n            pred = self.forward(mixed_x)[:, : len(LABELS)]\n\n            y_a = batch_a[\"target\"].clamp(min=0, max=self.clip_max)\n            y_b = batch_b[\"target\"].clamp(min=0, max=self.clip_max)\n\n            loss_a = WeightedSmoothL1Loss(beta=self.beta)(pred, y_a)\n            loss_b = WeightedSmoothL1Loss(beta=self.beta)(pred, y_b)\n            loss = lam * loss_a + (1 - lam) * loss_b\n\n            if \"grass-clover\" in batch:\n                input, target = (\n                    batch[\"grass-clover\"][\"input\"],\n                    batch[\"grass-clover\"][\"target\"],\n                )\n                batch_dict = {\"input\": input, \"target\": target}\n                mixed_x, batch_a, batch_b, lam = mixup(input, batch_dict, alpha=0.1)\n\n                with bn_eval_only(self):\n                    # set_model_domain(self, \"grass\")\n                    pred = self.forward(mixed_x)[:, len(LABELS) : len(LABELS) + 3]\n\n                    y_a = batch_a[\"target\"].clamp(min=0, max=self.clip_max)\n                    y_b = batch_b[\"target\"].clamp(min=0, max=self.clip_max)\n\n                    loss_a = WeightedSmoothL1Loss(beta=self.beta, weights=[1.0])(\n                        pred, y_a\n                    )\n                    loss_b = WeightedSmoothL1Loss(beta=self.beta, weights=[1.0])(\n                        pred, y_b\n                    )\n                    loss += 0.04 * (lam * loss_a + (1 - lam) * loss_b)\n\n            if \"irish\" in batch:\n                input, target = (\n                    batch[\"irish\"][\"input\"],\n                    batch[\"irish\"][\"target\"],\n                )\n                batch_dict = {\"input\": input, \"target\": target}\n                mixed_x, batch_a, batch_b, lam = mixup(input, batch_dict, alpha=0.1)\n\n                with bn_eval_only(self):\n                    # set_model_domain(self, \"irish\")\n                    pred = self.forward(mixed_x)[:, -2:]\n\n                    y_a = batch_a[\"target\"].clamp(min=0, max=self.clip_max)\n                    y_b = batch_b[\"target\"].clamp(min=0, max=self.clip_max)\n\n                    loss_a = WeightedSmoothL1Loss(beta=self.beta, weights=[1.0])(\n                        pred, y_a\n                    )\n                    loss_b = WeightedSmoothL1Loss(beta=self.beta, weights=[1.0])(\n                        pred, y_b\n                    )\n                    loss += 0.02 * (lam * loss_a + (1 - lam) * loss_b)\n\n            # # Predict ratio\n            # total, logit = self.forward(mixed_x)\n            # ratio = logit.softmax(dim=1)\n            # grams = ratio * total\n            # gdm = grams[:, 0:1] + grams[:, 2:3]\n            # pred = torch.hstack([grams, gdm, total])\n\n            # # Ratio loss\n            # ratio_a = batch_a[\"target\"][:, :3] / batch_a[\"target\"][:, 4:5]\n            # ratio_b = batch_b[\"target\"][:, :3] / batch_b[\"target\"][:, 4:5]\n            # log_soft = F.log_softmax(logit, dim=1)\n            # loss_a = 1.0 * F.kl_div(log_soft, ratio_a, reduction=\"batchmean\")\n            # loss_b = 1.0 * F.kl_div(log_soft, ratio_b, reduction=\"batchmean\")\n            # loss = lam * loss_a + (1 - lam) * loss_b\n            # # print(\"ratio\", loss_a, loss_b)\n\n            # # Total loss\n            # y_a = batch_a[\"target\"].clamp(min=0, max=self.clip_max)\n            # y_b = batch_b[\"target\"].clamp(min=0, max=self.clip_max)\n            # loss_a = nn.SmoothL1Loss(beta=self.beta)(total, y_a[:, 3:4])\n            # loss_b = nn.SmoothL1Loss(beta=self.beta)(total, y_b[:, 3:4])\n            # loss += lam * loss_a + (1 - lam) * loss_b\n            # # print(\"total\", loss_a, loss_b)\n\n            # # Other loss\n            # loss_a = 0.2 * nn.SmoothL1Loss(beta=self.beta)(grams, y_a[:, :3])\n            # loss_b = 0.2 * nn.SmoothL1Loss(beta=self.beta)(grams, y_b[:, :3])\n            # loss += lam * loss_a + (1 - lam) * loss_b\n            # loss_a = 0.4 * nn.SmoothL1Loss(beta=self.beta)(gdm, y_a[:, 3:4])\n            # loss_b = 0.4 * nn.SmoothL1Loss(beta=self.beta)(gdm, y_b[:, 3:4])\n            # loss += lam * loss_a + (1 - lam) * loss_b\n            # # print(\"other\", loss_a, loss_b)\n        else:\n            # set_model_domain(self, \"csiro\")\n            pred = self.forward(input)[:, : len(LABELS)]\n            # loss = nn.MSELoss()(\n            # loss = nn.SmoothL1Loss(beta=self.beta)(\n            # loss = WeightedMSELoss()(\n            loss = WeightedSmoothL1Loss(beta=self.beta)(\n                pred,\n                target.clamp(min=0, max=self.clip_max),\n                # torch.sqrt(target.clamp(min=0, max=self.clip_max)),\n            )\n\n            self.train_r2.update(pred.clamp(min=0), target)\n            self.log(\n                \"train_r2\", self.train_r2, on_step=False, on_epoch=True, prog_bar=True\n            )\n\n        # if \"grass-clover\" in batch:\n        #     input, target = batch[\"grass-clover\"]\n\n        #     with bn_eval_only(self):\n        #         # set_model_domain(self, \"grass\")\n        #         pred = self.forward(input)[:, len(LABELS) :]\n        #         loss += 0.04 * WeightedSmoothL1Loss(beta=self.beta, weights=[1.0])(\n        #             pred, target.clamp(min=0, max=self.clip_max)\n        #         )\n\n        self.log(\"train_loss\", loss, on_step=False, on_epoch=True, prog_bar=True)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        input, target = batch[\"input\"], batch[\"target\"]\n        # set_model_domain(self, \"csiro\")\n        pred = self.forward(input)[:, : len(LABELS)]\n\n        # # Predict ratio\n        # total, logit = self.forward(input)\n        # ratio = logit.softmax(dim=1)\n        # grams = ratio * total\n        # gdm = grams[:, 0:1] + grams[:, 2:3]\n        # pred = torch.hstack([grams, gdm, total])\n\n        # loss = WeightedSmoothL1Loss(beta=self.beta)(\n        loss = WeightedSmoothL1Loss()(\n            pred,\n            target.clamp(min=0, max=self.clip_max),\n            # torch.sqrt(target.clamp(min=0, max=self.clip_max)),\n        )\n        self.log(\"val_loss\", loss, on_step=False, on_epoch=True, prog_bar=True)\n\n        self.validation_step_outputs.append(\n            {\n                \"pred\": pred,\n                \"label\": target,\n            }\n        )\n\n    def on_validation_epoch_end(self):\n        preds = torch.cat([x[\"pred\"] for x in self.validation_step_outputs], dim=0)\n        labels = torch.cat([x[\"label\"] for x in self.validation_step_outputs], dim=0)\n\n        # Log\n        r2_score = self.val_r2(preds.clamp(min=0), labels)\n        self.val_r2.reset()\n        self.log(\"val_r2\", r2_score, on_epoch=True, prog_bar=True)\n\n        is_best = r2_score > self.best_score\n        if is_best:\n            self.best_score = r2_score\n\n        # Save results\n        if is_best and self.logger and hasattr(self.logger, \"log_dir\"):\n            results = {\n                \"preds\": preds.cpu().numpy().tolist(),\n                \"labels\": labels.cpu().numpy().tolist(),\n            }\n            with open(self.logger.log_dir + \"/results.json\", \"w\") as f:\n                json.dump(results, f, indent=4)\n\n        # Clear outputs\n        self.validation_step_outputs.clear()\n\n    def predict_step(self, batch, batch_idx: int = 0):\n        # set_model_domain(self, \"csiro\")\n        if isinstance(batch, dict):\n            input = batch[\"input\"]\n        else:\n            input = batch\n        return self.forward(input)[:, : len(LABELS)].clamp(min=0)\n\n        # # Predict ratio\n        # total, logit = self.forward(batch)\n        # ratio = logit.softmax(dim=1)\n        # grams = ratio * total\n        # gdm = grams[:, 0:1] + grams[:, 2:3]\n        # return torch.hstack([grams, gdm, total]).clamp(min=0)\n\n    def configure_optimizers(self):\n        params = [p for p in self.parameters() if p.requires_grad]\n\n        if self.disasble_bias_wd:\n            # Disable weight_decay for bias and BatchNorm\n            decay_params = []\n            no_decay_params = []\n            for name, param in self.named_parameters():\n                if not param.requires_grad:\n                    continue\n                if param.dim() == 1 or name.endswith(\"bias\"):\n                    # BatchNorm の weight/bias (1D) や Linear の bias\n                    no_decay_params.append(param)\n                else:\n                    decay_params.append(param)\n            params = [\n                {\"params\": decay_params, \"weight_decay\": self.weight_decay},\n                {\"params\": no_decay_params, \"weight_decay\": 0.0},\n            ]\n\n        optimizer = hydra.utils.instantiate(\n            self.optim_cfg, params=params, _convert_=\"all\"\n        )\n\n        scheduler = torch.optim.lr_scheduler.OneCycleLR(\n            optimizer,\n            max_lr=self.base_lr,\n            total_steps=self.trainer.estimated_stepping_batches,\n            **self.scheduler_params,\n        )\n        # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        #     optimizer,\n        #     T_max=self.trainer.estimated_stepping_batches,\n        #     eta_min=self.base_lr * 0.1,\n        # )\n        # scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n        #     optimizer,\n        #     T_0=self.trainer.estimated_stepping_batches // 10,\n        #     T_mult=1,\n        #     eta_min=self.base_lr * 0.01,\n        # )\n        # scheduler = transformers.get_cosine_schedule_with_warmup(\n        #     optimizer,\n        #     num_warmup_steps=int(self.trainer.estimated_stepping_batches * 0.05),\n        #     num_training_steps=self.trainer.estimated_stepping_batches,\n        # )\n\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": scheduler,\n                \"interval\": \"step\",\n                \"frequency\": 1,\n            },\n        }\n\n\nclass AugDataset(Dataset):\n    \"\"\"Dataset that handles dictionary of augmented features.\"\"\"\n\n    def __init__(\n        self,\n        aug_features: dict,\n        targets: np.ndarray | None = None,\n        aug_keys: list | None = None,\n        mixup: bool = False,\n        image_paths: np.ndarray | None = None,\n        label_replacements: dict | None = None,\n        spatial: bool = False,\n    ):\n        self.aug_features = aug_features\n        self.targets = targets\n        self.mixup = mixup\n        self.image_paths = image_paths\n        self.label_replacements = label_replacements if label_replacements else {}\n        self.spatial = spatial\n        self.aug_keys = aug_keys if aug_keys else [\"identity\"]\n\n    def __len__(self):\n        if self.targets is not None:\n            return len(self.targets)\n        return len(self.image_paths)\n\n    def __getitem__(self, idx):\n        import random\n\n        if self.mixup:\n            # Randomly select two augmentations for mixup\n            # aug_key1, aug_key2 = random.sample(self.aug_keys, 2)\n            aug_key1 = \"identity\"\n            aug_key2 = random.choice(self.aug_keys)\n\n            # Mixup parameter\n            alpha = 0.05\n            lam = np.random.beta(alpha, alpha)\n            # lam = np.random.beta(1e6, 1e-6)\n            # lam = np.random.uniform(0, 1)\n\n            # Mix features\n            feature1 = self.aug_features[aug_key1][idx]\n            feature2 = self.aug_features[aug_key2][idx]\n            mixed_feature = lam * feature1 + (1 - lam) * feature2\n\n            # Get target, applying label replacement if specified\n            target = self.targets[idx].copy()\n            if self.image_paths is not None and len(self.label_replacements) > 0:\n                image_path = self.image_paths[idx]\n                if image_path in self.label_replacements:\n                    target = self.label_replacements[image_path]\n\n            return {\"input\": mixed_feature, \"target\": target}\n        else:\n            # Randomly select one augmentation\n            aug_key = random.choice(self.aug_keys)\n\n            # Get target, applying label replacement if specified\n            if self.targets is None:\n                target = None\n            else:\n                target = self.targets[idx].copy()\n                if self.image_paths is not None and len(self.label_replacements) > 0:\n                    image_path = self.image_paths[idx]\n                    if image_path in self.label_replacements:\n                        target = self.label_replacements[image_path]\n\n            if self.spatial:\n                image_path = self.image_paths[idx]\n                if \"test\" in image_path:\n                    feature_maps = np.load(\n                        f\"data/features/feature_maps_test_pe_spatial_gigantic-250_{Path(image_path).stem}.npy\"\n                    )\n                    # feature_maps = np.load(\n                    #     f\"data/features/feature_maps_test_dinov3-vit7b16-200_{Path(image_path).stem}.npy\"\n                    # )\n                else:\n                    feature_maps = np.load(\n                        f\"data/features/feature_maps_pe_spatial_gigantic-250_{Path(image_path).stem}.npy\"\n                    )\n                    # feature_maps = np.load(\n                    #     f\"data/features/feature_maps_dinov3-vit7b16-200_{Path(image_path).stem}.npy\"\n                    # )\n                feats = self.aug_features[aug_key][idx][:4096]\n                feats = np.repeat(\n                    feats[:, np.newaxis, np.newaxis], 32 * 32, axis=1\n                ).reshape(4096, 32, 32)\n                feature_maps = np.concatenate([feats, feature_maps])\n\n                item = {\"input\": feature_maps}\n            else:\n                item = {\"input\": self.aug_features[aug_key][idx]}\n            if target is not None:\n                item[\"target\"] = target\n            return item\n\n\nclass BiomassDataModule(LightningDataModule):\n    def __init__(\n        self,\n        X_train: np.ndarray,\n        y_train: np.ndarray,\n        X_val: np.ndarray,\n        y_val: np.ndarray,\n        X_test: np.ndarray,\n        features_grass_clover: np.ndarray | None = None,\n        y_train_grass_clover: np.ndarray | None = None,\n        features_irish: np.ndarray | None = None,\n        y_train_irish: np.ndarray | None = None,\n        batch_size: int = 16,\n        num_workers: int = 4,\n        image_paths_train: np.ndarray | None = None,\n        image_paths_val: np.ndarray | None = None,\n        image_paths_test: np.ndarray | None = None,\n        label_replacements: dict | None = None,\n        spatial: bool = False,\n    ):\n        super().__init__()\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n\n        label_replacements = {\n            \"train/ID230058600.jpg\": np.array(\n                [19.57037, 6.107088, 1.535656, 21.10603, 27.21311]\n            ),\n            \"train/ID1403107574.jpg\": np.array(\n                [10.86399737, 0.86185927, 0.92111741, 11.78511478, 12.64697405]\n            ),\n            \"train/ID1337107565.jpg\": np.array(\n                [1.21730281, 0.06680434, 0.87944882, 2.09675163, 2.16355597]\n            ),\n            \"train/ID1761544403.jpg\": np.array(\n                [4.2394537, 0.15564277, 5.72451089, 9.96396459, 10.11960736]\n            ),\n            \"train/ID681680726.jpg\": np.array(\n                [2.74912401, 2.38283418, 16.97968112, 19.72880514, 22.11163931]\n            ),\n            # \"train/ID1357758282.jpg\": np.array(\n            #     [42.53141942, 1.40089515, 0.0, 42.53141942, 43.93231457]\n            # ),\n        }\n\n        self.train_dataset = AugDataset(\n            X_train,\n            y_train,\n            image_paths=image_paths_train,\n            label_replacements=label_replacements,\n            aug_keys=[\"identity\"]\n            if spatial\n            else [\n                # \"identity\",\n                \"hflip\",\n                \"vflip\",\n                \"hvflip\",\n                # \"rot90\",\n            ],\n            mixup=False if spatial else True,\n            spatial=spatial,\n        )\n        self.val_dataset = AugDataset(\n            X_val,\n            y_val,\n            image_paths=image_paths_val,\n            spatial=spatial,\n        )\n        self.test_dataset = AugDataset(\n            X_test,\n            image_paths=image_paths_test,\n            spatial=spatial,\n        )\n\n        self.grass_clover_dataset = None\n        if features_grass_clover is not None:\n            self.grass_clover_dataset = AugDataset(\n                features_grass_clover,\n                y_train_grass_clover,\n                # aug_keys=features_grass_clover.keys(),\n                aug_keys=[\n                    \"identity\",\n                    # \"hflip\",\n                    # \"vflip\",\n                    # \"hvflip\",\n                    # \"rot90\",\n                ],\n                mixup=False,\n                spatial=spatial,\n            )\n\n        self.irish_dataset = None\n        if features_irish is not None:\n            self.irish_dataset = AugDataset(\n                features_irish,\n                y_train_irish,\n                # aug_keys=features_irish.keys(),\n                aug_keys=[\n                    \"identity\",\n                    \"hflip\",\n                    # \"vflip\",\n                    # \"hvflip\",\n                    # \"rot90\",\n                ],\n                mixup=False,\n                spatial=spatial,\n            )\n\n    def train_dataloader(self):\n        iterables = {\n            \"csiro-biomass\": DataLoader(\n                self.train_dataset,\n                batch_size=self.batch_size,\n                shuffle=True,\n                drop_last=True,\n                num_workers=self.num_workers,\n                persistent_workers=self.num_workers > 0,\n                pin_memory=False,\n                # prefetch_factor=2,\n            )\n        }\n        if self.grass_clover_dataset is not None:\n            iterables[\"grass-clover\"] = DataLoader(\n                self.grass_clover_dataset,\n                batch_size=16,\n                shuffle=True,\n                drop_last=True,\n                num_workers=self.num_workers,\n                persistent_workers=self.num_workers > 0,\n                pin_memory=False,\n                # prefetch_factor=2,\n            )\n        if self.irish_dataset is not None:\n            iterables[\"irish\"] = DataLoader(\n                self.irish_dataset,\n                batch_size=16,\n                shuffle=True,\n                drop_last=True,\n                num_workers=self.num_workers,\n                persistent_workers=self.num_workers > 0,\n                pin_memory=False,\n                # prefetch_factor=2,\n            )\n        return CombinedLoader(\n            iterables,\n            mode=\"max_size_cycle\",  # \"max_size_cycle\", \"min_size\"\n        )\n\n    def val_dataloader(self):\n        return DataLoader(\n            self.val_dataset,\n            batch_size=self.batch_size,\n            shuffle=False,\n            drop_last=False,\n            num_workers=self.num_workers,\n            persistent_workers=self.num_workers > 0,\n            pin_memory=False,\n            # prefetch_factor=2,\n        )\n\n    def predict_dataloader(self):\n        return DataLoader(\n            self.test_dataset,\n            batch_size=self.batch_size,\n            shuffle=False,\n            drop_last=False,\n            num_workers=self.num_workers,\n            persistent_workers=self.num_workers > 0,\n            pin_memory=False,\n            # prefetch_factor=2,\n        )\n\n\ndef cross_validate_mlp(\n    df: pl.DataFrame,\n    df_test: pl.DataFrame,\n    X_train_full: np.ndarray,\n    X_test: np.ndarray,\n    df_grass_clover: pl.DataFrame | None = None,\n    features_grass_clover: np.ndarray | None = None,\n    df_irish: pl.DataFrame | None = None,\n    features_irish: np.ndarray | None = None,\n    l2_norm: Literal[\"pre\", \"post\", \"none\"] | None = \"pre\",\n    hidden_dim: int = cfg.mlp.hidden_dim,\n    dropout1: float = cfg.mlp.dropout1,\n    dropout2: float = cfg.mlp.dropout2,\n    act_layer: DictConfig = cfg.mlp.act_layer,\n    norm_layer: DictConfig = cfg.mlp.norm_layer,\n    num_layers: int = cfg.mlp.num_layers,\n    drop_path_rate: float | None = cfg.mlp.drop_path_rate,\n    out_act: DictConfig = cfg.mlp.out_act,\n    mixup_alpha: float = cfg.mlp.mixup_alpha,\n    clip_max: float = cfg.mlp.clip_max,\n    beta: float = cfg.mlp.beta,\n    base_lr: float = cfg.mlp.base_lr,\n    weight_decay: float = cfg.mlp.weight_decay,\n    disasble_bias_wd: bool = cfg.mlp.disasble_bias_wd,\n    optim_cfg: DictConfig = cfg.mlp.optimizer,\n    scheduler_params: DictConfig = cfg.mlp.scheduler_params,\n    batch_size: int = cfg.mlp.batch_size,\n    max_epochs: int = cfg.mlp.max_epochs,\n    ema_decay: float | None = None,\n    patience: int = cfg.mlp.patience,\n    random_state: int | None = None,\n    enable_progress_bar: bool = True,\n    log_dir: str | None = None,\n    val_fold: int | None = None,\n    infer: bool = False,\n    ckpt_dir: str | None = None,\n    label_replacements: dict | None = None,\n    spatial: bool = False,\n) -> tuple:\n    \"\"\"Function to perform cross-validation with PyTorch Lightning MLP.\"\"\"\n    # 入力データを破壊しないようにコピーを作る\n    X_train_full = deepcopy(X_train_full)\n    X_test = deepcopy(X_test)\n    if features_grass_clover is not None:\n        features_grass_clover = deepcopy(features_grass_clover)\n    if features_irish is not None:\n        features_irish = deepcopy(features_irish)\n\n    # Set random seed\n    if random_state is not None:\n        seed_everything(random_state)\n\n    n_folds = df[\"fold\"].max() + 1\n\n    # Prepare data\n    # spatial = False\n    # if isinstance(X_train_full, defaultdict) and X_train_full[\"identity\"].ndim == 4:\n    #     spatial = True\n    print(f\"{spatial=}\")\n\n    y_train_full = df.select(LABELS).to_numpy()\n    image_paths_full = df.select(\"image_path\").to_numpy().flatten()\n    image_paths_test = df_test.select(\"image_path\").to_numpy().flatten()\n    pred_oof = np.zeros_like(y_train_full)\n    pred_test = np.zeros([len(X_test[\"identity\"]), len(LABELS)])\n\n    assert l2_norm is None or l2_norm in [\"pre\", \"post\", \"none\"]\n    if l2_norm == \"pre\":\n        # for aug_key in [\"identity\"]:\n        for aug_key in X_train_full.keys():\n            X_train_full[aug_key] = F.normalize(\n                torch.tensor(X_train_full[aug_key]), dim=1\n            ).numpy()\n            if aug_key in X_test:\n                X_test[aug_key] = F.normalize(\n                    torch.tensor(X_test[aug_key]), dim=1\n                ).numpy()\n            if features_grass_clover is not None and aug_key in features_grass_clover:\n                features_grass_clover[aug_key] = F.normalize(\n                    torch.tensor(features_grass_clover[aug_key]), dim=1\n                ).numpy()\n            if features_irish is not None and aug_key in features_irish:\n                features_irish[aug_key] = F.normalize(\n                    torch.tensor(features_irish[aug_key]), dim=1\n                ).numpy()\n\n    # Apply StandardScaler\n    if not spatial:\n        # scaler = StandardScaler()\n        # scaler.fit(X_train_full[\"identity\"])\n        # for aug_key in [\"identity\"]:\n        for aug_key in X_train_full.keys():\n            scaler = StandardScaler()\n            scaler.fit(X_train_full[aug_key])\n            X_train_full[aug_key] = scaler.transform(X_train_full[aug_key])\n            if aug_key in X_test:\n                X_test[aug_key] = scaler.transform(X_test[aug_key])\n\n            # X_train_full[aug_key][:, :4916] = scaler.fit_transform(\n            #     X_train_full[aug_key][:, :4916]\n            # )\n            # X_train_full[aug_key][:, 4916:] = scaler.fit_transform(\n            #     X_train_full[aug_key][:, 4916:]\n            # )\n            # scaler = StandardScaler()\n            # scaler.fit(X_train_full[\"identity\"])\n            if features_grass_clover is not None and aug_key in features_grass_clover:\n                features_grass_clover[aug_key] = scaler.transform(\n                    features_grass_clover[aug_key]\n                )\n                # features_grass_clover[aug_key] = scaler.fit_transform(features_grass_clover[aug_key])\n            if features_irish is not None and aug_key in features_irish:\n                features_irish[aug_key] = scaler.transform(features_irish[aug_key])\n                # features_irish[aug_key] = scaler.fit_transform(features_irish[aug_key])\n    # else:\n    #     mean = X_train_full.mean(axis=(0, 2, 3), keepdims=True)\n    #     std = X_train_full.std(axis=(0, 2, 3), keepdims=True)\n    #     X_train_full = (X_train_full - mean) / std\n    #     X_test = (X_test - mean) / std\n\n    # grass-clover\n    y_train_grass_clover = None\n    if features_grass_clover is not None:\n        y_train_grass_clover = df_grass_clover.select(\n            # \"fresh_grass\",\n            \"dry_grass\",\n            # \"fresh_white_clover\",\n            # \"dry_white_clover\",\n            # \"fresh_red_clover\",\n            # \"dry_red_clover\",\n            # \"fresh_clover\",\n            \"dry_clover\",\n            # \"fresh_weeds\",\n            # \"dry_weeds\",\n            \"dry_total\",\n            # \"dry_clover_fraction\",\n            # \"dry_red_clover_fraction\",\n            # \"dry_white_clover_fraction\",\n            # \"dry_grass_fraction\",\n            # \"dry_weeds_fraction\",\n        ).to_numpy()\n\n    # Irish Grass Clover\n    y_train_irish = None\n    if features_irish is not None:\n        y_train_irish = df_irish.select(\n            \"Grass Dried g\",\n            \"Clover Dried g\",\n            # \"Weeds Dried g\",\n        ).to_numpy()\n\n    if l2_norm == \"post\":\n        for aug_key in X_train_full.keys():\n            X_train_full[aug_key] = F.normalize(\n                torch.tensor(X_train_full[aug_key]), dim=1\n            ).numpy()\n            if aug_key in X_test:\n                X_test[aug_key] = F.normalize(\n                    torch.tensor(X_test[aug_key]), dim=1\n                ).numpy()\n            if features_grass_clover is not None and aug_key in features_grass_clover:\n                features_grass_clover[aug_key] = F.normalize(\n                    torch.tensor(features_grass_clover[aug_key]), dim=1\n                ).numpy()\n            if features_irish is not None and aug_key in features_irish:\n                features_irish[aug_key] = F.normalize(\n                    torch.tensor(features_irish[aug_key]), dim=1\n                ).numpy()\n\n    if infer:\n        # Perform inference with pre-trained models\n        ckpt_paths = list(Path(ckpt_dir).rglob(\"*.ckpt\"))\n        # ckpt_paths = [\n        #     x for x in ckpt_paths if int(x.parent.parent.stem.split(\"version_\")[1]) < 19\n        # ]  # test\n        print(f\"\\n[Inference mode with {len(ckpt_paths)} pre-trained checkpoints]\")\n\n        folds = [int(x.stem.split(\"fold\")[1].split(\"-epoch\")[0]) for x in ckpt_paths]\n        assert set(folds) == set(range(n_folds))\n        assert len(ckpt_paths) % n_folds == 0\n        for fold in range(n_folds):\n            # Split data\n            val_idxs = (df[\"fold\"] == fold).to_numpy()\n            # X_val = X_train_full[val_idxs]\n            X_val = {k: v[val_idxs] for k, v in X_train_full.items()}\n            y_val = y_train_full[val_idxs]\n\n            datamodule = BiomassDataModule(\n                X_train=X_train_full,\n                y_train=y_train_full,\n                X_val=X_val,\n                y_val=y_val,\n                X_test=X_test,\n                # features_grass_clover=features_grass_clover,\n                # y_train_grass_clover=y_train_grass_clover,\n                # features_irish=features_irish,\n                # y_train_irish=y_train_irish,\n                batch_size=batch_size,\n                image_paths_train=image_paths_full,\n                image_paths_val=image_paths_full[val_idxs],\n                image_paths_test=image_paths_test,\n                label_replacements=label_replacements,\n                spatial=spatial,\n            )\n            trainer = Trainer(\n                accelerator=\"auto\",\n                devices=1,\n                precision=32,\n                enable_progress_bar=False,\n                enable_model_summary=False,\n            )\n\n            ckpt_idxs = [i for i, x in enumerate(folds) if x == fold]\n\n            val_preds = []\n            test_preds = []\n            for idx in ckpt_idxs:\n                # Load model\n                print(ckpt_paths[idx])\n                model = MLPRegressor.load_from_checkpoint(ckpt_paths[idx])\n                model.eval()\n                model.to(DEVICE)\n\n                # Predict on validation set\n                with torch.no_grad():\n                    # X_val_tensor = torch.FloatTensor(X_val[\"identity\"]).to(DEVICE)\n                    # preds_ = model.predict_step(X_val_tensor).cpu().numpy()\n\n                    preds = trainer.predict(\n                        model, dataloaders=datamodule.val_dataloader()\n                    )\n                    preds = torch.cat(preds).cpu().numpy()\n\n                    # assert np.allclose(preds, preds_)\n\n                    val_preds.append(preds)\n\n                # Predict on test set\n                with torch.no_grad():\n                    # model.to(DEVICE)\n                    # X_test_tensor = torch.FloatTensor(X_test[\"identity\"]).to(DEVICE)\n                    # preds_ = model.predict_step(X_test_tensor).cpu().numpy()\n\n                    preds = trainer.predict(model, datamodule)\n                    preds = torch.cat(preds).cpu().numpy()\n\n                    # assert np.allclose(preds, preds_)\n\n                    test_preds.append(preds)\n\n            # Simple Average\n            pred_ensemble = np.mean(val_preds, axis=0)\n            score = competition_metric(y_val, pred_ensemble)\n            print(f\"Fold {fold} (avg): {score:.4f}\")\n\n            # # NNLS\n            # W, _ = nnls(np.stack([x.ravel() for x in val_preds]).T, y_val.ravel())\n            # W /= W.sum()  # Constrain weights to sum to 1\n            # # W = [1.0, 0.0, 0.0, 0.0]\n            # pred_ensemble = sum(W[i] * val_preds[i] for i in range(len(val_preds)))\n            # score = competition_metric(y_val, pred_ensemble)\n            # print(f\"Fold {fold} (nnls): {score:.4f}, weights: {W}\")\n\n            # Assign ensemble results\n            pred_oof[val_idxs] = pred_ensemble\n            pred_test += np.mean(test_preds, axis=0) / n_folds\n            # pred_test += (\n            #     sum(W[i] * test_preds[i] for i in range(len(test_preds))) / n_folds\n            # )\n    else:\n        if val_fold is not None:\n            # Do hold-out validation\n            assert 0 <= val_fold < n_folds\n            folds = [val_fold]\n        else:\n            # Do cross-validation\n            folds = list(range(n_folds))\n\n        for fold in folds:\n            print(f\"\\n[Training Fold {fold}]\")\n\n            # Split data\n            train_idxs = (df[\"fold\"] != fold).to_numpy()\n            val_idxs = (df[\"fold\"] == fold).to_numpy()\n\n            # X_train = X_train_full[train_idxs]\n            X_train = {k: v[train_idxs] for k, v in X_train_full.items()}\n            y_train = y_train_full[train_idxs]\n            # X_val = X_train_full[val_idxs]\n            X_val = {k: v[val_idxs] for k, v in X_train_full.items()}\n            y_val = y_train_full[val_idxs]\n\n            image_paths_train = image_paths_full[train_idxs]\n            image_paths_val = image_paths_full[val_idxs]\n\n            datamodule = BiomassDataModule(\n                X_train=X_train,\n                y_train=y_train,\n                X_val=X_val,\n                y_val=y_val,\n                X_test=X_test,\n                features_grass_clover=features_grass_clover,\n                y_train_grass_clover=y_train_grass_clover,\n                features_irish=features_irish,\n                y_train_irish=y_train_irish,\n                batch_size=batch_size,\n                image_paths_train=image_paths_train,\n                image_paths_val=image_paths_val,\n                image_paths_test=image_paths_test,\n                label_replacements=label_replacements,\n                spatial=spatial,\n            )\n\n            output_dim = len(LABELS)\n            if df_grass_clover is not None:\n                output_dim += y_train_grass_clover.shape[1]\n            if df_irish is not None:\n                output_dim += y_train_irish.shape[1]\n\n            model = MLPRegressor(\n                input_dim=X_train[\"identity\"].shape[1],\n                output_dim=output_dim,\n                # output_dim=1,\n                act_layer=act_layer,\n                norm_layer=norm_layer,\n                num_layers=num_layers,\n                drop_path_rate=drop_path_rate,\n                out_act=out_act,\n                hidden_dim=hidden_dim,\n                dropout1=dropout1,\n                dropout2=dropout2,\n                base_lr=base_lr,\n                weight_decay=weight_decay,\n                disasble_bias_wd=disasble_bias_wd,\n                optim_cfg=optim_cfg,\n                scheduler_params=scheduler_params,\n                mixup_alpha=mixup_alpha,\n                spatial=spatial,\n                clip_max=clip_max,\n                beta=beta,\n            )\n\n            # print(\"Load pre-trained checkpoint\")\n            # model = MLPRegressor.load_from_checkpoint(\n            #     \"logs/mlp/multi-model/lightning_logs/version_2/checkpoints/mlp-fold1-epoch=00-val_r2=-1.1557.ckpt\",\n            #     strict=True,\n            # )\n\n            # model = torch.compile(model)\n\n            # Trainer\n            early_stop_callback = EarlyStopping(\n                monitor=\"val_r2\", patience=patience, mode=\"max\", verbose=False\n            )\n            ckpt_callback = ModelCheckpoint(\n                filename=f\"mlp-fold{fold}-{{epoch:02d}}-{{val_r2:.4f}}\",\n                monitor=\"val_r2\",\n                mode=\"max\",\n                save_top_k=1,\n                save_weights_only=True,\n            )\n            callbacks = [early_stop_callback, ckpt_callback]\n            # if ema_decay is not None:\n            #     ema_callback = EMACallback(decay=ema_decay)\n            trainer = Trainer(\n                accelerator=\"auto\",\n                devices=1,\n                precision=16,\n                max_epochs=max_epochs,\n                enable_progress_bar=enable_progress_bar,\n                enable_model_summary=False,\n                # num_sanity_val_steps=0,\n                # gradient_clip_val=1.0,\n                callbacks=callbacks,\n                default_root_dir=log_dir,\n            )\n\n            # Train\n            trainer.fit(model, datamodule)\n\n            # Load best model\n            model = MLPRegressor.load_from_checkpoint(ckpt_callback.best_model_path)\n            model.eval()\n            model.to(DEVICE)\n\n            # Predict on validation set\n            with torch.no_grad():\n                preds = trainer.predict(model, dataloaders=datamodule.val_dataloader())\n                preds = torch.cat(preds).cpu().numpy()\n                pred_oof[val_idxs] = preds\n\n            # Predict on test set\n            with torch.no_grad():\n                preds = trainer.predict(model, datamodule)\n                preds = torch.cat(preds).cpu().numpy()\n                pred_test += preds / n_folds\n\n            # Calculate fold score\n            score = competition_metric(y_val, pred_oof[val_idxs])\n            print(f\"Fold {fold}: {score:.4f}\")\n\n            del datamodule, model, trainer\n            gc.collect()\n            torch.cuda.empty_cache()\n\n            if val_fold is not None:\n                return pred_oof[val_idxs], pred_test, score\n\n    # Calculate overall CV score\n    score = competition_metric(y_train_full, pred_oof)\n    print(f\"\\nCV Score: {score:.4f}\")\n\n    # Apply projection to enforce constraints\n    pred_oof = project_nonneg_equalities(pred_oof)\n    pred_test = project_nonneg_equalities(pred_test)\n    score = competition_metric(y_train_full, pred_oof)\n    print(f\"PP Score: {score:.4f}\")\n\n    return pred_oof, pred_test, score\n\n\ndef cross_validate(\n    pipe: Pipeline,\n    df: pl.DataFrame,\n    df_test: pl.DataFrame,\n    features: np.ndarray,\n    features_test: np.ndarray,\n) -> tuple:\n    \"\"\"Function to perform GroupKFold cross-validation with a given model.\"\"\"\n    # Prepare data\n    X_train = features\n    y_train = df.select(LABELS).to_numpy()\n    pred_oof = np.zeros_like(y_train)\n\n    X_test = features_test\n    pred_test = np.zeros([len(X_test), len(LABELS)])\n\n    model = pipe.named_steps[\"model\"]\n\n    # CV\n    for fold in range(df[\"fold\"].max() + 1):\n        train_idxs = df[\"fold\"] != fold\n        val_idxs = df[\"fold\"] == fold\n        if type(model) is Ridge:\n            # weight = [weights[x] for x in LABELS]\n            # sample_weight = (y_train[train_idxs] * weight).sum(axis=1) / y_train[\n            #     train_idxs\n            # ].mean()\n            m = deepcopy(pipe)\n            m.fit(\n                X_train[train_idxs],\n                # np.sqrt(y_train[train_idxs].clip(0, 112.1)),\n                np.stack(\n                    [\n                        y_train[train_idxs, 0] / y_train[train_idxs, 4],\n                        y_train[train_idxs, 1] / y_train[train_idxs, 4],\n                        y_train[train_idxs, 2] / y_train[train_idxs, 4],\n                        np.sqrt(y_train[train_idxs, 3].clip(0, 112.1)),\n                        np.sqrt(y_train[train_idxs, 4].clip(0, 112.1)),\n                    ]\n                ).T,\n                # model__sample_weight=sample_weight,\n            )\n            # m.fit(X_train, np.sqrt(y_train.clip(0, 112.1)))  # overfitting\n            # pred_oof[val_idxs] = m.predict(X_train[val_idxs]).clip(0) ** 2\n            outs = m.predict(X_train[val_idxs])\n            total = outs[:, 4].clip(0) ** 2\n            sum_ = outs[:, 0] + outs[:, 1] + outs[:, 2]\n            pred_oof[val_idxs] = np.stack(\n                [\n                    (total * outs[:, 0]).clip(0) / sum_,\n                    (total * outs[:, 1]).clip(0) / sum_,\n                    (total * outs[:, 2]).clip(0) / sum_,\n                    # outs[:, 3].clip(0) ** 2,\n                    (total * outs[:, 0]).clip(0) / sum_\n                    + (total * outs[:, 2]).clip(0) / sum_,\n                    total,\n                ]\n            ).T\n            # pred_test[:, :] += m.predict(X_test).clip(0) ** 2 / n_folds\n            outs = m.predict(X_test)\n            total = outs[:, 4].clip(0) ** 2\n            sum_ = outs[:, 0] + outs[:, 1] + outs[:, 2]\n            pred_test[:, :] += (\n                np.vstack(\n                    [\n                        (total * outs[:, 0]).clip(0) / sum_,\n                        (total * outs[:, 1]).clip(0) / sum_,\n                        (total * outs[:, 2]).clip(0) / sum_,\n                        # outs[:, 3].clip(0) ** 2,\n                        (total * outs[:, 0]).clip(0) / sum_\n                        + (total * outs[:, 2]).clip(0) / sum_,\n                        total,\n                    ]\n                ).T\n                / n_folds\n            )\n        elif type(model) is SVR or type(model) is LGBMRegressor:\n            for i in range(len(LABELS)):\n                m = deepcopy(pipe)\n                m.fit(X_train[train_idxs], y_train[train_idxs, i].clip(0, 112.1))\n                # m.fit(X_train, y_train[:, i].clip(0, 112.1))  # overfitting\n                pred_oof[val_idxs, i] = m.predict(X_train[val_idxs]).clip(0)\n                pred_test[:, i] += m.predict(X_test).clip(0) / n_folds\n\n                # # Pseudo-labeling\n                # X_train2 = np.vstack([X_train[train_idxs], X_train[val_idxs]])\n                # y_train2 = np.hstack([y_train[train_idxs, i], pred_oof[val_idxs, i]])\n                # m = deepcopy(pipe)\n                # m.fit(X_train2, y_train2.clip(0, 112.1))\n                # pred_oof[val_idxs, i] = m.predict(X_train[val_idxs]).clip(0)\n        else:\n            m = deepcopy(pipe)\n            m.fit(X_train[train_idxs], y_train[train_idxs].clip(0, 112.1))\n            # m.fit(X_train, y_train.clip(0, 112.1))  # overfitting\n            pred_oof[val_idxs] = m.predict(X_train[val_idxs]).clip(0)\n            pred_test[:, :] += m.predict(X_test).clip(0) / n_folds\n\n        score = competition_metric(y_train[val_idxs], pred_oof[val_idxs])\n        print(f\"Fold {fold}: {score:.4f}\")\n\n    score = competition_metric(y_train, pred_oof)\n    print(f\"CV Score: {score:.4f}\")\n\n    # # --- Calculate remaining biomass fractions ---\n    # # Dry_Dead_g = Dry_Total_g - GDM_g\n    # pred_oof[:, LABELS.index(\"Dry_Dead_g\")] = (\n    #     pred_oof[:, LABELS.index(\"Dry_Total_g\")]\n    #     - pred_oof[:, LABELS.index(\"GDM_g\")]\n    # )\n    # # Dry_Green_g = GDM_g - Dry_Clover_g\n    # pred_oof[:, LABELS.index(\"Dry_Green_g\")] = (\n    #     pred_oof[:, LABELS.index(\"GDM_g\")]\n    #     - pred_oof[:, LABELS.index(\"Dry_Clover_g\")]\n    # )\n    # # --- Apply projection to enforce constraints ---\n    # pred_oof = pred_oof @ make_projection().T\n    # --- Apply projection to enforce constraints ---\n    pred_oof = project_nonneg_equalities(pred_oof)\n    pred_test = project_nonneg_equalities(pred_test)\n    score = competition_metric(y_train, pred_oof)\n    print(f\"PP Score: {score:.4f}\")\n\n    return pred_oof, pred_test\n\n\n# ==============================================================================\n# Apply Regressors\n# ==============================================================================\n\n# Compare Performance of Multiple Models\noof_preds = []\ntest_preds = []\n\n# PyTorch Lightning MLP\nif cfg.opt:\n    # Hyperparameter optimization\n    def objective(trial):\n        # df_grass_clover_trial = features_grass_clover_trial = None\n        # use_grass_clover = trial.suggest_categorical(\"use_grass_clover\", [False, True])\n        # if use_grass_clover:\n        #     df_grass_clover_trial = df_grass_clover\n        #     features_grass_clover_trial = features_grass_clover\n\n        hidden_dim = trial.suggest_int(\"hidden_dim\", 512, 1536, step=128)\n        dropout1 = trial.suggest_float(\"dropout1\", 0.0, 0.35)\n        dropout2 = trial.suggest_float(\"dropout2\", 0.0, 0.35)\n        num_layers = trial.suggest_int(\"num_layers\", 1, 3)\n        drop_path_rate = trial.suggest_float(\"drop_path_rate\", 0.0, 0.5)\n        mixup_alpha = trial.suggest_float(\"mixup_alpha\", 0.0, 0.6)\n        clip_max = trial.suggest_float(\"clip_max\", 110.0, 180.0)\n        beta = trial.suggest_float(\"beta\", 0.0, 10.0)\n        base_lr = trial.suggest_float(\"base_lr\", 2e-3, 2e-2, log=True)\n        weight_decay = trial.suggest_float(\"weight_decay\", 1e-7, 1e-2, log=True)\n        # disasble_bias_wd = trial.suggest_categorical(\"disasble_bias_wd\", [False, True])\n        # batch_size = trial.suggest_int(\"batch_size\", 4, 16, step=4)\n        # max_epochs = trial.suggest_int(\"max_epochs\", 60, 100, step=10)\n        # # ema_decay = trial.suggest_float(\"ema_decay\", 0.01, 0.99999, log=True)\n\n        act_layer = OmegaConf.create(\n            trial.suggest_categorical(\n                \"act_layer\",\n                [\n                    {\"_target_\": \"torch.nn.ReLU\"},\n                    {\"_target_\": \"torch.nn.SiLU\"},\n                    # {\"_target_\": \"torch.nn.Mish\"},\n                    {\"_target_\": \"torch.nn.LeakyReLU\"},\n                    # {\"_target_\": \"torch.nn.GELU\"},\n                ],\n            )\n        )\n        # norm_layer = OmegaConf.create(\n        #     trial.suggest_categorical(\n        #         \"norm_layer\",\n        #         [\n        #             {\"_target_\": \"torch.nn.BatchNorm1d\", \"affine\": True},\n        #             # {\"_target_\": \"torch.nn.BatchNorm1d\", \"affine\": False},\n        #             # {\"_target_\": \"torch.nn.LayerNorm\"},\n        #             # {\"_target_\": \"torch.nn.Identity\"},\n        #         ],\n        #     )\n        # )\n        out_act = OmegaConf.create(\n            trial.suggest_categorical(\n                \"out_act\",\n                [\n                    {\"_target_\": \"torch.nn.Identity\"},\n                    {\"_target_\": \"torch.nn.Softplus\"},\n                ],\n            )\n        )\n        optim = trial.suggest_categorical(\n            \"optim\",\n            [\n                \"torch.optim.AdamW\",\n                \"timm.optim.lion.Lion\",\n                # \"torch.optim.SGD\",\n            ],\n        )\n        if optim == \"timm.optim.lion.Lion\":\n            optim_cfg = OmegaConf.create(\n                {\n                    \"_target_\": optim,\n                    \"lr\": base_lr,\n                    \"weight_decay\": weight_decay,\n                    \"betas\": (\n                        trial.suggest_float(\"lion_beta1\", 0.84, 0.92),\n                        trial.suggest_float(\"lion_beta2\", 0.9, 0.99, log=True),\n                    ),\n                }\n            )\n        elif optim == \"torch.optim.AdamW\":\n            optim_cfg = OmegaConf.create(\n                {\n                    \"_target_\": optim,\n                    \"lr\": base_lr,\n                    \"weight_decay\": weight_decay,\n                    \"betas\": (\n                        trial.suggest_float(\"adamw_beta1\", 0.85, 0.95),\n                        trial.suggest_float(\"adamw_beta2\", 0.996, 0.999, log=True),\n                    ),\n                    \"eps\": trial.suggest_float(\"eps\", 3e-7, 1e-5, log=True),\n                }\n            )\n        elif optim == \"torch.optim.SGD\":\n            optim_cfg = OmegaConf.create(\n                {\n                    \"_target_\": optim,\n                    \"lr\": base_lr,\n                    \"momentum\": trial.suggest_float(\"momentum\", 0.8, 0.95),\n                    \"weight_decay\": weight_decay,\n                }\n            )\n        else:\n            raise ValueError(f\"Unknown optimizer: {optim}\")\n        scheduler_params = OmegaConf.create(\n            {\n                \"pct_start\": trial.suggest_float(\"pct_start\", 0.0, 0.2),\n                \"div_factor\": 1e7,\n                \"final_div_factor\": trial.suggest_float(\"final_div_factor\", 1.0, 10.0)\n                / 1e7,\n            }\n        )\n\n        # Save trials log\n        trials_df = trial.study.trials_dataframe()\n        name = \"optuna_trials\"\n        if cfg.val_fold is not None:\n            name += f\"-fold={cfg.val_fold}\"\n        trials_df.to_csv(Path(cfg.log_dir, name + \".csv\"), index=None)\n\n        pred_oof, _, score = cross_validate_mlp(\n            df,\n            df_test,\n            features,\n            features_test,\n            df_grass_clover=df_grass_clover,\n            features_grass_clover=features_grass_clover,\n            # df_grass_clover=df_grass_clover_trial,\n            # features_grass_clover=features_grass_clover_trial,\n            df_irish=df_irish,\n            features_irish=features_irish,\n            hidden_dim=hidden_dim,\n            dropout1=dropout1,\n            dropout2=dropout2,\n            act_layer=act_layer,\n            # norm_layer=norm_layer,\n            num_layers=num_layers,\n            drop_path_rate=drop_path_rate,\n            out_act=out_act,\n            mixup_alpha=mixup_alpha,\n            clip_max=clip_max,\n            beta=beta,\n            base_lr=base_lr,\n            weight_decay=weight_decay,\n            # disasble_bias_wd=disasble_bias_wd,\n            optim_cfg=optim_cfg,\n            scheduler_params=scheduler_params,\n            # batch_size=batch_size,\n            # max_epochs=max_epochs,\n            # ema_decay=ema_decay,\n            patience=cfg.mlp.patience,\n            log_dir=cfg.log_dir,\n            val_fold=cfg.val_fold,\n            enable_progress_bar=False,\n            spatial=cfg.mlp.spatial,\n        )\n        return score\n\n    study = optuna.create_study(\n        storage=cfg.storage,\n        study_name=cfg.study_name,\n        sampler=hydra.utils.instantiate(cfg.sampler),\n        direction=\"maximize\",\n        load_if_exists=cfg.load_if_exists,\n    )\n    study.optimize(objective, n_trials=cfg.n_trials)\n\n    exit()\nelse:\n    # Normal training\n    if hasattr(cfg, \"mlp\") and cfg.mlp:\n        print(\"\\n--- PyTorch Lightning MLP ---\")\n        pred_oof, pred_test, _ = cross_validate_mlp(\n            df,\n            df_test,\n            features,\n            features_test,\n            df_grass_clover=df_grass_clover,\n            features_grass_clover=features_grass_clover,\n            df_irish=df_irish,\n            features_irish=features_irish,\n            hidden_dim=cfg.mlp.hidden_dim,\n            dropout1=cfg.mlp.dropout1,\n            dropout2=cfg.mlp.dropout2,\n            act_layer=cfg.mlp.act_layer,\n            norm_layer=cfg.mlp.norm_layer,\n            num_layers=cfg.mlp.num_layers,\n            drop_path_rate=cfg.mlp.drop_path_rate,\n            out_act=cfg.mlp.out_act,\n            mixup_alpha=cfg.mlp.mixup_alpha,\n            clip_max=cfg.mlp.clip_max,\n            beta=cfg.mlp.beta,\n            base_lr=cfg.mlp.base_lr,\n            weight_decay=cfg.mlp.weight_decay,\n            disasble_bias_wd=cfg.mlp.disasble_bias_wd,\n            optim_cfg=cfg.mlp.optimizer,\n            scheduler_params=cfg.mlp.scheduler_params,\n            batch_size=cfg.mlp.batch_size,\n            max_epochs=cfg.mlp.max_epochs,\n            ema_decay=cfg.mlp.ema_decay,\n            patience=cfg.mlp.patience,\n            log_dir=cfg.log_dir,\n            val_fold=cfg.val_fold,\n            enable_progress_bar=False if is_kaggle() else True,\n            infer=cfg.infer,\n            ckpt_dir=cfg.ckpt_dir_kaggle if is_kaggle() else cfg.ckpt_dir,\n            spatial=cfg.mlp.spatial,\n        )\n        oof_preds.append(pred_oof)\n        test_preds.append(pred_test)\n\n# exit()\n\n# features = np.load(f\"{FEATS_DIR}/feature_maps_{key}.npy\")\n# features = features.mean(axis=(2, 3))\n\n\n# Other regressors\nif hasattr(cfg, \"regressors\") and cfg.regressors:\n    pipelines = hydra.utils.instantiate(cfg.regressors, _convert_=\"all\")\n    for pipe in pipelines:\n        print(f\"\\n--- {pipe.named_steps['model']} ---\")\n        pred_oof, pred_test = cross_validate(\n            pipe, df, df_test, features[\"identity\"], features_test[\"identity\"]\n        )\n        oof_preds.append(pred_oof)\n        test_preds.append(pred_test)\n\n# ==============================================================================\n# Ensemble\n# ==============================================================================\n\n\ndef min_func(W):\n    ensembled = sum(W[i] * oof_preds[i] for i in range(len(oof_preds)))\n    return -competition_metric(y_train, ensembled)\n\n\nprint(\"\\nEnsembling and saving OOF predictions...\")\ny_train = df.select(LABELS).to_numpy()\n\n# Simple Average\n# oof_pred_ensemble = (oof_preds[0] + oof_preds[1] + oof_preds[2]) / 3\noof_pred_ensemble = np.mean(oof_preds, axis=0)\nscore = competition_metric(y_train, oof_pred_ensemble)\nprint(f\"Final CV Score (avg): {score:.4f}\")\n\n# minimize\nres = minimize(\n    min_func,\n    [1 / len(oof_preds)] * len(oof_preds),\n    method=\"Nelder-Mead\",\n    tol=1e-6,\n    bounds=[(0, None)] * len(oof_preds),\n)\nW = res.x\nprint(\"Ensemble weights (minimize):\", W)\noof_pred_ensemble_ = sum(W[i] * oof_preds[i] for i in range(len(oof_preds)))\nscore = competition_metric(y_train, oof_pred_ensemble_)\nprint(f\"Final CV Score (minimize): {score:.4f}\")\n\n# NNLS\nW, _ = nnls(np.stack([x.ravel() for x in oof_preds]).T, y_train.ravel())\nprint(\"Ensemble weights (nnls):\", W)\noof_pred_ensemble = sum(W[i] * oof_preds[i] for i in range(len(oof_preds)))\nscore = competition_metric(y_train, oof_pred_ensemble)\nprint(f\"Final CV Score (nnls): {score:.4f}\")\n\n\"\"\"OPT\nfor a in np.linspace(200, 2000, 10):\n    print(f\"\\n{a}\")\n    _ = cross_validate(\n        Pipeline(\n            steps=[\n                [\"scaler\", StandardScaler()],\n                [\"model\", Ridge(alpha=a)],\n            ]\n        ),\n        df, df_test, features, features_test,\n    )\n\"\"\"\n\n# exit()\n\n# Save OOF predictions\noof_df = df.select([\"image_path\"]).to_pandas()\noof_df[LABELS] = oof_pred_ensemble\noutpath = f\"oof_{cfg_name}.csv\"\noof_df.to_csv(outpath, index=False)\nprint(f\"\\nOOF predictions saved to {outpath}\")\n\nW = [1.0, 0.0, 0.0]\n\n# Ensemble Test Predictions\nprint(\"\\nEnsembling test predictions...\")\n# pred_test = (test_preds[0] + test_preds[1] + test_preds[2]) / 3\n# pred_test = np.mean(test_preds, axis=0)\npred_test = sum(W[i] * test_preds[i] for i in range(len(test_preds)))\n\n# ==============================================================================\n# Post-processing\n# ==============================================================================\n\nif True:\n    from sklearn.linear_model import LogisticRegression\n    from sklearn.model_selection import StratifiedKFold, cross_val_score\n    from sklearn.pipeline import make_pipeline\n    from sklearn.preprocessing import StandardScaler\n    from sklearn.svm import LinearSVC\n\n    states = [\"Tas\", \"Vic\", \"NSW\", \"WA\"]\n    y_train_ = np.array([states.index(x) for x in list(df[\"State\"])])[:, np.newaxis]\n    prob_oof = np.zeros([len(features[\"identity\"]), len(states)])\n    prob_test = np.zeros([len(features_test[\"identity\"]), len(states)])\n    clf = make_pipeline(\n        StandardScaler(),\n        LogisticRegression(\n            solver=\"saga\",\n            penalty=\"l2\",\n            C=1.0,\n            max_iter=1000,\n            n_jobs=-1,\n            class_weight=\"balanced\",\n        ),\n    )\n    # clf = make_pipeline(\n    #     StandardScaler(),  # 埋め込みがL2正規化済みなら不要な場合もあり\n    #     LinearSVC(\n    #         C=1.0,\n    #         class_weight=\"balanced\",  # 不均衡なら\n    #         max_iter=20000,\n    #         random_state=42,\n    #     ),\n    # )\n    n_splits_ = 5\n    cv = StratifiedKFold(n_splits=n_splits_, shuffle=True, random_state=42)\n    # scores = cross_val_score(clf, features['identity'], y_train_, cv=cv, scoring=\"accuracy\")\n    # print(scores)\n\n    for train_idxs, val_idxs in cv.split(features[\"identity\"], y_train_):\n        clf.fit(features[\"identity\"][train_idxs], y_train_[train_idxs])\n        # clf.fit(features['identity'], y_train_)  # overfitting\n        prob_oof[val_idxs] = clf.predict_proba(features[\"identity\"][val_idxs])\n        prob_test[:, :] += clf.predict_proba(features_test[\"identity\"]) / n_splits_\n\n    post = oof_pred_ensemble.copy()\n\n    thr = 0.99\n\n    # NSW Dry_Clover_g = 0\n    post[prob_oof[:, states.index(\"NSW\")] > thr, LABELS.index(\"Dry_Clover_g\")] = 0\n    pred_test[prob_test[:, states.index(\"NSW\")] > thr, LABELS.index(\"Dry_Clover_g\")] = 0\n\n    # WA Dry_Dead_g = 0\n    post[prob_oof[:, states.index(\"WA\")] > thr, LABELS.index(\"Dry_Dead_g\")] = 0\n    pred_test[prob_test[:, states.index(\"WA\")] > thr, LABELS.index(\"Dry_Dead_g\")] = 0\n\n    score = competition_metric(y_train, post)\n    print(f\"PP Score w/ State: {score:.4f}\")\n    post = project_nonneg_equalities(post)\n    pred_test = project_nonneg_equalities(pred_test)\n    score = competition_metric(y_train, post)\n    print(f\"PP Score: {score:.4f}\")\n\n\n# ==============================================================================\n# Create Submission File\n# ==============================================================================\n\n# Format for submission\nsubmission_wide = pl.concat(\n    [\n        df_test.select(\"image_path\"),\n        pl.DataFrame(pred_test, schema=LABELS),\n    ],\n    how=\"horizontal\",\n)\nsubmission = (\n    test.join(submission_wide, on=\"image_path\")\n    .with_columns(\n        pl.coalesce(\n            *[pl.when(pl.col(\"target_name\") == col).then(pl.col(col)) for col in LABELS]\n        ).alias(\"target\")\n    )\n    .select(\"sample_id\", \"target\")\n)\n\n# Save\nsubmission.write_csv(f\"submission_{cfg_name}.csv\")\nprint(f\"\\nCreated submission_{cfg_name}.csv\")\nprint(submission.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-01T15:43:53.583575Z","iopub.execute_input":"2025-12-01T15:43:53.583801Z","iopub.status.idle":"2025-12-01T15:43:53.618831Z","shell.execute_reply.started":"2025-12-01T15:43:53.583784Z","shell.execute_reply":"2025-12-01T15:43:53.618224Z"},"jupyter":{"outputs_hidden":true},"collapsed":true,"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Run Scripts","metadata":{}},{"cell_type":"code","source":"# !python model_3_embed.py siglip2-giant.yml compute_features=true infer=true\n# !python model_3_embed.py pe_spatial_gigantic.yml compute_features=true infer=true\n# !python model_3_embed.py dinov3-vit7b16.yml compute_features=true infer=true\n\n# !python model_3_embed.py multi-model.yml compute_features=true infer=false\n!python model_3_embed.py multi-model.yml compute_features=true infer=true\n\n# !python model_3_embed.py spatial.yml compute_features=false infer=true","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-01T15:43:53.619623Z","iopub.execute_input":"2025-12-01T15:43:53.619877Z","iopub.status.idle":"2025-12-01T15:47:53.714500Z","shell.execute_reply.started":"2025-12-01T15:43:53.619860Z","shell.execute_reply":"2025-12-01T15:47:53.713466Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Weighted Average Ensemble\n\nEnsemble the prediction results of all models.","metadata":{}},{"cell_type":"code","source":"import os\n\nimport pandas as pd\n\n\ndef ensemble():\n    print(\"--- [Final Step] Performing a weighted average of all model predictions ---\")\n\n    # Model definition\n    models = {\n        # \"ConvnextTiny\": {  # LB 0.61\n        #     \"path\": \"/kaggle/working/submission_ConvnextTiny.csv\",\n        #     \"weight\": 0.35\n        # },\n        # \"siglip2-giant\": {\n        #     \"path\": \"/kaggle/working/submission_siglip2-giant.csv\",\n        #     \"weight\": 0.25\n        # },\n        # \"tulip-so400m\": {\n        #     \"path\": \"/kaggle/working/submission_tulip-so400m.csv\",\n        #     \"weight\": 0.20\n        # },\n        # \"pe_spatial_gigantic\": {  # LB ~0.74\n        #     \"path\": \"/kaggle/working/submission_pe_spatial_gigantic.csv\",\n        #     \"weight\": 0.5\n        # },\n        # \"dinov3-vit7b16\": {  # LB ~0.74\n        #     \"path\": \"/kaggle/working/submission_dinov3-vit7b16.csv\",\n        #     \"weight\": 1.0,\n        # },\n        \"multi-model\": {  # LB 0.76~0.77\n            \"path\": \"/kaggle/working/submission_multi-model.csv\",\n            \"weight\": 1.0,\n        },\n        # \"spatial\": {  # LB 0.76~0.77\n        #     \"path\": \"/kaggle/working/submission_spatial.csv\",\n        #     \"weight\": 1.0,\n        # },\n    }\n\n    # Load the submission file from each model\n    dfs = {}\n    for key, config in models.items():\n        filepath = config[\"path\"]\n        if not os.path.exists(filepath):\n            raise FileNotFoundError(f\"{filepath} not found.\")\n        dfs[key] = pd.read_csv(filepath).set_index(\"sample_id\")\n\n    # Perform the weighted average\n    submission = sum(config[\"weight\"] * dfs[key] for key, config in models.items())\n\n    # Save the final submission file\n    submission.to_csv(\"submission.csv\")\n    print(\"\\nThe final submission file 'submission.csv' has been created.\")\n    print(submission.head())\n\n\nif __name__ == \"__main__\":\n    ensemble()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-01T15:47:53.715909Z","iopub.execute_input":"2025-12-01T15:47:53.716198Z","iopub.status.idle":"2025-12-01T15:47:54.331757Z","shell.execute_reply.started":"2025-12-01T15:47:53.716173Z","shell.execute_reply":"2025-12-01T15:47:54.330678Z"}},"outputs":[],"execution_count":null}]}