viz_utils.py

import json 
import numpy as np

from matplotlib import cm
import matplotlib
from PIL import Image, ImageColor, ImageFont, ImageDraw 
import numpy as np
import pdb
from datetime import date
today = date.today()

#########################################
#  Draw keypoints on image
def draw_keypoints_on_image(image,
                            keypoints,
                            map_label_id_to_str,
                            flag_show_str_labels,
                            use_normalized_coordinates=True,                           
                            font_size=8,
                            keypt_color="#ff0000",
                            marker_size=2,
                            ):
    """Draws keypoints on an image.
    Modified from:
        https://www.programcreek.com/python/?code=fjchange%2Fobject_centric_VAD%2Fobject_centric_VAD-master%2Fobject_detection%2Futils%2Fvisualization_utils.py
    Args:
    image: a PIL.Image object.
    keypoints: a numpy array with shape [num_keypoints, 2].
    map_label_id_to_str: dict with keys=label number and values= label string
    flag_show_str_labels: boolean to select whether or not to show string labels
    color: color to draw the keypoints with. Default is red.
    radius: keypoint radius. Default value is 2.
    use_normalized_coordinates: if True (default), treat keypoint values as
        relative to the image.  Otherwise treat them as absolute.

    
    """
    # get a drawing context
    draw = ImageDraw.Draw(image,"RGBA")  

    im_width, im_height = image.size
    keypoints_x = [k[0] for k in keypoints]
    keypoints_y = [k[1] for k in keypoints]
    alpha = [k[2] for k in keypoints]
    norm = matplotlib.colors.Normalize(vmin=0, vmax=255)

    names_for_color = [i for i in map_label_id_to_str.keys()]
    colores = np.linspace(0, 255, num=len(names_for_color),dtype= int)

    # adjust keypoints coords if required
    if use_normalized_coordinates:
        keypoints_x = tuple([im_width * x for x in keypoints_x])
        keypoints_y = tuple([im_height * y for y in keypoints_y])
    
    #cmap = matplotlib.cm.get_cmap('hsv')
    cmap2 = matplotlib.cm.get_cmap('Greys')
    # draw ellipses around keypoints
    for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y)):
        round_fill = list(cm.viridis(norm(colores[i]),bytes=True))#[round(num*255) for num in list(cmap(i))[:3]] #check!
        if np.isnan(alpha[i]) == False : 
            round_fill[3] = round(alpha[i] *255)
        #print(round_fill)
        #round_outline = [round(num*255) for num in list(cmap2(alpha[i]))[:3]]
        draw.ellipse([(keypoint_x - marker_size, keypoint_y - marker_size),
                      (keypoint_x + marker_size, keypoint_y + marker_size)],
                       fill=tuple(round_fill), outline= 'black', width=1) #fill and outline: [0,255]

        # add string labels around keypoints
        if flag_show_str_labels:
            draw.text((keypoint_x + marker_size, keypoint_y + marker_size),#(0.5*im_width, 0.5*im_height), #-------
                      map_label_id_to_str[i],
                      ImageColor.getcolor(keypt_color, "RGB") # rgb #
                      )

#########################################
#  Draw bboxes on image
def draw_bbox_w_text(img,
                     results,
                     font_size=8): #TODO: select color too?
    #pdb.set_trace()
    bbxyxy = results
    w, h = bbxyxy[2], bbxyxy[3]
    shape = [(bbxyxy[0], bbxyxy[1]), (w , h)]
    imgR = ImageDraw.Draw(img)  
    imgR.rectangle(shape,  outline ="red",width=5) ##bb for animal

    confidence = bbxyxy[4]
    string_bb = 'animal ' + str(round(confidence, 2))

    text_size = font.getsize(string_bb) # (h,w)
    position = (bbxyxy[0],bbxyxy[1] - text_size[1] -2 )
    left, top, right, bottom = imgR.textbbox(position, string_bb, font=font)
    imgR.rectangle((left, top-5, right+5, bottom+5), fill="red")
    imgR.text((bbxyxy[0] + 3 ,bbxyxy[1] - text_size[1] -2 ), string_bb, font=font, fill="black")

    return imgR

###########################################
def save_results_as_json(md_results, dlc_outputs, map_dlc_label_id_to_str, thr,model,mega_model_input, path_to_output_file = 'download_predictions.json'):

    """
    Output detections as json file

    """
    # initialise dict to save to json
    info = {}
    info['date'] = str(today)
    info['MD_model'] = str(mega_model_input)
    # info from megaDetector
    info['file']= md_results.files[0]
    number_bb = len(md_results.xyxy[0].tolist())
    info['number_of_bb'] = number_bb
    # info from DLC
    number_bb_thr = len(dlc_outputs)
    labels = [n for n in map_dlc_label_id_to_str.values()]
    
    # create list of bboxes above th
    new_index = []
    for i in range(number_bb):
        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[i]

        if confidence > thr:
            new_index.append(i)

    # define aux dict for every bounding box above threshold
    for i in range(number_bb_thr):
        aux={}
        # MD output
        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[new_index[i]]
        aux['corner_1'] = (corner_x1,corner_y1)
        aux['corner_2'] = (corner_x2,corner_y2)
        aux['predict MD'] = md_results.names[0]
        aux['confidence MD'] = confidence
        
        # DLC output
        info['dlc_model'] = model
        kypts = []
        for s in dlc_outputs[i]:
            aux1 = []
            for j in s:
                aux1.append(float(j))

            kypts.append(aux1)
        aux['dlc_pred']  = dict(zip(labels,kypts))
        info['bb_' + str(new_index[i]) ]=aux

    # save dict as json
    with open(path_to_output_file, 'w') as f:
        json.dump(info, f, indent=1)
        print('Output file saved at {}'.format(path_to_output_file))

    return path_to_output_file


def save_results_only_dlc(dlc_outputs,map_label_id_to_str,model,output_file = 'dowload_predictions_dlc.json'):

    """
    write json dlc output
    """
    info = {}
    info['date'] = str(today)
    labels = [n for n in map_label_id_to_str.values()]
    info['dlc_model'] = model
    kypts = []
    for s in dlc_outputs:
        aux1 = []
        for j in s:
            aux1.append(float(j))

        kypts.append(aux1)
    info['dlc_pred']  = dict(zip(labels,kypts))

    with open(output_file, 'w') as f:
        json.dump(info, f, indent=1)
        print('Output file saved at {}'.format(output_file))

    return output_file


###########################################