Update app.py

app.py

@@ -3,34 +3,32 @@ import torch
 import numpy as np
 from PIL import Image
 import pandas as pd
-from pathlib import Path
 from sklearn.metrics.pairwise import cosine_similarity
 
-# Import your model and necessary functions
-from src.config import ConfigManager
-from src.token_classifier import load_token_classifier, predict
-from your_model_file import YourModel  # Replace with your actual model import
+
+from token_classifier import load_token_classifier, predict
+from model import Model
+from dataset import RetrievalDataset
+from generate_embeds import encode_database
+
 
 # Load model and configurations
 def load_model():
-    model = YourModel()  # Initialize your model
+    model = Model(model_name="ViTamin-L-384", pretrained=None)
+    model.load("weights.pth")
     model.eval()
     return model
 
-def load_dataset():
-    # Load your default dataset
-    database_df = pd.read_csv('database.csv')  # Adjust path as needed
-    return database_df
 
 def process_single_query(model, query_image_path, query_text, database_embeddings, database_df):
-    device = ConfigManager().get("training")["device"]
+    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)
     
     # Get token classifier
     token_classifier, token_classifier_tokenizer = load_token_classifier(
-        ConfigManager().get("paths")["pretrained_token_classifier_path"],
+        "trained_distil_bert_base",
         device
     )
     
@@ -87,8 +85,16 @@ def process_single_query(model, query_image_path, query_text, database_embedding
 
 # Initialize model and database
 model = load_model()
-database_df = load_dataset()
-database_embeddings = encode_database(model, database_df)  # Using your existing function
+
+test_dataset = RetrievalDataset(
+    img_dir_path="sample_evaluation/images",
+    annotations_file_path="sample_evaluation/data.csv",
+    split='test',
+    transform=model.processor,
+    tokenizer=model.tokenizer
+)
+
+database_embeddings = encode_database(model, test_dataset.load_database())  # Using your existing function
 
 def interface_fn(selected_image, query_text):
     result_image_path = process_single_query(
@@ -96,7 +102,7 @@ def interface_fn(selected_image, query_text):
         selected_image, 
         query_text, 
         database_embeddings, 
-        database_df
+        test_dataset.load_database()
     )
     return Image.open(result_image_path)