import cv2 import numpy as np import os import json import time import threading import torch from PIL import Image from ultralytics import YOLO import customtkinter as ctk from tkinter import messagebox # =============================================== # SMART OBJECT IDENTIFIER v2 # - YOLOv8 real-time detection (80 classes) # - Teach unknown objects (type on camera) # - Internet recognition via CLIP # - Persistent learning (saves to disk) # =============================================== # ------------------------- # Config # ------------------------- WEBCAM_INDEX = 0 FRAME_WIDTH = 1280 FRAME_HEIGHT = 720 YOLO_CONFIDENCE = 0.40 CUSTOM_MATCH_THRESHOLD = 0.82 CLIP_MODEL_NAME = "openai/clip-vit-base-patch32" # ------------------------- # Paths # ------------------------- OBJECTS_FOLDER = "learned_objects" OBJECTS_JSON = "objects_db.json" os.makedirs(OBJECTS_FOLDER, exist_ok=True) if not os.path.exists(OBJECTS_JSON): with open(OBJECTS_JSON, "w") as f: json.dump([], f) # ------------------------- # Device # ------------------------- device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") # ------------------------- # Load YOLOv8 # ------------------------- print("Loading YOLOv8 model...") yolo_model = YOLO("yolov8n.pt") print("YOLOv8 loaded!\n") # ------------------------- # CLIP (Lazy Load) # ------------------------- clip_model = None clip_processor = None clip_loaded = False def load_clip(): global clip_model, clip_processor, clip_loaded if clip_loaded: return True try: from transformers import CLIPProcessor, CLIPModel print("Loading CLIP model (first time downloads ~350MB)...") clip_model = CLIPModel.from_pretrained(CLIP_MODEL_NAME).to(device) clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME) clip_model.eval() clip_loaded = True print("CLIP loaded!\n") return True except Exception as e: print(f"Failed to load CLIP: {e}") return False # ------------------------- # CLIP Vocabulary # ------------------------- CLIP_VOCABULARY = [ "smartphone", "laptop", "tablet", "keyboard", "mouse", "monitor", "headphones", "earbuds", "charger", "USB cable", "power bank", "speaker", "microphone", "webcam", "router", "remote control", "game controller", "smartwatch", "calculator", "hard drive", "flash drive", "printer", "camera", "pen", "pencil", "eraser", "ruler", "notebook", "book", "paper", "stapler", "scissors", "tape", "glue stick", "marker", "highlighter", "folder", "envelope", "paperclip", "sticky note", "chair", "table", "desk", "sofa", "bed", "shelf", "cabinet", "lamp", "fan", "mirror", "curtain", "pillow", "blanket", "bottle", "cup", "mug", "glass", "plate", "bowl", "spoon", "fork", "knife", "pan", "pot", "kettle", "toaster", "blender", "microwave", "refrigerator", "stove", "cutting board", "thermos", "apple", "banana", "orange", "bread", "sandwich", "pizza", "burger", "rice", "chips", "biscuit", "chocolate", "candy", "cake", "water bottle", "juice box", "soda can", "coffee cup", "shirt", "jacket", "hoodie", "jeans", "shoe", "sneaker", "hat", "cap", "sunglasses", "watch", "wallet", "bag", "backpack", "belt", "tie", "scarf", "gloves", "toothbrush", "comb", "soap", "towel", "tissue box", "umbrella", "key", "keychain", "ID card", "coin", "ball", "football", "basketball", "cricket bat", "tennis racket", "badminton racket", "yoga mat", "dumbbell", "toy car", "stuffed animal", "rubik's cube", "playing cards", "hammer", "screwdriver", "wrench", "pliers", "drill", "measuring tape", "nail", "screw", "bicycle", "helmet", "plant", "flower", "leaf", "rock", "clock", "calendar", "photo frame", "vase", "candle", "lighter", "dustbin", "broom", "bucket", "hanger", "iron", "battery", "bulb", "medicine box", "bandage", "thermometer", "mask", "box", "container", "jar", "tin can", "plastic bag", "rope", "wire", "chain", "magnet", "compass", "globe", "trophy", "medal", "airpods", "earphone case", "phone case", "screen protector", "mouse pad", "laptop stand", "ring light", "tripod", ] # ------------------------- # Threaded Camera # ------------------------- class FastCapture: def __init__(self, src=0, width=1280, height=720): self.cap = cv2.VideoCapture(src) self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width) self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height) self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) self.ret = False self.frame = None self.lock = threading.Lock() self.running = True self.ret, self.frame = self.cap.read() self.thread = threading.Thread(target=self._update, daemon=True) self.thread.start() def _update(self): while self.running: ret, frame = self.cap.read() with self.lock: self.ret = ret self.frame = frame def read(self): with self.lock: if self.frame is not None: return self.ret, self.frame.copy() return False, None def release(self): self.running = False self.thread.join(timeout=2) self.cap.release() # =================================================== # LEARNED OBJECTS DATABASE (Persistent) # =================================================== # Each object: {name, embedding (.npy file), image} # Embeddings saved as .npy files so they survive restarts # without needing CLIP to reload them. # =================================================== learned_objects = [] # [{name, embedding (numpy), image_path}] def load_learned_objects(): """Load learned objects from disk. Embeddings stored as .npy files.""" global learned_objects learned_objects = [] if not os.path.exists(OBJECTS_JSON): return with open(OBJECTS_JSON, "r") as f: entries = json.load(f) for entry in entries: name = entry["name"] image_path = entry.get("image", "") embedding_path = entry.get("embedding_file", "") # Load saved embedding from .npy file if embedding_path and os.path.exists(embedding_path): embedding = np.load(embedding_path) learned_objects.append({ "name": name, "embedding": embedding, "image_path": image_path }) print(f" [OK] Loaded: {name}") else: print(f" [SKIP] No embedding file for: {name}") print(f"Loaded {len(learned_objects)} custom objects.\n") def get_clip_embedding(pil_image): """Get CLIP embedding for a PIL image.""" if not clip_loaded: return None try: inputs = clip_processor(images=pil_image, return_tensors="pt").to(device) with torch.no_grad(): emb = clip_model.get_image_features(**inputs) emb = emb / emb.norm(dim=-1, keepdim=True) return emb.cpu().numpy()[0] except Exception as e: print(f"CLIP error: {e}") return None def save_learned_object(name, frame, bbox): """ Save a new learned object: - Crop image saved as .jpg - CLIP embedding saved as .npy - Entry added to objects_db.json """ x1, y1, x2, y2 = bbox h, w = frame.shape[:2] pad = 10 x1, y1 = max(0, x1 - pad), max(0, y1 - pad) x2, y2 = min(w, x2 + pad), min(h, y2 + pad) crop = frame[y1:y2, x1:x2] if crop.size == 0: return False timestamp = int(time.time()) image_path = os.path.join(OBJECTS_FOLDER, f"{name}_{timestamp}.jpg") embedding_path = os.path.join(OBJECTS_FOLDER, f"{name}_{timestamp}.npy") # Save crop image cv2.imwrite(image_path, crop) # Compute and save CLIP embedding crop_rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB) pil_crop = Image.fromarray(crop_rgb) if clip_loaded: embedding = get_clip_embedding(pil_crop) if embedding is not None: np.save(embedding_path, embedding) # Add to memory learned_objects.append({ "name": name, "embedding": embedding, "image_path": image_path }) else: # Save without embedding (image only) embedding_path = "" else: # No CLIP loaded — save image only, embedding computed later embedding_path = "" # Update JSON with open(OBJECTS_JSON, "r") as f: entries = json.load(f) entries.append({ "name": name, "image": image_path, "embedding_file": embedding_path }) with open(OBJECTS_JSON, "w") as f: json.dump(entries, f, indent=4) print(f"Saved object: {name} -> {image_path}") return True def match_custom_object(pil_image): """Match image against learned objects using CLIP embeddings.""" if len(learned_objects) == 0 or not clip_loaded: return None, 0.0 embedding = get_clip_embedding(pil_image) if embedding is None: return None, 0.0 best_name = None best_score = -1 for obj in learned_objects: if obj["embedding"] is not None: sim = np.dot(embedding, obj["embedding"]) if sim > best_score: best_score = sim best_name = obj["name"] if best_score >= CUSTOM_MATCH_THRESHOLD: return best_name, float(best_score) return None, float(best_score) def identify_from_internet(pil_image, top_k=5): """CLIP zero-shot classification against vocabulary.""" if not clip_loaded: return [] try: texts = [f"a photo of a {obj}" for obj in CLIP_VOCABULARY] inputs_img = clip_processor(images=pil_image, return_tensors="pt").to(device) with torch.no_grad(): img_feat = clip_model.get_image_features(**inputs_img) img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True) all_scores = [] batch_size = 50 for i in range(0, len(texts), batch_size): batch = texts[i:i + batch_size] inputs_txt = clip_processor(text=batch, return_tensors="pt", padding=True, truncation=True).to(device) with torch.no_grad(): txt_feat = clip_model.get_text_features(**inputs_txt) txt_feat = txt_feat / txt_feat.norm(dim=-1, keepdim=True) sims = (img_feat @ txt_feat.T).squeeze(0) all_scores.extend(sims.cpu().numpy().tolist()) results = sorted(zip(CLIP_VOCABULARY, all_scores), key=lambda x: x[1], reverse=True) return results[:top_k] except Exception as e: print(f"CLIP identify error: {e}") return [] # =================================================== # DRAWING FUNCTIONS # =================================================== def draw_detections(frame, detections, fps, selected_idx): for i, det in enumerate(detections): name = det["name"] score = det["score"] x1, y1, x2, y2 = det["bbox"] source = det["source"] if source == "custom": color = (255, 180, 0) # Blue for learned label = f"[L] {name} {score:.0%}" else: color = (0, 255, 0) # Green for YOLO label = f"{name} {score:.0%}" if i == selected_idx: color = (0, 255, 255) # Yellow highlight thickness = 3 else: thickness = 2 cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness) (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2) cv2.rectangle(frame, (x1, y1 - th - 10), (x1 + tw + 6, y1), color, cv2.FILLED) cv2.putText(frame, label, (x1 + 3, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2) # FPS cv2.putText(frame, f"FPS: {fps:.1f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2) # Stats cv2.putText(frame, f"Objects: {len(detections)} | Learned: {len(learned_objects)}", (10, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1) # Controls help_text = "Tab=Select | T=Teach | I=Auto-ID | ESC=Quit" cv2.putText(frame, help_text, (10, frame.shape[0] - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1) return frame def draw_input_mode(frame, prompt, typed_text, suggestions=None): """Draw text input overlay on the camera frame.""" overlay = frame.copy() h, w = frame.shape[:2] # Dark overlay at bottom cv2.rectangle(overlay, (0, h - 120), (w, h), (30, 30, 30), cv2.FILLED) frame = cv2.addWeighted(overlay, 0.85, frame, 0.15, 0) # Prompt cv2.putText(frame, prompt, (15, h - 90), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 200, 255), 1) # Typed text with cursor display_text = typed_text + "|" cv2.putText(frame, display_text, (15, h - 55), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2) # Hint cv2.putText(frame, "Enter=Confirm | Esc=Cancel", (15, h - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (150, 150, 150), 1) # Show suggestions if available if suggestions: for i, (name, score) in enumerate(suggestions): y = h - 130 - (len(suggestions) - i) * 25 text = f"{i+1}. {name} ({score:.1%})" cv2.putText(frame, text, (15, y), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (100, 255, 100), 1) label_y = h - 130 - len(suggestions) * 25 - 10 cv2.putText(frame, "Type number to pick, or type custom name:", (15, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 200, 255), 1) return frame # =================================================== # MAIN DETECTION LOOP # =================================================== def run_detection(): cap = FastCapture(WEBCAM_INDEX, FRAME_WIDTH, FRAME_HEIGHT) selected_idx = -1 fps = 0.0 prev_time = time.time() frame_count = 0 cached_detections = [] # Input mode state input_mode = None # None, "teach", or "identify" typed_text = "" frozen_frame = None frozen_bbox = None clip_suggestions = None # For auto-identify results print("\nDetection started!") print(" Tab = Select | T = Teach | I = Auto-ID | ESC = Quit\n") while True: # --- INPUT MODE: typing on screen --- if input_mode is not None: display = frozen_frame.copy() # Highlight the selected object if frozen_bbox: x1, y1, x2, y2 = frozen_bbox cv2.rectangle(display, (x1, y1), (x2, y2), (0, 255, 255), 3) if input_mode == "teach": display = draw_input_mode(display, "What is this object? Type name:", typed_text) elif input_mode == "identify": display = draw_input_mode(display, "Pick a number or type custom name:", typed_text, clip_suggestions) cv2.imshow("Smart Object Identifier", display) key = cv2.waitKey(30) & 0xFF if key == 27: # ESC — cancel input_mode = None typed_text = "" clip_suggestions = None continue elif key == 13: # ENTER — confirm final_name = typed_text.strip() # For identify mode, check if number was typed if input_mode == "identify" and clip_suggestions: try: idx = int(final_name) - 1 if 0 <= idx < len(clip_suggestions): final_name = clip_suggestions[idx][0] except ValueError: pass if final_name and frozen_bbox: success = save_learned_object(final_name, frozen_frame, frozen_bbox) if success: print(f" >> Learned: {final_name}") input_mode = None typed_text = "" clip_suggestions = None continue elif key == 8: # BACKSPACE typed_text = typed_text[:-1] elif 32 <= key <= 126: # Printable characters typed_text += chr(key) continue # Skip normal detection while typing # --- NORMAL MODE: detect objects --- ret, frame = cap.read() if not ret or frame is None: print("Camera read failed.") break frame_count += 1 # Run YOLO every 2 frames if frame_count % 2 == 0: results = yolo_model(frame, conf=YOLO_CONFIDENCE, verbose=False) detections = [] for r in results: boxes = r.boxes for box in boxes: x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int) conf = float(box.conf[0]) cls_id = int(box.cls[0]) cls_name = yolo_model.names[cls_id] # Try custom match if CLIP is loaded custom_name = None crop = frame[y1:y2, x1:x2] if crop.size > 0 and len(learned_objects) > 0 and clip_loaded: crop_rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB) pil_crop = Image.fromarray(crop_rgb) custom_name, custom_score = match_custom_object(pil_crop) if custom_name: detections.append({ "name": custom_name, "score": custom_score, "bbox": [x1, y1, x2, y2], "source": "custom" }) else: detections.append({ "name": cls_name, "score": conf, "bbox": [x1, y1, x2, y2], "source": "yolo" }) cached_detections = detections if selected_idx >= len(cached_detections): selected_idx = max(0, len(cached_detections) - 1) if cached_detections else -1 # FPS now = time.time() dt = now - prev_time if dt > 0: fps = 1.0 / dt prev_time = now # Draw display = draw_detections(frame, cached_detections, fps, selected_idx) cv2.imshow("Smart Object Identifier", display) key = cv2.waitKey(1) & 0xFF if key == 27: # ESC break elif key == 9: # Tab if len(cached_detections) > 0: selected_idx = (selected_idx + 1) % len(cached_detections) else: selected_idx = -1 elif key == ord('t') or key == ord('T'): if 0 <= selected_idx < len(cached_detections): # Freeze frame and enter teach mode frozen_frame = frame.copy() frozen_bbox = cached_detections[selected_idx]["bbox"] input_mode = "teach" typed_text = "" print(" [TEACH MODE] Type the object name on screen, press Enter") else: print(" Press Tab first to select an object!") elif key == ord('i') or key == ord('I'): if 0 <= selected_idx < len(cached_detections): if not clip_loaded: print(" CLIP not loaded! Click '🌐 Load CLIP' in dashboard first.") else: # Freeze and run CLIP frozen_frame = frame.copy() frozen_bbox = cached_detections[selected_idx]["bbox"] bx1, by1, bx2, by2 = frozen_bbox crop = frame[by1:by2, bx1:bx2] if crop.size > 0: crop_rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB) pil_crop = Image.fromarray(crop_rgb) print(" [AUTO-ID] Running CLIP analysis...") clip_suggestions = identify_from_internet(pil_crop, top_k=5) if clip_suggestions: input_mode = "identify" typed_text = "" print(" CLIP suggestions ready — pick a number or type name") else: print(" CLIP returned no results.") frozen_frame = None frozen_bbox = None else: print(" Press Tab first to select an object!") cap.release() cv2.destroyAllWindows() print("Detection stopped.\n") # =============================================== # GUI DASHBOARD # =============================================== ctk.set_appearance_mode("dark") ctk.set_default_color_theme("blue") app_ui = ctk.CTk() app_ui.title("Smart Object Identifier") app_ui.geometry("700x700") app_ui.resizable(False, False) def refresh_objects_list(): objects_box.configure(state="normal") objects_box.delete("0.0", "end") if not os.path.exists(OBJECTS_JSON): objects_box.insert("end", " No objects learned yet.\n") objects_box.configure(state="disabled") return with open(OBJECTS_JSON, "r") as f: entries = json.load(f) if len(entries) == 0: objects_box.insert("end", " No objects learned yet.\n") else: counts = {} for e in entries: n = e["name"] counts[n] = counts.get(n, 0) + 1 for name, count in counts.items(): if count > 1: objects_box.insert("end", f" • {name} [{count} samples]\n") else: objects_box.insert("end", f" • {name}\n") objects_box.configure(state="disabled") def delete_object(): name = delete_entry.get().strip() if not name: messagebox.showerror("Error", "Enter an object name.") return with open(OBJECTS_JSON, "r") as f: entries = json.load(f) new_entries = [] removed = False for entry in entries: if entry["name"].lower() == name.lower(): for key in ["image", "embedding_file"]: path = entry.get(key, "") if path and os.path.exists(path): os.remove(path) removed = True else: new_entries.append(entry) if not removed: messagebox.showerror("Error", f"'{name}' not found.") return with open(OBJECTS_JSON, "w") as f: json.dump(new_entries, f, indent=4) global learned_objects learned_objects = [o for o in learned_objects if o["name"].lower() != name.lower()] messagebox.showinfo("Deleted", f"'{name}' removed!") refresh_objects_list() delete_entry.delete(0, "end") def on_start(): thread = threading.Thread(target=run_detection, daemon=True) thread.start() def on_load_clip(): def _load(): if load_clip(): load_learned_objects() app_ui.after(0, refresh_objects_list) app_ui.after(0, lambda: messagebox.showinfo("CLIP Ready", "CLIP loaded! Internet recognition & custom matching enabled.")) else: app_ui.after(0, lambda: messagebox.showerror("Error", "Failed to load CLIP. Check internet connection.")) thread = threading.Thread(target=_load, daemon=True) thread.start() def on_rebuild_embeddings(): """Recompute CLIP embeddings for all saved objects that don't have one.""" if not clip_loaded: messagebox.showerror("Error", "Load CLIP first!") return with open(OBJECTS_JSON, "r") as f: entries = json.load(f) updated = 0 for entry in entries: emb_file = entry.get("embedding_file", "") if not emb_file or not os.path.exists(emb_file): img_path = entry.get("image", "") if img_path and os.path.exists(img_path): img = Image.open(img_path).convert("RGB") emb = get_clip_embedding(img) if emb is not None: emb_path = img_path.replace(".jpg", ".npy").replace(".png", ".npy") np.save(emb_path, emb) entry["embedding_file"] = emb_path updated += 1 print(f" Rebuilt embedding for: {entry['name']}") with open(OBJECTS_JSON, "w") as f: json.dump(entries, f, indent=4) load_learned_objects() refresh_objects_list() messagebox.showinfo("Done", f"Rebuilt {updated} embeddings!") # =============================================== # BUILD GUI # =============================================== # Title ctk.CTkLabel(app_ui, text="🔍 Smart Object Identifier", font=("Arial", 26, "bold")).pack(pady=(20, 5)) ctk.CTkLabel(app_ui, text="YOLOv8 + CLIP Learning + Internet Recognition", font=("Arial", 12), text_color="gray").pack(pady=(0, 15)) # --- Controls --- ctrl = ctk.CTkFrame(app_ui, corner_radius=10) ctrl.pack(padx=20, pady=5, fill="x") ctk.CTkLabel(ctrl, text="Controls", font=("Arial", 16, "bold")).pack(pady=(10, 5)) ctk.CTkLabel(ctrl, text="Tab=Select | T=Teach | I=Auto-Identify | ESC=Quit", font=("Arial", 12), text_color="lightgray").pack(pady=(0, 5)) btn_row = ctk.CTkFrame(ctrl, fg_color="transparent") btn_row.pack(pady=(5, 10), fill="x", padx=20) ctk.CTkButton(btn_row, text="🚀 Start Detection", width=200, height=45, font=("Arial", 14, "bold"), fg_color="#3498db", hover_color="#2980b9", command=on_start).pack(side="left", padx=10, expand=True) ctk.CTkButton(btn_row, text="🌐 Load CLIP (Internet)", width=200, height=45, font=("Arial", 14, "bold"), fg_color="#8e44ad", hover_color="#7d3c98", command=on_load_clip).pack(side="left", padx=10, expand=True) # --- How to Use --- howto = ctk.CTkFrame(app_ui, corner_radius=10) howto.pack(padx=20, pady=5, fill="x") ctk.CTkLabel(howto, text="How It Works", font=("Arial", 16, "bold")).pack(pady=(10, 5)) ctk.CTkLabel(howto, text=( "1. Click 'Start Detection' — camera opens, YOLO detects objects\n" "2. Press Tab to highlight/select an object (yellow box)\n" "3. Press T — type the object name ON THE CAMERA, press Enter\n" "4. Press I — AI guesses from internet, pick a number or type name\n" "5. Green = YOLO known | Blue = Your taught objects" ), font=("Arial", 11), justify="left", text_color="lightgray").pack(padx=15, pady=(0, 10)) # --- Delete --- del_frame = ctk.CTkFrame(app_ui, corner_radius=10) del_frame.pack(padx=20, pady=5, fill="x") delete_entry = ctk.CTkEntry(del_frame, width=300, height=35, placeholder_text="Enter Object Name to Delete") delete_entry.pack(side="left", padx=(15, 5), pady=10) ctk.CTkButton(del_frame, text="🗑️ Delete", width=100, height=35, fg_color="#e74c3c", hover_color="#c0392b", command=delete_object).pack(side="left", padx=(5, 10), pady=10) ctk.CTkButton(del_frame, text="🔄 Rebuild", width=100, height=35, fg_color="#f39c12", hover_color="#e67e22", command=on_rebuild_embeddings).pack(side="left", padx=(0, 15), pady=10) # --- Learned Objects List --- ctk.CTkLabel(app_ui, text="Learned Objects", font=("Arial", 16, "bold")).pack(pady=(15, 5)) objects_box = ctk.CTkTextbox(app_ui, width=400, height=140, corner_radius=8) objects_box.pack(pady=5) # Info ctk.CTkLabel(app_ui, text=f"YOLOv8: {len(yolo_model.names)} classes | CLIP: {len(CLIP_VOCABULARY)} categories", font=("Arial", 11), text_color="gray").pack(pady=(10, 5)) # ------------------------- # Startup # ------------------------- # Load objects from disk (embeddings) load_learned_objects() refresh_objects_list() print("Dashboard ready!\n") app_ui.mainloop()