correct pred

yolov5/utils/dataloaders.py

@@ -304,9 +304,16 @@ class LoadSpectros:
                                 cmap = self.cmap, vmin=vmin)
         plt.subplots_adjust(top=1, bottom=0, left=0, right=1)
         im0 = axim.make_image(fig.canvas)[0][:,:,:-1][:,:,::-1]
+        # If impossible to write the image in the same folder, write in current path
+        if cv2.imwrite(path, im0) == False:
+            cv2.imwrite(path.split('/')[-1], im0)
+            im0 = cv2.imread(path.split('/')[-1])
+            os.system('rm '+ path.split('/')[-1])
+        else : 
             cv2.imwrite(path, im0)
             im0 = cv2.imread(path)
             os.system('rm '+path)
+
         plt.close()
         im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0]  # padded resize
         im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW