feat: ReId 객체의 embedding dim을 faiss 초기화에 이용

Peachypie98 · Jul 30, 2023 · 4958483 · 4958483
1 parent 24fe523
commit 4958483
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/re_id/module/reid.py b/re_id/module/reid.py
@@ -45,15 +45,15 @@ def __init__(self, model, checkpoint, gallery_path=None, person_thr=0.6, cosine_
         self.cluster = None
 
         random_tensor = torch.randn(1, 3, 224, 224).to(self.device)
-        embedding_dim = self.model(random_tensor).shape[-1]
+        self.embedding_dim = self.model(random_tensor).shape[-1]
 
         self.player_dict = dict()
 
         self.person_thr = person_thr
         self.cosine_thr = cosine_thr
 
         res = faiss.StandardGpuResources()
-        self.faiss_index = faiss.GpuIndexFlatIP(res, embedding_dim)
+        self.faiss_index = faiss.GpuIndexFlatIP(res, self.embedding_dim)
         self.faiss_index = faiss.IndexIDMap(self.faiss_index)
 
         self.gallery_path = gallery_path
@@ -180,7 +180,7 @@ def init_gallery(self, frames):
             frames (list):   (3, W, H) * (person) * (frames num)
         """
 
-        accumulated_query_vector = np.empty((0, 1000), dtype=np.float32)
+        accumulated_query_vector = np.empty((0, self.embedding_dim), dtype=np.float32)
 
         for frame in frames:
             person_img_lst = []