Spaces:
Build error
Build error
File size: 6,358 Bytes
b6f51cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 |
import json
import numpy as np
from matplotlib import cm
import matplotlib
from PIL import Image, ImageColor, ImageFont, ImageDraw
import numpy as np
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,
path_to_output_file = 'download_predictions.json'):
"""
Output detections as json file
"""
# initialise dict to save to json
info = {}
# 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
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
########################################### |