mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
feat(compressors): add OWLv2 bbox crop to HashPipeline and refactor image utilities
- Add Owlv2ForObjectDetection and Owlv2Processor imports to model_loader - Refactor load_dino_model to return tuple of processor and model - Rewrite generate_proposals_batch to group images by bbox count for efficient batching - Add _normalize_single_bbox_list helper for bbox normalization - Update verification.py to use new pipeline architecture with detect/segment/filter/crop steps
This commit is contained in:
@@ -42,12 +42,13 @@ def generate_proposals(
|
||||
- predicted_iou: Model's confidence in the mask
|
||||
- stability_score: Stability score for the mask
|
||||
"""
|
||||
if not bboxes:
|
||||
normalized_bboxes = _normalize_single_bbox_list(bboxes)
|
||||
if not normalized_bboxes:
|
||||
return []
|
||||
|
||||
device = get_device()
|
||||
image_rgb = image.convert("RGB")
|
||||
input_boxes = [bboxes]
|
||||
input_boxes = [normalized_bboxes]
|
||||
inputs = processor(
|
||||
images=image_rgb,
|
||||
input_boxes=input_boxes,
|
||||
@@ -84,22 +85,68 @@ def generate_proposals_batch(
|
||||
if not image_list:
|
||||
return []
|
||||
|
||||
if len(image_list) != len(bboxes_per_image):
|
||||
raise ValueError(
|
||||
f"Length mismatch: {len(image_list)} images, {len(bboxes_per_image)} bbox groups"
|
||||
)
|
||||
|
||||
normalized_bboxes_per_image = [
|
||||
_normalize_single_bbox_list(bboxes) for bboxes in bboxes_per_image
|
||||
]
|
||||
|
||||
proposals_per_image: list[list[dict[str, Any]]] = [[] for _ in image_list]
|
||||
valid_indices = [
|
||||
index for index, bboxes in enumerate(normalized_bboxes_per_image) if bboxes
|
||||
]
|
||||
if not valid_indices:
|
||||
return proposals_per_image
|
||||
|
||||
device = get_device()
|
||||
image_rgb_list = [img.convert("RGB") for img in image_list]
|
||||
grouped_indices: dict[int, list[int]] = {}
|
||||
for image_index in valid_indices:
|
||||
box_count = len(normalized_bboxes_per_image[image_index])
|
||||
grouped_indices.setdefault(box_count, []).append(image_index)
|
||||
|
||||
inputs = processor(
|
||||
images=image_rgb_list,
|
||||
input_boxes=bboxes_per_image,
|
||||
return_tensors="pt",
|
||||
).to(device)
|
||||
for grouped_image_indices in grouped_indices.values():
|
||||
image_rgb_list = [
|
||||
image_list[image_index].convert("RGB")
|
||||
for image_index in grouped_image_indices
|
||||
]
|
||||
batched_input_boxes = [
|
||||
normalized_bboxes_per_image[image_index]
|
||||
for image_index in grouped_image_indices
|
||||
]
|
||||
|
||||
outputs = model(**inputs, multimask_output=False)
|
||||
all_masks = processor.post_process_masks(
|
||||
outputs.pred_masks.cpu(),
|
||||
inputs["original_sizes"],
|
||||
)
|
||||
inputs = processor(
|
||||
images=image_rgb_list,
|
||||
input_boxes=batched_input_boxes,
|
||||
return_tensors="pt",
|
||||
).to(device)
|
||||
|
||||
return [_masks_to_proposals(image_masks) for image_masks in all_masks]
|
||||
outputs = model(**inputs, multimask_output=False)
|
||||
all_masks = processor.post_process_masks(
|
||||
outputs.pred_masks.cpu(),
|
||||
inputs["original_sizes"],
|
||||
)
|
||||
|
||||
batched_proposals = [
|
||||
_masks_to_proposals(image_masks) for image_masks in all_masks
|
||||
]
|
||||
for output_index, image_index in enumerate(grouped_image_indices):
|
||||
proposals_per_image[image_index] = batched_proposals[output_index]
|
||||
|
||||
return proposals_per_image
|
||||
|
||||
|
||||
def _normalize_single_bbox_list(bboxes: Sequence[Sequence[float]]) -> list[list[float]]:
|
||||
normalized: list[list[float]] = []
|
||||
for bbox in bboxes:
|
||||
bbox_list = list(bbox)
|
||||
if len(bbox_list) != 4:
|
||||
continue
|
||||
normalized.append([float(value) for value in bbox_list])
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def detect_objects(
|
||||
|
||||
Reference in New Issue
Block a user