Try saving heatmap outside clip_seg_tta

Taxabind/TaxaBind/SatBind/clip_seg_tta.py

@@ -40,6 +40,7 @@ from types import SimpleNamespace
 import torch.nn as nn
 import spaces   # integration with ZeroGPU on hf
 from torch.autograd import enable_grad           # handy alias
+import copy
 
 # import matplotlib
 # matplotlib.use("Agg")  # <-- key line to avoid tkinter dependency
@@ -416,6 +417,9 @@ class ClipSegTTA:
                 self.tta_time = time.time() - start_time
                 # print("self.tta_time: ", self.tta_time)
 
+                # Make deep copy of self.model_local.imo_encoder
+                self.model_local.imo_encoder = copy.deepcopy(self.model_local.imo_encoder)
+
                 # Visualization every 'num_viz_steps' steps (if enabled)
                 if (step + 1) % num_viz_steps == 0 and viz_heatmap:
                     # Visualize only the first sample in the batch 
@@ -430,6 +434,9 @@ class ClipSegTTA:
                         species_name=self.species_name
                     )
 
+            # Save final heatmap after TTA steps 
+            self.generate_heatmap(img, imo, txt, sound=sound, modality=modality)
+
             ## NOTE: Added due to app.py (to allocate to GPU only when needed on HF)
             # if self.device.type == "cuda":
             print("Deallocating models from GPU...")
@@ -438,13 +445,8 @@ class ClipSegTTA:
             self.model_local.imo_encoder.to(self.device)
             self.model_local.bio_model.to(self.device)
 
-            # Move tensors to CPU before generating heatmap to avoid dtype/device mismatches
-            img_cpu = img.to(self.device) if isinstance(img, torch.Tensor) else img
-            imo_cpu = imo.to(self.device)
-            sound_cpu = sound.to(self.device) if (sound is not None and isinstance(sound, torch.Tensor)) else sound
+            return self.heatmap
 
-            # Save final heatmap after TTA steps 
-            self.generate_heatmap(img_cpu, imo_cpu, txt, sound=sound_cpu, modality=modality)
 
 
     def generate_query_embeds(self, img, imo, txt, sound=None, modality="image"):

test_multi_robot_worker.py

@@ -653,7 +653,7 @@ class TestWorker:
             # print("!!! num_tta_steps", num_tta_steps)
                 
             # TTA Update
-            self.clip_seg_tta.execute_tta(
+            heatmap = self.clip_seg_tta.execute_tta(
                 filt_traj_coords, 
                 filt_targets_found_on_path, 
                 tta_steps=NUM_TTA_STEPS, 
@@ -665,6 +665,9 @@ class TestWorker:
                 target_found_idxs=self.env.target_found_idxs,
                 reset_weights=RESET_WEIGHTS
             )
+
+            self.clip_seg_tta.heatmap = heatmap
+
             self.env.segmentation_info_mask = np.expand_dims(self.clip_seg_tta.heatmap.T.flatten(), axis=1)
             self.env.segmentation_info_mask_unnormalized = np.expand_dims(self.clip_seg_tta.heatmap_unnormalized.T.flatten(), axis=1)
             self.step_since_tta = 0