File size: 9,109 Bytes
46d2a02 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 |
import os
import numpy as np
from PIL import Image
from torchvision import transforms, models
import torch
import torch.nn.functional as F
import streamlit as st
import pickle
from sklearn.neighbors import NearestNeighbors
import faiss
# Set up the image transformation
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# Data augmentation transforms
augment_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(20),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
transforms.RandomResizedCrop(224, scale=(0.8, 1.0), ratio=(0.75, 1.33)),
])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@st.cache_resource
def load_model():
model = models.efficientnet_b0(pretrained=True)
model.classifier = torch.nn.Identity() # Remove the final classification layer
model = model.to(device)
model.eval()
return model
model = load_model()
def extract_features(img):
img_t = transform(img)
batch_t = torch.unsqueeze(img_t, 0).to(device)
with torch.no_grad():
features = model(batch_t)
features = F.normalize(features, p=2, dim=1)
return features.cpu().squeeze().numpy()
def generate_augmented_images(img, num_augmented=5):
augmented_images = []
for _ in range(num_augmented):
augmented = augment_transform(img)
augmented_images.append(augmented)
return augmented_images
# def load_and_index_images(root_dir): #without adding data augmented images
# image_paths = []
# features_list = []
# categories = []
# for category in os.listdir(root_dir):
# category_path = os.path.join(root_dir, category)
# if os.path.isdir(category_path):
# for img_name in os.listdir(category_path):
# img_path = os.path.join(category_path, img_name)
# img = Image.open(img_path).convert('RGB')
# features = extract_features(img)
# image_paths.append(img_path)
# features_list.append(features)
# categories.append(category)
# features_array = np.array(features_list).astype('float32')
# d = features_array.shape[1] # dimension
# index = faiss.IndexFlatIP(d) # use inner product (cosine similarity on normalized vectors)
# index.add(features_array)
# return index, image_paths, categories
def load_and_index_images(root_dir):
image_paths = []
features_list = []
categories = []
for category in os.listdir(root_dir):
category_path = os.path.join(root_dir, category)
if os.path.isdir(category_path):
for img_name in os.listdir(category_path):
img_path = os.path.join(category_path, img_name)
img = Image.open(img_path).convert('RGB')
# Generate augmented images
augmented_images = generate_augmented_images(img)
features = extract_features(img)
image_paths.append(img_path)
features_list.append(features)
categories.append(category)
for aug_img in augmented_images:
aug_features = extract_features(aug_img)
features_list.append(aug_features)
image_paths.append(img_path) # Use original path for augmented images
categories.append(category)
features_array = np.array(features_list).astype('float32')
d = features_array.shape[1] # dimension
index = faiss.IndexFlatIP(d) # use inner product (cosine similarity on normalized vectors)
index.add(features_array)
return index, image_paths, categories
def save_index_and_metadata(nn, image_paths, categories, index_file, metadata_file):
with open(index_file, 'wb') as f:
pickle.dump(nn, f)
with open(metadata_file, 'wb') as f:
pickle.dump((image_paths, categories), f)
def load_index_and_metadata(index_file, metadata_file):
with open(index_file, 'rb') as f:
nn = pickle.load(f)
with open(metadata_file, 'rb') as f:
image_paths, categories = pickle.load(f)
return nn, image_paths, categories
def search_similar_images(index, image_paths, categories, query_features, k=20):
query_features = query_features.reshape(1, -1).astype('float32')
similarities, indices = index.search(query_features, k)
similar_images = [image_paths[i] for i in indices[0]]
similarity_scores = similarities[0]
similar_categories = [categories[i] for i in indices[0]]
return similar_images, similarity_scores, similar_categories
def index_files_exist(index_file, metadata_file):
return os.path.exists(index_file) and os.path.exists(metadata_file)
def main():
st.title("Image Classification and Similarity Search")
index_file = "faiss-d2-nn_index.pkl"
metadata_file = "faiss-d2-image_metadata.pkl"
if not index_files_exist(index_file, metadata_file):
st.warning("Index files not found. Creating new index...")
root_dir = "Dataset2" # Replace with your dataset path
index, image_paths, categories = load_and_index_images(root_dir)
save_index_and_metadata(index, image_paths, categories, index_file, metadata_file)
st.success("Index created and saved successfully!")
else:
index, image_paths, categories = load_index_and_metadata(index_file, metadata_file)
st.success("Index loaded successfully!")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file).convert('RGB')
query_features = extract_features(image)
# Search for similar images
similar_images, similarities, similar_categories = search_similar_images(index, image_paths, categories, query_features, k=50)
# Get the predicted class (most common category among top 5 similar images)
predicted_class = max(set(similar_categories[:5]), key=similar_categories[:5].count)
# Display query and matched image
col1, col2 = st.columns(2)
with col1:
st.subheader("Query Image")
st.image(image, caption="Uploaded Image", use_column_width=True)
st.write(f"Image ID: {uploaded_file.name}")
with col2:
if similar_images:
st.subheader("Matched Image")
matched_image_path = similar_images[0]
st.image(Image.open(matched_image_path),
caption=f"Matched Image (Similarity: {similarities[0]:.2f})",
use_column_width=True)
st.write(f"Image ID: {os.path.basename(matched_image_path)}")
else:
st.write("No matched image found")
st.subheader(f"Product Category: {predicted_class}")
similarity_threshold = st.slider("Similarity threshold", min_value=0.0, max_value=1.0, value=0.5, step=0.05)
# Filter results based on similarity threshold and predicted class, and remove duplicates
query_file_name = uploaded_file.name
seen_file_names = set([query_file_name]) # Add query image to seen set
filtered_results = []
for img, sim, cat in zip(similar_images[1:], similarities[1:], similar_categories[1:]): # Start from index 1
file_name = os.path.basename(img)
if sim >= similarity_threshold and cat == predicted_class and file_name not in seen_file_names:
filtered_results.append((img, sim))
seen_file_names.add(file_name)
# Rest of the code remains the same
if filtered_results:
max_images = len(filtered_results)
num_display = st.slider("Number of similar images to display", min_value=0, max_value=max_images, value=min(20, max_images))
st.subheader("Similar Images")
st.info(f"Displaying {num_display} out of {max_images} unique similar images found for the uploaded query image.")
# Create a grid for displaying similar images
num_cols = 5
num_rows = (num_display + num_cols - 1) // num_cols
for row in range(num_rows):
cols = st.columns(num_cols)
for col in range(num_cols):
idx = row * num_cols + col
if idx < num_display:
img_path, sim = filtered_results[idx]
with cols[col]:
st.image(Image.open(img_path), use_column_width=True)
st.write(f"Similarity: {sim:.2f}")
st.write(f"Image ID: {os.path.basename(img_path)}")
else:
st.info("No similar images found above the similarity threshold in the predicted class.")
if __name__ == "__main__":
main() |