PERF : Searching Optimized

Author: IcodeG00D

BFSearch.py

@@ -1,71 +1,67 @@
-#BF Search Project
 import cv2
+import torch
+import numpy as np
+from PIL import Image
+import os
+import clip
+  # Using openai-clip package
 
-# Define the paths of the target images
-target_image_paths = [
-'C:\\Users\\Kushagra pathak\\Desktop\\python\\testimg\\765f28e4-33d3-4a9c-be0e-4bc9e6cee4fe.jfif',
-    'C:\\Users\\Kushagra pathak\\Desktop\\python\\testimg\\93c4ec94-5e2e-4e44-9baf-6f7f43eeb5c7.jfif',
-    'C:\\Users\\Kushagra pathak\\Desktop\\python\\testimg\\d612eb6b-089c-4d91-b435-a03d51e01adf.jfif',
-     # Add more image paths as needed
-]
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, preprocess = clip.load("ViT-B/32", device=device)
 
-# Create a list to store the target images and their descriptors
-target_images = []
-target_keypoints = []
-target_descriptors = []
+# Folder with your 5 images
+dataset_folder = r"D:\Projects\Brute-Force-main\testimg"
+dataset_embeddings = []
+image_files = []
 
-# Create a feature detector
-orb = cv2.ORB_create()
+print("Loading dataset images...")
+for img_name in os.listdir(dataset_folder):
+    path = os.path.join(dataset_folder, img_name)
+    try:
+        image = preprocess(Image.open(path)).unsqueeze(0).to(device)
+        with torch.no_grad():
+            embedding = model.encode_image(image)
+            embedding /= embedding.norm(dim=-1, keepdim=True)
+            dataset_embeddings.append(embedding)
+            image_files.append(path)
+    except Exception as e:
+        print(f"Failed loading {path}: {e}")
 
-# Load the target images and compute descriptors
-for path in target_image_paths:
-    target_image = cv2.imread(path)
-    keypoints_target, descriptors_target = orb.detectAndCompute(target_image, None)
-    if descriptors_target is not None:
-        target_images.append(target_image)
-        target_keypoints.append(keypoints_target)
-        target_descriptors.append(descriptors_target)
+print(f"Loaded {len(image_files)} images.")
 
-# Create a feature matcher
-bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
-
-# Initialize the camera
-camera = cv2.VideoCapture(0)
+# Start webcam
+cap = cv2.VideoCapture(0)
 
 while True:
-    # Capture frame-by-frame
-    ret, frame = camera.read()
-
-    # Detect features in the frame
-    keypoints_frame, descriptors_frame = orb.detectAndCompute(frame, None)
+    ret, frame = cap.read()
+    if not ret:
+        break
 
-    best_match_idx = None  # Index of the best matching target image
-    best_match_distance = float('inf')  # Initial distance set to infinity
+    small_frame = cv2.resize(frame, (320, 240))
+    rgb = cv2.cvtColor(small_frame, cv2.COLOR_BGR2RGB)
+    pil_img = Image.fromarray(rgb)
+    image = preprocess(pil_img).unsqueeze(0).to(device)
 
-    # Match the features between the frame and each target image
-    for i, descriptors_target in enumerate(target_descriptors):
-        matches = bf.match(descriptors_frame, descriptors_target)
-        matches = sorted(matches, key=lambda x: x.distance)
+    with torch.no_grad():
+        query_embedding = model.encode_image(image)
+        query_embedding /= query_embedding.norm(dim=-1, keepdim=True)
 
-        # Check if the current target image has a better match
-        if matches[0].distance < best_match_distance:
-            best_match_idx = i
-            best_match_distance = matches[0].distance
+    similarities = [torch.cosine_similarity(query_embedding, emb).item() for emb in dataset_embeddings]
+    best_idx = np.argmax(similarities)
+    best_score = similarities[best_idx]
 
-    if best_match_idx is not None:
-        # Draw the best match on the frame
-        matched_frame = cv2.drawMatches(
-            frame, keypoints_frame,
-            target_images[best_match_idx], target_keypoints[best_match_idx], [matches[0]], None, flags=2
-        )
-        cv2.imshow('Object Detection', matched_frame)
+    if best_score > 0.60:
+        matched_img = cv2.imread(image_files[best_idx])
+        matched_img = cv2.resize(matched_img, (small_frame.shape[1], small_frame.shape[0]))
+        combined = np.hstack((small_frame, matched_img))
+        label = f"Match: {os.path.basename(image_files[best_idx])} ({best_score:.2f})"
+        cv2.putText(combined, label, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
+        cv2.imshow("Object Match", combined)
     else:
-        cv2.imshow('Object Detection', frame)
+        cv2.imshow("Object Match", small_frame)
 
-    # Check for the 'q' key to exit the program
     if cv2.waitKey(1) & 0xFF == ord('q'):
         break
 
-# Release the camera and close all windows
-camera.release()
+cap.release()
 cv2.destroyAllWindows()