diff --git a/get_train_annot_YOLO.py b/get_train_annot_YOLO.py index b637accd0f09f3f40401a81561e4b61c839efd81..0980f0a1236750d97f9cafad488071c07580a8b6 100644 --- a/get_train_annot_YOLO.py +++ b/get_train_annot_YOLO.py @@ -8,8 +8,10 @@ import ipdb import random from datetime import date import argparse +import cv2 import matplotlib.patches as patches from matplotlib.patches import Rectangle +from random import randrange from PIL import Image from mycolorpy import colorlist as mcp @@ -78,7 +80,18 @@ liste_espece.to_csv(str(directory+'liste_especes.csv'),index = False) print('\n',data) -color = mcp.gen_color(cmap = "Wistia", n= len(list_espece)) +#color = mcp.gen_color(cmap = "Wistia", n= len(list_espece)) + +colors = pd.DataFrame(columns = ['color', 'species']) + +for i in range (30): + r = randrange(255) + g = randrange(255) + b = randrange(255) + rand_color = (r, g, b) + new = pd.DataFrame([[rand_color, i]], columns = ['color', 'species']) + colors = pd.concat([colors, new]) + def process(x): count, (f, grp) = x @@ -163,20 +176,46 @@ def process(x): plt.savefig(os.path.join(directory,str('images_'+str(today.day)+'_'+str(today.month)),row.Code,str(name+'.jpg'))) plt.savefig(os.path.join(directory,str('images_'+str(today.day)+'_'+str(today.month)),'all',str(name+'.jpg'))) + plt.close() + if args.export != None: + + im = cv2.imread(os.path.join(directory,str('images_'+str(today.day)+'_'+str(today.month)),'all',str(name+'.jpg'))) + im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) + H,W = im.shape[0], im.shape[1] for l in range(len(fin)): - try : - plt.gca().add_patch(Rectangle(((fin.x.iloc[l]*len(time))-(0.5*fin.width.iloc[l]*len(time)), ((fin.y.iloc[l]*len(freq))+0.5*fin.height.iloc[l]*len(freq))), fin.width.iloc[l]*len(time), fin.height.iloc[l]*len(freq), linewidth=3, edgecolor=color[int(fin.id.iloc[l])], facecolor='none')) + x, y, w, h = fin.x.iloc[l]*W , fin.y.iloc[l]*H , fin.w.iloc[l]*W , fin.h.iloc[l]*H + + shape1 = (int(x-(0.5*w)), int(y+(0.5*h))) + shape2 = (int(x+(0.5*w)), int(y+(0.5*h))) + shape3 = (int(x+-(0.5*w)), int(y-(0.5*h))) + shape4 = (int(x+(0.5*w)), int(y-(0.5*h))) + + #rectangle text shape + + shp1 = shape4[0]-10, shape4[1]+20 + shp2 = shape4[0], shape4[1]+20 + shp3 = shape4[0]-10, shape4[1] + shp4 = shape4[0], shape4[1] + + #text placement + + text_shape = shp1[0], shp1[1]-5 + + label = str(fin.id.iloc[l]) + cv2.rectangle(im, pt1=shape1, pt2=shape4, color= colors[colors.species == label].color, thickness=1) - except IndexError: - ipdb.set_trace() - try: - plt.savefig(os.path.join(directory, str('images_annotes_'+str(today.day)+'_'+str(today.month)),str(name+'.jpg'))) + cv2.rectangle(im, pt1=shp1 , pt2= shp4, color= colors[colors.species == label].color, thickness= -1) + cv2.putText(im, label, text_shape, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1) + + plt.imshow(im) + try: + plt.savefig(os.path.join(directory, str('images_annotes_'+str(today.day)+'_'+str(today.month)),str(name+'.jpg'))) - except Exception: - os.mkdir(os.path.join(directory,str('images_annotes_'+str(today.day)+'_'+str(today.month)))) - plt.savefig(os.path.join(directory, str('images_annotes_'+str(today.day)+'_'+str(today.month)),str(name+'.jpg'))) + except Exception: + os.mkdir(os.path.join(directory,str('images_annotes_'+str(today.day)+'_'+str(today.month)))) + plt.savefig(os.path.join(directory, str('images_annotes_'+str(today.day)+'_'+str(today.month)),str(name+'.jpg'))) plt.close()