mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
feat(scenegraph): refactor image scene graph query into reusable function
- Export ImageSceneGraphQueryResult and query_image_against_scene_graph from scenegraph module - Replace inline hamming-distance-based image matching with dedicated query_image_against_scene_graph function - Improve top_matches structure by extracting similarity scores and hash_bytes from matches - Add .codegraph/ to gitignore (machine-local data, should not be committed) - Add CodeGraph configuration for multi-language indexing
This commit is contained in:
57
mini-nav/scenegraph/query.py
Normal file
57
mini-nav/scenegraph/query.py
Normal file
@@ -0,0 +1,57 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from .hash_codec import bits_tensor_to_hash_bytes
|
||||
from .scenegraph import SceneGraphMatch, SimpleSceneGraph
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageSceneGraphQueryResult:
|
||||
query_crop_index: int
|
||||
query_hash: bytes
|
||||
query_crop: Image.Image
|
||||
matches: list[SceneGraphMatch]
|
||||
|
||||
|
||||
def query_image_against_scene_graph(
|
||||
image: Image.Image,
|
||||
pipeline: Any,
|
||||
scene_graph: SimpleSceneGraph,
|
||||
text_labels: list[str],
|
||||
*,
|
||||
top_k: int = 1,
|
||||
batch_size: int = 1,
|
||||
) -> list[ImageSceneGraphQueryResult]:
|
||||
output = pipeline.process_batch([image], text_labels, batch_size=batch_size)
|
||||
hash_bits = output.hash_bits
|
||||
cropped_images = list(output.cropped_images)
|
||||
|
||||
if hash_bits.numel() == 0:
|
||||
if cropped_images:
|
||||
raise ValueError("hash_bits and cropped_images must align")
|
||||
return []
|
||||
|
||||
if hash_bits.dim() == 1:
|
||||
hash_bits = hash_bits.unsqueeze(0)
|
||||
|
||||
if hash_bits.shape[0] != len(cropped_images):
|
||||
raise ValueError("hash_bits and cropped_images must align")
|
||||
|
||||
results: list[ImageSceneGraphQueryResult] = []
|
||||
for crop_index, query_bits in enumerate(hash_bits):
|
||||
query_hash = bits_tensor_to_hash_bytes(query_bits)
|
||||
matches = scene_graph.query_by_visual_hash(query_hash, top_k=top_k)
|
||||
results.append(
|
||||
ImageSceneGraphQueryResult(
|
||||
query_crop_index=crop_index,
|
||||
query_hash=query_hash,
|
||||
query_crop=cropped_images[crop_index],
|
||||
matches=matches,
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
Reference in New Issue
Block a user