Merge pull request kerrj#25 from shengjie-lin/multi-phrases-fix

Multi phrases fix

Author: kerrj

lerf/data/utils/feature_dataloader.py

@@ -42,6 +42,7 @@ def load(self):
         self.data = torch.from_numpy(np.load(self.cache_path)).to(self.device)
 
     def save(self):
+        os.makedirs(self.cache_path.parent, exist_ok=True)
         cache_info_path = self.cache_path.with_suffix(".info")
         with open(cache_info_path, "w") as f:
             f.write(json.dumps(self.cfg))

lerf/lerf.py

@@ -90,7 +90,7 @@ def get_max_across(self, ray_samples, weights, hashgrid_field, scales_shape, pre
         n_phrases = len(self.image_encoder.positives)
         n_phrases_maxs = [None for _ in range(n_phrases)]
         n_phrases_sims = [None for _ in range(n_phrases)]
-        for _, scale in enumerate(scales_list):
+        for i, scale in enumerate(scales_list):
             scale = scale.item()
             with torch.no_grad():
                 clip_output = self.lerf_field.get_output_from_hashgrid(
@@ -100,12 +100,13 @@ def get_max_across(self, ray_samples, weights, hashgrid_field, scales_shape, pre
                 )
             clip_output = self.renderer_clip(embeds=clip_output, weights=weights.detach())
 
-            for i in range(n_phrases):
-                probs = self.image_encoder.get_relevancy(clip_output, i)
-                pos_prob = probs[..., 0:1]
-                if n_phrases_maxs[i] is None or pos_prob.max() > n_phrases_sims[i].max():
-                    n_phrases_maxs[i] = scale
-                    n_phrases_sims[i] = pos_prob
+            for j in range(n_phrases):
+                if preset_scales is None or j == i:
+                    probs = self.image_encoder.get_relevancy(clip_output, j)
+                    pos_prob = probs[..., 0:1]
+                    if n_phrases_maxs[j] is None or pos_prob.max() > n_phrases_sims[j].max():
+                        n_phrases_maxs[j] = scale
+                        n_phrases_sims[j] = pos_prob
         return torch.stack(n_phrases_sims), torch.Tensor(n_phrases_maxs)
 
     def get_outputs(self, ray_bundle: RayBundle):