Update app.py

app.py

@@ -4,12 +4,38 @@ import numpy as np
 from PIL import Image
 import pandas as pd
 from sklearn.metrics.pairwise import cosine_similarity
+from tqdm import tqdm
 
 
 from token_classifier import load_token_classifier, predict
 from model import Model
 from dataset import RetrievalDataset
-from generate_embeds import encode_database
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+batch_size = 512
+
+
+def encode_database(model, df: pd.DataFrame) -> np.ndarray :
+    """
+    Process database images and generate embeddings.
+
+    Args:
+    df (pd. DataFrame ): DataFrame with column:
+    - target_image: str, paths to database images
+
+    Returns:
+    np.ndarray: Embeddings array (num_images, embedding_dim)
+    """
+    model.eval()
+    all_embeddings = []
+    for i in tqdm(range(0, len(df), batch_size)):
+        target_imgs = torch.stack([model.processor(Image.open(target_image_path)) for target_image_path in df['target_image'][i:i+batch_size]]).to(device)
+        with torch.no_grad():
+            # target_imgs_embedding = model.encode_database_image(target_imgs)
+            target_imgs_embedding = model.feature_extractor.encode_image(target_imgs)
+        target_imgs_embedding = torch.nn.functional.normalize(target_imgs_embedding, dim=1, p=2)
+        all_embeddings.append(target_imgs_embedding.detach().cpu().numpy())
+    return np.concatenate(all_embeddings)
 
 
 # Load model and configurations
@@ -21,7 +47,6 @@ def load_model():
 
 
 def process_single_query(model, query_image_path, query_text, database_embeddings, database_df):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
     # Process query image
     query_img = model.processor(Image.open(query_image_path)).unsqueeze(0).to(device)