def __init__(self, root, transform=None, pre_transform=None, fold="train", small=False, pool_size=1): assert fold in [ "train", "val", "test_images" ], "Input fold={} should be in [\"train\", \"val\", \"test_images\"]".format( fold) if fold == "test_images": print_utils.print_error( "ERROR: fold {} not yet implemented!".format(fold)) exit() self.root = root self.fold = fold makedirs(self.processed_dir) self.small = small if self.small: print_utils.print_info( "INFO: Using small version of the Mapping challenge dataset.") self.pool_size = pool_size self.coco = None self.image_id_list = self.load_image_ids() self.stats_filepath = os.path.join(self.processed_dir, "stats.pt") self.stats = None if os.path.exists(self.stats_filepath): self.stats = torch.load(self.stats_filepath) self.processed_flag_filepath = os.path.join( self.processed_dir, "processed-flag-small" if self.small else "processed-flag") super(MappingChallenge, self).__init__(root, transform, pre_transform)
def eval_process(gpu, config, shared_dict, barrier): from frame_field_learning.evaluate import evaluate torch.manual_seed(0) # Ensure same seed for all processes # --- Find data directory --- # root_dir_candidates = [ os.path.join(data_dirpath, config["dataset_params"]["root_dirname"]) for data_dirpath in config["data_dir_candidates"] ] root_dir, paths_tried = python_utils.choose_first_existing_path( root_dir_candidates, return_tried_paths=True) if root_dir is None: print_utils.print_error( "GPU {} -> ERROR: Data root directory amongst \"{}\" not found!". format(gpu, paths_tried)) raise NotADirectoryError( f"Couldn't find a directory in {paths_tried} (gpu:{gpu})") print_utils.print_info("GPU {} -> Using data from {}".format( gpu, root_dir)) config["data_root_dir"] = root_dir # --- Get dataset # - CHANGE HERE TO ADD YOUR OWN DATASET eval_ds, = get_folds( config, root_dir, folds=config["fold"]) # config["fold"] is already a list (of length 1) # --- Instantiate backbone network (its backbone will be used to extract features) backbone = get_backbone(config["backbone_params"]) evaluate(gpu, config, shared_dict, barrier, eval_ds, backbone)
def inference_no_patching(config, model, tile_data): with torch.no_grad(): batch = { "image": tile_data["image"], "image_mean": tile_data["image_mean"], "image_std": tile_data["image_std"] } try: pred, batch = network_inference(config, model, batch) except RuntimeError as e: print_utils.print_error("ERROR: " + str(e)) if 1 < config["optim_params"]["eval_batch_size"]: print_utils.print_info( "INFO: Try lowering the effective batch_size (which is {} currently). " "Note that in eval mode, the effective bath_size is equal to double the batch_size " "because gradients do not need to " "be computed so double the memory is available. " "You can override the effective batch_size with the eval_batch_size parameter." .format(config["optim_params"]["eval_batch_size"])) else: print_utils.print_info( "INFO: The effective batch_size is 1 but the GPU still ran out of memory." "You can specify parameters to split the image into patches for inference:\n" "--eval_patch_size is the size of the patch and should be chosen as big as memory allows.\n" "--eval_patch_overlap (optional, default=200) adds overlaps between patches to avoid border artifacts." .format(config["optim_params"]["eval_batch_size"])) raise e tile_data["seg"] = pred["seg"] if "crossfield" in pred: tile_data["crossfield"] = pred["crossfield"] return tile_data
def get_folds(config, root_dir, folds): assert set(folds).issubset({"train", "val", "train_val", "test"}), \ 'fold in folds should be in ["train", "val", "train_val", "test"]' if config["dataset_params"]["root_dirname"] == "AerialImageDataset": return get_inria_aerial_folds(config, root_dir, folds) elif config["dataset_params"][ "root_dirname"] == "luxcarta_precise_buildings": return get_luxcarta_buildings(config, root_dir, folds) elif config["dataset_params"][ "root_dirname"] == "mapping_challenge_dataset": return get_mapping_challenge(config, root_dir, folds) elif config["dataset_params"]["root_dirname"] == "segbuildings": return get_opencities_competition(config, root_dir, folds) elif config["dataset_params"]["root_dirname"] == "xview2_xbd_dataset": return get_xview2_dataset(config, root_dir, folds) else: print_utils.print_error( "ERROR: config[\"data_root_partial_dirpath\"] = \"{}\" is an unknown dataset! " "If it is a new dataset, add it in dataset_folds.py's get_folds() function." .format(config["dataset_params"]["root_dirname"])) exit()
def launch_train(args): assert args.config is not None, "Argument --config must be specified. Run 'python main.py --help' for help on arguments." config = run_utils.load_config(args.config) if config is None: print_utils.print_error( "ERROR: cannot continue without a config file. Exiting now...") sys.exit() config["runs_dirpath"] = args.runs_dirpath if args.run_name is not None: config["run_name"] = args.run_name config["new_run"] = args.new_run config["init_run_name"] = args.init_run_name if args.samples is not None: config["samples"] = args.samples if args.batch_size is not None: config["optim_params"]["batch_size"] = args.batch_size if args.max_epoch is not None: config["optim_params"]["max_epoch"] = args.max_epoch if args.fold is None: if "fold" in config: fold = set(config["fold"]) else: fold = {"train"} # Default values for train else: fold = set(args.fold) assert fold == {"train"} or fold == {"train", "val"}, \ "Argument fold when training should be either: ['train'] or ['train', 'val']" config["fold"] = list(fold) print_utils.print_info("Training on fold(s): {}".format(config["fold"])) config["nodes"] = args.nodes config["gpus"] = args.gpus config["nr"] = args.nr config["world_size"] = args.gpus * args.nodes # --- Load params in config set as relative path to another JSON file config = run_utils.load_defaults_in_config( config, filepath_key="defaults_filepath") # Setup num_workers per process: if config["num_workers"] is None: config["num_workers"] = int(torch.multiprocessing.cpu_count() / config["gpus"]) # --- Distributed init: os.environ['MASTER_ADDR'] = args.master_addr os.environ['MASTER_PORT'] = args.master_port manager = torch.multiprocessing.Manager() shared_dict = manager.dict() shared_dict["run_dirpath"] = None shared_dict["init_checkpoints_dirpath"] = None barrier = manager.Barrier(args.gpus) torch.multiprocessing.spawn(train_process, nprocs=args.gpus, args=(config, shared_dict, barrier))
def launch_inference_from_filepath(args): from frame_field_learning.inference_from_filepath import inference_from_filepath # --- First step: figure out what run (experiment) is to be evaluated # Option 1: the run_name argument is given in which case that's our run run_name = None config = None if args.run_name is not None: run_name = args.run_name # Else option 2: Check if a config has been given to look for the run_name if args.config is not None: config = run_utils.load_config(args.config) if config is not None and "run_name" in config and run_name is None: run_name = config["run_name"] # Else abort... if run_name is None: print_utils.print_error( "ERROR: the run to evaluate could no be identified with the given arguments. " "Please specify either the --run_name argument or the --config argument " "linking to a config file that has a 'run_name' field filled with the name of " "the run name to evaluate.") sys.exit() # --- Second step: get path to the run and if --config was not specified, load the config from the run's folder run_dirpath = frame_field_learning.local_utils.get_run_dirpath( args.runs_dirpath, run_name) if config is None: config = run_utils.load_config(config_dirpath=run_dirpath) if config is None: print_utils.print_error( f"ERROR: the default run's config file at {run_dirpath} could not be loaded. " f"Exiting now...") sys.exit() # --- Add command-line arguments if args.batch_size is not None: config["optim_params"]["batch_size"] = args.batch_size if args.eval_batch_size is not None: config["optim_params"]["eval_batch_size"] = args.eval_batch_size else: config["optim_params"][ "eval_batch_size"] = 2 * config["optim_params"]["batch_size"] # --- Load params in config set as relative path to another JSON file config = run_utils.load_defaults_in_config( config, filepath_key="defaults_filepath") config["eval_params"]["run_dirpath"] = run_dirpath if args.eval_patch_size is not None: config["eval_params"]["patch_size"] = args.eval_patch_size if args.eval_patch_overlap is not None: config["eval_params"]["patch_overlap"] = args.eval_patch_overlap backbone = get_backbone(config["backbone_params"]) inference_from_filepath(config, args.in_filepath, backbone, args.out_dirpath)
def get_run_dirpath(runs_dirpath, run_name): working_dir = os.path.dirname(os.path.abspath(__file__)) runs_dir = os.path.join(working_dir, runs_dirpath) try: run_dirpath = run_utils.setup_run_dir(runs_dir, run_name, check_exists=True) except FileNotFoundError as e: print_utils.print_error(f"ERROR: {e}") sys.exit() return run_dirpath
def load_json(filepath): if not os.path.exists(filepath): return False try: with open(filepath, 'r') as f: minified = jsmin(f.read()) data = json.loads(minified) except json.decoder.JSONDecodeError as e: print_utils.print_error( f"ERROR in load_json(filepath): {e} from JSON at {filepath}") exit() return data
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 launch_eval_coco(args): # --- Init: fills mode-specific default command-line arguments if args.fold is None: fold = {"test"} # Default value for eval_coco else: fold = set(args.fold) assert len(fold) == 1, \ "Argument fold when evaluating with COCO should be a single fold" # --- Find which run and which config file to evaluate the run with if args.run_name is None and args.config is None: print_utils.print_error( "ERROR: At least of one --run_name or --config has to be specified." ) sys.exit() elif args.run_name is None and args.config is not None: # Load config config = run_utils.load_config(args.config) # Verify it has a run_name specified if "run_name" not in config: print_utils.print_error( "ERROR: run_name was not found in the provided config file, you can specify it with --run_name" ) sys.exit() run_name = config["run_name"] elif args.run_name is not None and args.config is None: # Load run_name's config run_dirpath = frame_field_learning.local_utils.get_run_dirpath( args.runs_dirpath, args.run_name) config = run_utils.load_config(config_dirpath=run_dirpath) run_name = args.run_name else: # Load specified config and use specified run_name config = run_utils.load_config(args.config) run_name = args.run_name # --- Load params in config set as relative path to another JSON file config = run_utils.load_defaults_in_config( config, filepath_key="defaults_filepath") # --- Second step: Replace parameters in config file from command-line arguments config["eval_params"]["run_name"] = run_name if args.samples is not None: config["samples"] = args.samples config["fold"] = list(fold) # Setup num_workers per process: if config["num_workers"] is None: config["num_workers"] = torch.multiprocessing.cpu_count() eval_coco(config)
def get_mapping_challenge(config, root_dir, folds): from torch_lydorn.torchvision.datasets import MappingChallenge if "train" in folds or "val" in folds or "train_val" in folds: train_online_cpu_transform = data_transforms.get_online_cpu_transform( config, augmentations=config["data_aug_params"]["enable"]) ds = MappingChallenge( root_dir, transform=train_online_cpu_transform, pre_transform=data_transforms.get_offline_transform_patch(), small=config["dataset_params"]["small"], fold="train", pool_size=config["num_workers"]) torch.manual_seed( config["dataset_params"]["seed"]) # Ensure a seed is set train_split_length = int( round(config["dataset_params"]["train_fraction"] * len(ds))) val_split_length = len(ds) - train_split_length train_ds, val_ds = torch.utils.data.random_split( ds, [train_split_length, val_split_length]) ds_list = [] for fold in folds: if fold == "train": ds_list.append(train_ds) elif fold == "val": ds_list.append(val_ds) elif fold == "train_val": ds_list.append(ds) elif fold == "test": # The val fold from the original challenge is used as test here # because we don't have the ground truth for the test_images fold of the challenge: test_online_cpu_transform = data_transforms.get_eval_online_cpu_transform( ) test_ds = MappingChallenge( root_dir, transform=test_online_cpu_transform, pre_transform=data_transforms.get_offline_transform_patch(), small=config["dataset_params"]["small"], fold="val", pool_size=config["num_workers"]) ds_list.append(test_ds) else: print_utils.print_error( "ERROR: fold \"{}\" not recognized, implement it in dataset_folds.py." .format(fold)) exit() return ds_list
def train_process(gpu, config, shared_dict, barrier): from frame_field_learning.train import train print_utils.print_info( "GPU {} -> Ready. There are {} GPU(s) available on this node.".format( gpu, torch.cuda.device_count())) torch.manual_seed(0) # Ensure same seed for all processes # --- Find data directory --- # root_dir_candidates = [ os.path.join(data_dirpath, config["dataset_params"]["root_dirname"]) for data_dirpath in config["data_dir_candidates"] ] root_dir, paths_tried = python_utils.choose_first_existing_path( root_dir_candidates, return_tried_paths=True) if root_dir is None: print_utils.print_error( "GPU {} -> ERROR: Data root directory amongst \"{}\" not found!". format(gpu, paths_tried)) exit() print_utils.print_info("GPU {} -> Using data from {}".format( gpu, root_dir)) # --- Get dataset splits # - CHANGE HERE TO ADD YOUR OWN DATASET # We have to adapt the config["fold"] param to the folds argument of the get_folds function fold = set(config["fold"]) if fold == {"train"}: # Val will be used for evaluating the model after each epoch: train_ds, val_ds = get_folds(config, root_dir, folds=["train", "val"]) elif fold == {"train", "val"}: # Both train and val are meant to be used for training train_ds, = get_folds(config, root_dir, folds=["train_val"]) val_ds = None else: # Should not arrive here since main makes sure config["fold"] is either one of the above print_utils.print_error("ERROR: specified folds not recognized!") raise NotImplementedError # --- Instantiate backbone network if config["backbone_params"]["name"] in ["deeplab50", "deeplab101"]: assert 1 < config["optim_params"]["batch_size"], \ "When using backbone {}, batch_size has to be at least 2 for the batchnorm of the ASPPPooling to work."\ .format(config["backbone_params"]["name"]) backbone = get_backbone(config["backbone_params"]) # --- Launch training train(gpu, config, shared_dict, barrier, train_ds, val_ds, backbone)
def get_luxcarta_buildings(config, root_dir, folds): from torch_lydorn.torchvision.datasets import LuxcartaBuildings # --- Online transform done on the host (CPU): online_cpu_transform = data_transforms.get_online_cpu_transform( config, augmentations=config["data_aug_params"]["enable"]) data_patch_size = config["dataset_params"]["data_patch_size"] if config[ "data_aug_params"]["enable"] else config["input_patch_size"] ds = LuxcartaBuildings( root_dir, transform=online_cpu_transform, patch_size=data_patch_size, patch_stride=config["dataset_params"]["input_patch_size"], pre_transform=data_transforms.get_offline_transform_patch(), fold="train", pool_size=config["num_workers"]) torch.manual_seed(config["dataset_params"]["seed"]) # Ensure a seed is set train_split_length = int( round(config["dataset_params"]["train_fraction"] * len(ds))) val_split_length = len(ds) - train_split_length train_ds, val_ds = torch.utils.data.random_split( ds, [train_split_length, val_split_length]) ds_list = [] for fold in folds: if fold == "train": ds_list.append(train_ds) elif fold == "val": ds_list.append(val_ds) elif fold == "test": # TODO: handle patching with multi-GPU processing print_utils.print_error( "WARNING: handle patching with multi-GPU processing") ds = LuxcartaBuildings( root_dir, transform=online_cpu_transform, pre_transform=data_transforms.get_offline_transform_patch(), fold="test", pool_size=config["num_workers"]) ds_list.append(ds) else: print_utils.print_error( "ERROR: fold \"{}\" not recognized, implement it in dataset_folds.py." .format(fold)) return ds_list
def eval_coco(config): assert len(config["fold"]) == 1, "There should be only one specified fold" fold = config["fold"][0] if fold != "test": raise NotImplementedError pool = Pool(processes=config["num_workers"]) # Find data dir root_dir_candidates = [os.path.join(data_dirpath, config["dataset_params"]["root_dirname"]) for data_dirpath in config["data_dir_candidates"]] root_dir, paths_tried = python_utils.choose_first_existing_path(root_dir_candidates, return_tried_paths=True) if root_dir is None: print_utils.print_error( "ERROR: Data root directory amongst \"{}\" not found!".format(paths_tried)) exit() print_utils.print_info("Using data from {}".format(root_dir)) raw_dir = os.path.join(root_dir, "raw") # Get run's eval results dir results_dirpath = os.path.join(root_dir, config["eval_params"]["results_dirname"]) run_results_dirpath = run_utils.setup_run_dir(results_dirpath, config["eval_params"]["run_name"], check_exists=True) # Setup coco annType = 'segm' # initialize COCO ground truth api gt_annotation_filename = "annotation-small.json" if config["dataset_params"]["small"] else "annotation.json" gt_annotation_filepath = os.path.join(raw_dir, "val", gt_annotation_filename) # We are using the original val fold as our test fold print_utils.print_info("INFO: Load gt from " + gt_annotation_filepath) cocoGt = COCO(gt_annotation_filepath) # image_id = 0 # annotation_ids = cocoGt.getAnnIds(imgIds=image_id) # annotation_list = cocoGt.loadAnns(annotation_ids) # print(annotation_list) # initialize COCO detections api annotation_filename_list = fnmatch.filter(os.listdir(run_results_dirpath), fold + ".annotation.*.json") eval_one_partial = partial(eval_one, run_results_dirpath=run_results_dirpath, cocoGt=cocoGt, config=config, annType=annType, pool=pool) # with Pool(8) as p: # r = list(tqdm(p.imap(eval_one_partial, annotation_filename_list), total=len(annotation_filename_list))) for annotation_filename in annotation_filename_list: eval_one_partial(annotation_filename)
def get_xview2_dataset(config, root_dir, folds): from torch_lydorn.torchvision.datasets import xView2Dataset if "train" in folds or "val" in folds or "train_val" in folds: train_online_cpu_transform = data_transforms.get_online_cpu_transform( config, augmentations=config["data_aug_params"]["enable"]) ds = xView2Dataset( root_dir, fold="train", pre_process=True, patch_size=config["dataset_params"]["data_patch_size"], pre_transform=data_transforms.get_offline_transform_patch(), transform=train_online_cpu_transform, small=config["dataset_params"]["small"], pool_size=config["num_workers"]) torch.manual_seed( config["dataset_params"]["seed"]) # Ensure a seed is set train_split_length = int( round(config["dataset_params"]["train_fraction"] * len(ds))) val_split_length = len(ds) - train_split_length train_ds, val_ds = torch.utils.data.random_split( ds, [train_split_length, val_split_length]) ds_list = [] for fold in folds: if fold == "train": ds_list.append(train_ds) elif fold == "val": ds_list.append(val_ds) elif fold == "train_val": ds_list.append(ds) elif fold == "test": raise NotImplementedError( "Test fold not yet implemented (skip pre-processing?)") elif fold == "hold": raise NotImplementedError( "Hold fold not yet implemented (skip pre-processing?)") else: print_utils.print_error( "ERROR: fold \"{}\" not recognized, implement it in dataset_folds.py." .format(fold)) exit() return ds_list
def get_opencities_competition(config, root_dir, folds): from torch_lydorn.torchvision.datasets import RasterizedOpenCities, OpenCitiesTestDataset data_patch_size = config["dataset_params"]["data_patch_size"] if config[ "data_aug_params"]["enable"] else config["input_patch_size"] ds_list = [] for fold in folds: if fold == "train": train_ds = RasterizedOpenCities( tier=1, augment=False, small_subset=False, resize_size=data_patch_size, data_dir=root_dir, baseline_mode=False, val=False, val_split=config["dataset_params"]["val_fraction"]) ds_list.append(train_ds) elif fold == "val": val_ds = RasterizedOpenCities( tier=1, augment=False, small_subset=False, resize_size=data_patch_size, data_dir=root_dir, baseline_mode=False, val=True, val_split=config["dataset_params"]["val_fraction"]) ds_list.append(val_ds) elif fold == "test": test_ds = OpenCitiesTestDataset(root_dir + "/test/", 1024) ds_list.append(test_ds) else: print_utils.print_error( "ERROR: fold \"{}\" not recognized, implement it in dataset_folds.py." .format(fold)) return ds_list
def main(): torch.manual_seed(0) # --- Process args --- # args = get_args() # --- Setup run --- # run_dirpath = local_utils.get_run_dirpath(args.runs_dirpath, args.run_name) # 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["eval_params"]["run_dirpath"] = run_dirpath if args.eval_patch_size is not None: config["eval_params"]["patch_size"] = args.eval_patch_size if args.eval_patch_overlap is not None: config["eval_params"]["patch_overlap"] = args.eval_patch_overlap backbone = get_backbone(config["backbone_params"]) polygonize_mask(config, args.filepath, backbone, args.out_ext)
def get_backbone(backbone_params): set_download_dir() if backbone_params["name"] == "unet": from torchvision.models.segmentation._utils import _SimpleSegmentationModel from frame_field_learning.unet import UNetBackbone backbone = UNetBackbone(backbone_params["input_features"], backbone_params["features"]) backbone = _SimpleSegmentationModel(backbone, classifier=torch.nn.Identity()) elif backbone_params["name"] == "fcn50": backbone = torchvision.models.segmentation.fcn_resnet50( pretrained=backbone_params["pretrained"], num_classes=21) backbone.classifier = torch.nn.Sequential( *list(backbone.classifier.children())[:-1], torch.nn.Conv2d(512, backbone_params["features"], kernel_size=(1, 1), stride=(1, 1))) elif backbone_params["name"] == "fcn101": backbone = torchvision.models.segmentation.fcn_resnet101( pretrained=backbone_params["pretrained"], num_classes=21) backbone.classifier = torch.nn.Sequential( *list(backbone.classifier.children())[:-1], torch.nn.Conv2d(512, backbone_params["features"], kernel_size=(1, 1), stride=(1, 1))) elif backbone_params["name"] == "deeplab50": backbone = torchvision.models.segmentation.deeplabv3_resnet50( pretrained=backbone_params["pretrained"], num_classes=21) backbone.classifier = torch.nn.Sequential( *list(backbone.classifier.children())[:-1], torch.nn.Conv2d(256, backbone_params["features"], kernel_size=(1, 1), stride=(1, 1))) elif backbone_params["name"] == "deeplab101": backbone = torchvision.models.segmentation.deeplabv3_resnet101( pretrained=backbone_params["pretrained"], num_classes=21) backbone.classifier = torch.nn.Sequential( *list(backbone.classifier.children())[:-1], torch.nn.Conv2d(256, backbone_params["features"], kernel_size=(1, 1), stride=(1, 1))) elif backbone_params["name"] == "unet_resnet": from torchvision.models.segmentation._utils import _SimpleSegmentationModel from frame_field_learning.unet_resnet import UNetResNetBackbone backbone = UNetResNetBackbone( backbone_params["encoder_depth"], num_filters=backbone_params["num_filters"], dropout_2d=backbone_params["dropout_2d"], pretrained=backbone_params["pretrained"], is_deconv=backbone_params["is_deconv"]) backbone = _SimpleSegmentationModel(backbone, classifier=torch.nn.Identity()) elif backbone_params["name"] == "ictnet": from torchvision.models.segmentation._utils import _SimpleSegmentationModel from frame_field_learning.ictnet import ICTNetBackbone backbone = ICTNetBackbone(in_channels=backbone_params["in_channels"], out_channels=backbone_params["out_channels"], preset_model=backbone_params["preset_model"], dropout_2d=backbone_params["dropout_2d"], efficient=backbone_params["efficient"]) backbone = _SimpleSegmentationModel(backbone, classifier=torch.nn.Identity()) else: print_utils.print_error( "ERROR: config[\"backbone_params\"][\"name\"] = \"{}\" is an unknown backbone!" "If it is a new backbone you want to use, " "add it in backbone.py's get_backbone() function.".format( backbone_params["name"])) raise RuntimeError("Specified backbone {} unknown".format( backbone_params["name"])) return backbone
def __call__(self, polygonize_params, seg_batch, crossfield_batch=None, pre_computed=None): """ :param polygonize_params: :param seg_batch: (N, C, H, W) :param crossfield_batch: (N, 4, H, W) :param pre_computed: None o a Dictionary of pre-computed values used for various methods :return: """ assert len(seg_batch.shape) == 4, "seg_batch should be (N, C, H, W)" assert pre_computed is None or isinstance( pre_computed, dict), "pre_computed should be either None or a dict" batch_size = seg_batch.shape[0] # Check if polygonize_params["method"] is a list or a string: if type(polygonize_params["method"]) == list: # --- For speed up, pre-compute anything that is used by multiple methods: if pre_computed is None: pre_computed = {} if ("simple" in polygonize_params["method"] or "acm" in polygonize_params["method"] ) and "init_contours_batch" not in pre_computed: indicator_batch = seg_batch[:, 0, :, :] np_indicator_batch = indicator_batch.cpu().numpy() init_contours_batch = polygonize_utils.compute_init_contours_batch( np_indicator_batch, polygonize_params["common_params"]["init_data_level"], pool=self.pool) pre_computed["init_contours_batch"] = init_contours_batch # --- # Run one method after the other: out_polygons_dict_batch = [{} for _ in range(batch_size)] out_probs_dict_batch = [{} for _ in range(batch_size)] for method_name in polygonize_params["method"]: new_polygonize_params = polygonize_params.copy() new_polygonize_params["method"] = method_name polygons_batch, probs_batch = self( new_polygonize_params, seg_batch, crossfield_batch=crossfield_batch, pre_computed=pre_computed) if polygons_batch is not None: for i, (polygons, probs) in enumerate( zip(polygons_batch, probs_batch)): out_polygons_dict_batch[i][method_name] = polygons out_probs_dict_batch[i][method_name] = probs return out_polygons_dict_batch, out_probs_dict_batch # --- Else: run the one method if polygonize_params["method"] == "acm": if crossfield_batch is None: # Cannot run the ACM method return None, None polygons_batch, probs_batch = polygonize_acm.polygonize( seg_batch, crossfield_batch, polygonize_params["acm_method"], pool=self.pool, pre_computed=pre_computed) elif polygonize_params["method"] == "asm": if crossfield_batch is None: # Cannot run the ASM method return None, None polygons_batch, probs_batch = self.polygonizer_asm( seg_batch, crossfield_batch, pre_computed=pre_computed) elif polygonize_params["method"] == "simple": polygons_batch, probs_batch = polygonize_simple.polygonize( seg_batch, polygonize_params["simple_method"], pool=self.pool, pre_computed=pre_computed) else: print_utils.print_error( "ERROR: polygonize method {} not recognized!".format( polygonize_params["method"])) raise NotImplementedError return polygons_batch, probs_batch
def inference_with_patching(config, model, tile_data): assert len(tile_data["image"].shape) == 4 and tile_data["image"].shape[0] == 1, \ f"When using inference with patching, tile_data should have a batch size of 1, " \ f"with image's shape being (1, C, H, W), not {tile_data['image'].shape}" with torch.no_grad(): # Init tile outputs (image is (N, C, H, W)): height = tile_data["image"].shape[2] width = tile_data["image"].shape[3] seg_channels = config["seg_params"]["compute_interior"] \ + config["seg_params"]["compute_edge"] \ + config["seg_params"]["compute_vertex"] if config["compute_seg"]: tile_data["seg"] = torch.zeros((1, seg_channels, height, width), device=config["device"]) if config["compute_crossfield"]: tile_data["crossfield"] = torch.zeros((1, 4, height, width), device=config["device"]) weight_map = torch.zeros( (1, 1, height, width), device=config["device"] ) # Count number of patches on top of each pixel # Split tile in patches: stride = config["eval_params"]["patch_size"] - config["eval_params"][ "patch_overlap"] patch_boundingboxes = image_utils.compute_patch_boundingboxes( (height, width), stride=stride, patch_res=config["eval_params"]["patch_size"]) # Compute patch pixel weights to merge overlapping patches back together smoothly: patch_weights = np.ones((config["eval_params"]["patch_size"] + 2, config["eval_params"]["patch_size"] + 2), dtype=np.float) patch_weights[0, :] = 0 patch_weights[-1, :] = 0 patch_weights[:, 0] = 0 patch_weights[:, -1] = 0 patch_weights = scipy.ndimage.distance_transform_edt(patch_weights) patch_weights = patch_weights[1:-1, 1:-1] patch_weights = torch.tensor(patch_weights, device=config["device"]).float() patch_weights = patch_weights[ None, None, :, :] # Adding batch and channels dims # Predict on each patch and save in outputs: for bbox in tqdm(patch_boundingboxes, desc="Running model on patches", leave=False): # Crop data batch = { "image": tile_data["image"][:, :, bbox[0]:bbox[2], bbox[1]:bbox[3]], "image_mean": tile_data["image_mean"], "image_std": tile_data["image_std"], } # Send batch to device try: pred, batch = network_inference(config, model, batch) except RuntimeError as e: print_utils.print_error("ERROR: " + str(e)) print_utils.print_info( "INFO: Reduce --eval_patch_size until the patch fits in memory." ) raise e if config["compute_seg"]: tile_data[ "seg"][:, :, bbox[0]:bbox[2], bbox[1]:bbox[3]] += patch_weights * pred["seg"] if config["compute_crossfield"]: tile_data["crossfield"][:, :, bbox[0]:bbox[2], bbox[1]:bbox[ 3]] += patch_weights * pred["crossfield"] weight_map[:, :, bbox[0]:bbox[2], bbox[1]:bbox[3]] += patch_weights # Take care of overlapping parts if config["compute_seg"]: tile_data["seg"] /= weight_map if config["compute_crossfield"]: tile_data["crossfield"] /= weight_map return tile_data
def main(): # Test using transforms from the frame_field_learning project: from frame_field_learning import data_transforms config = { "data_dir_candidates": ["/data/titane/user/nigirard/data", "~/data", "/data"], "dataset_params": { "root_dirname": "xview2_xbd_dataset", "pre_process": True, "small": False, "data_patch_size": 725, "input_patch_size": 512, "train_fraction": 0.75 }, "num_workers": 8, "data_aug_params": { "enable": True, "vflip": True, "affine": True, "scaling": [0.9, 1.1], "color_jitter": True, "device": "cuda" } } # Find data_dir data_dir = python_utils.choose_first_existing_path( config["data_dir_candidates"]) if data_dir is None: print_utils.print_error("ERROR: Data directory not found!") exit() else: print_utils.print_info("Using data from {}".format(data_dir)) root_dir = os.path.join(data_dir, config["dataset_params"]["root_dirname"]) # --- Transforms: --- # # --- pre-processing transform (done once then saved on disk): # --- Online transform done on the host (CPU): online_cpu_transform = data_transforms.get_online_cpu_transform( config, augmentations=config["data_aug_params"]["enable"]) train_online_cuda_transform = data_transforms.get_online_cuda_transform( config, augmentations=config["data_aug_params"]["enable"]) kwargs = { "pre_process": config["dataset_params"]["pre_process"], "transform": online_cpu_transform, "patch_size": config["dataset_params"]["data_patch_size"], "pre_transform": data_transforms.get_offline_transform_patch(), "small": config["dataset_params"]["small"], "pool_size": config["num_workers"], } # --- --- # fold = "train" if fold == "train": dataset = xView2Dataset(root_dir, fold="train", **kwargs) elif fold == "val": dataset = xView2Dataset(root_dir, fold="train", **kwargs) elif fold == "test": dataset = xView2Dataset(root_dir, fold="test", **kwargs) else: raise NotImplementedError print(f"dataset has {len(dataset)} samples.") print("# --- Sample 0 --- #") sample = dataset[0] for key, item in sample.items(): print("{}: {}".format(key, type(item))) print("# --- Samples --- #") # for data in tqdm(dataset): # pass data_loader = torch.utils.data.DataLoader( dataset, batch_size=1, shuffle=False, num_workers=config["num_workers"]) print("# --- Batches --- #") for batch in tqdm(data_loader): # batch["distances"] = batch["distances"].float() # batch["sizes"] = batch["sizes"].float() # im = np.array(batch["image"][0]) # im = np.moveaxis(im, 0, -1) # skimage.io.imsave('im_before_transform.png', im) # # distances = np.array(batch["distances"][0]) # distances = np.moveaxis(distances, 0, -1) # skimage.io.imsave('distances_before_transform.png', distances) # # sizes = np.array(batch["sizes"][0]) # sizes = np.moveaxis(sizes, 0, -1) # skimage.io.imsave('sizes_before_transform.png', sizes) print("----") print(batch["name"]) print("image:", batch["image"].shape, batch["image"].min().item(), batch["image"].max().item()) im = np.array(batch["image"][0]) im = np.moveaxis(im, 0, -1) skimage.io.imsave('im.png', im) if "gt_polygons_image" in batch: print("gt_polygons_image:", batch["gt_polygons_image"].shape, batch["gt_polygons_image"].min().item(), batch["gt_polygons_image"].max().item()) seg = np.array(batch["gt_polygons_image"][0]) / 255 seg = np.moveaxis(seg, 0, -1) seg_display = utils.get_seg_display(seg) seg_display = (seg_display * 255).astype(np.uint8) skimage.io.imsave("gt_seg.png", seg_display) if "gt_crossfield_angle" in batch: print("gt_crossfield_angle:", batch["gt_crossfield_angle"].shape, batch["gt_crossfield_angle"].min().item(), batch["gt_crossfield_angle"].max().item()) gt_crossfield_angle = np.array(batch["gt_crossfield_angle"][0]) gt_crossfield_angle = np.moveaxis(gt_crossfield_angle, 0, -1) skimage.io.imsave('gt_crossfield_angle.png', gt_crossfield_angle) if "distances" in batch: print("distances:", batch["distances"].shape, batch["distances"].float().min().item(), batch["distances"].float().max().item()) distances = np.array(batch["distances"][0]) distances = np.moveaxis(distances, 0, -1) skimage.io.imsave('distances.png', distances) if "sizes" in batch: print("sizes:", batch["sizes"].shape, batch["sizes"].float().min().item(), batch["sizes"].float().max().item()) sizes = np.array(batch["sizes"][0]) sizes = np.moveaxis(sizes, 0, -1) skimage.io.imsave('sizes.png', sizes) # valid_mask = np.array(batch["valid_mask"][0]) # valid_mask = np.moveaxis(valid_mask, 0, -1) # skimage.io.imsave('valid_mask.png', valid_mask) input("Press enter to continue...") print("Apply online tranform:") batch = utils.batch_to_cuda(batch) batch = train_online_cuda_transform(batch) batch = utils.batch_to_cpu(batch) print("image:", batch["image"].shape, batch["image"].min().item(), batch["image"].max().item()) print("gt_polygons_image:", batch["gt_polygons_image"].shape, batch["gt_polygons_image"].min().item(), batch["gt_polygons_image"].max().item()) print("gt_crossfield_angle:", batch["gt_crossfield_angle"].shape, batch["gt_crossfield_angle"].min().item(), batch["gt_crossfield_angle"].max().item()) # print("distances:", batch["distances"].shape, batch["distances"].min().item(), batch["distances"].max().item()) # print("sizes:", batch["sizes"].shape, batch["sizes"].min().item(), batch["sizes"].max().item()) # Save output to visualize seg = np.array(batch["gt_polygons_image"][0]) seg = np.moveaxis(seg, 0, -1) seg_display = utils.get_seg_display(seg) seg_display = (seg_display * 255).astype(np.uint8) skimage.io.imsave("gt_seg.png", seg_display) im = np.array(batch["image"][0]) im = np.moveaxis(im, 0, -1) skimage.io.imsave('im.png', im) gt_crossfield_angle = np.array(batch["gt_crossfield_angle"][0]) gt_crossfield_angle = np.moveaxis(gt_crossfield_angle, 0, -1) skimage.io.imsave('gt_crossfield_angle.png', gt_crossfield_angle) distances = np.array(batch["distances"][0]) distances = np.moveaxis(distances, 0, -1) skimage.io.imsave('distances.png', distances) sizes = np.array(batch["sizes"][0]) sizes = np.moveaxis(sizes, 0, -1) skimage.io.imsave('sizes.png', sizes) # valid_mask = np.array(batch["valid_mask"][0]) # valid_mask = np.moveaxis(valid_mask, 0, -1) # skimage.io.imsave('valid_mask.png', valid_mask) input("Press enter to continue...")
def eval_one(annotation_filename, run_results_dirpath, cocoGt, config, annType, pool=None): print("---eval_one") annotation_name = os.path.splitext(annotation_filename)[0] if "samples" in config: stats_filepath = os.path.join(run_results_dirpath, "{}.stats.{}.{}.json".format("test", annotation_name, config["samples"])) metrics_filepath = os.path.join(run_results_dirpath, "{}.metrics.{}.{}.json".format("test", annotation_name, config["samples"])) else: stats_filepath = os.path.join(run_results_dirpath, "{}.stats.{}.json".format("test", annotation_name)) metrics_filepath = os.path.join(run_results_dirpath, "{}.metrics.{}.json".format("test", annotation_name)) res_filepath = os.path.join(run_results_dirpath, annotation_filename) if not os.path.exists(res_filepath): print_utils.print_warning("WARNING: result not found at filepath {}".format(res_filepath)) return print_utils.print_info("Evaluate {} annotations:".format(annotation_filename)) try: cocoDt = cocoGt.loadRes(res_filepath) except AssertionError as e: print_utils.print_error("ERROR: {}".format(e)) print_utils.print_info("INFO: continuing by removing unrecognised images") res = json.load(open(res_filepath)) print("Initial res length:", len(res)) annsImgIds = [ann["image_id"] for ann in res] image_id_rm = set(annsImgIds) - set(cocoGt.getImgIds()) print_utils.print_warning("Remove {} image ids!".format(len(image_id_rm))) new_res = [ann for ann in res if ann["image_id"] not in image_id_rm] print("New res length:", len(new_res)) cocoDt = cocoGt.loadRes(new_res) # {4601886185638229705, 4602408603195004682, 4597274499619802317, 4600985465712755606, 4597238470822783353, # 4597418614807878173} # image_id = 0 # annotation_ids = cocoDt.getAnnIds(imgIds=image_id) # annotation_list = cocoDt.loadAnns(annotation_ids) # print(annotation_list) if not os.path.exists(stats_filepath): # Run COCOeval cocoEval = COCOeval(cocoGt, cocoDt, annType) cocoEval.evaluate() cocoEval.accumulate() cocoEval.summarize() # Save stats stats = {} stat_names = ["AP", "AP_50", "AP_75", "AP_S", "AP_M", "AP_L", "AR", "AR_50", "AR_75", "AR_S", "AR_M", "AR_L"] assert len(stat_names) == cocoEval.stats.shape[0] for i, stat_name in enumerate(stat_names): stats[stat_name] = cocoEval.stats[i] python_utils.save_json(stats_filepath, stats) else: print("COCO stats already computed, skipping...") if not os.path.exists(metrics_filepath): # Verify that cocoDt has polygonal segmentation masks and not raster masks: if isinstance(cocoDt.loadAnns(cocoDt.getAnnIds(imgIds=cocoDt.getImgIds()[0]))[0]["segmentation"], list): metrics = {} # Run additionnal metrics print_utils.print_info("INFO: Running contour metrics") contour_eval = ContourEval(cocoGt, cocoDt) max_angle_diffs = contour_eval.evaluate(pool=pool) metrics["max_angle_diffs"] = list(max_angle_diffs) python_utils.save_json(metrics_filepath, metrics) else: print("Contour metrics already computed, skipping...")
def get_inria_aerial_folds(config, root_dir, folds): from torch_lydorn.torchvision.datasets import InriaAerial # --- Online transform done on the host (CPU): online_cpu_transform = data_transforms.get_online_cpu_transform( config, augmentations=config["data_aug_params"]["enable"]) mask_only = config["dataset_params"]["mask_only"] kwargs = { "pre_process": config["dataset_params"]["pre_process"], "transform": online_cpu_transform, "patch_size": config["dataset_params"]["data_patch_size"], "patch_stride": config["dataset_params"]["input_patch_size"], "pre_transform": data_transforms.get_offline_transform_patch(distances=not mask_only, sizes=not mask_only), "small": config["dataset_params"]["small"], "pool_size": config["num_workers"], "gt_source": config["dataset_params"]["gt_source"], "gt_type": config["dataset_params"]["gt_type"], "gt_dirname": config["dataset_params"]["gt_dirname"], "mask_only": mask_only, } train_val_split_point = config["dataset_params"]["train_fraction"] * 36 partial_train_tile_filter = functools.partial( inria_aerial_train_tile_filter, train_val_split_point=train_val_split_point) partial_val_tile_filter = functools.partial( inria_aerial_val_tile_filter, train_val_split_point=train_val_split_point) ds_list = [] for fold in folds: if fold == "train": ds = InriaAerial(root_dir, fold="train", tile_filter=partial_train_tile_filter, **kwargs) ds_list.append(ds) elif fold == "val": ds = InriaAerial(root_dir, fold="train", tile_filter=partial_val_tile_filter, **kwargs) ds_list.append(ds) elif fold == "train_val": ds = InriaAerial(root_dir, fold="train", **kwargs) ds_list.append(ds) elif fold == "test": ds = InriaAerial(root_dir, fold="test", **kwargs) ds_list.append(ds) else: print_utils.print_error( "ERROR: fold \"{}\" not recognized, implement it in dataset_folds.py." .format(fold)) return ds_list
def get_online_cpu_transform(config, augmentations=False): if augmentations and config["data_aug_params"]["device"] == "cpu": print_utils.print_error( "ERROR: CPU augmentations is not supported anymore. " "Look at CudaDataAugmentation to see what additional augs would need to be implemented." ) raise NotImplementedError online_transform_list = [] # Convert to PIL images if not augmentations \ or (augmentations and config["data_aug_params"]["device"] == "cpu"): online_transform_list.extend([ torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.ToPILImage(), key="image"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.ToPILImage(), key="gt_polygons_image"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.ToPILImage(), key="gt_crossfield_angle"), ]) # Add rotation data augmentation: if augmentations and config["data_aug_params"]["device"] == "cpu" and \ config["data_aug_params"]["affine"]: online_transform_list.extend([ torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.SampleUniform( -180, 180), outkey="rand_angle"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.functional.rotate, key=["image", "rand_angle"], outkey="image", resample=PIL.Image.BILINEAR), torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.functional.rotate, key=["gt_polygons_image", "rand_angle"], outkey="gt_polygons_image", resample=PIL.Image.BILINEAR), torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.functional.rotate, key=["gt_crossfield_angle", "rand_angle"], outkey="gt_crossfield_angle", resample=PIL.Image.NEAREST), ]) # Crop to final size if not augmentations \ or (augmentations and config["data_aug_params"]["device"] == "cpu"): if "input_patch_size" in config["dataset_params"]: online_transform_list.extend([ torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.CenterCrop( config["dataset_params"]["input_patch_size"]), key="image"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.CenterCrop( config["dataset_params"]["input_patch_size"]), key="gt_polygons_image"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.CenterCrop( config["dataset_params"]["input_patch_size"]), key="gt_crossfield_angle"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.CenterCrop( config["dataset_params"]["input_patch_size"]), key="distances"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.CenterCrop( config["dataset_params"]["input_patch_size"]), key="sizes"), ]) # Random Horizontal flip: if augmentations and config["data_aug_params"]["device"] == "cpu" and \ config["data_aug_params"]["vflip"]: online_transform_list.extend([ torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.RandomBool( p=0.5), outkey="rand_flip"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.ConditionApply( transform=torchvision.transforms.functional.vflip), key=["image", "rand_flip"], outkey="image"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.ConditionApply( transform=torchvision.transforms.functional.vflip), key=["gt_polygons_image", "rand_flip"], outkey="gt_polygons_image"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.ConditionApply( transform=torchvision.transforms.functional.vflip), key=["gt_crossfield_angle", "rand_flip"], outkey="gt_crossfield_angle"), ]) # Other augs: if augmentations and config["data_aug_params"]["device"] == "cpu" and \ config["data_aug_params"]["color_jitter"]: online_transform_list.append( torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.ColorJitter(brightness=0.05, contrast=0.05, saturation=.5, hue=.1), key="image")) # Convert to PyTorch tensors: online_transform_list.extend([ # Print(), torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.ToTensor(), key="image"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.Lambda(torch.from_numpy), key="image_mean"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torchvision.transforms.Lambda(torch.from_numpy), key="image_std"), torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.ToTensor(), key="gt_polygons_image", ignore_key_error=True), torch_lydorn.torchvision.transforms.TransformByKey( torch_lydorn.torchvision.transforms.ToTensor(), key="gt_crossfield_angle", ignore_key_error=True), torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.ToTensor(), key="distances", ignore_key_error=True), torch_lydorn.torchvision.transforms.TransformByKey( transform=torch_lydorn.torchvision.transforms.ToTensor(), key="sizes", ignore_key_error=True), ]) online_transform_list.append( torch_lydorn.torchvision.transforms.RemoveKeys(keys=["gt_polygons"])) online_transform = torchvision.transforms.Compose(online_transform_list) return online_transform
def main(): # Test using transforms from the frame_field_learning project: from frame_field_learning import data_transforms config = { "data_dir_candidates": [ "/data/titane/user/nigirard/data", "~/data", "/data", "/home/krishna/building-footprints-custom/frameField/data" ], "dataset_params": { "small": True, "root_dirname": "mapping_challenge_dataset", "seed": 0, "train_fraction": 0.75 }, "num_workers": 8, "data_aug_params": { "enable": False, "vflip": True, "affine": True, "color_jitter": True, "device": "cuda" } } # Find data_dir data_dir = python_utils.choose_first_existing_path( config["data_dir_candidates"]) if data_dir is None: print_utils.print_error("ERROR: Data directory not found!") exit() else: print_utils.print_info("Using data from {}".format(data_dir)) root_dir = os.path.join(data_dir, config["dataset_params"]["root_dirname"]) # --- Transforms: --- # # --- pre-processing transform (done once then saved on disk): # --- Online transform done on the host (CPU): train_online_cpu_transform = data_transforms.get_online_cpu_transform( config, augmentations=config["data_aug_params"]["enable"]) test_online_cpu_transform = data_transforms.get_eval_online_cpu_transform() train_online_cuda_transform = data_transforms.get_online_cuda_transform( config, augmentations=config["data_aug_params"]["enable"]) # --- --- # dataset = MappingChallenge( root_dir, transform=test_online_cpu_transform, pre_transform=data_transforms.get_offline_transform_patch(), fold="train", small=config["dataset_params"]["small"], pool_size=config["num_workers"]) print("# --- Sample 0 --- #") sample = dataset[0] print(sample.keys()) for key, item in sample.items(): print("{}: {}".format(key, type(item))) print(sample["image"].shape) print(len(sample["gt_polygons_image"])) print("# --- Samples --- #") # for data in tqdm(dataset): # pass data_loader = torch.utils.data.DataLoader( dataset, batch_size=10, shuffle=True, num_workers=config["num_workers"]) print("# --- Batches --- #") for batch in tqdm(data_loader): print("Images:") print(batch["image_relative_filepath"]) print(batch["image"].shape) print(batch["gt_polygons_image"].shape) print("Apply online tranform:") batch = utils.batch_to_cuda(batch) batch = train_online_cuda_transform(batch) batch = utils.batch_to_cpu(batch) print(batch["image"].shape) print(batch["gt_polygons_image"].shape) # Save output to visualize seg = np.array(batch["gt_polygons_image"][0]) seg = np.moveaxis(seg, 0, -1) seg_display = utils.get_seg_display(seg) seg_display = (seg_display * 255).astype(np.uint8) skimage.io.imsave("gt_seg.png", seg_display) skimage.io.imsave("gt_seg_edge.png", seg[:, :, 1]) im = np.array(batch["image"][0]) im = np.moveaxis(im, 0, -1) skimage.io.imsave('im.png', im) gt_crossfield_angle = np.array(batch["gt_crossfield_angle"][0]) gt_crossfield_angle = np.moveaxis(gt_crossfield_angle, 0, -1) skimage.io.imsave('gt_crossfield_angle.png', gt_crossfield_angle) distances = np.array(batch["distances"][0]) distances = np.moveaxis(distances, 0, -1) skimage.io.imsave('distances.png', distances) sizes = np.array(batch["sizes"][0]) sizes = np.moveaxis(sizes, 0, -1) skimage.io.imsave('sizes.png', sizes) # valid_mask = np.array(batch["valid_mask"][0]) # valid_mask = np.moveaxis(valid_mask, 0, -1) # skimage.io.imsave('valid_mask.png', valid_mask) input("Press enter to continue...")
def launch_eval(args): # --- Init: fills mode-specific default command-line arguments if args.fold is None: fold = {"test"} # Default value for eval mode else: fold = set(args.fold) assert len(fold) == 1, "Argument 'fold' must be a single fold in eval mode" # --- First step: figure out what run (experiment) is to be evaluated # Option 1: the run_name argument is given in which case that's our run run_name = None config = None if args.run_name is not None: run_name = args.run_name # Else option 2: Check if a config has been given to look for the run_name if args.config is not None: config = run_utils.load_config(args.config) if config is not None and "run_name" in config and run_name is None: run_name = config["run_name"] # Else abort... if run_name is None: print_utils.print_error( "ERROR: the run to evaluate could no be identified with the given arguments. " "Please specify either the --run_name argument or the --config argument " "linking to a config file that has a 'run_name' field filled with the name of " "the run name to evaluate.") sys.exit() # --- Second step: get path to the run and if --config was not specified, load the config from the run's folder run_dirpath = frame_field_learning.local_utils.get_run_dirpath( args.runs_dirpath, run_name) if config is None: config = run_utils.load_config(config_dirpath=run_dirpath) if config is None: print_utils.print_error( f"ERROR: the default run's config file at {run_dirpath} could not be loaded. " f"Exiting now...") sys.exit() # --- Third step: Replace parameters in config file from command-line arguments if args.dataset_params is not None: config["dataset_params"] = python_utils.load_json(args.dataset_params) if args.samples is not None: config["samples"] = args.samples if args.batch_size is not None: config["optim_params"]["batch_size"] = args.batch_size if args.eval_batch_size is not None: config["optim_params"]["eval_batch_size"] = args.eval_batch_size else: config["optim_params"][ "eval_batch_size"] = 2 * config["optim_params"]["batch_size"] config["fold"] = list(fold) config["nodes"] = args.nodes config["gpus"] = args.gpus config["nr"] = args.nr config["world_size"] = args.gpus * args.nodes # --- Load params in config set as relative path to another JSON file config = run_utils.load_defaults_in_config( config, filepath_key="defaults_filepath") config["eval_params"]["run_dirpath"] = run_dirpath if args.eval_patch_size is not None: config["eval_params"]["patch_size"] = args.eval_patch_size if args.eval_patch_overlap is not None: config["eval_params"]["patch_overlap"] = args.eval_patch_overlap # Setup num_workers per process: if config["num_workers"] is None: config["num_workers"] = int(torch.multiprocessing.cpu_count() / config["gpus"]) # --- Distributed init: os.environ['MASTER_ADDR'] = args.master_addr os.environ['MASTER_PORT'] = args.master_port manager = torch.multiprocessing.Manager() shared_dict = manager.dict() shared_dict["name_list"] = manager.list() shared_dict["iou_list"] = manager.list() shared_dict["seg_coco_list"] = manager.list() shared_dict["poly_coco_list"] = manager.list() barrier = manager.Barrier(args.gpus) torch.multiprocessing.spawn(eval_process, nprocs=args.gpus, args=(config, shared_dict, barrier))
def compute_seg_loss_weigths(pred_batch, gt_batch, config): """ Combines distances (from U-Net paper) with sizes (from https://github.com/neptune-ai/open-solution-mapping-challenge). @param pred_batch: @param gt_batch: @return: """ device = gt_batch["distances"].device use_dist = config["loss_params"]["seg_loss_params"]["use_dist"] use_size = config["loss_params"]["seg_loss_params"]["use_size"] w0 = config["loss_params"]["seg_loss_params"]["w0"] sigma = config["loss_params"]["seg_loss_params"]["sigma"] height = gt_batch["image"].shape[2] width = gt_batch["image"].shape[3] im_radius = math.sqrt(height * width) / 2 # --- Class imbalance weight (not forgetting background): gt_polygons_mask = (0 < gt_batch["gt_polygons_image"]).float() background_freq = 1 - torch.sum(gt_batch["class_freq"], dim=1) pixel_class_freq = gt_polygons_mask * gt_batch["class_freq"][:, :, None, None] + \ (1 - gt_polygons_mask) * background_freq[:, None, None, None] if pixel_class_freq.min() == 0: print_utils.print_error( "ERROR: pixel_class_freq has some zero values, can't divide by zero!" ) raise ZeroDivisionError freq_weights = 1 / pixel_class_freq # print("freq_weights:", freq_weights.min().item(), freq_weights.max().item()) # Compute size weights # print("sizes:", gt_batch["sizes"].min().item(), gt_batch["sizes"].max().item()) # print("distances:", gt_batch["distances"].min().item(), gt_batch["distances"].max().item()) # print("im_radius:", im_radius) size_weights = None if use_size: if gt_batch["sizes"].min() == 0: print_utils.print_error( ("ERROR: sizes tensor has zero values, can't divide by zero!")) raise ZeroDivisionError size_weights = 1 + 1 / (im_radius * gt_batch["sizes"]) distance_weights = None if use_dist: # print("distances:", gt_batch["distances"].min().item(), gt_batch["distances"].max().item()) distance_weights = gt_batch["distances"] * (height + width ) # Denormalize distances distance_weights = w0 * torch.exp(-(distance_weights**2) / (sigma**2)) # print("sum(distances == 0):", torch.sum(gt_batch["distances"] == 0).item()) # print("distance_weights:", distance_weights.min().item(), distance_weights.max().item()) # print(distance_weights.shape, distance_weights.min().item(), distance_weights.max().item()) # print(size_weights.shape, size_weights.min().item(), size_weights.max().item()) # print(freq_weights.shape, freq_weights.min().item(), freq_weights.max().item()) gt_batch["seg_loss_weights"] = freq_weights if use_dist: gt_batch["seg_loss_weights"] += distance_weights if use_size: gt_batch["seg_loss_weights"] *= size_weights # print(gt_batch["seg_loss_weights"].shape, gt_batch["seg_loss_weights"].min().item(), gt_batch["seg_loss_weights"].max().item()) # print("seg_loss_weights:", size_weights.min().item(), size_weights.max().item()) # print("freq_weights:", freq_weights.min().item(), freq_weights.max().item()) # print("size_weights:", size_weights.min().item(), size_weights.max().item()) # print("distance_weights:", distance_weights.min().item(), distance_weights.max().item()) # Display: # display_seg_loss_weights = gt_batch["seg_loss_weights"][0].cpu().detach().numpy() # display_distance_weights = distance_weights[0].cpu().detach().numpy() # skimage.io.imsave("seg_loss_dist_weights.png", display_distance_weights[0]) # display_size_weights = size_weights[0].cpu().detach().numpy() # skimage.io.imsave("seg_loss_size_weights.png", display_size_weights[0]) # display_freq_weights = freq_weights[0].cpu().detach().numpy() # display_freq_weights = display_freq_weights - display_freq_weights.min() # display_freq_weights /= display_freq_weights.max() # skimage.io.imsave("seg_loss_freq_weights.png", np.moveaxis(display_freq_weights, 0, -1)) # for i in range(3): # skimage.io.imsave(f"seg_loss_weights_{i}.png", display_seg_loss_weights[i]) # skimage.io.imsave(f"freq_weights_{i}.png", display_freq_weights[i]) return pred_batch, gt_batch
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)