added anchor_self to self_action

gauenk · Nov 21, 2023 · 3176baf · 3176baf
1 parent 5a63c0b
commit 3176baf
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 4 deletions.
diff --git a/lib/stnls/nn/non_local_attn_stack.py b/lib/stnls/nn/non_local_attn_stack.py
@@ -134,7 +134,6 @@ def forward(self, vid, flows=None, state=None):
 
         # -- search --
         dists,inds = self.run_search(q_vid,k_vid,flows,state)
-        # print(inds.shape)
 
         # -- normalize --
         weights,inds = self.run_normalize(dists,inds)

diff --git a/lib/stnls/search/impl/non_local_search.py b/lib/stnls/search/impl/non_local_search.py
@@ -67,11 +67,11 @@ def forward(vid0, vid1, flows,
             reflect_bounds, full_ws, patch_offset,  dist_type_i)
 
     # -- anchor --
-    menu = [None,"anchor","anchor_each","remove",]
+    menu = [None,"anchor","anchor_self","anchor_each","remove",]
     menu += ["remove_ref_frame","anchor_and_remove_ref_frame"]
     assert self_action in menu
     anchor_self = False if self_action is None else "anchor" in self_action
-    if self_action == "anchor":
+    if self_action in ["anchor","anchor_self"]:
         stnls.nn.anchor_self(dists,inds,stride0,H,W)
     elif self_action == "anchor_each":
         stnls.nn.anchor_self_time(dists,inds,flows,wt,stride0,H,W)

diff --git a/lib/stnls/search/impl/refinement.py b/lib/stnls/search/impl/refinement.py
@@ -73,7 +73,7 @@ def forward(vid0, vid1, flows,
         assert self_action == None
 
     # -- topk --
-    assert self_action in [None,"anchor","anchor_each"]
+    assert self_action in [None,"anchor","anchor_self","anchor_each"]
     anchor_self = False if self_action is None else "anchor" in self_action
     if topk_mode == "all":
         dim = 3