File size: 3,060 Bytes
6e73cd3 4390f6c 6e73cd3 4390f6c 6e73cd3 38057e4 6e73cd3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
from CircumSpect.object_detection.groundingdino.util.inference import load_model, predict
import CircumSpect.object_detection.groundingdino.datasets.transforms as T
from torchvision.ops import box_convert
from utils import setup_device
from typing import Tuple, List
import supervision as sv
from io import BytesIO
from PIL import Image
import numpy as np
import requests
import torch
import cv2
model = load_model("./CircumSpect/object_detection/groundingdino/config/GroundingDINO_SwinT_OGC.py",
"./CircumSpect/object_detection/weights/groundingdino_swint_ogc.pth")
def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]:
transform = T.Compose(
[
T.RandomResize([800], max_size=1333),
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
image_source = Image.open(BytesIO(requests.get(
image_path).content) if image_path.startswith("http") else image_path).convert("RGB")
image = np.asarray(image_source)
image_transformed, _ = transform(image_source, None)
return image, image_transformed
def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str]) -> np.ndarray:
h, w, _ = image_source.shape
boxes = boxes * torch.Tensor([w, h, w, h])
xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
detections = sv.Detections(xyxy=xyxy)
labels = [
f"{phrase} {logit:.2f}"
for phrase, logit
in zip(phrases, logits)
]
coordinates = []
for i in list(detections):
left, top, right, bottom = list(list(i)[0])
coordinates.append((int((left+right)/2), int((top+bottom)/2)))
object_coordinates = [[" ".join(object_name.split()[:-1])+": "+object_name.split(
)[-1]+"%", coordinate] for object_name, coordinate in zip(labels, coordinates)]
box_annotator = sv.BoxAnnotator()
annotated_frame = cv2.cvtColor(image_source, cv2.COLOR_RGB2BGR)
annotated_frame = box_annotator.annotate(
scene=annotated_frame, detections=detections, labels=labels)
return annotated_frame, object_coordinates
BOX_TRESHOLD = 0.35
TEXT_TRESHOLD = 0.25
device = setup_device()
def locate_object(objects, image):
image_source, image = load_image(image)
boxes, logits, phrases = predict(
model=model,
image=image,
device=device,
caption=objects,
box_threshold=BOX_TRESHOLD,
text_threshold=TEXT_TRESHOLD
)
annotated_frame, object_coordinates = annotate(
image_source=image_source, boxes=boxes, logits=logits, phrases=phrases)
cv2.circle(annotated_frame, object_coordinates[0][1], 2, (255, 0, 0), 2)
cv2.imwrite("detected_objects.png", annotated_frame)
return annotated_frame, object_coordinates
if __name__ == "__main__":
frame, coord = locate_object(
"drawer", "https://images.nationalgeographic.org/image/upload/v1638890052/EducationHub/photos/robots-3d-landing-page.jpg")
print(coord)
|