File size: 2,557 Bytes
50ddf24 |
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 |
import numpy as np
from sklearn.metrics.pairwise import pairwise_distances
from typing import List, Dict
from utils.config import Config
from PIL import Image
import pandas as pd
import tensorflow as tf
import io
import os
# Load the dataset (replace with the actual path to your dataset)
dataset_path = Config.read('app', 'dataset')
# Ensure the dataset exists
if not os.path.exists(dataset_path):
raise FileNotFoundError(f"The dataset file at {dataset_path} was not found.")
# Load the dataset
data = pd.read_pickle(dataset_path)
# Ensure the dataset has the necessary columns: 'asin', 'title', 'brand', 'medium_image_url'
required_columns = ['asin', 'title', 'brand', 'medium_image_url']
for col in required_columns:
if col not in data.columns:
raise ValueError(f"Missing required column: {col} in the dataset")
# Load the pre-trained CNN features and corresponding ASINs
bottleneck_features_train = np.load(Config.read('app', 'cnnmodel'))
bottleneck_features_train = bottleneck_features_train.astype(np.float64)
asins = np.load(Config.read('app', 'cssasins'))
asins = list(asins)
# Helper function to extract features from the uploaded image using a pre-trained model
def extract_features_from_image(image_bytes):
image = Image.open(io.BytesIO(image_bytes))
image = image.resize((224, 224))
image_array = np.array(image) / 255.0
image_array = np.expand_dims(image_array, axis=0)
# Load the VGG16 model for feature extraction
model = tf.keras.applications.VGG16(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
features = model.predict(image_array)
features = features.flatten()
return features
# Function to get similar products based on CNN features
def get_similar_products_cnn(image_features, num_results: int) -> List[Dict]:
pairwise_dist = pairwise_distances(bottleneck_features_train, image_features.reshape(1, -1))
# Get the indices of the closest products
indices = np.argsort(pairwise_dist.flatten())[0:num_results]
results = []
for i in range(len(indices)):
# Get the product details for each closest match
product_details = data[['asin', 'brand', 'title', 'medium_image_url']].loc[data['asin'] == asins[indices[i]]]
for indx, row in product_details.iterrows():
result = {
'asin': row['asin'],
'brand': row['brand'],
'title': row['title'],
'url': row['medium_image_url']
}
results.append(result)
return results
|