diff --git a/cocoviewer.py b/cocoviewer.py index 470a4be..033dd0d 100644 --- a/cocoviewer.py +++ b/cocoviewer.py @@ -3,6 +3,7 @@ View images with bboxes from the COCO dataset. """ +import math import argparse import colorsys import json @@ -148,6 +149,41 @@ def get_categories(instances: dict) -> dict: return categories +def show_label(label_size, c, b, draw): + text = c[0] + + try: + try: + # Should work for Linux + font = ImageFont.truetype("DejaVuSans.ttf", size=label_size) + except OSError: + # Should work for Windows + font = ImageFont.truetype("Arial.ttf", size=label_size) + except OSError: + # Load default, note no resize option + # TODO: Implement notification message as popup window + font = ImageFont.load_default() + + tw, th = draw.textsize(text, font) + tx0 = b[0] + ty0 = b[1] - th + + # TODO: Looks weird! We need image dims to make it right + tx0 = max(b[0], max(b[0], tx0)) if tx0 < 0 else tx0 + ty0 = max(b[1], max(0, ty0)) if ty0 < 0 else ty0 + + tx1 = tx0 + tw + ty1 = ty0 + th + + # TODO: The same here + if tx1 > b[2]: + tx0 = max(0, tx0 - (tx1 - b[2])) + tx1 = tw if tx0 == 0 else b[2] + + draw.rectangle((tx0, ty0, tx1, ty1), fill=c[-1]) + draw.text((tx0, ty0), text, (255, 255, 255), font=font) + + def draw_bboxes(draw, objects, labels, obj_categories, ignore, width, label_size): """Puts rectangles on the image.""" # Extracting bbox coordinates @@ -166,38 +202,43 @@ def draw_bboxes(draw, objects, labels, obj_categories, ignore, width, label_size draw.rectangle(b, outline=c[-1], width=width) if labels: - text = c[0] + show_label(label_size, c, b, draw) + - try: - try: - # Should work for Linux - font = ImageFont.truetype("DejaVuSans.ttf", size=label_size) - except OSError: - # Should work for Windows - font = ImageFont.truetype("Arial.ttf", size=label_size) - except OSError: - # Load default, note no resize option - # TODO: Implement notification message as popup window - font = ImageFont.load_default() +def draw_rotated_bboxes(draw, objects, labels, obj_categories, ignore, width, label_size): + for i, (c, obj) in enumerate(zip(obj_categories, objects)): + if i not in ignore: + x = int(obj['bbox'][0]) + y = int(obj['bbox'][1]) + box_width = int(obj['bbox'][2]) + box_height = int(obj['bbox'][3]) + box = [x, y, x+box_width, y+box_height] - tw, th = draw.textsize(text, font) - tx0 = b[0] - ty0 = b[1] - th + rotation = obj['attributes']['rotation'] - # TODO: Looks weird! We need image dims to make it right - tx0 = max(b[0], max(b[0], tx0)) if tx0 < 0 else tx0 - ty0 = max(b[1], max(0, ty0)) if ty0 < 0 else ty0 + points = [(int(x), int(y)), + (int(x + box_width), int(y)), + (int(x + box_width), int(y + box_height)), + (int(x), int(y + box_height))] - tx1 = tx0 + tw - ty1 = ty0 + th + center_point = (int(x + box_width / 2), int(y + box_height / 2)) - # TODO: The same here - if tx1 > b[2]: - tx0 = max(0, tx0 - (tx1 - b[2])) - tx1 = tw if tx0 == 0 else b[2] + angle_rad = rotation * (3.14159 / 180) + rotated_points = [] - draw.rectangle((tx0, ty0, tx1, ty1), fill=c[-1]) - draw.text((tx0, ty0), text, (255, 255, 255), font=font) + for point in points: + translated_x = point[0] - center_point[0] + translated_y = point[1] - center_point[1] + rotated_x = translated_x * math.cos(angle_rad) - translated_y * math.sin(angle_rad) + rotated_y = translated_x * math.sin(angle_rad) + translated_y * math.cos(angle_rad) + final_x = rotated_x + center_point[0] + final_y = rotated_y + center_point[1] + rotated_points.append((final_x, final_y)) + + draw.polygon(rotated_points, outline=c[-1], width=width) + + if labels: + show_label(label_size, c, box, draw) def draw_masks(draw, objects, obj_categories, ignore, alpha): @@ -616,7 +657,11 @@ def compose_image( draw_masks(draw, objects, names_colors, ignore, alpha) # Draw bounding boxes if bboxes_on: - draw_bboxes(draw, objects, labels_on, names_colors, ignore, width, label_size) + if 'attributes' in objects[0].keys(): + if 'rotation' in objects[0]['attributes']: + draw_rotated_bboxes(draw, objects, labels_on, names_colors, ignore, width, label_size) + else: + draw_bboxes(draw, objects, labels_on, names_colors, ignore, width, label_size) del draw # Resulting image self.current_composed_image = Image.alpha_composite(img_open, draw_layer) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..6c1215d --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +cffi~=1.15.1 +future~=0.18.3 +Pillow~=9.0.1 +pip~=22.0.4 +wheel~=0.37.1 +cryptography~=3.4.8 +numpy~=1.21.6 +setuptools~=59.6.0 +certifi~=2020.6.20 +pycparser~=2.21 \ No newline at end of file