med_rpg/utils/visual_bbox.py

import math
from PIL import Image, ImageDraw, ImageFont
import numpy as np

class DashedImageDraw(ImageDraw.ImageDraw):

    def thick_line(self, xy, direction, fill=None, width=0):
        #xy – Sequence of 2-tuples like [(x, y), (x, y), ...]
        #direction – Sequence of 2-tuples like [(x, y), (x, y), ...]
        if xy[0] != xy[1]:
            self.line(xy, fill = fill, width = width)
        else:
            x1, y1 = xy[0]            
            dx1, dy1 = direction[0]
            dx2, dy2 = direction[1]
            if dy2 - dy1 < 0:
                x1 -= 1
            if dx2 - dx1 < 0:
                y1 -= 1
            if dy2 - dy1 != 0:
                if dx2 - dx1 != 0:
                    k = - (dx2 - dx1)/(dy2 - dy1)
                    a = 1/math.sqrt(1 + k**2)
                    b = (width*a - 1) /2
                else:
                    k = 0
                    b = (width - 1)/2
                x3 = x1 - math.floor(b)
                y3 = y1 - int(k*b)
                x4 = x1 + math.ceil(b)
                y4 = y1 + int(k*b)
            else:
                x3 = x1
                y3 = y1 - math.floor((width - 1)/2)
                x4 = x1
                y4 = y1 + math.ceil((width - 1)/2)
            self.line([(x3, y3), (x4, y4)], fill = fill, width = 1)
        return   
        
    def dashed_line(self, xy, dash=(2,2), fill=None, width=0):
        #xy – Sequence of 2-tuples like [(x, y), (x, y), ...]
        for i in range(len(xy) - 1):
            x1, y1 = xy[i]
            x2, y2 = xy[i + 1]
            x_length = x2 - x1
            y_length = y2 - y1
            length = math.sqrt(x_length**2 + y_length**2)
            dash_enabled = True
            postion = 0
            while postion <= length:
                for dash_step in dash:
                    if postion > length:
                        break
                    if dash_enabled:
                        start = postion/length
                        end = min((postion + dash_step - 1) / length, 1)
                        self.thick_line([(round(x1 + start*x_length),
                                          round(y1 + start*y_length)),
                                         (round(x1 + end*x_length),
                                          round(y1 + end*y_length))],
                                        xy, fill, width)
                    dash_enabled = not dash_enabled
                    postion += dash_step
        return

    def dashed_rectangle(self, xy, dash=(2,2), outline=None, width=0):
        #xy - Sequence of [(x1, y1), (x2, y2)] where (x1, y1) is top left corner and (x2, y2) is bottom right corner
        x1, y1 = xy[0]
        x2, y2 = xy[1]
        halfwidth1 = math.floor((width - 1)/2)
        halfwidth2 = math.ceil((width - 1)/2)
        min_dash_gap = min(dash[1::2])
        end_change1 = halfwidth1 + min_dash_gap + 1
        end_change2 = halfwidth2 + min_dash_gap + 1
        odd_width_change = (width - 1)%2        
        self.dashed_line([(x1 - halfwidth1, y1), (x2 - end_change1, y1)],
                         dash, outline, width)       
        self.dashed_line([(x2, y1 - halfwidth1), (x2, y2 - end_change1)],
                         dash, outline, width)
        self.dashed_line([(x2 + halfwidth2, y2 + odd_width_change),
                          (x1 + end_change2, y2 + odd_width_change)],
                         dash, outline, width)
        self.dashed_line([(x1 + odd_width_change, y2 + halfwidth2),
                          (x1 + odd_width_change, y1 + end_change2)],
                         dash, outline, width)
        return

def visualBBox(real_im, pred_box, bbox):
    # real_im = Image.open(image_path).convert('RGB')
    # real_im = np.asarray(real_im)
    draw = ImageDraw.Draw(real_im)
    dash_draw = DashedImageDraw(real_im)
    # color_gt = (255, 0, 0)
    # color_predict = (0, 255, 0)
    color_gt = (230, 209, 57)
    # color_predict = (169, 60, 189)
    case_id = bbox[0]
    bbox = bbox[1:]

    if case_id == 1:
        color_predict = (255, 0, 0)
        color_gt = (255, 0, 0)
    elif case_id == 2:
        color_predict = (0, 255, 0)
        color_gt = (0, 255, 0)
    elif case_id == 3:
        color_predict = (0, 0, 255)
        color_gt = (0, 0, 255)
    else:
        color_predict = (169, 60, 189)

    if bbox is not None:
        # draw.rectangle(bbox, outline=color_gt, width=3)
        dash_draw.dashed_rectangle([(bbox[0], bbox[1]), (bbox[2], bbox[3])], dash = (5, 3), outline = color_gt, width = 3)
        # x1, y1, x2, y2 = bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]
        # draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], fill='yellow', width=3, dash=(15, 10))

    draw.rectangle(pred_box, outline=color_predict, width=3)
    # save_path = os.path.basename(image_path)
    # real_im.save('outputs/{}.png'.format(save_path))

    # try:
    #     font = ImageFont.truetype('arial.ttf', 24)
    # except IOError:
    #     font = ImageFont.load_default()

    # display_str = "GT"
    # text_width, text_height = font.getsize(display_str)
    # margin = np.ceil(0.05 * text_height)
    # draw.rectangle(
    #     [(left, text_bottom - text_height - 2 * margin), (left + text_width,
    #                                                     text_bottom)],
    #     fill=color)
    # draw.text(
    #     (left + margin, text_bottom - text_height - margin),
    #     display_str,
    #     fill='black',
    #     font=font)
    # text_bottom -= text_height - 2 * margin

    del draw, dash_draw
    return real_im



if __name__ == '__main__':

    pass