File size: 5,374 Bytes
4346a4c 9cfe107 4346a4c 9cfe107 4346a4c | 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 | import logging
import cv2
import torch
import torchvision.transforms as T
import torchvision.models as models
from PIL import Image
import numpy as np
from typing import Optional, List, Dict, Any, Tuple
from backend.config import DEVICE, SPECIFIC_SIMILARITY_THRESHOLD
from backend.database.storage import db
logger = logging.getLogger(__name__)
class EmbeddingRecognizer:
def __init__(self, model_name: str = "resnet18"):
self.model_name = model_name
self.model = None
self.transform = T.Compose([
T.Resize((224, 224)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
self.load_model()
def load_model(self):
try:
logger.info("Loading PyTorch vision backbone for Specific Object Recognition...")
backbone = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
# Remove classification layer to get 512-dim feature embedding vector
self.model = torch.nn.Sequential(*list(backbone.children())[:-1])
self.model.eval()
self.model.to(DEVICE)
logger.info("Specific Object Embedding backbone loaded successfully.")
except Exception as e:
logger.error(f"Error loading embedding backbone model: {e}")
self.model = None
def extract_embedding_from_crop(self, crop_np: np.ndarray) -> Optional[np.ndarray]:
if self.model is None or crop_np is None or crop_np.size == 0:
return None
try:
pil_img = Image.fromarray(crop_np).convert("RGB")
tensor = self.transform(pil_img).unsqueeze(0).to(DEVICE)
with torch.no_grad():
feat = self.model(tensor)
feat = feat.squeeze().cpu().numpy()
# L2 normalize vector
norm = np.linalg.norm(feat)
if norm > 0:
feat = feat / norm
return feat
except Exception as e:
logger.error(f"Failed to extract embedding from crop: {e}")
return None
def extract_embedding_from_image(self, img_np: np.ndarray) -> Optional[np.ndarray]:
if self.model is None or img_np is None or img_np.size == 0:
return None
try:
rgb_img = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB) if len(img_np.shape) == 3 and img_np.shape[2] == 3 else img_np
pil_img = Image.fromarray(rgb_img).convert("RGB")
tensor = self.transform(pil_img).unsqueeze(0).to(DEVICE)
with torch.no_grad():
feat = self.model(tensor)
feat = feat.squeeze().cpu().numpy()
norm = np.linalg.norm(feat)
if norm > 0:
feat = feat / norm
return feat
except Exception as e:
logger.error(f"Failed to extract embedding from image: {e}")
return None
def extract_embedding_from_image_path(self, image_path: str) -> Optional[np.ndarray]:
try:
pil_img = Image.open(image_path).convert("RGB")
tensor = self.transform(pil_img).unsqueeze(0).to(DEVICE)
with torch.no_grad():
feat = self.model(tensor)
feat = feat.squeeze().cpu().numpy()
norm = np.linalg.norm(feat)
if norm > 0:
feat = feat / norm
return feat
except Exception as e:
logger.error(f"Failed to extract embedding from file {image_path}: {e}")
return None
def recognize_crop(self, crop_np: np.ndarray, threshold: float = SPECIFIC_SIMILARITY_THRESHOLD) -> Tuple[Optional[str], float]:
emb = self.extract_embedding_from_crop(crop_np)
if emb is None:
return None, 0.0
matched_name, similarity = db.find_best_matching_object(emb, min_threshold=threshold)
return matched_name, similarity
def refine_detections_with_specific_objects(
self, image_np: np.ndarray, candidate_detections: List[Dict[str, Any]], threshold: float = SPECIFIC_SIMILARITY_THRESHOLD
) -> List[Dict[str, Any]]:
"""Scans candidate bounding boxes and overrides label if matched to a specific user-taught object."""
if image_np is None or not candidate_detections:
return candidate_detections
h, w = image_np.shape[:2]
refined = []
for det in candidate_detections:
bbox = det["bbox"]
x1, y1, x2, y2 = max(0, bbox[0]), max(0, bbox[1]), min(w, bbox[2]), min(h, bbox[3])
if (x2 - x1) > 10 and (y2 - y1) > 10:
crop = image_np[y1:y2, x1:x2]
matched_name, similarity = self.recognize_crop(crop, threshold=threshold)
if matched_name is not None:
# Specific Object Identity Matched!
det_copy = dict(det)
det_copy["label"] = matched_name
det_copy["specific_identity"] = matched_name
det_copy["type"] = "specific"
det_copy["source"] = "embedding"
det_copy["confidence"] = round(similarity, 4)
refined.append(det_copy)
continue
refined.append(det)
return refined
|