基于Python、scikit-learn与Keras实现实拍手表图像相似检索

Hey there! Let's walk through building this watch image matching system using Python, scikit-learn, and Keras. We'll tackle those tricky real-world photos (like messy backgrounds or dim lighting) head-on with a robust pipeline:

First, we need to get all images (your 3000 samples and input photos) into a consistent format, and fix common issues with real shots:

• Resize everything:统一所有图像尺寸到224x224 (matches most pre-trained model inputs). Use PIL or OpenCV for this—here's a quick PIL example:

  from PIL import Image
  
  def resize_image(img_path, target_size=(224,224)):
      img = Image.open(img_path).convert('RGB')  # Force RGB to avoid grayscale issues
      return img.resize(target_size)
  

• Fix lighting & noise: For dark photos, use histogram equalization to boost contrast. If the background is distracting, try simple color thresholding to isolate the watch (or use a lightweight segmentation model if backgrounds are super messy):

  import cv2
  import numpy as np
  
  def enhance_low_light(img_path):
      img = cv2.imread(img_path)
      img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
      # Apply histogram equalization to each channel
      for channel in range(3):
          img_rgb[:,:,channel] = cv2.equalizeHist(img_rgb[:,:,channel])
      return img_rgb
  

• Normalize pixels: Scale pixel values to [0,1] or the range expected by your pre-trained model (like [-1,1] for some Keras models) to keep feature extraction consistent.

Training a CNN from scratch on 3000 images won't give great results—instead, leverage pre-trained models (like ResNet50 or VGG16) to pull out high-level visual features (things like dial shape, strap color, texture):

from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import Model
import numpy as np

# Load ResNet50 without the top classification layer
base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(224,224,3))
# Create a feature extractor that outputs a flat vector (2048 dimensions for ResNet50)
feature_extractor = Model(inputs=base_model.input, outputs=base_model.output)

def get_image_features(img_path):
    # Load and preprocess the image
    img = image.load_img(img_path, target_size=(224,224))
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)
    
    # Extract features and flatten to a 1D vector
    features = feature_extractor.predict(img_array)
    return features.flatten()

# Precompute features for all 3000 watch samples and store them
sample_features = []
sample_names = []

# Replace this loop with your actual file list (e.g., using os.listdir)
for watch_filename in ['Watch_1.jpg', 'Watch_2.jpg', ..., 'Watch_3000.jpg']:
    feat = get_image_features(watch_filename)
    sample_features.append(feat)
    sample_names.append(watch_filename)

# Convert to numpy array for fast calculations
sample_features = np.array(sample_features)

We'll use cosine similarity (better than Euclidean distance for high-dimensional features) to measure how close the input photo's features are to your sample library:

from sklearn.metrics.pairwise import cosine_similarity

def find_top_match(input_img_path):
    # Extract features from the input photo
    input_features = get_image_features(input_img_path).reshape(1, -1)
    
    # Calculate similarity scores against all samples
    similarity_scores = cosine_similarity(input_features, sample_features)[0]
    
    # Sort samples by similarity (highest first)
    sorted_indices = similarity_scores.argsort()[::-1]
    
    # Grab the top match
    top_match_name = sample_names[sorted_indices[0]]
    top_match_score = similarity_scores[sorted_indices[0]]
    
    return top_match_name, top_match_score

# Test it out!
input_photo_path = 'my_real_watch_photo.jpg'
match, score = find_top_match(input_photo_path)
print(f"Your closest match is: {match} (similarity score: {score:.4f})")

• Reduce feature dimensionality: Use PCA to shrink your 2048-dim features to something smaller (like 256-dim) to speed up calculations without losing too much info:

  from sklearn.decomposition import PCA
  
  # Fit PCA on your sample features
  pca = PCA(n_components=256)
  sample_features_pca = pca.fit_transform(sample_features)
  
  # When processing input photos, transform their features with the same PCA
  input_features_pca = pca.transform(input_features)
  

• Speed up large-scale search: If you ever expand beyond 3000 samples, use the FAISS library for approximate nearest neighbor search—it's way faster than brute-force cosine similarity.
• Semantic matching: If you need to prioritize things like "round dial" or "brown strap" over pixel-perfect matches, add a small classification model (trained on labeled data) to tag each sample with these attributes. Then combine visual features with these semantic tags when calculating similarity.

Hope this gives you a solid starting point! If you hit snags with background removal or fine-tuning features, feel free to dig into those specific parts.

内容的提问来源于stack exchange，提问作者Outcast