diff --git a/extensions/chaeo/annotators.py b/extensions/chaeo/annotators.py new file mode 100644 index 0000000000000000000000000000000000000000..f951551a24b560ab068e8b4b41ea6eec82271791 --- /dev/null +++ b/extensions/chaeo/annotators.py @@ -0,0 +1,44 @@ +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +def draw_boxes_on_2d_image(img, boxes, **kwargs): + pilimg = Image.fromarray(np.copy(img)) # drawing modifies array in-place + draw = ImageDraw.Draw(pilimg) + font_size = kwargs.get('font_size', 18) + linewidth = kwargs.get('linewidth', 4) + + draw.font = ImageFont.truetype(font="arial.ttf", size=font_size) + + for box in boxes: + y0 = box['info'].y0 + y1 = box['info'].y1 + x0 = box['info'].x0 + x1 = box['info'].x1 + xm = round((x0 + x1) / 2) + + la = box['info'].label + zi = box['info'].zi + + draw.rectangle([(x0, y0), (x1, y1)], outline='white', width=linewidth) + + if kwargs.get('add_label') is True: + draw.text((xm, y0), f'Z{zi:04d}-L{la:04d}', fill='white', anchor='mb') + + return pilimg + + +def draw_boxes_on_3d_image(img, boxes, draw_full_depth=False, **kwargs): + annotated = np.zeros(img.shape, dtype=img.dtype) + + for zi in range(0, img.shape[0]): + if draw_full_depth: + zi_boxes = boxes + else: + zi_boxes = [bb for bb in boxes if bb['info'].zi == zi] + annotated[zi, :, :] = draw_boxes_on_2d_image(img[zi, :, :], zi_boxes, **kwargs) + + if clip := kwargs.get('rescale_clip'): + assert clip >= 0.0 and clip <= 1.0 + annotated = rescale(annotated, clip=clip) + + return annotated \ No newline at end of file diff --git a/extensions/chaeo/tests/__init__.py b/extensions/chaeo/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/extensions/runners.py b/extensions/runners.py new file mode 100644 index 0000000000000000000000000000000000000000..2a4ebdec95bf07e4a3ffa1bf427194ec04ecda36 --- /dev/null +++ b/extensions/runners.py @@ -0,0 +1,3 @@ +""" +Batch runners specific to generating chaetoceros annotations +""" \ No newline at end of file