Support rotation

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)

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