wehub-resource-sync
119 lines
3.8 KiB
Python
119 lines
3.8 KiB
Python
import cv2
|
|
import numpy as np
|
|
import supervision as sv
|
|
|
|
import torch
|
|
import torchvision
|
|
from torchvision.transforms import ToTensor
|
|
|
|
from groundingdino.util.inference import Model
|
|
|
|
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
# GroundingDINO config and checkpoint
|
|
GROUNDING_DINO_CONFIG_PATH = "GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py"
|
|
GROUNDING_DINO_CHECKPOINT_PATH = "./groundingdino_swint_ogc.pth"
|
|
|
|
# Building GroundingDINO inference model
|
|
grounding_dino_model = Model(model_config_path=GROUNDING_DINO_CONFIG_PATH, model_checkpoint_path=GROUNDING_DINO_CHECKPOINT_PATH)
|
|
|
|
# Building MobileSAM predictor
|
|
EFFICIENT_SAM_CHECHPOINT_PATH = "./EfficientSAM/efficientsam_s_gpu.jit"
|
|
efficientsam = torch.jit.load(EFFICIENT_SAM_CHECHPOINT_PATH)
|
|
|
|
|
|
# Predict classes and hyper-param for GroundingDINO
|
|
SOURCE_IMAGE_PATH = "./EfficientSAM/LightHQSAM/example_light_hqsam.png"
|
|
CLASSES = ["bench"]
|
|
BOX_THRESHOLD = 0.25
|
|
TEXT_THRESHOLD = 0.25
|
|
NMS_THRESHOLD = 0.8
|
|
|
|
|
|
# load image
|
|
image = cv2.imread(SOURCE_IMAGE_PATH)
|
|
|
|
# detect objects
|
|
detections = grounding_dino_model.predict_with_classes(
|
|
image=image,
|
|
classes=CLASSES,
|
|
box_threshold=BOX_THRESHOLD,
|
|
text_threshold=TEXT_THRESHOLD
|
|
)
|
|
|
|
# annotate image with detections
|
|
box_annotator = sv.BoxAnnotator()
|
|
labels = [
|
|
f"{CLASSES[class_id]} {confidence:0.2f}"
|
|
for _, _, confidence, class_id, _, _
|
|
in detections]
|
|
annotated_frame = box_annotator.annotate(scene=image.copy(), detections=detections, labels=labels)
|
|
|
|
# save the annotated grounding dino image
|
|
cv2.imwrite("EfficientSAM/LightHQSAM/groundingdino_annotated_image.jpg", annotated_frame)
|
|
|
|
|
|
# NMS post process
|
|
print(f"Before NMS: {len(detections.xyxy)} boxes")
|
|
nms_idx = torchvision.ops.nms(
|
|
torch.from_numpy(detections.xyxy),
|
|
torch.from_numpy(detections.confidence),
|
|
NMS_THRESHOLD
|
|
).numpy().tolist()
|
|
|
|
detections.xyxy = detections.xyxy[nms_idx]
|
|
detections.confidence = detections.confidence[nms_idx]
|
|
detections.class_id = detections.class_id[nms_idx]
|
|
|
|
print(f"After NMS: {len(detections.xyxy)} boxes")
|
|
|
|
|
|
def efficient_sam_box_prompt_segment(image, pts_sampled, model):
|
|
bbox = torch.reshape(torch.tensor(pts_sampled), [1, 1, 2, 2])
|
|
bbox_labels = torch.reshape(torch.tensor([2, 3]), [1, 1, 2])
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
img_tensor = ToTensor()(image)
|
|
|
|
predicted_logits, predicted_iou = model(
|
|
img_tensor[None, ...].cuda(),
|
|
bbox.cuda(),
|
|
bbox_labels.cuda(),
|
|
)
|
|
predicted_logits = predicted_logits.cpu()
|
|
all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy()
|
|
predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy()
|
|
|
|
max_predicted_iou = -1
|
|
selected_mask_using_predicted_iou = None
|
|
for m in range(all_masks.shape[0]):
|
|
curr_predicted_iou = predicted_iou[m]
|
|
if (
|
|
curr_predicted_iou > max_predicted_iou
|
|
or selected_mask_using_predicted_iou is None
|
|
):
|
|
max_predicted_iou = curr_predicted_iou
|
|
selected_mask_using_predicted_iou = all_masks[m]
|
|
return selected_mask_using_predicted_iou
|
|
|
|
|
|
# collect segment results from EfficientSAM
|
|
result_masks = []
|
|
for box in detections.xyxy:
|
|
mask = efficient_sam_box_prompt_segment(image, box, efficientsam)
|
|
result_masks.append(mask)
|
|
|
|
detections.mask = np.array(result_masks)
|
|
|
|
# annotate image with detections
|
|
box_annotator = sv.BoxAnnotator()
|
|
mask_annotator = sv.MaskAnnotator()
|
|
labels = [
|
|
f"{CLASSES[class_id]} {confidence:0.2f}"
|
|
for _, _, confidence, class_id, _, _
|
|
in detections]
|
|
annotated_image = mask_annotator.annotate(scene=image.copy(), detections=detections)
|
|
annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
|
|
|
|
# save the annotated grounded-sam image
|
|
cv2.imwrite("EfficientSAM/gronded_efficient_sam_anontated_image.jpg", annotated_image)
|