""" """ import sys import cv2 as cv from PIL import Image from legacy.utls.utls import * def cv_imshow(title='default',img=None): print(img.shape) cv.imshow(title,img) cv.waitKey(0) cv.destroyAllWindows() def rotated_img_extractor(x=None, gt=None,img_width=None, img_height=None,i=None, two_data=False): if two_data: if img_width==img_height: # for images whose sizes are the same if i % 90 == 0: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height, img_width)) elif i % 19 == 0: if i == 57: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (100, 100), (720 - 100, 720 - 100), (0, 0, 255), (2)) rot_x = rot_x[100:720 - 100, 100:720 - 100, :] rot_gt = rot_gt[100:720 - 100, 100:720 - 100] # print("just for check 19: ", i, rot_x.shape) elif i == 285: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height, img_width)) rot_x = rot_x[75:720 - 75, 75:720 - 75, :] rot_gt = rot_gt[75:720 - 75, 75:720 - 75] else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height, img_width)) rot_x = rot_x[95:720 - 95, 95:720 - 95, :] rot_gt = rot_gt[95:720 - 95, 95:720 - 95] elif i % 23 == 0: if i == 161: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height, img_width)) rot_x = rot_x[85:720 - 85, 85:720 - 85, :] rot_gt = rot_gt[85:720 - 85, 85:720 - 85] else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (105, 105), (720 - 105, 720 - 105), (0, 0, 255), (2)) rot_x = rot_x[105:720 - 105, 105:720 - 105, :] rot_gt = rot_gt[105:720 - 105, 105:720 - 105] # print("just for check 23:", i, rot_x.shape) return rot_x, rot_gt else: # # for images whose sizes are ***not*** the same ********************************* img_size = img_width if img_width < img_height else img_height if i % 90 == 0: if i==180: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height+250, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height+250, img_width)) # a = np.copy(rot_x) rot_x = rot_x[10:img_size-90, 10:img_size+110, :] rot_gt = rot_gt[10:img_size-90, 10:img_size+110] else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + 450, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height + 450, img_width)) # a = np.copy(rot_x) rot_x = rot_x[100:img_size + 200, 300:img_size + 200, :] rot_gt = rot_gt[100:img_size + 200, 300:img_size + 200] elif i % 19 == 0: if i == 57: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height+i+5, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height+i+5, img_width)) # a = np.copy(rot_x) # # x y x y # cv.rectangle(a, (275, 275), (img_size+55, img_size+55), (0, 0, 255), (2)) # y x rot_x = rot_x[275:img_size+55, 275:img_size+55, :] rot_gt = rot_gt[275:img_size+55, 275:img_size+55] # print("just for check 19: ", i, rot_x.shape) elif i == 285: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height+i, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height+i, img_width)) rot_x = rot_x[100:img_size-50,355:img_size+205, :] rot_gt = rot_gt[100:img_size-50,355:img_size+205] # print("just for check 19: ", i, rot_x.shape) elif i==19: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height+200, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height+200, img_width)) # a = np.copy(rot_x) # cv.rectangle(a, (150, 150), (img_size+30, img_size-70), (0, 0, 255), (2)) rot_x = rot_x[150:img_size-70, 150:img_size+30, :] rot_gt = rot_gt[150:img_size-70, 150:img_size+30] # print("just for check 19: ", i, rot_x.shape) else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height+250, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height+250, img_width)) # a = np.copy(rot_x) # cv.rectangle(a, (400, 115), (img_size+180, img_size-105), (0, 0, 255), (2)) rot_x = rot_x[115:img_size-105, 400:img_size+180, :] rot_gt = rot_gt[115:img_size-105, 400:img_size+180] # print("just for check 19: ", i, rot_x.shape) elif i % 23 == 0: if i == 161: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height+i+200, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height+i, img_width)) # a = rot_x # cv.rectangle(a, (95, 50), (img_size+75, img_size-170), (0, 0, 255), (2)) rot_x = rot_x[50:img_size-170, 95:img_size+75, :] rot_gt = rot_gt[50:img_size-170, 95:img_size+75] # print("just for check 23: ", i, rot_x.shape) elif i==207: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + 250, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height + 250, img_width)) # a = np.copy(rot_x) # cv.rectangle(a, (200, 185), (img_size + 160, img_size - 95), (0, 0, 255), (2)) rot_x = rot_x[185:img_size - 95, 200:img_size + 160, :] rot_gt = rot_gt[185:img_size - 95, 200:img_size + 160] else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height+250, img_width)) rot_gt = cv.warpAffine(gt, adjus_M, (img_height+250, img_width)) # a = np.copy(rot_x) # cv.rectangle(a, (390, 115), (img_size+170, img_size-105), (0, 0, 255), (2)) rot_x = rot_x[115:img_size-105, 390:img_size+170, :] rot_gt = rot_gt[115:img_size-105, 390:img_size+170] return rot_x,rot_gt else: # For NIR imagel but just NIR (ONE data) if img_height==img_width: if i % 90 == 0: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) # print("just for check 90: ", i) elif i % 19 == 0: if i == 57: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (100, 100), (720 - 100, 720 - 100), (0, 0, 255), (2)) rot_x = rot_x[100:720 - 100, 100:720 - 100, :] # print("just for check 19: ", i, rot_x.shape) elif i == 285: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (75, 75), (720 - 75, 720 - 75), (0, 0, 255), (2)) rot_x = rot_x[75:720 - 75, 75:720 - 75, :] # print("just for check 19: ", i, rot_x.shape) else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x,(95,95),(720-95,720-95),(0,0,255),(2) ) rot_x = rot_x[95:720 - 95, 95:720 - 95, :] # print("just for check 19: ", i, rot_x.shape) elif i % 23 == 0: if i == 161: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x,(85,85),(720-85,720-85),(0,0,255),(2) ) rot_x = rot_x[85:720 - 85, 85:720 - 85, :] # print("just for check 23: ", i, rot_x.shape) elif i==207: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (105, 105), (720 - 105, 720 - 105), (0, 0, 255), (2)) rot_x = rot_x[105:720 - 105, 105:720 - 105, :] # print("just for check 23:", i, rot_x.shape) else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (105, 105), (720 - 105, 720 - 105), (0, 0, 255), (2)) rot_x = rot_x[105:720 - 105, 105:720 - 105, :] # print("just for check 23:", i, rot_x.shape) else: print("Error line 221 in dataset_manager") sys.exit() else: img_size = img_width if img_width < img_height else img_height if i % 90 == 0: if i == 180: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + 250, img_width)) rot_x = rot_x[10:img_size - 90, 10:img_size + 110, :] else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + 450, img_width)) rot_x = rot_x[100:img_size + 200, 300:img_size + 200, :] elif i % 19 == 0: if i == 57: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + i + 5, img_width)) # # x y x y # cv.rectangle(a, (275, 275), (img_size+55, img_size+55), (0, 0, 255), (2)) # y x rot_x = rot_x[275:img_size + 55, 275:img_size + 55, :] elif i == 285: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + i, img_width)) # a = np.copy(rot_x) # cv.rectangle(a, (355, 100), (img_size+205, img_size-50), (0, 0, 255), (2)) rot_x = rot_x[100:img_size - 50, 355:img_size + 205, :] elif i == 19: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + 200, img_width)) # a = np.copy(rot_x) # cv.rectangle(a, (150, 150), (img_size+30, img_size-70), (0, 0, 255), (2)) rot_x = rot_x[150:img_size - 70, 150:img_size + 30, :] else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + 250, img_width)) # a = np.copy(rot_x) # cv.rectangle(a, (400, 115), (img_size+180, img_size-105), (0, 0, 255), (2)) rot_x = rot_x[115:img_size - 105, 400:img_size + 180, :] elif i % 23 == 0: if i == 161: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + i + 200, img_width)) # a = rot_x # cv.rectangle(a, (95, 50), (img_size+75, img_size-170), (0, 0, 255), (2)) rot_x = rot_x[50:img_size - 170, 95:img_size + 75, :] elif i == 207: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + 250, img_width)) # a = np.copy(rot_x) # cv.rectangle(a, (200, 185), (img_size + 160, img_size - 95), (0, 0, 255), (2)) rot_x = rot_x[185:img_size - 95, 200:img_size + 160, :] else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(x, adjus_M, (img_height + 250, img_width)) # a = np.copy(rot_x) # cv.rectangle(a, (390, 115), (img_size+170, img_size-105), (0, 0, 255), (2)) rot_x = rot_x[115:img_size - 105, 390:img_size + 170, :] return rot_x def augment_data(args): data_for = 'train' # choice [train validation, test] imgs_splitted = True imgs_rotated = True imgs_flipped = True imgs_gamma_corrected = True # degrees # [19, 46, 57, 90, 114, 138, 161, 180, 207, 230, 247 270, 285, 322, 342] degrees = [19, 23*2, 19*3, 90, 19*6, 23*6,23*7,180, 23*9, 23*10,19*13, 270, 19*15,23*14,19*18] # if data_for=='train': # **************** training data ************************# if data_for=='train' and not args.use_nir: base_dataset_dir = args.dataset_dir.lower() + args.train_dataset + '/edges' GT_dir = os.path.join(base_dataset_dir, 'edge_maps/train/rgbr/aug') X_dir = os.path.join(base_dataset_dir, 'imgs/train/rgbr/aug') # this implementation is just for BIPED dataset gt_list = os.listdir(os.path.join(GT_dir,'real')) # gt_list.sort() x_list = os.listdir(os.path.join(X_dir,'real')) x_list.sort() n = len(gt_list) if len(x_list)==len(gt_list) else 0 if n==0: print('there is some inconsistence in the size of lists') sys.exit() # making 720x720 image size (splitting) ****************** # ******************************************************** if not imgs_splitted: tmp_img = cv.imread(os.path.join( os.path.join(X_dir, 'real'), x_list[0])) img_width = tmp_img.shape[1] img_height = tmp_img.shape[0] if not os.path.exists(os.path.join(GT_dir, 'p1')): os.makedirs(os.path.join(GT_dir, 'p1')) if not os.path.exists(os.path.join(GT_dir, 'p2')): os.makedirs(os.path.join(GT_dir, 'p2')) if not os.path.exists(os.path.join(X_dir, 'p1')): os.makedirs(os.path.join(X_dir, 'p1')) if not os.path.exists(os.path.join(X_dir, 'p2')): os.makedirs(os.path.join(X_dir, 'p2')) for i in range(n): x_tmp = cv.imread(os.path.join( os.path.join(X_dir,'real'), x_list[i])) gt_tmp = cv.imread(os.path.join( os.path.join(GT_dir,'real'),gt_list[i])) x_tmp1 = x_tmp[:,0:img_height,:] x_tmp2 = x_tmp[:,img_width-img_height:img_width,:] gt_tmp1 = gt_tmp[:,0:img_height] gt_tmp2 = gt_tmp[:,img_width-img_height:img_width] cv.imwrite(os.path.join(X_dir,os.path.join('p1',x_list[i])), x_tmp1) cv.imwrite(os.path.join(X_dir, os.path.join('p2', x_list[i])), x_tmp2) cv.imwrite(os.path.join(GT_dir, os.path.join('p1', gt_list[i])), gt_tmp1) cv.imwrite(os.path.join(GT_dir, os.path.join('p2', gt_list[i])), gt_tmp2) print('saved image: ', x_list[i], gt_list[i]) print('...split done') # *************** image rotation ******************************* if not imgs_rotated: # for p1 *********** folder_name='real' # choice [p1,p2,real] which are the source of files previously prepared # folder_name_x = 'real' if folder_name=='p1': x_aug_list = os.listdir(os.path.join(X_dir, 'p1')) x_aug_list.sort() gt_aug_list = os.listdir(os.path.join(GT_dir, 'p1')) gt_aug_list.sort() elif folder_name == 'p2': x_aug_list = os.listdir(os.path.join(X_dir, 'p2')) x_aug_list.sort() gt_aug_list = os.listdir(os.path.join(GT_dir, 'p2')) gt_aug_list.sort() elif folder_name=='real': x_aug_list = os.listdir(os.path.join(X_dir, 'real')) x_aug_list.sort() gt_aug_list = os.listdir(os.path.join(GT_dir, 'real')) gt_aug_list.sort() else: print("error reading folder name") sys.exit() if n==len(x_aug_list) and n== len(gt_aug_list): pass else: print("Error reading data. The is an inconsistency in the data ") sys.exit() tmp_img = cv.imread(os.path.join(X_dir, os.path.join(folder_name,x_aug_list[1]))) img_width = tmp_img.shape[1] img_height = tmp_img.shape[0] for i in (degrees): if folder_name=='p1': current_X_dir = X_dir+'/p1_rot_'+str(i) current_GT_dir = GT_dir+'/p1_rot_'+str(i) elif folder_name=='p2': current_X_dir = X_dir + '/p2_rot_' + str(i) current_GT_dir = GT_dir + '/p2_rot_' + str(i) elif folder_name=='real': current_X_dir = X_dir + '/real_rot_' + str(i) current_GT_dir = GT_dir + '/real_rot_' + str(i) else: print('error') sys.exit() if not os.path.exists(current_X_dir): os.makedirs(current_X_dir) if not os.path.exists(current_GT_dir): os.makedirs(current_GT_dir) for j in range(n): # i = degrees[j] if folder_name=='real': tmp_x = cv.imread(os.path.join(X_dir, os.path.join(folder_name, x_aug_list[j]))) tmp_gt = cv.imread(os.path.join(GT_dir, os.path.join(folder_name, gt_aug_list[j]))) rot_x,rot_gt=rotated_img_extractor(tmp_x, tmp_gt,img_width, img_height,i, True) # [19, 46, 90, 114, 138, 161, 180, 207, 230, 247, 270, 285, 322, 342] else: tmp_x = cv.imread(os.path.join(X_dir, os.path.join(folder_name,x_aug_list[j]))) tmp_gt = cv.imread(os.path.join(GT_dir, os.path.join(folder_name,gt_aug_list[j]))) rot_x, rot_gt = rotated_img_extractor(tmp_x, tmp_gt, img_width, img_height, i, True) cv.imwrite(os.path.join(current_GT_dir,gt_aug_list[j]),rot_gt) cv.imwrite(os.path.join(current_X_dir, x_aug_list[j]),rot_x) tmp_imgs = np.concatenate((rot_x,rot_gt), axis=1) cv.imshow("rotated", tmp_imgs) cv.waitKey(400) # 1000= 1 seg print("rotation with {} degrees fullfiled ".format(i)) cv.destroyAllWindows() print("... rotation done in ", folder_name) # **************** flipping horizontally *********** if not imgs_flipped: type_aug= '_flip' dir_list = os.listdir(X_dir) dir_list.sort() for i in (dir_list): X_list = os.listdir(os.path.join(X_dir,i)) X_list.sort() GT_list = os.listdir(os.path.join(GT_dir,i)) GT_list.sort() save_dir_x = X_dir+'/'+str(i)+type_aug save_dir_gt = GT_dir+'/'+str(i)+type_aug if not os.path.exists(save_dir_x): os.makedirs(save_dir_x) if not os.path.exists(save_dir_gt): os.makedirs(save_dir_gt) print("Working on the dir: ", os.path.join(X_dir,i),os.path.join(GT_dir,i) ) for j in range(n): x_tmp = cv.imread(os.path.join(X_dir,os.path.join(i,X_list[j]))) gt_tmp = cv.imread(os.path.join(GT_dir, os.path.join(i, GT_list[j]))) flip_x = np.fliplr(x_tmp) flip_gt = np.fliplr(gt_tmp) tmp_imgs = np.concatenate((flip_x, flip_gt), axis=1) cv.imshow("rotated", tmp_imgs) cv.waitKey(350) cv.imwrite(os.path.join(save_dir_gt,GT_list[j]),flip_gt) cv.imwrite(os.path.join(save_dir_x, X_list[j]),flip_x) print("End flipping file in {}".format(os.path.join(X_dir,i))) cv.destroyAllWindows() print("... Flipping data augmentation finished") # ***********Data augmentation based on gamma correction ********** if not imgs_gamma_corrected: gamma30 = '_ga30' gamma60 = '_ga60' dir_list = os.listdir(X_dir) dir_list.sort() for i in (dir_list): X_list = os.listdir(os.path.join(X_dir,i)) X_list.sort() GT_list = os.listdir(os.path.join(GT_dir,i)) GT_list.sort() save_dir_x30 = X_dir+'/'+str(i)+gamma30 save_dir_gt30 = GT_dir+'/'+str(i)+gamma30 save_dir_x60 = X_dir + '/' + str(i)+ gamma60 save_dir_gt60 = GT_dir + '/' + str(i)+ gamma60 if not os.path.exists(save_dir_x30): os.makedirs(save_dir_x30) if not os.path.exists(save_dir_gt30): os.makedirs(save_dir_gt30) if not os.path.exists(save_dir_x60): os.makedirs(save_dir_x60) if not os.path.exists(save_dir_gt60): os.makedirs(save_dir_gt60) print("Working on the dir: ", os.path.join(X_dir, i), os.path.join(GT_dir, i)) for j in range(n): x_tmp = cv.imread(os.path.join(X_dir, os.path.join(i, X_list[j]))) gt_tmp = cv.imread(os.path.join(GT_dir, os.path.join(i, GT_list[j]))) x_tmp = normalization_data_01(x_tmp) x_tmp = gamma_correction(x_tmp,0.4040,False) gam30_x = gamma_correction(x_tmp,0.3030,True) gam60_x = gamma_correction(x_tmp,0.6060, True) gam30_x = np.uint8(normalization_data_0255(gam30_x)) gam60_x = np.uint8(normalization_data_0255(gam60_x)) tmp_imgs1 = np.concatenate((gam30_x, gt_tmp), axis=1) tmp_imgs2 = np.concatenate((gam60_x, gt_tmp), axis=1) tmp_imgs = np.concatenate((tmp_imgs2,tmp_imgs1),axis=0) cv.imshow("gamma ", tmp_imgs) cv.waitKey(350) cv.imwrite(os.path.join(save_dir_gt30, GT_list[j]), gt_tmp) cv.imwrite(os.path.join(save_dir_x30, X_list[j]), gam30_x) cv.imwrite(os.path.join(save_dir_gt60, GT_list[j]), gt_tmp) cv.imwrite(os.path.join(save_dir_x60, X_list[j]), gam60_x) print("End gamma correction, file in {}".format(os.path.join(X_dir, i))) cv.destroyAllWindows() print("... gamma correction data augmentation finished") # ************** for validation ******************************** elif data_for=='validation' and not args.use_nir: train_GT_dir = args.dataset_dir + args.train_dataset + '/valid' train_X_dir = args.dataset_dir + args.train_dataset + '/valid' gt_list = os.listdir(os.path.join(train_GT_dir, 'GT_un')) # gt_list.sort() x_list = os.listdir(os.path.join(train_X_dir, 'X_un')) x_list.sort() n = len(gt_list) if len(x_list) == len(gt_list) else 0 if n == 0: print('there is some inconsistence in the size of lists') sys.exit() # making 720x720 image size (splitting) ****************** if not imgs_splitted: tmp_img = cv.imread(os.path.join( os.path.join(train_X_dir, 'X_un'), x_list[0])) img_width = tmp_img.shape[1] img_height = tmp_img.shape[0] for i in range(n): x_tmp = cv.imread(os.path.join( os.path.join(train_X_dir, 'X_un'), x_list[i])) gt_tmp = cv.imread(os.path.join( os.path.join(train_GT_dir, 'GT_un'), gt_list[i])) x_tmp1 = x_tmp[:, 0:img_height, :] x_tmp2 = x_tmp[:, img_width - img_height:img_width, :] gt_tmp1 = gt_tmp[:, 0:img_height] gt_tmp2 = gt_tmp[:, img_width - img_height:img_width] cv.imwrite(os.path.join(train_X_dir, os.path.join('X/p1', x_list[i])), x_tmp1) cv.imwrite(os.path.join(train_X_dir, os.path.join('X/p2', x_list[i])), x_tmp2) cv.imwrite(os.path.join(train_GT_dir, os.path.join('GT/p1', gt_list[i])), gt_tmp1) cv.imwrite(os.path.join(train_GT_dir, os.path.join('GT/p2', gt_list[i])), gt_tmp2) print('saved image: ', x_list[i], gt_list[i]) tmp_imgs = np.concatenate((x_tmp,gt_tmp), axis=1) cv.imshow("rotated", tmp_imgs) cv.waitKey(500) cv.destroyAllWindows() print('...split for validation finished') # image rotation if not imgs_rotated: folder_name = 'p2' # choice [p1,p2,GT_u] which are the source of files previously prepared X_dir = args.dataset_dir + args.train_dataset + '/valid' GT_dir = args.dataset_dir + args.train_dataset + '/valid' folder_name_x = 'X_un' if folder_name == 'p1': x_aug_list = os.listdir(os.path.join(X_dir, 'X/p1')) x_aug_list.sort() gt_aug_list = os.listdir(os.path.join(GT_dir, 'GT/p1')) gt_aug_list.sort() elif folder_name == 'p2': x_aug_list = os.listdir(os.path.join(X_dir, 'X/p2')) x_aug_list.sort() gt_aug_list = os.listdir(os.path.join(GT_dir, 'GT/p2')) gt_aug_list.sort() elif folder_name == 'GT_u': x_aug_list = os.listdir(os.path.join(X_dir, 'X_un')) x_aug_list.sort() gt_aug_list = os.listdir(os.path.join(GT_dir, 'GT_un')) gt_aug_list.sort() else: print("error reading folder name") sys.exit() # image size tmp_img = cv.imread(os.path.join(X_dir, os.path.join('X/' + folder_name, x_aug_list[1]))) img_width = tmp_img.shape[1] img_height = tmp_img.shape[0] for i in (degrees): if folder_name == 'p1': current_X_dir = X_dir + '/X/p1_rot_' + str(i) current_GT_dir = GT_dir + '/GT/p1_rot_' + str(i) elif folder_name == 'p2': current_X_dir = X_dir + '/X/p2_rot_' + str(i) current_GT_dir = GT_dir + '/GT/p2_rot_' + str(i) elif folder_name == 'GT_u': current_X_dir = X_dir + 'X/un_rot_' + str(i) current_GT_dir = GT_dir + 'GT/un_rot_' + str(i) else: print('error') sys.exit() if not os.path.exists(current_X_dir): os.makedirs(current_X_dir) if not os.path.exists(current_GT_dir): os.makedirs(current_GT_dir) for j in range(n): if folder_name == 'GT_un': tmp_x = cv.imread(os.path.join(X_dir, os.path.join(folder_name_x, x_aug_list[j]))) tmp_gt = cv.imread(os.path.join(GT_dir, os.path.join(folder_name, gt_aug_list[j]))) else: tmp_x = cv.imread(os.path.join(X_dir, os.path.join('X/' + folder_name, x_aug_list[j]))) tmp_gt = cv.imread(os.path.join(GT_dir, os.path.join('GT/' + folder_name, gt_aug_list[j]))) if i % 90 == 0: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(tmp_gt, adjus_M, (img_height, img_width)) elif i % 19 == 0: if i == 57: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(tmp_gt, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (100, 100), (720 - 100, 720 - 100), (0, 0, 255), (2)) rot_x = rot_x[100:720 - 100, 100:720 - 100, :] rot_gt = rot_gt[100:720 - 100, 100:720 - 100] # print("just for check 19: ", i, rot_x.shape) elif i == 285: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(tmp_gt, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (75, 75), (720 - 75, 720 - 75), (0, 0, 255), (2)) rot_x = rot_x[75:720 - 75, 75:720 - 75, :] rot_gt = rot_gt[75:720 - 75, 75:720 - 75] else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(tmp_gt, adjus_M, (img_height, img_width)) rot_x = rot_x[95:720 - 95, 95:720 - 95, :] rot_gt = rot_gt[95:720 - 95, 95:720 - 95] elif i % 23 == 0: if i == 161: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(tmp_gt, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x,(85,85),(720-85,720-85),(0,0,255),(2) ) rot_x = rot_x[85:720 - 85, 85:720 - 85, :] rot_gt = rot_gt[85:720 - 85, 85:720 - 85] # print("just for check 23: ", i, rot_x.shape) else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) rot_gt = cv.warpAffine(tmp_gt, adjus_M, (img_height, img_width)) rot_x = rot_x[105:720 - 105, 105:720 - 105, :] rot_gt = rot_gt[105:720 - 105, 105:720 - 105] else: print('Error using degrees for rotation') sys.exit() cv.imwrite(os.path.join(current_GT_dir, gt_aug_list[j]), rot_gt) cv.imwrite(os.path.join(current_X_dir, x_aug_list[j]), rot_x) tmp_imgs = np.concatenate((rot_x, rot_gt), axis=1) cv.imshow("rotated", tmp_imgs) cv.waitKey(1000) print("rotation with {} degrees fullfiled ".format(i)) cv.destroyAllWindows() print("... data rotation for validation finished ", folder_name) # flipping horizontally if not imgs_flipped: type_aug = 'flip_' X_dir = args.dataset_dir + args.train_dataset + '/valid/X' GT_dir = args.dataset_dir + args.train_dataset + '/valid/GT' dir_list = os.listdir(X_dir) dir_list.sort() for i in (dir_list): X_list = os.listdir(os.path.join(X_dir, i)) X_list.sort() GT_list = os.listdir(os.path.join(GT_dir, i)) GT_list.sort() save_dir_x = X_dir + '/' + type_aug + str(i) save_dir_gt = GT_dir + '/' + type_aug + str(i) if not os.path.exists(save_dir_x): os.makedirs(save_dir_x) if not os.path.exists(save_dir_gt): os.makedirs(save_dir_gt) print("Working on the dir: ", os.path.join(X_dir, i), os.path.join(GT_dir, i)) for j in range(n): x_tmp = cv.imread(os.path.join(X_dir, os.path.join(i, X_list[j]))) gt_tmp = cv.imread(os.path.join(GT_dir, os.path.join(i, GT_list[j]))) flip_x = np.fliplr(x_tmp) flip_gt = np.fliplr(gt_tmp) tmp_imgs = np.concatenate((flip_x, flip_gt), axis=1) cv.imshow("rotated", tmp_imgs) cv.waitKey(1000) cv.imwrite(os.path.join(save_dir_gt, GT_list[j]), flip_gt) cv.imwrite(os.path.join(save_dir_x, X_list[j]), flip_x) print("End flipping file in {}".format(os.path.join(X_dir, i))) cv.destroyAllWindows() print("... Flipping validation stage data augmentation finished") # ======================================================== # =====================Just for NIR ====================== # ======================================================== elif args.use_nir: if data_for == 'train': base_dataset_dir = args.dataset_dir + args.train_dataset + '/edges' X_dir = os.path.join(base_dataset_dir, 'imgs/train/nir/aug') x_list = os.listdir(os.path.join(X_dir, 'real')) x_list.sort() n = len(x_list) if n == 0: print('there is some inconsistence in the size of lists') sys.exit() # making 720x720 image size (splitting) ****************** if not imgs_splitted: tmp_img = cv.imread(os.path.join( os.path.join(X_dir, 'real'), x_list[0])) img_width = tmp_img.shape[1] img_height = tmp_img.shape[0] for i in range(n): x_tmp = cv.imread(os.path.join( os.path.join(X_dir, 'real'), x_list[i])) x_tmp1 = x_tmp[:, 0:img_height, :] x_tmp2 = x_tmp[:, img_width - img_height:img_width, :] cv.imwrite(os.path.join(X_dir, os.path.join('p1', x_list[i])), x_tmp1) cv.imwrite(os.path.join(X_dir, os.path.join('p2', x_list[i])), x_tmp2) print('saved image: ', x_list[i]) print('...split done') # ***************** image rotation ****************** if not imgs_rotated: # for p1 *********** folder_name = 'real' # choice [p1,p2,real] which are the source of files previously prepared # firstly X_un is not used folder_name_x = 'real' if folder_name == 'p1': x_aug_list = os.listdir(os.path.join(X_dir, 'p1')) x_aug_list.sort() elif folder_name == 'p2': x_aug_list = os.listdir(os.path.join(X_dir, 'p2')) x_aug_list.sort() elif folder_name == 'real': x_aug_list = os.listdir(os.path.join(X_dir, 'real')) x_aug_list.sort() else: print("error reading folder name") sys.exit() # image size tmp_img = cv.imread(os.path.join( X_dir, os.path.join(folder_name, x_aug_list[1]))) img_width = tmp_img.shape[1] img_height = tmp_img.shape[0] for i in (degrees): if folder_name == 'p1': current_X_dir = X_dir + '/p1_rot_' + str(i) elif folder_name == 'p2': current_X_dir = X_dir + '/p2_rot_' + str(i) elif folder_name == 'real': current_X_dir = X_dir + '/real_rot_' + str(i) else: print('error') sys.exit() if not os.path.exists(current_X_dir): os.makedirs(current_X_dir) for j in range(n): # i = degrees[j] if folder_name == 'real': tmp_x = cv.imread(os.path.join(X_dir, os.path.join(folder_name_x, x_aug_list[j]))) rot_x = rotated_img_extractor(x=tmp_x, gt=None, img_width=img_width, img_height=img_height, i=i, two_data=False) else: tmp_x = cv.imread(os.path.join(X_dir, os.path.join(folder_name, x_aug_list[j]))) rot_x = rotated_img_extractor(x=tmp_x, gt=None, img_width=img_width, img_height=img_height, i=i, two_data=False) cv.imwrite(os.path.join(current_X_dir, x_aug_list[j]), rot_x) cv.imshow("rotated", rot_x) cv.waitKey(450) print("rotation with {} degrees fullfiled ".format(i)) cv.destroyAllWindows() print("... rotation done in ", folder_name) # flipping horizontally if not imgs_flipped: type_aug = '_flip' dir_list = os.listdir(X_dir) dir_list.sort() for i in (dir_list): X_list = os.listdir(os.path.join(X_dir, i)) X_list.sort() save_dir_x = X_dir + '/' + str(i)+ type_aug if not os.path.exists(save_dir_x): os.makedirs(save_dir_x) print("Working on the dir: ", os.path.join(X_dir, i), i) for j in range(n): x_tmp = cv.imread(os.path.join(X_dir, os.path.join(i, X_list[j]))) flip_x = np.fliplr(x_tmp) cv.imshow("Flipping", x_tmp) cv.waitKey(450) cv.imwrite(os.path.join(save_dir_x, X_list[j]), flip_x) print("End flipping file in {}".format(os.path.join(X_dir, i))) cv.destroyAllWindows() print("... Flipping data augmentation finished") if not imgs_gamma_corrected: gamma30 = '_ga30' gamma60 = '_ga60' dir_list = os.listdir(X_dir) dir_list.sort() for i in (dir_list): X_list = os.listdir(os.path.join(X_dir, i)) X_list.sort() save_dir_x30 = X_dir + '/' + str(i) + gamma30 save_dir_x60 = X_dir + '/' + str(i) + gamma60 if not os.path.exists(save_dir_x30): os.makedirs(save_dir_x30) if not os.path.exists(save_dir_x60): os.makedirs(save_dir_x60) print("Working on the dir: ", os.path.join(X_dir, i)) for j in range(n): x_tmp = cv.imread(os.path.join(X_dir, os.path.join(i, X_list[j]))) x_tmp = normalization_data_01(x_tmp) x_tmp = gamma_correction(x_tmp, 0.4040, False) gam30_x = gamma_correction(x_tmp, 0.3030, True) gam60_x = gamma_correction(x_tmp, 0.6060, True) gam30_x = np.uint8(normalization_data_0255(gam30_x)) gam60_x = np.uint8(normalization_data_0255(gam60_x)) tmp_imgs = np.concatenate((gam30_x, gam60_x), axis=1) cv.imshow("gamma ", tmp_imgs) cv.waitKey(450) cv.imwrite(os.path.join(save_dir_x30, X_list[j]), gam30_x) cv.imwrite(os.path.join(save_dir_x60, X_list[j]), gam60_x) print("End gamma correction, file in {}".format(os.path.join(X_dir, i))) cv.destroyAllWindows() print("... gamma correction data augmentation finished") elif data_for == 'validation': train_X_dir = args.dataset_dir + args.train_dataset + '/nir_valid' x_list = os.listdir(os.path.join(train_X_dir, 'X_un')) x_list.sort() n = len(x_list) if n == 0: print('there is some inconsistence in the size of lists') sys.exit() # making 720x720 image size (splitting) ****************** if not imgs_splitted: tmp_img = cv.imread(os.path.join( os.path.join(train_X_dir, 'X_un'), x_list[0])) img_width = tmp_img.shape[1] img_height = tmp_img.shape[0] for i in range(n): x_tmp = cv.imread(os.path.join( os.path.join(train_X_dir, 'X_un'), x_list[i])) x_tmp1 = x_tmp[:, 0:img_height, :] x_tmp2 = x_tmp[:, img_width - img_height:img_width, :] cv.imwrite(os.path.join(train_X_dir, os.path.join('X/p1', x_list[i])), x_tmp1) cv.imwrite(os.path.join(train_X_dir, os.path.join('X/p2', x_list[i])), x_tmp2) print('saved image: ', x_list[i]) print('... validation split done') # image rotation if not imgs_rotated: X_dir = args.dataset_dir + args.train_dataset + '/nir_valid' # for p1 *********** folder_name = 'p2' # choice [p1,p2,X_un] which are the source of files previously prepared # firstly X_un is not used folder_name_x = 'X_un' if folder_name == 'p1': x_aug_list = os.listdir(os.path.join(X_dir, 'X/p1')) x_aug_list.sort() elif folder_name == 'p2': x_aug_list = os.listdir(os.path.join(X_dir, 'X/p2')) x_aug_list.sort() elif folder_name == 'X_un': x_aug_list = os.listdir(os.path.join(X_dir, 'X_un')) x_aug_list.sort() else: print("error reading folder name") sys.exit() # image size tmp_img = cv.imread(os.path.join(X_dir, os.path.join('X/' + folder_name, x_aug_list[1]))) img_width = tmp_img.shape[1] img_height = tmp_img.shape[0] for i in (degrees): if folder_name == 'p1': current_X_dir = X_dir + '/X/p1_rot_' + str(i) elif folder_name == 'p2': current_X_dir = X_dir + '/X/p2_rot_' + str(i) elif folder_name == 'X_un': current_X_dir = X_dir + 'X/un_rot_' + str(i) else: print('error') sys.exit() if not os.path.exists(current_X_dir): os.makedirs(current_X_dir) for j in range(n): # i = degrees[j] if folder_name == 'X_un': tmp_x = cv.imread( os.path.join(X_dir, os.path.join(folder_name_x, x_aug_list[j]))) else: tmp_x = cv.imread( os.path.join(X_dir, os.path.join('X/' + folder_name, x_aug_list[j]))) if i % 90 == 0: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) elif i % 19 == 0: if i == 57: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (100, 100), (720 - 100, 720 - 100), (0, 0, 255), (2)) rot_x = rot_x[100:720 - 100, 100:720 - 100, :] # print("just for check 19: ", i, rot_x.shape) elif i == 285: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (75, 75), (720 - 75, 720 - 75), (0, 0, 255), (2)) rot_x = rot_x[75:720 - 75, 75:720 - 75, :] # print("just for check 19: ", i, rot_x.shape) else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x,(95,95),(720-95,720-95),(0,0,255),(2) ) rot_x = rot_x[95:720 - 95, 95:720 - 95, :] # print("just for check 19: ", i, rot_x.shape) elif i % 23 == 0: if i == 161: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x,(85,85),(720-85,720-85),(0,0,255),(2) ) rot_x = rot_x[85:720 - 85, 85:720 - 85, :] # print("just for check 23: ", i, rot_x.shape) else: adjus_M = cv.getRotationMatrix2D((img_width / 2, img_height / 2), i, 1) rot_x = cv.warpAffine(tmp_x, adjus_M, (img_height, img_width)) # cv.rectangle(rot_x, (105, 105), (720 - 105, 720 - 105), (0, 0, 255), (2)) rot_x = rot_x[105:720 - 105, 105:720 - 105, :] # print("just for check 23:", i, rot_x.shape) else: print('Error using degrees for rotation') sys.exit() cv.imwrite(os.path.join(current_X_dir, x_aug_list[j]), rot_x) cv.imshow("rotated", rot_x) cv.waitKey(350) print("rotation with {} degrees fullfiled ".format(i)) cv.destroyAllWindows() print("... rotation done in ", folder_name) # flipping horizontally if not imgs_flipped: type_aug = 'flip_' X_dir = args.dataset_dir + args.train_dataset + '/nir_valid/X' dir_list = os.listdir(X_dir) dir_list.sort() for i in (dir_list): X_list = os.listdir(os.path.join(X_dir, i)) X_list.sort() save_dir_x = X_dir + '/' + type_aug + str(i) if not os.path.exists(save_dir_x): os.makedirs(save_dir_x) print("Working on the dir: ", os.path.join(X_dir, i), i) for j in range(n): x_tmp = cv.imread(os.path.join(X_dir, os.path.join(i, X_list[j]))) flip_x = np.fliplr(x_tmp) cv.imshow("rotated", x_tmp) cv.waitKey(350) cv.imwrite(os.path.join(save_dir_x, X_list[j]), flip_x) print("End flipping file in {}".format(os.path.join(X_dir, i))) cv.destroyAllWindows() print("... Flipping validation data augmentation finished") else: print("This part is not finished yet") else: print("Error, just train and validation code have written") def data_parser(args): if args.model_state=='train': train_files_name = args.train_list # dataset base dir base_dir = os.path.join(args.dataset_dir, args.train_dataset,'edges') \ if args.train_dataset.lower() == 'biped' else os.path.join(args.dataset_dir, args.train_dataset) train_list_path = os.path.join(base_dir, train_files_name) train_list = read_files_list(train_list_path) train_samples = split_pair_names(args,train_list, base_dir=base_dir) n_train = len(train_samples) print_info(" Enterely training set-up from {}, size: {}".format(train_list_path, n_train)) all_train_ids = np.arange(n_train) np.random.shuffle(all_train_ids) train_ids = all_train_ids[:int(args.train_split * len(train_list))] valid_ids = all_train_ids[int(args.train_split * len(train_list)):] print_info("Training set-up from {}, size: {}".format(train_list_path, len(train_ids))) print_info("Validation set-up from {}, size: {}".format(train_list_path, len(valid_ids))) cache_info = { "files_path": train_samples, "n_files": n_train, "train_indices": train_ids, "validation_indices": valid_ids } return cache_info # ************** for testing ********************** elif args.model_state=='test': base_dir = os.path.join(args.dataset_dir,args.test_dataset,'edges')\ if args.test_dataset.upper()=='BIPED' else os.path.join(args.dataset_dir,args.test_dataset) if args.test_dataset.upper() == "BIPED": test_files_name = args.test_list test_list_path = os.path.join(base_dir,test_files_name) test_list = read_files_list(test_list_path) test_samples = split_pair_names(args, test_list, base_dir) n_test = len(test_samples) print_info(" Enterely testing set-up from {}, size: {}".format(test_list_path, n_test)) test_ids = np.arange(n_test) # np.random.shuffle(test_ids) print_info("testing set-up from {}, size: {}".format(test_list_path, len(test_ids))) cache_out = [test_samples, test_ids] return cache_out elif args.test_dataset == "BSDS": test_files_name = args.test_list test_list_path = os.path.join(args.dataset_dir + args.test_dataset, test_files_name) test_list = read_files_list(test_list_path) test_samples = split_pair_names(args,test_list, args.dataset_dir + args.test_dataset) n_test = len(test_samples) print_info(" Enterely testing set-up from {}, size: {}".format(test_list_path, n_test)) test_ids = np.arange(n_test) # np.random.shuffle(test_ids) print_info("testing set-up from {}, size: {}".format(test_list_path, len(test_ids))) cache_out = [test_samples, test_ids] return cache_out else: # for NYUD test_files_name = args.test_list test_list_path = os.path.join(args.dataset_dir + args.test_dataset, test_files_name) test_list = read_files_list(test_list_path) test_samples = split_pair_names(args,test_list, args.dataset_dir + args.test_dataset) n_test = len(test_samples) print_info(" Enterely testing set-up from {}, size: {}".format(test_list_path, n_test)) test_ids = np.arange(n_test) # np.random.shuffle(test_ids) print_info("testing set-up from {}, size: {}".format(test_list_path, len(test_ids))) cache_out = [test_samples, test_ids] return cache_out else: print_error("The model state is just train and test") sys.exit() # ___________batch management ___________ def get_batch(arg,file_list, batch=None, use_batch=True): if use_batch: file_names =[] images=[] edgemaps=[] for idx, b in enumerate(batch): x = cv.imread(file_list[b][0]) # Image.open(file_list[b][0]) y = cv.imread(file_list[b][1]) # Image.open(file_list[b][1]) if arg.model_state=='test': pass else: x = cv.resize(x, dsize=(arg.image_width,arg.image_height)) # x.resize((arg.image_width, arg.image_height)) y = cv.resize(y,dsize=(arg.image_width, arg.image_height))# y.resize((arg.image_width, arg.image_height)) # pay attention here x = np.array(x, dtype=np.float32) # x = x[:, :, arg.channel_swap] # while using opencv it is not necessary x -= arg.mean_pixel_values[0:3] y = cv.cvtColor(y, cv.COLOR_BGR2GRAY)# np.array(y.convert('L'), dtype=np.float32) y = np.array(y, dtype=np.float32) if arg.train_dataset.lower()!='biped': y[y < 107] = 0 # first nothing second <50 third <30 y[y >= 107] = 255.0 # else: # y[y < 51] = 0 # first 100 if arg.target_regression: bin_y = y/255.0 else: bin_y = np.zeros_like(y) bin_y[np.where(y)]=1 bin_y = bin_y if bin_y.ndim ==2 else bin_y[:,:,0] bin_y = np.expand_dims(bin_y,axis=2) images.append(x) edgemaps.append(bin_y) file_names.append(file_list[b]) return images, edgemaps, file_names else: # for testing re-coding is needed if arg.test_dataset=='BIPED' and (arg.model_state=='test' and arg.use_nir): x_nir = Image.open(file_list[0]) x_rgb = Image.open(file_list[1]) real_size = x_rgb.shape # y = Image.open(file_list[2]) if not arg.image_width % 6 == 0 and arg.image_width > 1000: x_nir = x_nir.resize((arg.image_width, arg.image_height)) #-- x_rgb = x_rgb.resize((arg.image_width, arg.image_height)) # -- else: x_nir = x_nir.resize((arg.image_width, arg.image_height)) x_rgb = x_rgb.resize((arg.image_width, arg.image_height)) # y = y.resize((arg.image_width, arg.image_height)) x_nir = x_nir.convert("L") # pay attention here x_nir = np.array(x_nir, dtype=np.float32) x_nir = np.expand_dims(x_nir, axis=2) x_rgb = np.array(x_rgb, dtype=np.float32) x_rgb = x_rgb[:, :, arg.channel_swap] x = np.concatenate((x_rgb, x_nir), axis=2) x -= arg.mean_pixel_values # y = np.array(y.convert('L'), dtype=np.float32) # if arg.target_regression: # bin_y = y / 255.0 # else: # bin_y = np.zeros_like(y) # bin_y[np.where(y)] = 1 # # bin_y = bin_y if bin_y.ndim == 2 else bin_y[:, :, 0] # bin_y = np.expand_dims(bin_y, axis=2) images =x edgemaps = None file_info = (file_list[2], real_size) else: x = cv.imread(file_list[0]) y = cv.imread(file_list[1]) real_size = x.shape if arg.test_dataset.lower()=='biped' or arg.test_dataset.lower()=='multicue': pass else: x = cv.resize(x, dsize=(arg.image_width, arg.image_height)) y = cv.resize(y, dsize=(arg.image_width, arg.image_height)) x = np.array(x, dtype=np.float32) # x = x[:, :, arg.channel_swap] # while using opencv it is not necessary x -= arg.mean_pixel_values[:-1] y = cv.cvtColor(y, cv.COLOR_BGR2GRAY) y = np.array(y, dtype=np.float32) if arg.target_regression: bin_y = y / 255.0 else: bin_y = np.zeros_like(y) bin_y[np.where(y)] = 1 bin_y = bin_y if bin_y.ndim == 2 else bin_y[:, :, 0] bin_y = np.expand_dims(bin_y, axis=2) images = x edgemaps = bin_y file_info = (file_list[1],real_size) return images, edgemaps, file_info def get_training_batch(arg, data_ids): train_ids = data_ids['train_indices'] file_list= data_ids['files_path'] batch_ids = np.random.choice(train_ids,arg.batch_size_train) return get_batch(arg, file_list, batch_ids) def get_validation_batch(arg, data_ids): if arg.use_nir: valid_ids = data_ids['validation_indices'] file_list = data_ids['files_path'] batch_ids = np.random.choice(valid_ids, arg.batch_size_val) else: valid_ids = data_ids['validation_indices'] file_list= data_ids['files_path'] batch_ids = np.random.choice(valid_ids,arg.batch_size_val) return get_batch(arg,file_list,batch_ids) def get_testing_batch(arg,list_ids, use_batch=True, i=None): if use_batch: test_ids = list_ids[1] file_list = list_ids[0] batch_ids = test_ids[i:i + arg.batch_size_test] return get_batch(arg,file_list,batch_ids) else: return get_batch(arg, list_ids[0],list_ids[1], use_batch=False) def open_images(file_list): if len(file_list)>2 and not len(file_list)==3: imgs=[] file_names = [] for i in range(len(file_list)): tmp = Image.open(file_list[i]) imgs.append(tmp) file_names.append(file_list[i]) elif len(file_list)>2 and len(file_list)==3: imgs = Image.open(file_list[2]) file_names = file_list else: imgs = Image.open(file_list[1]) file_names= file_list return imgs, file_names # for testing on single images def get_single_image(args,file_path=None): if file_path is None: imgs_name = ["CLASSIC", None] img_dir = 'data' if args.test_dataset in imgs_name else os.path.join(args.dataset_dir, args.test_dataset) file_list = os.listdir(img_dir) data =[] for i in file_list: data.append(os.path.join(img_dir,i)) return data else: img = cv.imread(file_path) h,w,c=img.shape x = np.array(img, dtype=np.float32) if h==args.image_height and w==args.image_width: pass else: x = cv.resize(x, dsize=(args.image_width, args.image_height)) # x = x[:, :, arg.channel_swap] # while using opencv it is not necessary x -= args.mean_pixel_values[:-1] img_info = (file_path,img.shape) return x, img_info def visualize_result(imgs_list, arg): """ function for tensorflow results :param imgs_list: a list of prediction, gt and input data :param arg: :return: one image with the whole of imgs_list data """ n_imgs = len(imgs_list) data_list =[] for i in range(n_imgs): tmp = imgs_list[i] if len(tmp.shape)==3 or len(tmp.shape)==2: if len(tmp.shape)==2: tmp = np.uint8(image_normalization(tmp)) tmp=cv.cvtColor(tmp, cv.COLOR_GRAY2BGR) tmp = cv.bitwise_not(tmp) else: if tmp.shape[-1]==3: tmp=restore_rgb([arg.channel_swap,arg.mean_pixel_values[:3]],tmp) tmp = np.uint8(image_normalization(tmp)) else: tmp = np.squeeze(tmp) tmp = np.uint8(image_normalization(tmp)) tmp = cv.cvtColor(tmp, cv.COLOR_GRAY2BGR) tmp = cv.bitwise_not(tmp) else: tmp= np.squeeze(tmp) if len(tmp.shape) == 2: tmp = np.uint8(image_normalization(tmp)) tmp = cv.bitwise_not(tmp) tmp = cv.cvtColor(tmp, cv.COLOR_GRAY2BGR) else: tmp = np.uint8(image_normalization(tmp)) data_list.append(tmp) img = data_list[0] if n_imgs % 2 == 0: imgs = np.zeros((img.shape[0] * 2 + 10, img.shape[1] * (n_imgs // 2) + ((n_imgs // 2 - 1) * 5), 3)) else: imgs = np.zeros((img.shape[0] * 2 + 10, img.shape[1] * ((1 + n_imgs) // 2) + ((n_imgs // 2) * 5), 3)) n_imgs += 1 k=0 imgs = np.uint8(imgs) i_step = img.shape[0]+10 j_step = img.shape[1]+5 for i in range(2): for j in range(n_imgs//2): if k