File size: 2,394 Bytes
c446951 |
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 89 90 91 92 93 |
import cv2
import requests
import base64
import supervision as sv
import argparse
import os
from typing import List
def parse_arguments() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Process image and annotate with detections."
)
parser.add_argument(
"--image_path", type=str, required=True, help="Path to the image file."
)
parser.add_argument(
"--class_list",
type=str,
nargs="+",
required=True,
help="List of classes to detect in the image.",
)
parser.add_argument(
"--dataset_id", type=str, required=True, help="Dataset ID for the API request."
)
parser.add_argument(
"--version_id", type=str, required=True, help="Version ID for the API request."
)
parser.add_argument(
"--confidence",
type=float,
default=0.5,
help="Confidence threshold for the detections.",
)
return parser.parse_args()
def annotate_image(
image_path: str,
class_list: List[str],
dataset_id: str,
version_id: str,
confidence: float,
api_key: str,
) -> None:
url = f"http://localhost:9001/{dataset_id}/{version_id}"
headers = {"Content-Type": "application/json"}
params = {
"api_key": api_key,
"confidence": confidence,
}
with open(image_path, "rb") as image_file:
encoded_image = base64.b64encode(image_file.read()).decode("utf-8")
image = cv2.imread(image_path)
response = requests.post(
url, headers=headers, params=params, data=encoded_image
).json()
detections = sv.Detections.from_roboflow(response, class_list=class_list)
box_annotator = sv.BoxAnnotator()
labels = [
f"{class_list[class_id]} {det_confidence:0.2f}"
for _, _, det_confidence, class_id, _ in detections
]
annotated_image = box_annotator.annotate(
image, detections=detections, labels=labels
)
cv2.imshow("Annotated image", annotated_image)
cv2.waitKey(0)
if __name__ == "__main__":
args = parse_arguments()
API_KEY = os.environ.get("API_KEY")
if not API_KEY:
raise ValueError("API_KEY not found in environment variables.")
annotate_image(
args.image_path,
args.class_list,
args.dataset_id,
args.version_id,
args.confidence,
API_KEY,
)
|