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()