mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
ba96cec406
- 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
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
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
|