Update app.py

app.py

@@ -1,50 +1,33 @@
-import torch
-import torch.nn as nn
-import torchvision.models as models
-import torchvision.transforms as transforms
-import os
-from PIL import Image
-
-# Define the ResNet-50 model
-model = models.resnet50(pretrained=True)
-
-# Remove the classification head (the fully connected layer)
-num_features = model.fc.in_features
-model.fc = nn.Identity()
-
-# Set the model to evaluation mode
-model.eval()
-
-# Define the preprocessing transforms
-preprocess = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225]
-    )
-])
-
-# Define the dictionary to store the feature vectors
-features = {}
-
-# Iterate over the images and extract the features
-image_dir = 'lfw'
-for root, dirs, files in os.walk(image_dir):
-    for file in files:
-        # Load the image
-        image_path = os.path.join(root, file)
-        image = Image.open(image_path).convert('RGB')
-
-        # Apply the preprocessing transforms
-        input_tensor = preprocess(image)
-        input_batch = input_tensor.unsqueeze(0)
-
-        # Extract the features from the penultimate layer
-        with torch.no_grad():
-            features_tensor = model(input_batch)
-            features_vector = torch.squeeze(features_tensor).numpy()
-
-        # Store the feature vector in the dictionary
-        features[file] = features_vector
+
+
+
+import streamlit as st
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import cv2
+
+# Load the test dataset
+test_data = pd.read_csv("test_dataset.csv")
+
+# Create a dropdown to select an image from the test dataset
+selected_image = st.sidebar.selectbox("Select an image", test_data["image"])
+
+# Create a file uploader to upload an image
+uploaded_file = st.sidebar.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+
+# Load the selected or uploaded image
+if uploaded_file is not None:
+    query_image = cv2.imread(uploaded_file.name)
+else:
+    query_image = cv2.imread(selected_image)
+
+# Display the query image
+st.image(query_image, caption="Query Image", use_column_width=True)
+
+# Use the similarity search system to find the most similar images
+similar_images = find_similar_images(query_image)
+
+# Display the most similar images
+for image in similar_images:
+    st.image(image, caption="Similar Image", use_column_width=True)