def inference_from_filepath(config, in_filepaths, backbone):
# --- Online transform performed on the device (GPU):
eval_online_cuda_transform = data_transforms.get_eval_online_cuda_transform(config)
print("Loading model...")
model = FrameFieldModel(config, backbone=backbone, eval_transform=eval_online_cuda_transform)
model.to(config["device"])
checkpoints_dirpath = run_utils.setup_run_subdir(config["eval_params"]["run_dirpath"], config["optim_params"]["checkpoints_dirname"])
model = inference.load_checkpoint(model, checkpoints_dirpath, config["device"])
model.eval()
# Read image
pbar = tqdm(in_filepaths, desc="Infer images")
for in_filepath in pbar:
pbar.set_postfix(status="Loading image")
image = skimage.io.imread(in_filepath)
if 3 < image.shape[2]:
print_utils.print_info(f"Image {in_filepath} has more than 3 channels. Keeping the first 3 channels and discarding the rest...")
image = image[:, :, :3]
elif image.shape[2] < 3:
print_utils.print_error(f"Image {in_filepath} has only {image.shape[2]} channels but the network expects 3 channels.")
raise ValueError
image_float = image / 255
mean = np.mean(image_float.reshape(-1, image_float.shape[-1]), axis=0)
std = np.std(image_float.reshape(-1, image_float.shape[-1]), axis=0)
sample = {
"image": torchvision.transforms.functional.to_tensor(image)[None, ...],
"image_mean": torch.from_numpy(mean)[None, ...],
"image_std": torch.from_numpy(std)[None, ...],
"image_filepath": [in_filepath],
}
pbar.set_postfix(status="Inference")
tile_data = inference.inference(config, model, sample, compute_polygonization=True)
tile_data = local_utils.batch_to_cpu(tile_data)
# Remove batch dim:
tile_data = local_utils.split_batch(tile_data)[0]
pbar.set_postfix(status="Saving output")
base_filepath = os.path.splitext(in_filepath)[0]
if config["compute_seg"]:
seg_mask = 0.5 < tile_data["seg"][0]
save_utils.save_seg_mask(seg_mask, base_filepath + ".mask", tile_data["image_filepath"])
save_utils.save_seg(tile_data["seg"], base_filepath, "seg", tile_data["image_filepath"])
save_utils.save_seg_luxcarta_format(tile_data["seg"], base_filepath, "seg_luxcarta_format", tile_data["image_filepath"])
if config["compute_crossfield"]:
save_utils.save_crossfield(tile_data["crossfield"], base_filepath, "crossfield")
if "poly_viz" in config["eval_params"]["save_individual_outputs"] and \
config["eval_params"]["save_individual_outputs"]["poly_viz"]:
save_utils.save_poly_viz(tile_data["image"], tile_data["polygons"], tile_data["polygon_probs"], base_filepath, "poly_viz")
def __init__(self, rank, gpu, config, model, optimizer, loss_func,
run_dirpath, init_checkpoints_dirpath=None, lr_scheduler=None):
self.rank = rank
self.gpu = gpu
self.config = config
self.model = model
self.optimizer = optimizer
self.lr_scheduler = lr_scheduler
self.loss_func = loss_func
self.init_checkpoints_dirpath = init_checkpoints_dirpath
logs_dirpath = run_utils.setup_run_subdir(run_dirpath, config["optim_params"]["logs_dirname"])
self.checkpoints_dirpath = run_utils.setup_run_subdir(run_dirpath, config["optim_params"]["checkpoints_dirname"])
if self.rank == 0:
self.logs_dirpath = logs_dirpath
train_logs_dirpath = os.path.join(self.logs_dirpath, "train")
val_logs_dirpath = os.path.join(self.logs_dirpath, "val")
self.train_writer = SummaryWriter(train_logs_dirpath)
self.val_writer = SummaryWriter(val_logs_dirpath)
else:
self.logs_dirpath = self.train_writer = self.val_writer = None
def __init__(self, gpu: int, config: dict, shared_dict, barrier, model,
run_dirpath):
self.gpu = gpu
self.config = config
assert 0 < self.config["eval_params"]["batch_size_mult"], \
"batch_size_mult in polygonize_params should be at least 1."
self.shared_dict = shared_dict
self.barrier = barrier
self.model = model
self.checkpoints_dirpath = run_utils.setup_run_subdir(
run_dirpath, config["optim_params"]["checkpoints_dirname"])
self.eval_dirpath = os.path.join(config["data_root_dir"], "eval_runs",
os.path.split(run_dirpath)[-1])
if self.gpu == 0:
os.makedirs(self.eval_dirpath, exist_ok=True)
print_utils.print_info("Saving eval outputs to {}".format(
self.eval_dirpath))
def polygonize_masks(run_dirpath, images_dirpath, gt_filepath, in_filepath,
out_filepath, batch_size, batch_size_mult):
coco_gt = COCO(gt_filepath)
coco_dt = coco_gt.loadRes(in_filepath)
# --- Load model --- #
# Load run's config file:
config = run_utils.load_config(config_dirpath=run_dirpath)
if config is None:
print_utils.print_error(
"ERROR: cannot continue without a config file. Exiting now...")
sys.exit()
config["backbone_params"][
"pretrained"] = False # Don't load pretrained model
backbone = get_backbone(config["backbone_params"])
eval_online_cuda_transform = data_transforms.get_eval_online_cuda_transform(
config)
model = FrameFieldModel(config,
backbone=backbone,
eval_transform=eval_online_cuda_transform)
model.to(config["device"])
checkpoints_dirpath = run_utils.setup_run_subdir(
run_dirpath, config["optim_params"]["checkpoints_dirname"])
model = inference.load_checkpoint(model, checkpoints_dirpath,
config["device"])
model.eval()
# --- Polygonize input COCO mask detections --- #
img_ids = coco_dt.getImgIds()
# img_ids = sorted(img_ids)[:1] # TODO: rm limit
output_annotations = []
model_data_list = [
] # Used to accumulate inputs and run model inference on it.
poly_data_list = [
] # Used to accumulate inputs and run polygonization on it.
for img_id in tqdm(img_ids, desc="Polygonizing"):
# Load image
img = coco_gt.loadImgs(img_id)[0]
image = skimage.io.imread(
os.path.join(images_dirpath, img["file_name"]))
# Draw mask from input COCO mask annotations
mask_image = np.zeros((img["height"], img["width"]))
score_image = np.zeros((img["height"], img["width"]))
dts = coco_dt.loadAnns(coco_dt.getAnnIds(imgIds=img_id))
for dt in dts:
dt_mask = cocomask.decode(dt["segmentation"])
mask_image = np.maximum(mask_image, dt_mask)
score_image = np.maximum(score_image, dt_mask * dt["score"])
# Accumulate inputs into the current batch
sample_data = {
"img_id": [img_id],
"mask_image":
torch_lydorn.torchvision.transforms.functional.to_tensor(
mask_image)[None, ...].float(),
"score_image":
torch_lydorn.torchvision.transforms.functional.to_tensor(
score_image)[None, ...].float(),
"image":
torch_lydorn.torchvision.transforms.functional.to_tensor(image)[
None, ...],
"image_mean":
torch.tensor(image_mean)[None, ...],
"image_std":
torch.tensor(image_std)[None, ...]
}
# Accumulate batch for running the model
model_data_list.append(sample_data)
if len(model_data_list) == batch_size:
# Run model
tile_data = run_model(config, model, model_data_list)
model_data_list = [] # Empty model batch
# Accumulate batch for running the polygonization
poly_data_list.append(tile_data)
if len(poly_data_list) == batch_size_mult:
coco_ann = run_polygonization(poly_data_list)
output_annotations.extend(coco_ann)
poly_data_list = []
# Finish with incomplete batches
if len(model_data_list):
tile_data = run_model(config, model, model_data_list)
poly_data_list.append(tile_data)
if len(poly_data_list):
coco_ann = run_polygonization(poly_data_list)
output_annotations.extend(coco_ann)
print("Saving output...")
with open(out_filepath, 'w') as outfile:
json.dump(output_annotations, outfile)
def polygonize_mask(config, mask_filepaths, backbone, out_ext):
"""
Reads
@param args:
@return:
"""
# --- Online transform performed on the device (GPU):
eval_online_cuda_transform = data_transforms.get_eval_online_cuda_transform(
config)
print("Loading model...")
model = FrameFieldModel(config,
backbone=backbone,
eval_transform=eval_online_cuda_transform)
model.to(config["device"])
checkpoints_dirpath = run_utils.setup_run_subdir(
config["eval_params"]["run_dirpath"],
config["optim_params"]["checkpoints_dirname"])
model = inference.load_checkpoint(model, checkpoints_dirpath,
config["device"])
model.eval()
rasterizer = torch_lydorn.torchvision.transforms.Rasterize(fill=True,
edges=False,
vertices=False)
# Read image
pbar = tqdm(mask_filepaths, desc="Infer images")
for mask_filepath in pbar:
pbar.set_postfix(status="Loading mask image")
mask_image = skimage.io.imread(mask_filepath)
input_image = mask_image
if len(input_image.shape) == 2:
# Make input_image shape (H, W, 1)
input_image = input_image[:, :, None]
if input_image.shape[2] == 1:
input_image = np.repeat(input_image, 3, axis=-1)
mean = np.array([0.5, 0.5, 0.5])
std = np.array([1, 1, 1])
tile_data = {
"image":
torchvision.transforms.functional.to_tensor(input_image)[None,
...],
"image_mean":
torch.from_numpy(mean)[None, ...],
"image_std":
torch.from_numpy(std)[None, ...],
"image_filepath": [mask_filepath],
}
pbar.set_postfix(status="Inference")
tile_data = inference.inference(config,
model,
tile_data,
compute_polygonization=False)
pbar.set_postfix(status="Polygonize")
seg_batch = torchvision.transforms.functional.to_tensor(mask_image)[
None, ...].float() / 255
crossfield_batch = tile_data["crossfield"]
polygons_batch, probs_batch = polygonize_acm.polygonize(
seg_batch, crossfield_batch, polygonize_config)
tile_data["polygons"] = polygons_batch
tile_data["polygon_probs"] = probs_batch
pbar.set_postfix(status="Saving output")
tile_data = local_utils.batch_to_cpu(tile_data)
tile_data = local_utils.split_batch(tile_data)[0]
base_filepath = os.path.splitext(mask_filepath)[0]
# save_utils.save_polygons(tile_data["polygons"], base_filepath, "polygons", tile_data["image_filepath"])
# save_utils.save_poly_viz(tile_data["image"], tile_data["polygons"], tile_data["polygon_probs"], base_filepath, name)
# geo_utils.save_shapefile_from_shapely_polygons(tile_data["polygons"], mask_filepath, base_filepath + "." + name + ".shp")
if out_ext == "geojson":
save_utils.save_geojson(tile_data["polygons"], base_filepath)
elif out_ext == "shp":
save_utils.save_shapefile(tile_data["polygons"], base_filepath,
"polygonized", mask_filepath)
else:
raise ValueError(f"out_ext '{out_ext}' invalid!")
# --- Compute IoU of mask image and extracted polygons
polygons_raster = rasterizer(mask_image, tile_data["polygons"])[:, :,
0]
mask = 128 < mask_image
polygons_mask = 128 < polygons_raster
iou = measures.iou(torch.tensor(polygons_mask).view(1, -1),
torch.tensor(mask).view(1, -1),
threshold=0.5)
print("IoU:", iou.item())
if iou < 0.9:
print(mask_filepath)