mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
SikongJueluo
101a675ece
- Introduce ImageHashPipeline Protocol for extensible hash computation - Rename query_crop_index to query_index for clarity - Make query_crop nullable to handle missing crop edge case - Add keyword-only arguments for better API clarity - Handle empty scene graph objects gracefully - Add comprehensive test coverage for query_image_against_scene_graph
77 lines
2.1 KiB
Python
77 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
from PIL import Image
|
|
|
|
from .hash_codec import bits_tensor_to_hash_bytes
|
|
from .scenegraph import SceneGraphMatch, SimpleSceneGraph
|
|
|
|
|
|
class ImageHashPipelineOutput(Protocol):
|
|
"""Protocol for the output of ImageHashPipeline.process_batch."""
|
|
|
|
hash_bits: Any # torch.Tensor
|
|
cropped_images: list[Image.Image]
|
|
|
|
|
|
@runtime_checkable
|
|
class ImageHashPipeline(Protocol):
|
|
"""Protocol for an image hash computation pipeline."""
|
|
|
|
def process_batch(
|
|
self,
|
|
images: list[Image.Image],
|
|
text_labels: list[str],
|
|
batch_size: int = 32,
|
|
) -> ImageHashPipelineOutput:
|
|
...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ImageSceneGraphQueryResult:
|
|
query_index: int
|
|
query_hash: bytes
|
|
query_crop: Image.Image | None
|
|
matches: list[SceneGraphMatch]
|
|
|
|
|
|
def query_image_against_scene_graph(
|
|
image: Image.Image,
|
|
*,
|
|
pipeline: ImageHashPipeline,
|
|
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:
|
|
return []
|
|
|
|
if hash_bits.dim() == 1:
|
|
hash_bits = hash_bits.unsqueeze(0)
|
|
|
|
results: list[ImageSceneGraphQueryResult] = []
|
|
for query_index, query_bits in enumerate(hash_bits):
|
|
query_hash = bits_tensor_to_hash_bytes(query_bits)
|
|
if scene_graph.objects:
|
|
matches = scene_graph.query_by_visual_hash(query_hash, top_k=top_k)
|
|
else:
|
|
matches = []
|
|
crop = cropped_images[query_index] if query_index < len(cropped_images) else None
|
|
results.append(
|
|
ImageSceneGraphQueryResult(
|
|
query_index=query_index,
|
|
query_hash=query_hash,
|
|
query_crop=crop,
|
|
matches=matches,
|
|
)
|
|
)
|
|
|
|
return results
|