From 22f82c54481ba928251bd89b1a1cbe18bb91da2f Mon Sep 17 00:00:00 2001 From: vahidrezanezhad Date: Wed, 16 Apr 2025 23:36:41 +0200 Subject: [PATCH] rotation augmentation is provided for machine based reading order --- train.py | 7 +++++-- utils.py | 23 ++++++++++++++++++++++- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/train.py b/train.py index 7e3e390..130c7f4 100644 --- a/train.py +++ b/train.py @@ -380,7 +380,10 @@ def run(_config, n_classes, n_epochs, input_height, dir_flow_train_labels = os.path.join(dir_train, 'labels') classes = os.listdir(dir_flow_train_labels) - num_rows =len(classes) + if augmentation: + num_rows = len(classes)*(len(thetha) + 1) + else: + num_rows = len(classes) #ls_test = os.listdir(dir_flow_train_labels) #f1score_tot = [0] @@ -390,7 +393,7 @@ def run(_config, n_classes, n_epochs, input_height, model.compile(loss="binary_crossentropy", optimizer = opt_adam,metrics=['accuracy']) for i in range(n_epochs): - history = model.fit(generate_arrays_from_folder_reading_order(dir_flow_train_labels, dir_flow_train_imgs, n_batch, input_height, input_width, n_classes), steps_per_epoch=num_rows / n_batch, verbose=1) + history = model.fit(generate_arrays_from_folder_reading_order(dir_flow_train_labels, dir_flow_train_imgs, n_batch, input_height, input_width, n_classes, thetha, augmentation), steps_per_epoch=num_rows / n_batch, verbose=1) model.save( os.path.join(dir_output,'model_'+str(i+indexer_start) )) with open(os.path.join(os.path.join(dir_output,'model_'+str(i)),"config.json"), "w") as fp: diff --git a/utils.py b/utils.py index d7ddb99..50c21af 100644 --- a/utils.py +++ b/utils.py @@ -363,6 +363,11 @@ def rotation_not_90_func(img, label, thetha): return rotate_max_area(img, rotated, rotated_label, thetha) +def rotation_not_90_func_single_image(img, thetha): + rotated = imutils.rotate(img, thetha) + return rotate_max_area(img, rotated, thetha) + + def color_images(seg, n_classes): ann_u = range(n_classes) if len(np.shape(seg)) == 3: @@ -410,7 +415,7 @@ def IoU(Yi, y_predi): #print("Mean IoU: {:4.3f}".format(mIoU)) return mIoU -def generate_arrays_from_folder_reading_order(classes_file_dir, modal_dir, batchsize, height, width, n_classes): +def generate_arrays_from_folder_reading_order(classes_file_dir, modal_dir, batchsize, height, width, n_classes, thetha, augmentation=False): all_labels_files = os.listdir(classes_file_dir) ret_x= np.zeros((batchsize, height, width, 3))#.astype(np.int16) ret_y= np.zeros((batchsize, n_classes)).astype(np.int16) @@ -433,6 +438,22 @@ def generate_arrays_from_folder_reading_order(classes_file_dir, modal_dir, batch ret_x= np.zeros((batchsize, height, width, 3))#.astype(np.int16) ret_y= np.zeros((batchsize, n_classes)).astype(np.int16) batchcount = 0 + + if augmentation: + for thetha_i in thetha: + img_rot = rotation_not_90_func_single_image(img, thetha_i) + + ret_x[batchcount, :,:,0] = img_rot[:,:,0]/3.0 + ret_x[batchcount, :,:,2] = img_rot[:,:,2]/3.0 + ret_x[batchcount, :,:,1] = img_rot[:,:,1]/5.0 + + ret_y[batchcount, :] = label_class + batchcount+=1 + if batchcount>=batchsize: + yield (ret_x, ret_y) + ret_x= np.zeros((batchsize, height, width, 3))#.astype(np.int16) + ret_y= np.zeros((batchsize, n_classes)).astype(np.int16) + batchcount = 0 def data_gen(img_folder, mask_folder, batch_size, input_height, input_width, n_classes, task='segmentation'): c = 0