support more demos with florence-2
This commit is contained in:
rentainhe
@@ -1,6 +1,7 @@
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
import argparse
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from PIL import Image
|
||||
@@ -127,7 +128,7 @@ def object_detection_and_segmentation(
|
||||
}
|
||||
}
|
||||
"""
|
||||
results = results["<OD>"]
|
||||
results = results[task_prompt]
|
||||
# parse florence-2 detection results
|
||||
input_boxes = np.array(results["bboxes"])
|
||||
class_names = results["labels"]
|
||||
@@ -163,22 +164,286 @@ def object_detection_and_segmentation(
|
||||
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
|
||||
cv2.imwrite(os.path.join(OUTPUT_DIR, "grounded_sam2_florence2_det_annotated_image.jpg"), annotated_frame)
|
||||
cv2.imwrite(os.path.join(output_dir, "grounded_sam2_florence2_det_annotated_image.jpg"), annotated_frame)
|
||||
|
||||
mask_annotator = sv.MaskAnnotator()
|
||||
annotated_frame = mask_annotator.annotate(scene=annotated_frame, detections=detections)
|
||||
cv2.imwrite(os.path.join(OUTPUT_DIR, "grounded_sam2_florence2_det_image_with_mask.jpg"), annotated_frame)
|
||||
cv2.imwrite(os.path.join(output_dir, "grounded_sam2_florence2_det_image_with_mask.jpg"), annotated_frame)
|
||||
|
||||
print(f'Successfully save annotated image to "{output_dir}"')
|
||||
|
||||
"""
|
||||
Pipeline 2: Dense Region Caption + Segmentation
|
||||
"""
|
||||
def dense_region_caption_and_segmentation(
|
||||
florence2_model,
|
||||
florence2_processor,
|
||||
sam2_predictor,
|
||||
image_path,
|
||||
task_prompt="<DENSE_REGION_CAPTION>",
|
||||
text_input=None,
|
||||
output_dir=OUTPUT_DIR
|
||||
):
|
||||
# run florence-2 object detection in demo
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
results = run_florence2(task_prompt, text_input, florence2_model, florence2_processor, image)
|
||||
|
||||
""" Florence-2 Object Detection Output Format
|
||||
{'<DENSE_REGION_CAPTION>':
|
||||
{
|
||||
'bboxes':
|
||||
[
|
||||
[33.599998474121094, 159.59999084472656, 596.7999877929688, 371.7599792480469],
|
||||
[454.0799865722656, 96.23999786376953, 580.7999877929688, 261.8399963378906],
|
||||
[224.95999145507812, 86.15999603271484, 333.7599792480469, 164.39999389648438],
|
||||
[449.5999755859375, 276.239990234375, 554.5599975585938, 370.3199768066406],
|
||||
[91.19999694824219, 280.0799865722656, 198.0800018310547, 370.3199768066406]
|
||||
],
|
||||
'labels': ['turquoise Volkswagen Beetle', 'wooden double doors with metal handles', 'wheel', 'wheel', 'door']
|
||||
}
|
||||
}
|
||||
"""
|
||||
results = results[task_prompt]
|
||||
# parse florence-2 detection results
|
||||
input_boxes = np.array(results["bboxes"])
|
||||
class_names = results["labels"]
|
||||
class_ids = np.array(list(range(len(class_names))))
|
||||
|
||||
# predict mask with SAM 2
|
||||
sam2_predictor.set_image(np.array(image))
|
||||
masks, scores, logits = sam2_predictor.predict(
|
||||
point_coords=None,
|
||||
point_labels=None,
|
||||
box=input_boxes,
|
||||
multimask_output=False,
|
||||
)
|
||||
|
||||
if masks.ndim == 4:
|
||||
masks = masks.squeeze(1)
|
||||
|
||||
# specify labels
|
||||
labels = [
|
||||
f"{class_name}" for class_name in class_names
|
||||
]
|
||||
|
||||
# visualization results
|
||||
img = cv2.imread(image_path)
|
||||
detections = sv.Detections(
|
||||
xyxy=input_boxes,
|
||||
mask=masks.astype(bool),
|
||||
class_id=class_ids
|
||||
)
|
||||
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
annotated_frame = box_annotator.annotate(scene=img.copy(), detections=detections)
|
||||
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
|
||||
cv2.imwrite(os.path.join(output_dir, "grounded_sam2_florence2_dense_region_cap_annotated_image.jpg"), annotated_frame)
|
||||
|
||||
mask_annotator = sv.MaskAnnotator()
|
||||
annotated_frame = mask_annotator.annotate(scene=annotated_frame, detections=detections)
|
||||
cv2.imwrite(os.path.join(output_dir, "grounded_sam2_florence2_dense_region_cap_image_with_mask.jpg"), annotated_frame)
|
||||
|
||||
print(f'Successfully save annotated image to "{output_dir}"')
|
||||
|
||||
|
||||
"""
|
||||
Pipeline 3: Region Proposal + Segmentation
|
||||
"""
|
||||
def region_proposal_and_segmentation(
|
||||
florence2_model,
|
||||
florence2_processor,
|
||||
sam2_predictor,
|
||||
image_path,
|
||||
task_prompt="<REGION_PROPOSAL>",
|
||||
text_input=None,
|
||||
output_dir=OUTPUT_DIR
|
||||
):
|
||||
# run florence-2 object detection in demo
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
results = run_florence2(task_prompt, text_input, florence2_model, florence2_processor, image)
|
||||
|
||||
""" Florence-2 Object Detection Output Format
|
||||
{'<REGION_PROPOSAL>':
|
||||
{
|
||||
'bboxes':
|
||||
[
|
||||
[33.599998474121094, 159.59999084472656, 596.7999877929688, 371.7599792480469],
|
||||
[454.0799865722656, 96.23999786376953, 580.7999877929688, 261.8399963378906],
|
||||
[224.95999145507812, 86.15999603271484, 333.7599792480469, 164.39999389648438],
|
||||
[449.5999755859375, 276.239990234375, 554.5599975585938, 370.3199768066406],
|
||||
[91.19999694824219, 280.0799865722656, 198.0800018310547, 370.3199768066406]
|
||||
],
|
||||
'labels': ['', '', '', '', '', '', '']
|
||||
}
|
||||
}
|
||||
"""
|
||||
results = results[task_prompt]
|
||||
# parse florence-2 detection results
|
||||
input_boxes = np.array(results["bboxes"])
|
||||
class_names = results["labels"]
|
||||
class_ids = np.array(list(range(len(class_names))))
|
||||
|
||||
# predict mask with SAM 2
|
||||
sam2_predictor.set_image(np.array(image))
|
||||
masks, scores, logits = sam2_predictor.predict(
|
||||
point_coords=None,
|
||||
point_labels=None,
|
||||
box=input_boxes,
|
||||
multimask_output=False,
|
||||
)
|
||||
|
||||
if masks.ndim == 4:
|
||||
masks = masks.squeeze(1)
|
||||
|
||||
# specify labels
|
||||
labels = [
|
||||
f"region_{idx}" for idx, class_name in enumerate(class_names)
|
||||
]
|
||||
|
||||
# visualization results
|
||||
img = cv2.imread(image_path)
|
||||
detections = sv.Detections(
|
||||
xyxy=input_boxes,
|
||||
mask=masks.astype(bool),
|
||||
class_id=class_ids
|
||||
)
|
||||
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
annotated_frame = box_annotator.annotate(scene=img.copy(), detections=detections)
|
||||
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
|
||||
cv2.imwrite(os.path.join(output_dir, "grounded_sam2_florence2_region_proposal.jpg"), annotated_frame)
|
||||
|
||||
mask_annotator = sv.MaskAnnotator()
|
||||
annotated_frame = mask_annotator.annotate(scene=annotated_frame, detections=detections)
|
||||
cv2.imwrite(os.path.join(output_dir, "grounded_sam2_florence2_region_proposal_with_mask.jpg"), annotated_frame)
|
||||
|
||||
print(f'Successfully save annotated image to "{output_dir}"')
|
||||
|
||||
|
||||
"""
|
||||
Pipeline 4: Phrase Grounding + Segmentation
|
||||
"""
|
||||
def phrase_grounding_and_segmentation(
|
||||
florence2_model,
|
||||
florence2_processor,
|
||||
sam2_predictor,
|
||||
image_path,
|
||||
task_prompt="<CAPTION_TO_PHRASE_GROUNDING>",
|
||||
text_input=None,
|
||||
output_dir=OUTPUT_DIR
|
||||
):
|
||||
# run florence-2 object detection in demo
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
results = run_florence2(task_prompt, text_input, florence2_model, florence2_processor, image)
|
||||
|
||||
""" Florence-2 Object Detection Output Format
|
||||
{'<CAPTION_TO_PHRASE_GROUNDING>':
|
||||
{
|
||||
'bboxes':
|
||||
[
|
||||
[34.23999786376953, 159.1199951171875, 582.0800170898438, 374.6399841308594],
|
||||
[1.5999999046325684, 4.079999923706055, 639.0399780273438, 305.03997802734375]
|
||||
],
|
||||
'labels': ['A green car', 'a yellow building']
|
||||
}
|
||||
}
|
||||
"""
|
||||
assert text_input is not None, "Text input should not be none when calling phrase grounding pipeline."
|
||||
results = results[task_prompt]
|
||||
# parse florence-2 detection results
|
||||
input_boxes = np.array(results["bboxes"])
|
||||
class_names = results["labels"]
|
||||
class_ids = np.array(list(range(len(class_names))))
|
||||
|
||||
# predict mask with SAM 2
|
||||
sam2_predictor.set_image(np.array(image))
|
||||
masks, scores, logits = sam2_predictor.predict(
|
||||
point_coords=None,
|
||||
point_labels=None,
|
||||
box=input_boxes,
|
||||
multimask_output=False,
|
||||
)
|
||||
|
||||
if masks.ndim == 4:
|
||||
masks = masks.squeeze(1)
|
||||
|
||||
# specify labels
|
||||
labels = [
|
||||
f"{class_name}" for class_name in class_names
|
||||
]
|
||||
|
||||
# visualization results
|
||||
img = cv2.imread(image_path)
|
||||
detections = sv.Detections(
|
||||
xyxy=input_boxes,
|
||||
mask=masks.astype(bool),
|
||||
class_id=class_ids
|
||||
)
|
||||
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
annotated_frame = box_annotator.annotate(scene=img.copy(), detections=detections)
|
||||
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
|
||||
cv2.imwrite(os.path.join(output_dir, "grounded_sam2_florence2_phrase_grounding.jpg"), annotated_frame)
|
||||
|
||||
mask_annotator = sv.MaskAnnotator()
|
||||
annotated_frame = mask_annotator.annotate(scene=annotated_frame, detections=detections)
|
||||
cv2.imwrite(os.path.join(output_dir, "grounded_sam2_florence2_phrase_grounding_with_mask.jpg"), annotated_frame)
|
||||
|
||||
print(f'Successfully save annotated image to "{output_dir}"')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
image_path = "./notebooks/images/groceries.jpg"
|
||||
parser = argparse.ArgumentParser("Grounded SAM 2 Florence-2 Demos", add_help=True)
|
||||
parser.add_argument("--image_path", type=str, default="./notebooks/images/cars.jpg", required=True, help="path to image file")
|
||||
parser.add_argument("--pipeline", type=str, default="object_detection_segmentation", required=True, help="path to image file")
|
||||
args = parser.parse_args()
|
||||
|
||||
# pipeline-1: detection + segmentation
|
||||
object_detection_and_segmentation(
|
||||
florence2_model=florence2_model,
|
||||
florence2_processor=florence2_processor,
|
||||
sam2_predictor=sam2_predictor,
|
||||
image_path=image_path
|
||||
)
|
||||
IMAGE_PATH = args.image_path
|
||||
PIPELINE = args.pipeline
|
||||
|
||||
print(f"Running pipeline: {PIPELINE} now.")
|
||||
|
||||
if PIPELINE == "object_detection_segmentation":
|
||||
# pipeline-1: detection + segmentation
|
||||
object_detection_and_segmentation(
|
||||
florence2_model=florence2_model,
|
||||
florence2_processor=florence2_processor,
|
||||
sam2_predictor=sam2_predictor,
|
||||
image_path=IMAGE_PATH
|
||||
)
|
||||
elif PIPELINE == "dense_region_caption_segmentation":
|
||||
# pipeline-2: dense region caption + segmentation
|
||||
dense_region_caption_and_segmentation(
|
||||
florence2_model=florence2_model,
|
||||
florence2_processor=florence2_processor,
|
||||
sam2_predictor=sam2_predictor,
|
||||
image_path=IMAGE_PATH
|
||||
)
|
||||
elif PIPELINE == "region_proposal_segmentation":
|
||||
# pipeline-3: dense region caption + segmentation
|
||||
region_proposal_and_segmentation(
|
||||
florence2_model=florence2_model,
|
||||
florence2_processor=florence2_processor,
|
||||
sam2_predictor=sam2_predictor,
|
||||
image_path=IMAGE_PATH
|
||||
)
|
||||
elif PIPELINE == "phrase_grounding_segmentation":
|
||||
# pipeline-4: phrase grounding + segmentation
|
||||
phrase_grounding_and_segmentation(
|
||||
florence2_model=florence2_model,
|
||||
florence2_processor=florence2_processor,
|
||||
sam2_predictor=sam2_predictor,
|
||||
image_path=IMAGE_PATH,
|
||||
text_input="The image shows two vintage Chevrolet cars parked side by side, with one being a red convertible and the other a pink sedan, \
|
||||
set against the backdrop of an urban area with a multi-story building and trees. \
|
||||
The cars have Cuban license plates, indicating a location likely in Cuba."
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Pipeline: {args.pipeline} is not implemented at this time")
|
||||
Reference in New Issue
Block a user