Spaces:
Build error
Build error
File size: 7,128 Bytes
b6f51cf 026ff32 b6f51cf db71c14 b6f51cf 026ff32 b6f51cf 026ff32 b6f51cf 026ff32 b6f51cf 026ff32 1f24f7d 026ff32 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 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 |
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)
if np.isnan(keypoint_x).any():
continue
#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
########################################### |