Spaces:
Build error
Build error
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() | |
FONTS = {'amiko': "fonts/Amiko-Regular.ttf", | |
'nature': "fonts/LoveNature.otf", | |
'painter':"fonts/PainterDecorator.otf", | |
'animals': "fonts/UncialAnimals.ttf", | |
'zen': "fonts/ZEN.TTF"} | |
######################################### | |
# 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_style='amiko', | |
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: | |
font = ImageFont.truetype(FONTS[font_style], | |
font_size) | |
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 # | |
font=font) | |
######################################### | |
# Draw bboxes on image | |
def draw_bbox_w_text(img, | |
results, | |
font_style='amiko', | |
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)) | |
font = ImageFont.truetype(FONTS[font_style], font_size) | |
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 | |
########################################### |