From 3176baff21ca65c6cee966e3bd6f1af6539b8ae1 Mon Sep 17 00:00:00 2001 From: gauenk Date: Tue, 21 Nov 2023 10:06:27 -0500 Subject: [PATCH] added anchor_self to self_action --- lib/stnls/nn/non_local_attn_stack.py | 1 - lib/stnls/search/impl/non_local_search.py | 4 ++-- lib/stnls/search/impl/refinement.py | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/lib/stnls/nn/non_local_attn_stack.py b/lib/stnls/nn/non_local_attn_stack.py index 300eebc..0a05df7 100644 --- 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 index 93f8fc0..53d726e 100644 --- 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 index fc0f625..2856f2d 100644 --- 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