From dcb26968f719acdd70610b6bb72d9ab013aecf6e Mon Sep 17 00:00:00 2001
From: Kartik Pradeepan <kartikspradeepan@gmail.com>
Date: Tue, 29 Oct 2024 05:05:16 -0400
Subject: [PATCH] Allow region_layer_map json to be loaded into BrainModel
 during commitment (#1314) (#1401)

* Add load_region_layer_map_from_json

* Cleaned up load_region_layer_map_json function and added error handling.

* Replaced print with logging

* Recommended fix:
- logger
- Path instead of OS
- Avoid cwd

* Recommended fix:
- logger
- Path instead of OS
- Avoid cwd
---
 .../brain_transformation/__init__.py          | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/brainscore_vision/model_helpers/brain_transformation/__init__.py b/brainscore_vision/model_helpers/brain_transformation/__init__.py
index a1a3c1dc7..553bf951d 100644
--- a/brainscore_vision/model_helpers/brain_transformation/__init__.py
+++ b/brainscore_vision/model_helpers/brain_transformation/__init__.py
@@ -1,3 +1,8 @@
+import json
+from pathlib import Path
+import logging
+from brainscore_vision.utils import fullname
+from brainscore_core.plugin_management import import_plugin
 from brainscore_vision import load_benchmark
 from brainscore_vision.model_helpers.brain_transformation.temporal import TemporalAligned
 from brainscore_vision.model_interface import BrainModel
@@ -22,6 +27,7 @@ class ModelCommitment(BrainModel):
     def __init__(self, identifier,
                  activations_model, layers, behavioral_readout_layer=None, region_layer_map=None,
                  visual_degrees=8):
+        self._logger = logging.getLogger(fullname(self))
         self.layers = layers
         self.activations_model = activations_model
         # We set the visual degrees of the ActivationsExtractorHelper here to avoid changing its signature.
@@ -30,12 +36,18 @@ def __init__(self, identifier,
         self.activations_model._extractor.set_visual_degrees(visual_degrees)  # for microsaccades
         self._visual_degrees = visual_degrees
         # region-layer mapping
+
+        # Attempt to load region_layer_map from JSON, if available
+        region_layer_map = self.load_region_layer_map_json(identifier) if region_layer_map is None else region_layer_map
+
+        # If region_layer_map is unavailable
         if region_layer_map is None:
             layer_selection = LayerSelection(model_identifier=identifier,
                                              activations_model=activations_model, layers=layers,
                                              visual_degrees=visual_degrees)
             region_layer_map = RegionLayerMap(layer_selection=layer_selection,
                                               region_benchmarks=STANDARD_REGION_BENCHMARKS)
+
         # neural
         layer_model = LayerMappedModel(identifier=identifier, activations_model=activations_model,
                                        region_layer_map=region_layer_map)
@@ -52,6 +64,27 @@ def __init__(self, identifier,
                                                })
         self.do_behavior = False
 
+    def load_region_layer_map_json(self, identifier):
+        '''
+        Attempts to load the region_layer_map from a JSON file in the model's directory
+        If file exists, load JSON. Otherwise, return None and proceed with legacy layer mapping
+        '''
+        try:
+            importer = import_plugin.ImportPlugin(library_root='brainscore_vision', plugin_type='models', identifier=identifier)
+            model_dir = importer.locate_plugin()
+            project_root = Path(__file__).resolve().parent.parent
+            region_layer_map_path = project_root / 'vision' / 'models' / model_dir / 'region_layer_map' / f'{identifier}.json'
+            if region_layer_map_path.exists():
+                with region_layer_map_path.open('r') as region_layer_map_file:
+                    self._logger.info(f"Successfully loaded region_layer_map for {identifier}")
+                    return json.load(region_layer_map_file)
+            else:
+                self._logger.info(f"No region_layer_map file found for {identifier}, proceeding with default layer mapping")
+                return None
+        except Exception as e:
+            self._logger.error(f"Error importing model to search for region_layer_map: {e}")
+            return None
+
     def visual_degrees(self) -> int:
         return self._visual_degrees