Spaces:
Sleeping
Sleeping
File size: 1,611 Bytes
350aa90 |
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 |
import argparse
from typing import Dict, List
import cv2
from mivolo.data.data_reader import PictureInfo, read_csv_annotation_file
from ultralytics.yolo.utils.plotting import Annotator, colors
def get_parser():
parser = argparse.ArgumentParser(description="Visualization")
parser.add_argument("--dataset_images", default="", type=str, required=True, help="path to images")
parser.add_argument("--annotation_file", default="", type=str, required=True, help="path to annotations")
return parser
def visualize(images_dir, new_annotation_file):
bboxes_per_image: Dict[str, List[PictureInfo]] = read_csv_annotation_file(new_annotation_file, images_dir)[0]
print(f"Found {len(bboxes_per_image)} unique images")
for image_path, bboxes in bboxes_per_image.items():
im_cv = cv2.imread(image_path)
annotator = Annotator(im_cv)
for i, bbox_info in enumerate(bboxes):
label = f"{bbox_info.gender} Age: {bbox_info.age}"
if any(coord != -1 for coord in bbox_info.bbox):
# draw face bbox if exist
annotator.box_label(bbox_info.bbox, label, color=colors(i, True))
if any(coord != -1 for coord in bbox_info.person_bbox):
# draw person bbox if exist
annotator.box_label(bbox_info.person_bbox, "p " + label, color=colors(i, True))
im_cv = annotator.result()
cv2.imshow("image", im_cv)
cv2.waitKey(0)
if __name__ == "__main__":
parser = get_parser()
args = parser.parse_args()
visualize(args.dataset_images, args.annotation_file)
|