app.py

import streamlit as st
import cv2
from PIL import Image
import numpy as np
import os
import matplotlib.pyplot as plt

# Define a function to search for similar images
def search_similar_images(img_path, num_results=10):
    # Load the query image
    query_image = cv2.imread(img_path)

    # Convert the query image to grayscale
    query_image_gray = cv2.cvtColor(query_image, cv2.COLOR_BGR2GRAY)

    # Resize the query image to a fixed size
    query_image_resized = cv2.resize(query_image_gray, (300, 300))

    # Calculate the histogram of the query image
    query_hist = cv2.calcHist([query_image_resized], [0], None, [256], [0, 256])

    # Normalize the histogram
    query_hist_norm = cv2.normalize(query_hist, query_hist, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)

    # Load all images from the 'images' directory
    images_dir = "images"
    image_files = os.listdir(images_dir)
    images = []
    for image_file in image_files:
        if image_file.endswith(".jpg") or image_file.endswith(".png"):
            image_path = os.path.join(images_dir, image_file)
            image = cv2.imread(image_path)
            images.append(image)

    # Calculate the histograms and similarities for each image
    similarities = []
    for image in images:
        # Convert the image to grayscale
        image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # Resize the image to a fixed size
        image_resized = cv2.resize(image_gray, (300, 300))

        # Calculate the histogram of the image
        image_hist = cv2.calcHist([image_resized], [0], None, [256], [0, 256])

        # Normalize the histogram
        image_hist_norm = cv2.normalize(image_hist, image_hist, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)

        # Calculate the correlation between the histograms
        similarity = cv2.compareHist(query_hist_norm, image_hist_norm, cv2.HISTCMP_CORREL)

        # Add the similarity to the list
        similarities.append(similarity)

    # Get the indices of the top num_results most similar images
    indices = np.argsort(similarities)[-num_results:]

    # Create a list of the top num_results most similar images
    results = []
    for index in indices:
        results.append(images[index])

    # Return the list of results
    return results

def display_results(similar_images):
    # Create a figure with a grid of 10 subplots
    fig, axs = plt.subplots(2, 5, figsize=(10, 4))

    # Loop through the similar images and display each one in a subplot
    for i,img in enumerate(similar_images):
        # Load the image using OpenCV
        # Convert to RGB for displaying with matplotlib
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # Calculate the row and column index for the current subplot
        row = i // 5
        col = i % 5

        # Display the image in the current subplot
        axs[row, col].imshow(img)
        axs[row, col].axis("off")
    
    # Adjust the spacing between subplots and display the figure
    plt.subplots_adjust(wspace=0.05, hspace=0.05)
    st.pyplot(fig)

def app():
    # Set the page title and icon
    st.set_page_config(page_title="Similar Images Search", page_icon=":mag:")
    # Add a title to the app
    st.title("Find Similar Images")

    # Add a description to the app
    st.markdown("This app allows you to search for similar images.")

    # Allow the user to choose between "Browse a single image" and "Select from test dataset"
    option = st.radio("Select an option:", ("Browse a single image", "Select from test dataset"))

    # If the user selects "Browse a single image"
    if option == "Browse a single image":
        # Allow the user to upload an image
        uploaded_file = st.file_uploader("Choose an image:", type=["jpg", "jpeg", "png"])

        # If the user has uploaded an image
        if uploaded_file is not None:
            # Load the image
            image = Image.open(uploaded_file).convert("RGB")

            # Display the uploaded image
            st.image(image, caption="Uploaded Image", use_column_width=True)

            # Search for similar images
            # print("Uploaded file ==> ",uploaded_file)
            # Save the image to a temporary file
            temp_file_path = "temp.jpg"
            image.save(temp_file_path)

            results = search_similar_images(temp_file_path)

            # Display the results as a grid of images
            st.markdown("---")
            st.subheader("Similar Images")
            display_results(results)

    # If the user selects "Select from test dataset"
    else:
        # Get a list of all test dataset image files
        test_dataset_dir = "images"
        test_dataset_files = os.listdir(test_dataset_dir)
        test_dataset_images = []

        # Display the list of test dataset image files in a sidebar
        st.sidebar.title("Test Dataset")
        for test_dataset_file in test_dataset_files:
            if test_dataset_file.endswith(".jpg") or test_dataset_file.endswith(".png"):
                test_dataset_image_path = os.path.join(test_dataset_dir, test_dataset_file)
                test_dataset_image = Image.open(test_dataset_image_path)
                test_dataset_images.append(test_dataset_image)

                # Display the file name as a text instead of an image
                st.sidebar.write(test_dataset_file)

        # Allow the user to select an image from the test dataset
        selected_image_path = st.sidebar.selectbox("Select an image:", test_dataset_files)

        # Load the selected image
        selected_image_path = os.path.join(test_dataset_dir, selected_image_path)
        selected_image = Image.open(selected_image_path)

        # Display the selected image
        st.image(selected_image, caption="Selected Image", use_column_width=True)

        # Search for similar images
        results = search_similar_images(selected_image_path)

        # Display the results as a grid of images
        if selected_image_path != os.path.join(test_dataset_dir, test_dataset_files[0]):
            st.markdown("---")
            st.subheader("Similar Images")
            display_results(results)
            # for result in results:
            #     st.image(result, caption="Similar Image", use_column_width=True)
            #     print("result==> ",result)

    # Add a footer to the app
    st.markdown("---")
    st.markdown("Created by Pruthul")


app()



####  == Image feature extraction using Resnet == (Failed when run using streamlit) ### 

# from sklearn.neighbors import NearestNeighbors
# from tensorflow.keras.preprocessing.image import load_img, img_to_array
# from tensorflow.keras.applications.resnet50 import ResNet50,preprocess_input

# def extract_features(image_path, model):
#     # Load and preprocess the image
#     image = load_img(image_path, target_size=(224, 224))
#     image_array = img_to_array(image)
#     image_array = preprocess_input(image_array)
#     # Extract the features using the ResNet50 model
#     features = model.predict(image_array.reshape(1, 224, 224, 3))
    
#     # Flatten the features and return them as a 1D array
#     features = features.flatten()
#     return features


# # Define a function to search for similar images
# def search_similar_images(img_path, num_results=10):
#     #load feature dictionary
#     with open("features.npy", "rb") as f:
#         image_dict = np.load(f, allow_pickle=True).item()

#     # Fit a nearest neighbor model on the features
#     nn_model = NearestNeighbors(n_neighbors=num_results, metric='cosine')

#     # Convert the dictionary to a matrix of feature vectors
#     features_list = np.array(list(image_dict.values()))

#     # Fit the model to the feature matrix
#     nn_model.fit(features_list)

#     # Define the file path to the test image
#     test_image_path = img_path

#     # Load all images from the 'images' directory
#     images_dir = "images"
#     image_files = os.listdir(images_dir)
#     images = []
#     for image_file in image_files:
#         if image_file.endswith(".jpg") or image_file.endswith(".png"):
#             image_path = os.path.join(images_dir, image_file)
#             image = cv2.imread(image_path)
#             images.append(image)

#     model = ResNet50(weights='imagenet', include_top=False)
#     # Extract features from the test image
#     test_image_features = extract_features(test_image_path, model)

#     # Reshape the test image features to match the shape of the feature vectors
#     test_image_features = test_image_features.reshape(1, -1)

#     # # Find the 10 most similar images to the test image
#     distances, indices = nn_model.kneighbors(test_image_features)

#     # Create a list of the top num_results most similar images
#     results = []
#     for index in indices[0]:
#         results.append(images[index])

#     # Return the list of results
#     return results