def run(src_dirs=constants.data_source_folders,
working_dir=constants.working_dir):
for src_dir_index, src_dir in enumerate(src_dirs):
shape_file_path = os.path.join(src_dir, "shapes/shapes.shp")
add_shapefile_classes_to_label_dictionary(shape_file_path)
print(str(len(classes)) + " classes present in dataset:")
print(classes)
utils.save_obj(classes, os.path.join(working_dir, "labelmap.pkl"))
(temp_dir, mask_tiles_dir, image_tiles_dir) = make_folders(working_dir)
for src_dir_index, src_dir in enumerate(src_dirs):
shape_file_path = os.path.join(src_dir, "shapes/shapes.shp")
images_folder = os.path.join(src_dir, "images")
print("Tiling all images in input folder: " + src_dir)
for image_path in progressbar.progressbar(
utils.get_all_image_paths_in_folder(images_folder)):
projected_image_path = os.path.join(temp_dir,
os.path.basename(image_path))
resize_image_and_change_coordinate_system(image_path,
projected_image_path)
image_path = projected_image_path
'''
#Change Coordinate System of Image if necessary
proj = osr.SpatialReference(wkt=gdal.Open(image_path).GetProjection())
epsg_code_of_image = proj.GetAttrValue('AUTHORITY',1)
if epsg_code_of_image != EPSG_TO_WORK_WITH:
projected_image_path = os.path.join(temp_dir,os.path.basename(image_path))
gdal.Warp(projected_image_path,image_path,dstSRS='EPSG:'+str(EPSG_TO_WORK_WITH))
image_path = projected_image_path
'''
mask_image_path = os.path.join(
temp_dir,
os.path.basename(image_path).replace(".tif", "_mask.tif"))
all_polygons = get_all_polygons_from_shapefile(shape_file_path)
all_polygons = convert_polygon_coords_to_pixel_coords(
all_polygons, image_path)
make_mask_image(image_path, mask_image_path, all_polygons)
tile_image(image_path, image_tiles_dir, src_dir_index)
tile_image(mask_image_path, mask_tiles_dir, src_dir_index)
#split_train_dir(image_tiles_dir, mask_tiles_dir, val_image_tiles_dir, val_mask_tiles_dir)
utils.delete_folder_contents(temp_dir)
shutil.rmtree(temp_dir)
def make_folders(working_dir):
temp_dir = os.path.join(working_dir,"temp")
os.makedirs(temp_dir,exist_ok=True)
utils.delete_folder_contents(temp_dir)
tiles_dir = os.path.join(temp_dir,"tiles/0")
os.makedirs(tiles_dir,exist_ok=True)
pred_tiles_dir = os.path.join(temp_dir,"pred_tiles/0")
os.makedirs(pred_tiles_dir,exist_ok=True)
return [temp_dir, tiles_dir, pred_tiles_dir]
def make_folders(project_dir):
training_data_dir = os.path.join(project_dir, "training_data")
folders = [
'train_frames/0', 'train_masks/0', 'val_frames/0', 'val_masks/0',
'test_frames/0', 'test_masks/0'
]
full_folder_paths = []
for folder in folders:
full_folder_path = os.path.join(training_data_dir, folder)
os.makedirs(full_folder_path, exist_ok=True)
utils.delete_folder_contents(full_folder_path)
full_folder_paths.append(full_folder_path)
return full_folder_paths
def make_folders(project_dir):
temp_dir = os.path.join(project_dir,"temp")
os.makedirs(temp_dir,exist_ok=True)
utils.delete_folder_contents(temp_dir)
training_data_dir = os.path.join(project_dir,"training_data")
os.makedirs(training_data_dir,exist_ok=True)
utils.delete_folder_contents(training_data_dir)
mask_tiles_dir = os.path.join(training_data_dir,"masks")
os.makedirs(mask_tiles_dir,exist_ok=True)
image_tiles_dir = os.path.join(training_data_dir,"images")
os.makedirs(image_tiles_dir,exist_ok=True)
return (temp_dir,mask_tiles_dir,image_tiles_dir)
def run(predict_folder,output_folder, working_dir=constants.working_dir, batch_size = constants.batch_size):
[temp_dir, tiles_dir, pred_tiles_dir] = make_folders(working_dir)
classes = utils.load_obj(constants.label_map)
utils.create_color_legend(classes, os.path.join(output_folder,"color_legend.png"))
print("Preparing image tiles...",flush=True)
images_to_predict = utils.get_all_image_paths_in_folder(predict_folder)
for image_path in progressbar.progressbar(images_to_predict):
utils.tile_image(image_path, tiles_dir, classes, tile_size=256, overlap=0)
print("Loading Trained Model...",flush=True)
all_tile_paths = utils.get_all_image_paths_in_folder(tiles_dir)
#test_frames_dir = "C:/Users/johan/Desktop/working_dir/training_data/train_frames/0"
#test_masks_dir = "C:/Users/johan/Desktop/working_dir/training_data/train_masks/0"
#print(len(classes))
model = unet_utils.get_small_unet(n_filters = 32,num_classes=len(classes),batch_size=batch_size)
model.compile(optimizer='adam', loss="categorical_crossentropy", metrics=[unet_utils.tversky_loss,unet_utils.dice_coef,'accuracy'])
#model.summary()
model_save_path = constants.trained_model
model.load_weights(model_save_path)
data_generator = unet_utils.DataGeneratorWithFileNames(tiles_dir,classes,batch_size=batch_size)
num_batches = int(np.ceil(float(len(all_tile_paths)) / float(batch_size)))
tile_index = 0
print("Predicting Tiles...",flush=True)
for batch_number in progressbar.progressbar(range(0,num_batches)):
img_batch,filenames=next(data_generator)
pred_all= model.predict(img_batch)
#pred = model.predict_generator(data_generator,steps=64)
#reassemble(image_path,mask_path,pred,classes)
for i in range(0,np.shape(pred_all)[0]):
if(tile_index == len(all_tile_paths)):
break
predicted_mask = utils.onehot_to_rgb(pred_all[i],classes)
img = img_batch[i] *255
mask_pred_path = os.path.join(pred_tiles_dir,os.path.basename(filenames[i]))
img_path = os.path.join(pred_tiles_dir,os.path.basename(filenames[i].replace(".png","a.png")))
save_array_as_image(mask_pred_path, predicted_mask)
save_array_as_image(img_path, img)
tile_index += 1
#mask = onehot_to_rgb(batch_mask[i],classes)
#img = batch_img[i] *255
#img_path = os.path.join("C:/Users/johan/Desktop/test", str(j) + "_" + str(i) + "_img.png")
#mask_path = os.path.join("C:/Users/johan/Desktop/test", str(j) + "_" + str(i) + "_mask.png")
#save_array_as_image(mask_path,mask)
#save_array_as_image(img_path, img)
print("Reassembling Tiles...",flush=True)
for image_path in progressbar.progressbar(images_to_predict):
dst_image_path = os.path.join(output_folder, os.path.basename(image_path))
reassemble(image_path,dst_image_path,pred_tiles_dir)
print("Cleaning up...",flush=True)
utils.delete_folder_contents(temp_dir)