diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..957b2579c6ef20995a09efd9a17f8fd90606f5ed --- /dev/null +++ b/.gitattributes @@ -0,0 +1,27 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bin.* filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zstandard filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000000000000000000000000000000000000..3b664107303df336bab8010caad42ddaed24550e --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "git.ignoreLimitWarning": true +} \ No newline at end of file diff --git a/LICENSE.txt b/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..c8b91007253607eb44b1b6e99185013a89798873 --- /dev/null +++ b/LICENSE.txt @@ -0,0 +1,94 @@ +*** +CONTRAT DE LICENCE DE LOGICIEL + +Logiciel PFFL ©Inria 2020, tout droit réservé, ci-après dénommé "le Logiciel". + +Le Logiciel a été conçu et réalisé par des chercheurs de l’équipe-projet TITANE d’Inria (Institut National de Recherche en Informatique et Automatique). + +Inria, Domaine de Voluceau, Rocquencourt - BP 105 +78153 Le Chesnay Cedex, France + +Inria détient tous les droits de propriété sur le Logiciel. + +Le Logiciel a été déposé auprès de l'Agence pour la Protection des Programmes (APP) sous le numéro . + +Le Logiciel est en cours de développement et Inria souhaite qu'il soit utilisé par la communauté scientifique de façon à le tester et l'évaluer, et afin qu’Inria puisse le cas échéant le faire évoluer. + +A cette fin, Inria a décidé de distribuer le Logiciel. + +Inria concède à l'utilisateur académique, gratuitement, sans droit de sous-licence, pour une période de un (1) an à compter du téléchargement du code source, le droit non-exclusif d'utiliser le Logiciel à fins de recherche. Toute autre utilisation sans l’accord préalable d’Inria est exclue. + +L’utilisateur académique reconnaît expressément avoir reçu d’Inria toutes les informations lui permettant d’apprécier l’adéquation du Logiciel à ses besoins et de prendre toutes les précautions utiles pour sa mise en œuvre et son utilisation. + +Le Logiciel est distribué sous forme d'un fichier source. + +Si le Logiciel est utilisé pour la publication de résultats, l’utilisateur devra citer le Logiciel de la façon suivante : + +@misc{girard2020polygonal, + title={Polygonal Building Segmentation by Frame Field Learning}, + author={Nicolas Girard and Dmitriy Smirnov and Justin Solomon and Yuliya Tarabalka}, + year={2020}, + eprint={2004.14875}, + archivePrefix={arXiv}, + primaryClass={cs.CV} +} + + +Tout utilisateur du Logiciel pourra communiquer ses remarques d'utilisation du Logiciel aux développeurs de PFFL : nicolas.girard@inria.fr + + +L'UTILISATEUR NE PEUT FAIRE NI UTILISATION NI EXPLOITATION NI DISTRIBUTION COMMERCIALE DU LOGICIEL SANS L'ACCORD EXPRÈS PRÉALABLE d’INRIA (stip-sam@inria.fr). +TOUT ACTE CONTRAIRE CONSTITUERAIT UNE CONTREFAÇON. + +LE LOGICIEL EST FOURNI "TEL QU'EN L'ÉTAT" SANS AUCUNE GARANTIE DE QUELQUE NATURE, IMPLICITE OU EXPLICITE, QUANT À SON UTILISATION COMMERCIALE, PROFESSIONNELLE, LÉGALE OU NON, OU AUTRE, SA COMMERCIALISATION OU SON ADAPTATION. + +SAUF LORSQU'EXPLICITEMENT PRÉVU PAR LA LOI, INRIA NE POURRA ÊTRE TENU POUR RESPONSABLE DE TOUT DOMMAGE OU PRÉJUDICE DIRECT,INDIRECT, (PERTES FINANCIÈRES DUES AU MANQUE À GAGNER, À L'INTERRUPTION D'ACTIVITÉS OU À LA PERTE DE DONNÉES, ETC...) DÉCOULANT DE L'UTILISATION DE TOUT OU PARTIE DU LOGICIEL OU DE L'IMPOSSIBILITÉ D'UTILISER CELUI-CI. +  +*** + +SOFTWARE LICENSE AGREEMENT + + +Software PFFL ©Inria – 2020, all rights reserved, hereinafter "the Software". + +This software has been developed by researchers of TITANE project team of Inria (Institut National de Recherche en Informatique et Automatique). + +Inria, Domaine de Voluceau, Rocquencourt - BP 105 +78153 Le Chesnay Cedex, FRANCE + + +Inria holds all the ownership rights on the Software. + +The Software has been registered with the Agence pour la Protection des Programmes (APP) under . + +The Software is still being currently developed. It is the Inria’s aim for the Software to be used by the scientific community so as to test it and, evaluate it so that Inria may improve it. + +For these reasons Inria has decided to distribute the Software. + +Inria grants to the academic user, a free of charge, without right to sublicense non-exclusive right to use the Software for research purposes for a period of one (1) year from the date of the download of the source code. Any other use without of prior consent of Inria is prohibited. + +The academic user explicitly acknowledges having received from Inria all information allowing him to appreciate the adequacy between of the Software and his needs and to undertake all necessary precautions for his execution and use. + +The Software is provided only as a source. + +In case of using the Software for a publication or other results obtained through the use of the Software, user should cite the Software as follows : + +@misc{girard2020polygonal, + title={Polygonal Building Segmentation by Frame Field Learning}, + author={Nicolas Girard and Dmitriy Smirnov and Justin Solomon and Yuliya Tarabalka}, + year={2020}, + eprint={2004.14875}, + archivePrefix={arXiv}, + primaryClass={cs.CV} +} + + +Every user of the Software could communicate to the developers of PFFL [nicolas.girard@inria.fr] his or her remarks as to the use of the Software. + +THE USER CANNOT USE, EXPLOIT OR COMMERCIALY DISTRIBUTE THE SOFTWARE WITHOUT PRIOR AND EXPLICIT CONSENT OF INRIA (stip-sam@inria.fr). ANY SUCH ACTION WILL CONSTITUTE A FORGERY. + +THIS SOFTWARE IS PROVIDED "AS IS" WITHOUT ANY WARRANTIES OF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES,WITH REGARDS TO COMMERCIAL USE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR ADAPTATION. + +UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL INRIA OR THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + + diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..0a06b85f87d313337f4e2f4d007a7f895ae9d487 --- /dev/null +++ b/README.md @@ -0,0 +1,13 @@ +--- +title: Luuu +emoji: 🌍 +colorFrom: red +colorTo: purple +sdk: gradio +sdk_version: 2.8.12 +app_file: app.py +pinned: false +license: apache-2.0 +--- + +Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..84e19b2adac127da674ca901f3ae6f6c994e3d7f --- /dev/null +++ b/app.py @@ -0,0 +1,48 @@ +''' +Author: Egrt +Date: 2022-03-19 10:23:48 +LastEditors: Egrt +LastEditTime: 2022-03-20 13:41:53 +FilePath: \Luuu\app.py +''' + +from gis import GIS +import gradio as gr +import os +os.system('pip install requirements.txt') +from zipfile import ZipFile +gis = GIS() + + +# --------模型推理---------- # +def inference(filepath): + filename, file_list = gis.detect_image(filepath) + with ZipFile("result.zip", "w") as zipObj: + zipObj.write(file_list[0], "{}.tif".format(filename+'mask')) + zipObj.write(file_list[1], "{}.tif".format(filename)) + zipObj.write(file_list[2], "{}.pdf".format(filename)) + zipObj.write(file_list[3], "{}.cpg".format(filename)) + zipObj.write(file_list[4], "{}.dbf".format(filename)) + zipObj.write(file_list[5], "{}.shx".format(filename)) + zipObj.write(file_list[6], "{}.shp".format(filename)) + zipObj.write(file_list[7], "{}.prj".format(filename)) + return "images/result.zip" + + +# --------网页信息---------- # +title = "基于帧场学习的多边形建筑提取" +description = "目前最先进图像分割模型通常以栅格形式输出分割，但地理信息系统中的应用通常需要矢量多边形。我们在遥感图像中提取建筑物的任务中，将帧场输出添加到深度分割模型中，将预测的帧场与地面实况轮廓对齐，帮助减少深度网络输出与下游任务中输出样式之间的差距。 @Luuuu🐋🐋" +article = "

Polygonization-by-Frame-Field-Learning | Github Repo

" +example_img_dir = 'images' +example_img_name = os.listdir(example_img_dir) +examples=[[os.path.join(example_img_dir, image_path)] for image_path in example_img_name if image_path.endswith('.png')] +gr.Interface( + inference, + [gr.inputs.Image(type="filepath", label="待检测图片")], + gr.outputs.File(label="检测结果"), + title=title, + description=description, + article=article, + enable_queue=True, + examples=examples + ).launch(debug=True) \ No newline at end of file diff --git a/backbone.py b/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..86a7a603dd465c0f8cd1b5e361d1c9ecb4fe18e7 --- /dev/null +++ b/backbone.py @@ -0,0 +1,71 @@ +import os + +import torch +import torchvision + +from lydorn_utils import print_utils + + +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 set_download_dir(): + os.environ['TORCH_HOME'] = 'models' # setting the environment variable diff --git a/child_processes.py b/child_processes.py new file mode 100644 index 0000000000000000000000000000000000000000..b6eaba8ca550a11b9bfc522fbe37b558e737a102 --- /dev/null +++ b/child_processes.py @@ -0,0 +1,76 @@ +import os + +import torch + +from lydorn_utils import python_utils +from lydorn_utils import print_utils + +from backbone import get_backbone +from dataset_folds import get_folds + + +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 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) + diff --git a/configs/README.md b/configs/README.md new file mode 100644 index 0000000000000000000000000000000000000000..330dca41efa1779159da234b00f26f64d8a82865 --- /dev/null +++ b/configs/README.md @@ -0,0 +1,68 @@ +This folder stores all configuration variables used for launching training runs (and evaluating the results from those runs) in the form on config files. + +The ```main.py``` script has a ```--config``` argument which can be the path to any of the "config.*.json" files in this folder. Of course you can also write your own config file. +Thus one "config.*.json" files points to all parameters used for the run. + +As we perform quite a lot of different experiments requiring only a few parameters to change, I have setup a hierarchical system of config files allowing default values that can be overwritten. +This way a single parameter can be shared by all experiments and can be changed for all of them by modifying it in only one location. +A config is stored as a JSON dictionary. This config dictionary can store nested dictionaries of parameters. +Whenever the "defaults_filepath" key is used in a config file, +its value is assumed to be the path to another config file whose dictionary is loaded and merged with the dictionary that had the "defaults_filepath" key. +The keys alongside the "defaults_filepath" key specify parameters that should overwrite the default values loaded from the "defaults_filepath" config file. + +To illustrate how this works, let's say the config folder looks like this: +``` +config +|-- config.defaults.json +`-- config.my_exp_1.json +`-- config.my_exp_2.json +``` + +Let's say config.defaults.json is: +```json +{ + "learning_rate": 0.1, + "batch_size": 16 +} +``` + +And config.my_exp_1.json is: +```json +{ + "defaults_filepath": "configs/config.defaults.json" +} +``` + +And config.my_exp_2.json is: +```json +{ + "defaults_filepath": "configs/config.defaults.json", + + "learning_rate": 0.01 + +} +``` + +When loaded by the ```main.py``` script, they will be expanded into the following. + +config.my_exp_1.json: +```json +{ + "learning_rate": 0.1, + "batch_size": 16 +} +``` + +config.my_exp_2.json: +```json +{ + "learning_rate": 0.01, + "batch_size": 16 + +} +``` + +When a lot of parameters are used by the actual config files we used, it is thus very easy to know that all "my_exp_2" does is change the learning rate. +Also if we want to change the batch size for all experiments, all we have to do is change its value in "config.defaults.json". +This principle of using the "defaults_filepath" key to point to another config file can be used in nested dictionary parameters as well. +A config file is thus the root of a config tree loaded recursively. \ No newline at end of file diff --git a/configs/backbone_params.deeplab101.json b/configs/backbone_params.deeplab101.json new file mode 100644 index 0000000000000000000000000000000000000000..cfdb2e730d0e694efd7bd7432d0f6a2fa24d97af --- /dev/null +++ b/configs/backbone_params.deeplab101.json @@ -0,0 +1,6 @@ +{ + "name": "deeplab101", + "input_features": 3, + "features": 128, + "pretrained": false +} \ No newline at end of file diff --git a/configs/backbone_params.deeplab50.json b/configs/backbone_params.deeplab50.json new file mode 100644 index 0000000000000000000000000000000000000000..bda1fd50bf37de8a88ad1a6b6c59566c196972e5 --- /dev/null +++ b/configs/backbone_params.deeplab50.json @@ -0,0 +1,6 @@ +{ + "name": "deeplab50", + "input_features": 3, + "features": 128, + "pretrained": false +} \ No newline at end of file diff --git a/configs/backbone_params.fcn101.json b/configs/backbone_params.fcn101.json new file mode 100644 index 0000000000000000000000000000000000000000..13ada28f5610151998f0e4c34acd07c94cd2ab9b --- /dev/null +++ b/configs/backbone_params.fcn101.json @@ -0,0 +1,6 @@ +{ + "name": "fcn101", + "input_features": 3, + "features": 256, + "pretrained": false +} \ No newline at end of file diff --git a/configs/backbone_params.fcn50.json b/configs/backbone_params.fcn50.json new file mode 100644 index 0000000000000000000000000000000000000000..ccc52e7d52ee418a91516bead26d912b9b19e467 --- /dev/null +++ b/configs/backbone_params.fcn50.json @@ -0,0 +1,6 @@ +{ + "name": "fcn50", + "input_features": 3, + "features": 256, + "pretrained": false +} \ No newline at end of file diff --git a/configs/backbone_params.ictnet.json b/configs/backbone_params.ictnet.json new file mode 100644 index 0000000000000000000000000000000000000000..0117f8121dfd25ff3aff4e32540d2e47eaf6c241 --- /dev/null +++ b/configs/backbone_params.ictnet.json @@ -0,0 +1,8 @@ +{ + "name": "ictnet", + "preset_model": "FC-DenseNet103", // FC-DenseNet56, FC-DenseNet67 and FC-DenseNet103 are possible + "in_channels": 3, + "out_channels": 32, + "dropout_2d": 0.0, // Default: 0.2 + "efficient": true // If true, use gradient checkpointing to reduce memory at the expense of some speed +} \ No newline at end of file diff --git a/configs/backbone_params.unet16.json b/configs/backbone_params.unet16.json new file mode 100644 index 0000000000000000000000000000000000000000..807930144cbb878db7f40242995bcc10c17d897a --- /dev/null +++ b/configs/backbone_params.unet16.json @@ -0,0 +1,5 @@ +{ + "name": "unet", + "input_features": 3, + "features": 16 +} \ No newline at end of file diff --git a/configs/backbone_params.unet_resnet101.json b/configs/backbone_params.unet_resnet101.json new file mode 100644 index 0000000000000000000000000000000000000000..2077f94dca6419379fedfe13e8ef639c885fa1f8 --- /dev/null +++ b/configs/backbone_params.unet_resnet101.json @@ -0,0 +1,9 @@ +{ + "name": "unet_resnet", + "encoder_depth": 101, // 34, 101 and 152 are possible + "input_features": 3, + "num_filters": 32, // Default: 32 + "pretrained": false, + "dropout_2d": 0.2, // Default: 0.2 + "is_deconv": false // Default: false +} \ No newline at end of file diff --git a/configs/config.defaults.inria_dataset_osm_aligned.json b/configs/config.defaults.inria_dataset_osm_aligned.json new file mode 100644 index 0000000000000000000000000000000000000000..b509226578fdc97524cfb0b7e8726acdc0251995 --- /dev/null +++ b/configs/config.defaults.inria_dataset_osm_aligned.json @@ -0,0 +1,15 @@ +{ + "defaults_filepath": "configs/config.defaults.json", + + "dataset_params": { + "defaults_filepath": "configs/dataset_params.inria_dataset_osm_aligned.json" // Path from the project's root to a JSON with default values for dataset_params + }, + + "eval_params" : { + "defaults_filepath": "configs/eval_params.inria_dataset.json" // Path from the project's root to a JSON with default values for eval_params + }, + + "optim_params": { + "gamma": 0.99 + } +} \ No newline at end of file diff --git a/configs/config.defaults.inria_dataset_osm_mask_only.json b/configs/config.defaults.inria_dataset_osm_mask_only.json new file mode 100644 index 0000000000000000000000000000000000000000..b3b54dfb3273834c9969095e43e2c614a54a2cac --- /dev/null +++ b/configs/config.defaults.inria_dataset_osm_mask_only.json @@ -0,0 +1,21 @@ +{ + "defaults_filepath": "configs/config.defaults.json", + + "dataset_params": { + "defaults_filepath": "configs/dataset_params.inria_dataset_osm_mask_only.json" // Path from the project's root to a JSON with default values for dataset_params + }, + + "eval_params" : { + "defaults_filepath": "configs/eval_params.inria_dataset.json" // Path from the project's root to a JSON with default values for eval_params + }, + + "optim_params": { + "gamma": 0.99 + }, + + "data_aug_params": { + "color_jitter": false + }, + + "compute_seg": false +} \ No newline at end of file diff --git a/configs/config.defaults.inria_dataset_polygonized.json b/configs/config.defaults.inria_dataset_polygonized.json new file mode 100644 index 0000000000000000000000000000000000000000..8c47f5234b3a9898aa9ca0451180ea34d9a656fa --- /dev/null +++ b/configs/config.defaults.inria_dataset_polygonized.json @@ -0,0 +1,15 @@ +{ + "defaults_filepath": "configs/config.defaults.json", + + "dataset_params": { + "defaults_filepath": "configs/dataset_params.inria_dataset_polygonized.json" // Path from the project's root to a JSON with default values for dataset_params + }, + + "eval_params" : { + "defaults_filepath": "configs/eval_params.inria_dataset.json" // Path from the project's root to a JSON with default values for eval_params + }, + + "optim_params": { + "gamma": 0.99 + } +} \ No newline at end of file diff --git a/configs/config.defaults.inria_dataset_polygonized_256.json b/configs/config.defaults.inria_dataset_polygonized_256.json new file mode 100644 index 0000000000000000000000000000000000000000..7064e894d8b177f5cf7b0b9056f98d2b5062d0fd --- /dev/null +++ b/configs/config.defaults.inria_dataset_polygonized_256.json @@ -0,0 +1,15 @@ +{ + "defaults_filepath": "configs/config.defaults.json", + + "dataset_params": { + "defaults_filepath": "configs/dataset_params.inria_dataset_polygonized_256.json" // Path from the project's root to a JSON with default values for dataset_params + }, + + "eval_params" : { + "defaults_filepath": "configs/eval_params.inria_dataset.json" // Path from the project's root to a JSON with default values for eval_params + }, + + "optim_params": { + "gamma": 0.99 + } +} \ No newline at end of file diff --git a/configs/config.defaults.json b/configs/config.defaults.json new file mode 100644 index 0000000000000000000000000000000000000000..dbe0a9bbd6ca2d03cd281b42461bcdd8bf4c0756 --- /dev/null +++ b/configs/config.defaults.json @@ -0,0 +1,40 @@ +{ + "data_dir_candidates": [ + "/data/titane/user/nigirard/data", + "~/data", + "/data" + ], + "num_workers": null, // If null, will use multiprocess.cpu_count() workers in total + "data_aug_params": { + "enable": true, + "vflip": true, + "affine": true, + "scaling": [0.75, 1.5], // Range of scaling factor to apply during affine transform. Set to None to not apply. + "color_jitter": true, + "device": "cuda" + }, + + "device": "cuda", // Only has effects when mode is val or test. When mode is train, always use CUDA + "use_amp": false, // Automatic Mixed Precision switch + + "compute_seg": true, + "compute_crossfield": true, + + "seg_params": { + "compute_interior": true, + "compute_edge": true, + "compute_vertex": false + }, + + "loss_params": { + "defaults_filepath": "configs/loss_params.json" // Path from the project's root to a JSON with default values for dataset_params + }, + + "optim_params": { + "defaults_filepath": "configs/optim_params.json" // Path from the project's root to a JSON with default values for optim_params + }, + + "polygonize_params": { + "defaults_filepath": "configs/polygonize_params.json" // Path from the project's root to a JSON with default values for polygonize_params + } +} diff --git a/configs/config.defaults.luxcarta_dataset.json b/configs/config.defaults.luxcarta_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..21ea3b97552ef921af769045f1b1c8f7c68a50aa --- /dev/null +++ b/configs/config.defaults.luxcarta_dataset.json @@ -0,0 +1,11 @@ +{ + "defaults_filepath": "configs/config.defaults.json", + + "dataset_params": { + "defaults_filepath": "configs/dataset_params.luxcarta_dataset.json" // Path from the project's root to a JSON with default values for dataset_params + }, + + "eval_params" : { + "defaults_filepath": "configs/eval_params.luxcarta_dataset.json" // Path from the project's root to a JSON with default values for eval_params + } +} \ No newline at end of file diff --git a/configs/config.defaults.mapping_dataset.json b/configs/config.defaults.mapping_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..2bb94ef62c10137ec593d181e315ad578e0ac7c2 --- /dev/null +++ b/configs/config.defaults.mapping_dataset.json @@ -0,0 +1,11 @@ +{ + "defaults_filepath": "configs/config.defaults.json", + + "dataset_params": { + "defaults_filepath": "configs/dataset_params.mapping_dataset.json" // Path from the project's root to a JSON with default values for dataset_params + }, + + "eval_params" : { + "defaults_filepath": "configs/eval_params.mapping_dataset.json" // Path from the project's root to a JSON with default values for eval_params + } +} \ No newline at end of file diff --git a/configs/config.defaults.xview2_dataset.json b/configs/config.defaults.xview2_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..db80dbbbe08e4d08dd8e69e75358ea090c92d6ac --- /dev/null +++ b/configs/config.defaults.xview2_dataset.json @@ -0,0 +1,15 @@ +{ + "defaults_filepath": "configs/config.defaults.json", + + "dataset_params": { + "defaults_filepath": "configs/dataset_params.xview2_dataset.json" // Path from the project's root to a JSON with default values for dataset_params + }, + + "eval_params" : { + "defaults_filepath": "configs/eval_params.xview2_dataset.json" // Path from the project's root to a JSON with default values for eval_params + }, + + "optim_params": { + "gamma": 0.99 + } +} \ No newline at end of file diff --git a/configs/config.inria_dataset_osm_aligned.unet_resnet101_pretrained.field_off.json b/configs/config.inria_dataset_osm_aligned.unet_resnet101_pretrained.field_off.json new file mode 100644 index 0000000000000000000000000000000000000000..9b4acc40fd021e902938f2a104cc46ad70bc2c79 --- /dev/null +++ b/configs/config.inria_dataset_osm_aligned.unet_resnet101_pretrained.field_off.json @@ -0,0 +1,24 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_osm_aligned.json", + + "run_name": "inria_dataset_osm_aligned.unet_resnet101_pretrained.field_off", + + "compute_crossfield": false, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "loss_params": { + "seg_loss_params": { + "use_size": false + } + }, + + "optim_params": { + "optimizer": "RMSProp", + "gamma": 0.99, + "batch_size": 3 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_osm_aligned.unet_resnet101_pretrained.json b/configs/config.inria_dataset_osm_aligned.unet_resnet101_pretrained.json new file mode 100644 index 0000000000000000000000000000000000000000..1e9cb3d79c21f56edf7398986b00c793d180c1de --- /dev/null +++ b/configs/config.inria_dataset_osm_aligned.unet_resnet101_pretrained.json @@ -0,0 +1,24 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_osm_aligned.json", + + "run_name": "inria_dataset_osm_aligned.unet_resnet101_pretrained", + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "loss_params": { + "seg_loss_params": { + "use_size": false + } + }, + + "optim_params": { + "optimizer": "RMSProp", + "gamma": 0.99, + "batch_size": 3 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_osm_mask_only.unet16.json b/configs/config.inria_dataset_osm_mask_only.unet16.json new file mode 100644 index 0000000000000000000000000000000000000000..c09a5f533cc48c2a7bf78181bcb9f3f39bb9f4d5 --- /dev/null +++ b/configs/config.inria_dataset_osm_mask_only.unet16.json @@ -0,0 +1,17 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_osm_mask_only.json", + + "run_name": "inria_dataset_osm_mask_only.unet16", + + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 16 // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_polygonized.ictnet.leaderboard.json b/configs/config.inria_dataset_polygonized.ictnet.leaderboard.json new file mode 100644 index 0000000000000000000000000000000000000000..24b94bd8bf3a8fc3d2903b800c043a66437e7056 --- /dev/null +++ b/configs/config.inria_dataset_polygonized.ictnet.leaderboard.json @@ -0,0 +1,38 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized.ictnet.leaderboard", + + + + + + + + "seg_params": { + "compute_interior": true, + "compute_edge": false, + "compute_vertex": false + }, + + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.ictnet.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "loss_params": { + "seg_loss_params": { + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, // Dist weights as in the original U-Net paper + "use_size": false // Size weights increasing importance of smaller buildings + } + }, + + "optim_params": { + "batch_size": 2 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_polygonized.ictnet.rmsprop.leaderboard.field_off.json b/configs/config.inria_dataset_polygonized.ictnet.rmsprop.leaderboard.field_off.json new file mode 100644 index 0000000000000000000000000000000000000000..760a1b896d31f5ec5f2113e51a0a721b90622b25 --- /dev/null +++ b/configs/config.inria_dataset_polygonized.ictnet.rmsprop.leaderboard.field_off.json @@ -0,0 +1,39 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized.ictnet.rmsprop.leaderboard.field_off", + + "compute_crossfield": false, + + + + + + "seg_params": { + "compute_interior": true, + "compute_edge": false, + "compute_vertex": false + }, + + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.ictnet.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "loss_params": { + "seg_loss_params": { + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, // Dist weights as in the original U-Net paper + "use_size": false // Size weights increasing importance of smaller buildings + } + }, + + "optim_params": { + "optimizer": "RMSProp", + "batch_size": 2 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_polygonized.ictnet.rmsprop.leaderboard.json b/configs/config.inria_dataset_polygonized.ictnet.rmsprop.leaderboard.json new file mode 100644 index 0000000000000000000000000000000000000000..3bebbd6375662a4835dbd02fc20e20d38dbbf8ad --- /dev/null +++ b/configs/config.inria_dataset_polygonized.ictnet.rmsprop.leaderboard.json @@ -0,0 +1,39 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized.ictnet.rmsprop.leaderboard", + + + + + + + + "seg_params": { + "compute_interior": true, + "compute_edge": false, + "compute_vertex": false + }, + + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.ictnet.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "loss_params": { + "seg_loss_params": { + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, // Dist weights as in the original U-Net paper + "use_size": false // Size weights increasing importance of smaller buildings + } + }, + + "optim_params": { + "optimizer": "RMSProp", + "batch_size": 2 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_polygonized.unet_resnet101.leaderboard.json b/configs/config.inria_dataset_polygonized.unet_resnet101.leaderboard.json new file mode 100644 index 0000000000000000000000000000000000000000..0a43e5fb9934f09395025f8912dd35578aabcb3f --- /dev/null +++ b/configs/config.inria_dataset_polygonized.unet_resnet101.leaderboard.json @@ -0,0 +1,33 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized.unet_resnet101.leaderboard", + + "seg_params": { + "compute_interior": true, + "compute_edge": false, + "compute_vertex": false + }, + + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": false + }, + + "loss_params": { + "seg_loss_params": { + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, // Dist weights as in the original U-Net paper + "use_size": false // Size weights increasing importance of smaller buildings + } + }, + + "optim_params": { + "batch_size": 4 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.field_off.json b/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.field_off.json new file mode 100644 index 0000000000000000000000000000000000000000..cd2529c42c9ef8f22b2051941d8fa5d7ce474595 --- /dev/null +++ b/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.field_off.json @@ -0,0 +1,18 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized.unet_resnet101_pretrained.field_off", + + "compute_crossfield": false, + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.json b/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.json new file mode 100644 index 0000000000000000000000000000000000000000..260d3d62bae55a8fb18eeb34c21e31a3991e82ee --- /dev/null +++ b/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.json @@ -0,0 +1,18 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized.unet_resnet101_pretrained", + + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.leaderboard.field_off.json b/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.leaderboard.field_off.json new file mode 100644 index 0000000000000000000000000000000000000000..0aeee724799d919e825d5bbb8ff43bb9eda84a90 --- /dev/null +++ b/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.leaderboard.field_off.json @@ -0,0 +1,35 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized.unet_resnet101_pretrained.leaderboard.field_off", + + "compute_crossfield": false, + + "seg_params": { + "compute_interior": true, + "compute_edge": false, + "compute_vertex": false + }, + + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "loss_params": { + "seg_loss_params": { + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, // Dist weights as in the original U-Net paper + "use_size": false // Size weights increasing importance of smaller buildings + } + }, + + "optim_params": { + "batch_size": 4 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.leaderboard.json b/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.leaderboard.json new file mode 100644 index 0000000000000000000000000000000000000000..82936d04e4b37bfa063e191a8108da0a78e5cc8c --- /dev/null +++ b/configs/config.inria_dataset_polygonized.unet_resnet101_pretrained.leaderboard.json @@ -0,0 +1,35 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized.unet_resnet101_pretrained.leaderboard", + + + + "seg_params": { + "compute_interior": true, + "compute_edge": false, + "compute_vertex": false + }, + + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "loss_params": { + "seg_loss_params": { + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, // Dist weights as in the original U-Net paper + "use_size": false // Size weights increasing importance of smaller buildings + } + }, + + "optim_params": { + "batch_size": 4 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.inria_dataset_polygonized_small.unet_resnet101_pretrained.json b/configs/config.inria_dataset_polygonized_small.unet_resnet101_pretrained.json new file mode 100644 index 0000000000000000000000000000000000000000..4043a6487364d1835ce3fb1b5c7179a699efeebe --- /dev/null +++ b/configs/config.inria_dataset_polygonized_small.unet_resnet101_pretrained.json @@ -0,0 +1,19 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized_small.unet_resnet101_pretrained", + + "dataset_params": { + "small": true + }, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "gamma": 1.0 + } +} diff --git a/configs/config.inria_dataset_polygonized_small.unet_resnet101_pretrained.no_aug.json b/configs/config.inria_dataset_polygonized_small.unet_resnet101_pretrained.no_aug.json new file mode 100644 index 0000000000000000000000000000000000000000..59acecae984441b8d97b13042b8badf8dad86a73 --- /dev/null +++ b/configs/config.inria_dataset_polygonized_small.unet_resnet101_pretrained.no_aug.json @@ -0,0 +1,23 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset_polygonized.json", + + "run_name": "inria_dataset_polygonized_small.unet_resnet101_pretrained.no_aug", + + "dataset_params": { + "small": true + }, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "gamma": 1.0, + "log_steps": 10 + }, + + "data_aug_params": { + "enable": false + } +} diff --git a/configs/config.inria_dataset_small.unet16.json b/configs/config.inria_dataset_small.unet16.json new file mode 100644 index 0000000000000000000000000000000000000000..d900d02112bb665f0a1e9ce9f0464ead8ecfdd6e --- /dev/null +++ b/configs/config.inria_dataset_small.unet16.json @@ -0,0 +1,19 @@ +{ + "defaults_filepath": "configs/config.defaults.inria_dataset.json", + + "run_name": "inria_dataset_small.unet16", + + "dataset_params": { + "small": true + }, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + + }, + + "optim_params": { + "batch_size": 12, // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "gamma": 0.99 + } +} diff --git a/configs/config.luxcarta_dataset.unet16.field_off.json b/configs/config.luxcarta_dataset.unet16.field_off.json new file mode 100644 index 0000000000000000000000000000000000000000..46ccdffe284db033f0b0f3b24d6c97b5f40b0943 --- /dev/null +++ b/configs/config.luxcarta_dataset.unet16.field_off.json @@ -0,0 +1,15 @@ +{ + "defaults_filepath": "configs/config.defaults.luxcarta_dataset.json", + + "run_name": "luxcarta_dataset.unet16.field_off", + + "compute_crossfield": false, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 16 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.luxcarta_dataset.unet16.json b/configs/config.luxcarta_dataset.unet16.json new file mode 100644 index 0000000000000000000000000000000000000000..374011d27957f6e529dbcb2a6cfcffe8cbc6601c --- /dev/null +++ b/configs/config.luxcarta_dataset.unet16.json @@ -0,0 +1,15 @@ +{ + "defaults_filepath": "configs/config.defaults.luxcarta_dataset.json", + + "run_name": "luxcarta_dataset.unet16", + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 16 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.asip.json b/configs/config.mapping_dataset.asip.json new file mode 100644 index 0000000000000000000000000000000000000000..62c7fb7c164cc33fa39f2f544b21695e02bd1e2c --- /dev/null +++ b/configs/config.mapping_dataset.asip.json @@ -0,0 +1,5 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.asip" +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.deeplab101.field_off.json b/configs/config.mapping_dataset.deeplab101.field_off.json new file mode 100644 index 0000000000000000000000000000000000000000..47aa41a6a4f93a50d4761c9b54d2cc5c035b17b6 --- /dev/null +++ b/configs/config.mapping_dataset.deeplab101.field_off.json @@ -0,0 +1,15 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.deeplab101.field_off", + + "compute_crossfield": false, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.deeplab101.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 8 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.mapping_dataset.deeplab101.field_off.train_val.json b/configs/config.mapping_dataset.deeplab101.field_off.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..190bd2bdc22e5b146e4f91a720ad396258900631 --- /dev/null +++ b/configs/config.mapping_dataset.deeplab101.field_off.train_val.json @@ -0,0 +1,16 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.deeplab101.field_off.train_val", + "fold": ["train", "val"], + + "compute_crossfield": false, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.deeplab101.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 8 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.deeplab101.json b/configs/config.mapping_dataset.deeplab101.json new file mode 100644 index 0000000000000000000000000000000000000000..c26b2028e9e1511cbb9328d4afd6f62e5bf88899 --- /dev/null +++ b/configs/config.mapping_dataset.deeplab101.json @@ -0,0 +1,16 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.deeplab101", + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.deeplab101.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 8 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.deeplab101.train_val.json b/configs/config.mapping_dataset.deeplab101.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..eaf5b34da772cb51a7bd926b1c2c6d516c188f95 --- /dev/null +++ b/configs/config.mapping_dataset.deeplab101.train_val.json @@ -0,0 +1,15 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.deeplab101.train_val", + "fold": ["train", "val"], + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.deeplab101.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 8 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.fcn101.json b/configs/config.mapping_dataset.fcn101.json new file mode 100644 index 0000000000000000000000000000000000000000..5cec638fa862b2626c18ecf28d88a3e32eef1ccc --- /dev/null +++ b/configs/config.mapping_dataset.fcn101.json @@ -0,0 +1,89 @@ +{ + "run_name": "mapping_dataset.fcn50", + + "data_dir_candidates": [ + "/local/shared/data", // try cluster local node first + "/data/titane/user/nigirard/data", // Try cluster /data directory + "~/data", // In home directory (docker) + "/data" // In landsat's /data volume (docker) + ], + "data_root_partial_dirpath": "mapping_challenge_dataset", + "dataset_params": { + "small": false + }, + "num_workers": 10, + "data_split_params": { + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.75, + "val_fraction": 0.25 // test_fraction is the rest + }, + "data_aug_params": { + "enable": true, + "vflip": true, + "rotate": true, + "color_jitter": true, + "device": "cuda" + }, + + "device": "cuda", // Only has effects when mode is val or test. When mode is train, always use CUDA + "use_amp": true, // Automatic Mixed Precision switch + + "backbone_params": { + "name": "fcn101", + "input_features": 3, + "features": 256, + "pretrained": false + }, + + "compute_seg": true, + "compute_crossfield": true, + + "seg_params": { + "compute_interior": true, + "compute_edge": true, + "compute_vertex": false + }, + + "loss_params": { + "multiloss": { + "normalization_params": { + "min_samples": 10, // Per GPU + "max_samples": 1000 // Per GPU + }, + "coefs": { + "seg_interior": 1, + "seg_edge": 1, + "seg_vertex": 0, + "crossfield_align": 1, + "crossfield_align90": 0.2, + "crossfield_smooth": 0.2, + "seg_interior_crossfield": 0.2, + "seg_edge_crossfield": 0.2, + "seg_edge_interior": 0.2 + } + }, + "seg_loss_params": { // https://github.com/neptune-ai/open-solution-mapping-challenge/blob/master/neptune.yaml + "bce_coef": 1.0, + "dice_coef": 0.2, + "w0": 50, // From original U-Net paper: distance weight to increase loss between objects + "sigma": 10 // From original U-Net paper: distance weight to increase loss between objects + } + }, + + "batch_size": 12, // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "base_lr": 1e-4, // Will be multiplied by the effective_batch_size=world_size*batch_size. + "max_lr": 1e-2, // Maximum learning rate + "warmup_epochs": 1, // Number of epochs for warmup (learning rate starts at lr*warmup_factor and gradually increases to lr) + "warmup_factor": 1e-3, + "weight_decay": 0, + "dropout_keep_prob": 1.0, // Not used for now + "max_epoch": 25, + "log_steps": 50, + "checkpoint_epoch": 1, + "checkpoints_to_keep": 10, // outputs + "logs_dirname": "logs", + "save_input_output": false, + "log_input_output": false, + "checkpoints_dirname": "checkpoints", + "eval_dirname": "eval" +} diff --git a/configs/config.mapping_dataset.fcn50.json b/configs/config.mapping_dataset.fcn50.json new file mode 100644 index 0000000000000000000000000000000000000000..4379649eb7bccca45851fe3a8d1be3276fee833c --- /dev/null +++ b/configs/config.mapping_dataset.fcn50.json @@ -0,0 +1,89 @@ +{ + "run_name": "mapping_dataset.fcn50", + + "data_dir_candidates": [ + "/local/shared/data", // try cluster local node first + "/data/titane/user/nigirard/data", // Try cluster /data directory + "~/data", // In home directory (docker) + "/data" // In landsat's /data volume (docker) + ], + "data_root_partial_dirpath": "mapping_challenge_dataset", + "dataset_params": { + "small": false + }, + "num_workers": 10, + "data_split_params": { + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.75, + "val_fraction": 0.25 // test_fraction is the rest + }, + "data_aug_params": { + "enable": true, + "vflip": true, + "rotate": true, + "color_jitter": true, + "device": "cuda" + }, + + "device": "cuda", // Only has effects when mode is val or test. When mode is train, always use CUDA + "use_amp": true, // Automatic Mixed Precision switch + + "backbone_params": { + "name": "fcn50", + "input_features": 3, + "features": 256, + "pretrained": false + }, + + "compute_seg": true, + "compute_crossfield": true, + + "seg_params": { + "compute_interior": true, + "compute_edge": true, + "compute_vertex": false + }, + + "loss_params": { + "multiloss": { + "normalization_params": { + "min_samples": 10, // Per GPU + "max_samples": 1000 // Per GPU + }, + "coefs": { + "seg_interior": 1, + "seg_edge": 1, + "seg_vertex": 0, + "crossfield_align": 1, + "crossfield_align90": 0.2, + "crossfield_smooth": 0.2, + "seg_interior_crossfield": 0.2, + "seg_edge_crossfield": 0.2, + "seg_edge_interior": 0.2 + } + }, + "seg_loss_params": { // https://github.com/neptune-ai/open-solution-mapping-challenge/blob/master/neptune.yaml + "bce_coef": 1.0, + "dice_coef": 0.2, + "w0": 50, // From original U-Net paper: distance weight to increase loss between objects + "sigma": 10 // From original U-Net paper: distance weight to increase loss between objects + } + }, + + "batch_size": 16, // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "base_lr": 1e-4, // Will be multiplied by the effective_batch_size=world_size*batch_size. + "max_lr": 1e-2, // Maximum learning rate + "warmup_epochs": 1, // Number of epochs for warmup (learning rate starts at lr*warmup_factor and gradually increases to lr) + "warmup_factor": 1e-3, + "weight_decay": 0, + "dropout_keep_prob": 1.0, // Not used for now + "max_epoch": 25, + "log_steps": 50, + "checkpoint_epoch": 1, + "checkpoints_to_keep": 10, // outputs + "logs_dirname": "logs", + "save_input_output": false, + "log_input_output": false, + "checkpoints_dirname": "checkpoints", + "eval_dirname": "eval" +} diff --git a/configs/config.mapping_dataset.open_solution.json b/configs/config.mapping_dataset.open_solution.json new file mode 100644 index 0000000000000000000000000000000000000000..779308eacb9ee1cce5037bc7b4bbff147d4ab0c1 --- /dev/null +++ b/configs/config.mapping_dataset.open_solution.json @@ -0,0 +1,5 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.open_solution" +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.open_solution_full.json b/configs/config.mapping_dataset.open_solution_full.json new file mode 100644 index 0000000000000000000000000000000000000000..c20bfef839c7854dc0f3179cc463eedca3efa7f8 --- /dev/null +++ b/configs/config.mapping_dataset.open_solution_full.json @@ -0,0 +1,5 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.open_solution_full" +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.polymapper.json b/configs/config.mapping_dataset.polymapper.json new file mode 100644 index 0000000000000000000000000000000000000000..7051cf0608479a5b32900f90a94fff225ba451d9 --- /dev/null +++ b/configs/config.mapping_dataset.polymapper.json @@ -0,0 +1,5 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.polymapper" +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet16.coupling_losses_0.4.json b/configs/config.mapping_dataset.unet16.coupling_losses_0.4.json new file mode 100644 index 0000000000000000000000000000000000000000..a20dec4e0c0b64f6888fcfb4bf495d5259740d65 --- /dev/null +++ b/configs/config.mapping_dataset.unet16.coupling_losses_0.4.json @@ -0,0 +1,88 @@ +{ + "run_name": "mapping_dataset.unet16.coupling_losses_0.4", + + "data_dir_candidates": [ + "/local/shared/data", // try cluster local node first + "/data/titane/user/nigirard/data", // Try cluster /data directory + "~/data", // In home directory (docker) + "/data" // In landsat's /data volume (docker) + ], + "data_root_partial_dirpath": "mapping_challenge_dataset", + "dataset_params": { + "small": false + }, + "num_workers": 10, + "data_split_params": { + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.75, + "val_fraction": 0.25 // test_fraction is the rest + }, + "data_aug_params": { + "enable": true, + "vflip": true, + "rotate": true, + "color_jitter": true, + "device": "cuda" + }, + + "device": "cuda", // Only has effects when mode is val or test. When mode is train, always use CUDA + "use_amp": true, // Automatic Mixed Precision switch + + "backbone_params": { + "name": "unet", + "input_features": 3, + "features": 16 + }, + + "compute_seg": true, + "compute_crossfield": true, + + "seg_params": { + "compute_interior": true, + "compute_edge": true, + "compute_vertex": false + }, + + "loss_params": { + "multiloss": { + "normalization_params": { + "min_samples": 10, // Per GPU + "max_samples": 1000 // Per GPU + }, + "coefs": { + "seg_interior": 1, + "seg_edge": 1, + "seg_vertex": 0, + "crossfield_align": 1, + "crossfield_align90": 0.4, + "crossfield_smooth": 0.4, + "seg_interior_crossfield": 0.4, + "seg_edge_crossfield": 0.4, + "seg_edge_interior": 0.4 + } + }, + "seg_loss_params": { // https://github.com/neptune-ai/open-solution-mapping-challenge/blob/master/neptune.yaml + "bce_coef": 1.0, + "dice_coef": 0.2, + "w0": 50, // From original U-Net paper: distance weight to increase loss between objects + "sigma": 10 // From original U-Net paper: distance weight to increase loss between objects + } + }, + + "batch_size": 32, // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "base_lr": 1e-4, // Will be multiplied by the effective_batch_size=world_size*batch_size. + "max_lr": 1e-2, // Maximum learning rate + "warmup_epochs": 1, // Number of epochs for warmup (learning rate starts at lr*warmup_factor and gradually increases to lr) + "warmup_factor": 1e-3, + "weight_decay": 0, + "dropout_keep_prob": 1.0, // Not used for now + "max_epoch": 25, + "log_steps": 50, + "checkpoint_epoch": 1, + "checkpoints_to_keep": 10, // outputs + "logs_dirname": "logs", + "save_input_output": false, + "log_input_output": false, + "checkpoints_dirname": "checkpoints", + "eval_dirname": "eval" +} diff --git a/configs/config.mapping_dataset.unet16.coupling_losses_off.json b/configs/config.mapping_dataset.unet16.coupling_losses_off.json new file mode 100644 index 0000000000000000000000000000000000000000..1e921f4973cc975879ea6d3f5f4845cc2f93a3dc --- /dev/null +++ b/configs/config.mapping_dataset.unet16.coupling_losses_off.json @@ -0,0 +1,26 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet16.coupling_losses_off", + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "loss_params": { + "multiloss": { + "coefs": { + "seg_interior_crossfield": 0.0, + "seg_edge_crossfield": 0.0, + "seg_edge_interior": 0.0 + } + } + }, + + "optim_params": { + "batch_size": 32 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.mapping_dataset.unet16.coupling_losses_off.train_val.json b/configs/config.mapping_dataset.unet16.coupling_losses_off.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..f6937f87b4c1324591833d2282638f6a8a0acdb0 --- /dev/null +++ b/configs/config.mapping_dataset.unet16.coupling_losses_off.train_val.json @@ -0,0 +1,26 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet16.coupling_losses_off.train_val", + "fold": ["train", "val"], + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "loss_params": { + "multiloss": { + "coefs": { + "seg_interior_crossfield": 0.0, + "seg_edge_crossfield": 0.0, + "seg_edge_interior": 0.0 + } + } + }, + + "optim_params": { + "batch_size": 32 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} diff --git a/configs/config.mapping_dataset.unet16.field_off.json b/configs/config.mapping_dataset.unet16.field_off.json new file mode 100644 index 0000000000000000000000000000000000000000..45984fc43fbe947ea8939520bfaa84f71fb7b4ef --- /dev/null +++ b/configs/config.mapping_dataset.unet16.field_off.json @@ -0,0 +1,16 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet16.field_off", + + + "compute_crossfield": false, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 32 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet16.field_off.train_val.json b/configs/config.mapping_dataset.unet16.field_off.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..bd9198a9e2b8bcec1bfe6e18df22f66c846cf0ca --- /dev/null +++ b/configs/config.mapping_dataset.unet16.field_off.train_val.json @@ -0,0 +1,16 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet16.field_off.train_val", + "fold": ["train", "val"], + + "compute_crossfield": false, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 32 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet16.json b/configs/config.mapping_dataset.unet16.json new file mode 100644 index 0000000000000000000000000000000000000000..4b12967914629b06c18646b8cddaa0f6d1a6fb91 --- /dev/null +++ b/configs/config.mapping_dataset.unet16.json @@ -0,0 +1,16 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet16", + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 32 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet16.train_val.json b/configs/config.mapping_dataset.unet16.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..e96321ce3b0a587fe7640986d966046c7676241b --- /dev/null +++ b/configs/config.mapping_dataset.unet16.train_val.json @@ -0,0 +1,16 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet16.train_val", + "fold": ["train", "val"], + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 32 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet_resnet101.json b/configs/config.mapping_dataset.unet_resnet101.json new file mode 100644 index 0000000000000000000000000000000000000000..497744bc0dbfaf90e7c1d4a54934dd66f7fae958 --- /dev/null +++ b/configs/config.mapping_dataset.unet_resnet101.json @@ -0,0 +1,16 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet_resnet101", + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 8 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet_resnet101_pretrained.align90_off.train_val.json b/configs/config.mapping_dataset.unet_resnet101_pretrained.align90_off.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..08d65cedc2af6c9bde343cce716b42342f799417 --- /dev/null +++ b/configs/config.mapping_dataset.unet_resnet101_pretrained.align90_off.train_val.json @@ -0,0 +1,23 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet_resnet101_pretrained.align90_off.train_val", + "fold": ["train", "val"], + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "max_epoch": 47 // Highest val was reached at epoch 47 + }, + + "loss_params": { + "coefs": { + "crossfield_align90": 0.0 + } + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet_resnet101_pretrained.edge_int_off.train_val.json b/configs/config.mapping_dataset.unet_resnet101_pretrained.edge_int_off.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..609b5965cd3faec2a7be52a886d8c241c704e735 --- /dev/null +++ b/configs/config.mapping_dataset.unet_resnet101_pretrained.edge_int_off.train_val.json @@ -0,0 +1,23 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet_resnet101_pretrained.edge_int_off.train_val", + "fold": ["train", "val"], + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "max_epoch": 47 // Highest val was reached at epoch 47 + }, + + "loss_params": { + "coefs": { + "seg_edge_interior": 0.0 + } + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet_resnet101_pretrained.field_off.json b/configs/config.mapping_dataset.unet_resnet101_pretrained.field_off.json new file mode 100644 index 0000000000000000000000000000000000000000..9742c60c20bc3b606bb1f0e448d2604efc1aa4ce --- /dev/null +++ b/configs/config.mapping_dataset.unet_resnet101_pretrained.field_off.json @@ -0,0 +1,17 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet_resnet101_pretrained.field_off", + + "compute_crossfield": false, + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet_resnet101_pretrained.field_off.train_val.json b/configs/config.mapping_dataset.unet_resnet101_pretrained.field_off.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..687d9c8b038373d7f40db4e9d6529e8344b3d3e5 --- /dev/null +++ b/configs/config.mapping_dataset.unet_resnet101_pretrained.field_off.train_val.json @@ -0,0 +1,19 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet_resnet101_pretrained.field_off.train_val", + "fold": ["train", "val"], + + "compute_crossfield": false, + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "max_epoch": 47 // Highest val was reached at epoch 47 + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet_resnet101_pretrained.json b/configs/config.mapping_dataset.unet_resnet101_pretrained.json new file mode 100644 index 0000000000000000000000000000000000000000..aa227d1b4840d4f5bc3cd6d3b96c59cae9a1949b --- /dev/null +++ b/configs/config.mapping_dataset.unet_resnet101_pretrained.json @@ -0,0 +1,17 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet_resnet101_pretrained", + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet_resnet101_pretrained.seg_framefield_off.train_val.json b/configs/config.mapping_dataset.unet_resnet101_pretrained.seg_framefield_off.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..169f89a0d5f7a4a86e2a73c877e30a1a30539052 --- /dev/null +++ b/configs/config.mapping_dataset.unet_resnet101_pretrained.seg_framefield_off.train_val.json @@ -0,0 +1,24 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet_resnet101_pretrained.seg_framefield_off.train_val", + "fold": ["train", "val"], + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "max_epoch": 47 // Highest val was reached at epoch 47 + }, + + "loss_params": { + "coefs": { + "seg_interior_crossfield": 0.0, + "seg_edge_crossfield": 0.0 + } + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet_resnet101_pretrained.smooth_off.train_val.json b/configs/config.mapping_dataset.unet_resnet101_pretrained.smooth_off.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..efad3f2e73f617ae01041db42c4463072bb5fb8f --- /dev/null +++ b/configs/config.mapping_dataset.unet_resnet101_pretrained.smooth_off.train_val.json @@ -0,0 +1,23 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet_resnet101_pretrained.smooth_off.train_val", + "fold": ["train", "val"], + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "max_epoch": 47 // Highest val was reached at epoch 47 + }, + + "loss_params": { + "coefs": { + "crossfield_smooth": 0.0 + } + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.unet_resnet101_pretrained.train_val.json b/configs/config.mapping_dataset.unet_resnet101_pretrained.train_val.json new file mode 100644 index 0000000000000000000000000000000000000000..07ad3e7ddfecc99db4797a78514ab71e3d4eeb21 --- /dev/null +++ b/configs/config.mapping_dataset.unet_resnet101_pretrained.train_val.json @@ -0,0 +1,19 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.unet_resnet101_pretrained.train_val", + "fold": ["train", "val"], + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 10, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "max_epoch": 47 // Highest val was reached at epoch 47 + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset.zorzi.json b/configs/config.mapping_dataset.zorzi.json new file mode 100644 index 0000000000000000000000000000000000000000..4bb845386455b31efd188fbe87058bb39b663040 --- /dev/null +++ b/configs/config.mapping_dataset.zorzi.json @@ -0,0 +1,5 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset.zorzi" +} \ No newline at end of file diff --git a/configs/config.mapping_dataset_small.deeplab101.json b/configs/config.mapping_dataset_small.deeplab101.json new file mode 100644 index 0000000000000000000000000000000000000000..1c6a5699c900bfff90fc8a04f18aabe6e9c8d0d7 --- /dev/null +++ b/configs/config.mapping_dataset_small.deeplab101.json @@ -0,0 +1,21 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset_small.deeplab101", + + + + + "dataset_params": { + "small": true + }, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.deeplab101.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 8, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "gamma": 0.99 + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset_small.unet16.json b/configs/config.mapping_dataset_small.unet16.json new file mode 100644 index 0000000000000000000000000000000000000000..b57a19c9ba3aa0ae42774ebb56b2fa06e3121abf --- /dev/null +++ b/configs/config.mapping_dataset_small.unet16.json @@ -0,0 +1,26 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset_small.unet16", + + + + + "dataset_params": { + "small": true + }, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + + }, + + "optim_params": { + "batch_size": 32, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "gamma": 0.99 + }, + + "eval_params" : { + "defaults_filepath": "configs/eval_params.mapping_dataset.json" // Path from the project's root to a JSON with default values for eval_params + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset_small.unet_resnet101.json b/configs/config.mapping_dataset_small.unet_resnet101.json new file mode 100644 index 0000000000000000000000000000000000000000..bf893075d74c91f88a4f69350811d55f2d98efcb --- /dev/null +++ b/configs/config.mapping_dataset_small.unet_resnet101.json @@ -0,0 +1,22 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset_small.unet_resnet101", + + + + + "dataset_params": { + "small": true + }, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + + "optim_params": { + "batch_size": 8, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "gamma": 0.99 + } +} \ No newline at end of file diff --git a/configs/config.mapping_dataset_small.unet_resnet101_pretrained.json b/configs/config.mapping_dataset_small.unet_resnet101_pretrained.json new file mode 100644 index 0000000000000000000000000000000000000000..7ea694a18834b9c3ce7d37c8a196f7b0652d251d --- /dev/null +++ b/configs/config.mapping_dataset_small.unet_resnet101_pretrained.json @@ -0,0 +1,31 @@ +{ + "defaults_filepath": "configs/config.defaults.mapping_dataset.json", + + "run_name": "mapping_dataset_small.unet_resnet101_pretrained", + + + + + "dataset_params": { + "small": true + }, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "optim_params": { + "batch_size": 8, // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "gamma": 1.0, + "log_steps": 10 + }, + + "eval_params" : { + "defaults_filepath": "configs/eval_params.mapping_dataset.json" // Path from the project's root to a JSON with default values for eval_params + }, + + "data_aug_params": { + "enable": false + } +} \ No newline at end of file diff --git a/configs/config.segbuildings_dataset.deeplab101.dev.json b/configs/config.segbuildings_dataset.deeplab101.dev.json new file mode 100644 index 0000000000000000000000000000000000000000..76a5c68dcee54708e56a15cfc4f8fc978737ffde --- /dev/null +++ b/configs/config.segbuildings_dataset.deeplab101.dev.json @@ -0,0 +1,89 @@ +{ + "run_name": null, + + "data_dir_candidates": [ + "/media/gaetan/data/drivendata/", + "/data/drivendata/" // try inside docker image + ], + "data_root_partial_dirpath": "segbuildings", + "dataset_params": { + "small": false + }, + "num_workers": 8, + "data_patch_size": 725, // Size of patch saved on disk if data aug is True (allows for rotating patches for the train split) + "input_patch_size": 512, // Size of patch fed to the model + "data_split_params": { + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.75, + "val_fraction": 0.25 // test_fraction is the rest + }, + "data_aug_params": { + "enable": true, + "vflip": true, + "rotate": true, + "color_jitter": true, + "device": "cuda" + }, + + "device": "cuda", // Only has effects when mode is val or test. When mode is train, always use CUDA + "use_amp": true, // Automatic Mixed Precision switch + + "backbone_params": { + "name": "deeplab101", + "input_features": 3, + "features": 256, + "pretrained": false + }, + + "compute_seg": true, + "compute_crossfield": true, + + "seg_params": { + "compute_interior": true, + "compute_edge": true, + "compute_vertex": false + }, + + "loss_params": { + "multiloss": { + "normalization_params": { + "min_samples": 10, // Per GPU + "max_samples": 1000 // Per GPU + }, + "coefs": { + "seg_interior": 1, + "seg_edge": 1, + "seg_vertex": 0, + "crossfield_align": 1, + "crossfield_align90": 0.2, + "crossfield_smooth": 0.2, + "seg_interior_crossfield": 0.2, + "seg_edge_crossfield": 0.2, + "seg_edge_interior": 0.2 + } + }, + "seg_loss_params": { // https://github.com/neptune-ai/open-solution-mapping-challenge/blob/master/neptune.yaml + "bce_coef": 1.0, + "dice_coef": 0.2, + "w0": 50, // From original U-Net paper: distance weight to increase loss between objects + "sigma": 10 // From original U-Net paper: distance weight to increase loss between objects + } + }, + + "batch_size": 1, // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "base_lr": 1e-4, // Will be multiplied by the effective_batch_size=world_size*batch_size. + "max_lr": 1e-2, // Maximum learning rate + "warmup_epochs": 1, // Number of epochs for warmup (learning rate starts at lr*warmup_factor and gradually increases to lr) + "warmup_factor": 1e-3, + "weight_decay": 0, + "dropout_keep_prob": 1.0, // Not used for now + "max_epoch": 10000, + "log_steps": 50, + "checkpoint_epoch": 1, + "checkpoints_to_keep": 10, // outputs + "logs_dirname": "logs", + "save_input_output": false, + "log_input_output": false, + "checkpoints_dirname": "checkpoints", + "eval_dirname": "eval" +} diff --git a/configs/config.segbuildings_dataset.deeplab50.dev.json b/configs/config.segbuildings_dataset.deeplab50.dev.json new file mode 100644 index 0000000000000000000000000000000000000000..63bb934d05ed9c138986ac8cd83b5d4d2d850dfc --- /dev/null +++ b/configs/config.segbuildings_dataset.deeplab50.dev.json @@ -0,0 +1,89 @@ +{ + "run_name": null, + + "data_dir_candidates": [ + "/media/gaetan/data/drivendata/", + "/data/drivendata/" // try inside docker image + ], + "data_root_partial_dirpath": "segbuildings", + "dataset_params": { + "small": false + }, + "num_workers": 8, + "data_patch_size": 725, // Size of patch saved on disk if data aug is True (allows for rotating patches for the train split) + "input_patch_size": 512, // Size of patch fed to the model + "data_split_params": { + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.75, + "val_fraction": 0.25 // test_fraction is the rest + }, + "data_aug_params": { + "enable": true, + "vflip": true, + "rotate": true, + "color_jitter": true, + "device": "cuda" + }, + + "device": "cuda", // Only has effects when mode is val or test. When mode is train, always use CUDA + "use_amp": false, // Automatic Mixed Precision switch + + "backbone_params": { + "name": "deeplab50", + "input_features": 3, + "features": 128, + "pretrained": false + }, + + "compute_seg": true, + "compute_crossfield": true, + + "seg_params": { + "compute_interior": true, + "compute_edge": true, + "compute_vertex": false + }, + + "loss_params": { + "multiloss": { + "normalization_params": { + "min_samples": 10, // Per GPU + "max_samples": 1000 // Per GPU + }, + "coefs": { + "seg_interior": 1, + "seg_edge": 1, + "seg_vertex": 0, + "crossfield_align": 1, + "crossfield_align90": 0.2, + "crossfield_smooth": 0.2, + "seg_interior_crossfield": 0.2, + "seg_edge_crossfield": 0.2, + "seg_edge_interior": 0.2 + } + }, + "seg_loss_params": { // https://github.com/neptune-ai/open-solution-mapping-challenge/blob/master/neptune.yaml + "bce_coef": 1.0, + "dice_coef": 0.2, + "w0": 50, // From original U-Net paper: distance weight to increase loss between objects + "sigma": 10 // From original U-Net paper: distance weight to increase loss between objects + } + }, + + "batch_size": 2, // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "base_lr": 1e-4, // Will be multiplied by the effective_batch_size=world_size*batch_size. + "max_lr": 1e-2, // Maximum learning rate + "warmup_epochs": 1, // Number of epochs for warmup (learning rate starts at lr*warmup_factor and gradually increases to lr) + "warmup_factor": 1e-3, + "weight_decay": 0, + "dropout_keep_prob": 1.0, // Not used for now + "max_epoch": 10000, + "log_steps": 50, + "checkpoint_epoch": 1, + "checkpoints_to_keep": 10, // outputs + "logs_dirname": "logs", + "save_input_output": false, + "log_input_output": false, + "checkpoints_dirname": "checkpoints", + "eval_dirname": "eval" +} diff --git a/configs/config.segbuildings_dataset.dev.json b/configs/config.segbuildings_dataset.dev.json new file mode 100644 index 0000000000000000000000000000000000000000..d74afb9400137486618bc4d491e7fdfe86e011bb --- /dev/null +++ b/configs/config.segbuildings_dataset.dev.json @@ -0,0 +1,87 @@ +{ + "run_name": null, + +"data_dir_candidates": [ + "/media/gaetan/data/drivendata/", + "/data/drivendata/" // try inside docker image + ], + "data_root_partial_dirpath": "segbuildings", + "num_workers": 10, + "data_patch_size": 725, // Size of patch saved on disk if data aug is True (allows for rotating patches for the train split) + "input_patch_size": 512, // Size of patch fed to the model + "data_split_params": { + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.9, + "val_fraction": 0.1 + // test_fraction is the rest + }, + "data_aug_params": { + "enable": true, + "vflip": true, + "rotate": true, + "color_jitter": true, + "device": "cuda" + }, + + "device": "cuda", // Only has effects when mode is val or test. When mode is train, always use CUDA + "use_amp": true, // Automatic Mixed Precision switch + + "backbone_params": { + "name": "unet", + "input_features": 3, + "features": 16 + }, + + "compute_seg": true, + "compute_crossfield": true, + + "seg_params": { + "compute_interior": true, + "compute_edge": true, + "compute_vertex": false + }, + + "loss_params": { + "multiloss": { + "normalization_params": { + "min_samples": 10, // Per GPU + "max_samples": 1000 // Per GPU + }, + "coefs": { + "seg_interior": 1, + "seg_edge": 1, + "seg_vertex": 0, + "crossfield_align": 1, + "crossfield_align90": 0.2, + "crossfield_smooth": 0.2, + "seg_interior_crossfield": 0.2, + "seg_edge_crossfield": 0.2, + "seg_edge_interior": 0.2 + } + }, + "seg_loss_params": { // https://github.com/neptune-ai/open-solution-mapping-challenge/blob/master/neptune.yaml + "bce_coef": 1.0, + "dice_coef": 0.2, + "w0": 50, // From original U-Net paper: distance weight to increase loss between objects + "sigma": 10 // From original U-Net paper: distance weight to increase loss between objects + } + }, + + "batch_size": 8, // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "base_lr": 1e-4, // Will be multiplied by the effective_batch_size=world_size*batch_size. + "max_lr": 1e-2, // Maximum learning rate + "warmup_epochs": 1, // Number of epochs for warmup (learning rate starts at lr*warmup_factor and gradually increases to lr) + "warmup_factor": 1e-3, + "weight_decay": 0, + "dropout_keep_prob": 1.0, // Not used for now + "max_epoch": 10000, + "log_steps": 50, + "checkpoint_epoch": 1, + "checkpoints_to_keep": 10, + // outputs + "logs_dirname": "logs", + "save_input_output": true, + "log_input_output": false, + "checkpoints_dirname": "checkpoints", + "eval_dirname": "eval" +} diff --git a/configs/config.segbuildings_dataset.unet16.json b/configs/config.segbuildings_dataset.unet16.json new file mode 100644 index 0000000000000000000000000000000000000000..3e51369e97963d1a59f808d0dc5f5f4bad3a8f1c --- /dev/null +++ b/configs/config.segbuildings_dataset.unet16.json @@ -0,0 +1,20 @@ +{ + "defaults_filepath": "configs/config.defaults.json", + + "run_name": "segbuildings_dataset.unet16", + + + + + "dataset_params": { + "defaults_filepath": "configs/dataset_params.segbuildings_dataset.json" // Path from the project's root to a JSON with default values for dataset_params + }, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet16.json" // Path from the project's root to a JSON with default values for backbone_params + }, + + "optim_params": { + "batch_size": 8 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + } +} \ No newline at end of file diff --git a/configs/config.segbuildings_dataset.unet_resnet.dev.json b/configs/config.segbuildings_dataset.unet_resnet.dev.json new file mode 100644 index 0000000000000000000000000000000000000000..23be7ad5b9fd6fe5f8bd1383238328160f161ea6 --- /dev/null +++ b/configs/config.segbuildings_dataset.unet_resnet.dev.json @@ -0,0 +1,92 @@ +{ + "run_name": null, + + "data_dir_candidates": [ + "/media/gaetan/data/drivendata/", + "/data/drivendata/" // try inside docker image + ], + "data_root_partial_dirpath": "segbuildings", + "dataset_params": { + "small": false + }, + "num_workers": 6, + "data_patch_size": 725, // Size of patch saved on disk if data aug is True (allows for rotating patches for the train split) + "input_patch_size": 512, // Size of patch fed to the model + "data_split_params": { + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.75, + "val_fraction": 0.25 // test_fraction is the rest + }, + "data_aug_params": { + "enable": true, + "vflip": true, + "rotate": true, + "color_jitter": true, + "device": "cuda" + }, + + "device": "cuda", // Only has effects when mode is val or test. When mode is train, always use CUDA + "use_amp": false, // Automatic Mixed Precision switch + + "backbone_params": { + "name": "unet_resnet", + "encoder_depth": 34, // 34, 101 and 152 are possible + "input_features": 3, + "num_filters": 32, // Default: 32 + "pretrained": true, + "dropout_2d": 0.2, // Default: 0.2 + "is_deconv": false // Default: false + }, + + "compute_seg": true, + "compute_crossfield": true, + + "seg_params": { + "compute_interior": true, + "compute_edge": true, + "compute_vertex": false + }, + + "loss_params": { + "multiloss": { + "normalization_params": { + "min_samples": 10, // Per GPU + "max_samples": 1000 // Per GPU + }, + "coefs": { + "seg_interior": 1, + "seg_edge": 1, + "seg_vertex": 0, + "crossfield_align": 1, + "crossfield_align90": 0.2, + "crossfield_smooth": 0.2, + "seg_interior_crossfield": 0.2, + "seg_edge_crossfield": 0.2, + "seg_edge_interior": 0.2 + } + }, + "seg_loss_params": { // https://github.com/neptune-ai/open-solution-mapping-challenge/blob/master/neptune.yaml + "bce_coef": 1.0, + "dice_coef": 0.2, + "w0": 50, // From original U-Net paper: distance weight to increase loss between objects + "sigma": 10 // From original U-Net paper: distance weight to increase loss between objects + } + }, + + "batch_size": 2, // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "base_lr": 1e-4, // Will be multiplied by the effective_batch_size=world_size*batch_size. + "max_lr": 1e-2, // Maximum learning rate + "warmup_epochs": 0, // Number of epochs for warmup (learning rate starts at lr*warmup_factor and gradually increases to lr) + "warmup_factor": 1e-3, + "weight_decay": 0, + "dropout_keep_prob": 1.0, // Not used for now + "max_epoch": 10000, + "log_steps": 50, + "checkpoint_epoch": 1, + "checkpoints_to_keep": 10, // outputs + "logs_dirname": "logs", + "save_input_output": false, + "log_input_output": false, + "checkpoints_dirname": "checkpoints", + "eval_dirname": "eval" +} diff --git a/configs/config.xview2_dataset.unet_resnet101_pretrained.json b/configs/config.xview2_dataset.unet_resnet101_pretrained.json new file mode 100644 index 0000000000000000000000000000000000000000..6992ff50110c362d49defe7b08c2d68a5db63a5b --- /dev/null +++ b/configs/config.xview2_dataset.unet_resnet101_pretrained.json @@ -0,0 +1,30 @@ +{ + "defaults_filepath": "configs/config.defaults.xview2_dataset.json", + + "run_name": "xview2_dataset.unet_resnet101_pretrained", + + + + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "loss_params": { + "seg_loss_params": { + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, // Dist weights as in the original U-Net paper + "use_size": false // Size weights increasing importance of smaller buildings + } + }, + + "optim_params": { + "batch_size": 4 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + }, + + "data_aug_params": { + "enable": false + } +} diff --git a/configs/config.xview2_dataset_small.unet_resnet101_pretrained.json b/configs/config.xview2_dataset_small.unet_resnet101_pretrained.json new file mode 100644 index 0000000000000000000000000000000000000000..de7be98fcaef99d503d797c75ffc07dbfea62b62 --- /dev/null +++ b/configs/config.xview2_dataset_small.unet_resnet101_pretrained.json @@ -0,0 +1,31 @@ +{ + "defaults_filepath": "configs/config.defaults.xview2_dataset.json", + + "run_name": "xview2_dataset_small.unet_resnet101_pretrained", + + "dataset_params": { + "small": true + }, + + "backbone_params": { + "defaults_filepath": "configs/backbone_params.unet_resnet101.json", // Path from the project's root to a JSON with default values for backbone_params + "pretrained": true + }, + + "loss_params": { + "seg_loss_params": { + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, // Dist weights as in the original U-Net paper + "use_size": false // Size weights increasing importance of smaller buildings + } + }, + + "optim_params": { + "batch_size": 4 // Overwrite default batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + }, + + "data_aug_params": { + "enable": false + } +} diff --git a/configs/dataset_params.inria_dataset.json b/configs/dataset_params.inria_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..2288273e1637b46ea4fdd59dbc1a11203c7d5769 --- /dev/null +++ b/configs/dataset_params.inria_dataset.json @@ -0,0 +1,12 @@ +{ + "root_dirname": "AerialImageDataset", + "gt_source": "disk", // Can be "disk" or "osm". "osm" will download annotations from OSM + "gt_type": "polygon", // If gt_source=="disk", reads polygon data from gt_dirname + "gt_dirname": "aligned_gt_polygons_2", // If null, will be replaced by default dirname depending on gt_type + "mask_only": false, // This is the default value to not use this option + "small": false, + "data_patch_size": 725, + "input_patch_size": 512, + + "train_fraction": 0.75 // Remaining is validation +} \ No newline at end of file diff --git a/configs/dataset_params.inria_dataset_ori_gt.json b/configs/dataset_params.inria_dataset_ori_gt.json new file mode 100644 index 0000000000000000000000000000000000000000..f309e536c95a013b10bd2a884505bec58f04d22b --- /dev/null +++ b/configs/dataset_params.inria_dataset_ori_gt.json @@ -0,0 +1,13 @@ +{ + "root_dirname": "AerialImageDataset", + "pre_process": false, + "gt_source": "disk", + "gt_type": "tif", + "gt_dirname": "gt", // gt_polygonized are annotations that are the polygonized version of the original gt masks on inria dataset + "mask_only": false, // This is the default value to not use this option + "small": false, + "data_patch_size": 425, + "input_patch_size": 300, + + "train_fraction": 0.75 // Remaining is validation +} \ No newline at end of file diff --git a/configs/dataset_params.inria_dataset_osm_aligned.json b/configs/dataset_params.inria_dataset_osm_aligned.json new file mode 100644 index 0000000000000000000000000000000000000000..dbd281b4cad892a9965f048b5b8c000114187746 --- /dev/null +++ b/configs/dataset_params.inria_dataset_osm_aligned.json @@ -0,0 +1,13 @@ +{ + "root_dirname": "AerialImageDataset", + "pre_process": true, + "gt_source": "disk", + "gt_type": "npy", + "gt_dirname": "aligned_gt_polygons_2", // aligned_gt_polygons_2 are annotations that are the aligned version of OSM stored in gt_polygons + "mask_only": false, // This is the default value to not use this option + "small": false, + "data_patch_size": 725, + "input_patch_size": 512, + + "train_fraction": 0.75 // Remaining is validation +} \ No newline at end of file diff --git a/configs/dataset_params.inria_dataset_osm_mask_only.json b/configs/dataset_params.inria_dataset_osm_mask_only.json new file mode 100644 index 0000000000000000000000000000000000000000..20e99f44dcaf8ff050ca82248fec6f069002e963 --- /dev/null +++ b/configs/dataset_params.inria_dataset_osm_mask_only.json @@ -0,0 +1,13 @@ +{ + "root_dirname": "AerialImageDataset", + "pre_process": true, + "gt_source": "disk", + "gt_type": "npy", + "gt_dirname": "gt_polygons", // gt_polygons are OSM annotations that were pre-downloaded and projected to the image's coordinate system + "mask_only": true, // Discard the RGB image, the input image to the network is a single-channel binary mask of the polygons. + "small": false, + "data_patch_size": 725, + "input_patch_size": 512, + + "train_fraction": 0.75 // Remaining is validation +} \ No newline at end of file diff --git a/configs/dataset_params.inria_dataset_polygonized.json b/configs/dataset_params.inria_dataset_polygonized.json new file mode 100644 index 0000000000000000000000000000000000000000..7eb0030ce48b4c03dd08dd408ca6ea6032e50cdf --- /dev/null +++ b/configs/dataset_params.inria_dataset_polygonized.json @@ -0,0 +1,13 @@ +{ + "root_dirname": "AerialImageDataset", + "pre_process": true, + "gt_source": "disk", + "gt_type": "geojson", + "gt_dirname": "gt_polygonized", // gt_polygonized are annotations that are the polygonized version of the original gt masks on inria dataset + "mask_only": false, // This is the default value to not use this option + "small": false, + "data_patch_size": 725, + "input_patch_size": 512, + + "train_fraction": 0.75 // Remaining is validation +} \ No newline at end of file diff --git a/configs/dataset_params.inria_dataset_polygonized_256.json b/configs/dataset_params.inria_dataset_polygonized_256.json new file mode 100644 index 0000000000000000000000000000000000000000..02a6702f7ae0e947b979037be2775c85613478d9 --- /dev/null +++ b/configs/dataset_params.inria_dataset_polygonized_256.json @@ -0,0 +1,13 @@ +{ + "root_dirname": "AerialImageDataset", + "pre_process": true, + "gt_source": "disk", + "gt_type": "geojson", + "gt_dirname": "gt_polygonized", // gt_polygonized are annotations that are the polygonized version of the original gt masks on inria dataset + "mask_only": false, // This is the default value to not use this option + "small": false, + "data_patch_size": 384, + "input_patch_size": 256, + + "train_fraction": 0.75 // Remaining is validation +} \ No newline at end of file diff --git a/configs/dataset_params.luxcarta_dataset.json b/configs/dataset_params.luxcarta_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..66cba6c53e94199627ab0f89757fc9a3f5a0c249 --- /dev/null +++ b/configs/dataset_params.luxcarta_dataset.json @@ -0,0 +1,8 @@ +{ + "root_dirname": "luxcarta_precise_buildings", + "data_patch_size": 725, + "input_patch_size": 512, + + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.9 // Remaining is validation +} \ No newline at end of file diff --git a/configs/dataset_params.mapping_dataset.json b/configs/dataset_params.mapping_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..7b0b0de9af007de7e9e3eef604565bc5308d2061 --- /dev/null +++ b/configs/dataset_params.mapping_dataset.json @@ -0,0 +1,7 @@ +{ + "root_dirname": "mapping_challenge_dataset", + "small": false, + + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.75 // Remaining is validation +} \ No newline at end of file diff --git a/configs/dataset_params.xview2_dataset.json b/configs/dataset_params.xview2_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..c2181b66de71223fff54e2b2fff1f4b908b9408b --- /dev/null +++ b/configs/dataset_params.xview2_dataset.json @@ -0,0 +1,14 @@ +{ + "root_dirname": "xview2_xbd_dataset", + "pre_process": true, + "gt_source": "disk", + "gt_type": "geojson", + "gt_dirname": "gt_polygonized", // gt_polygonized are annotations that are the polygonized version of the original gt masks on inria dataset + "mask_only": false, // This is the default value to not use this option + "small": false, + "data_patch_size": 725, + "input_patch_size": 512, + + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.75 // Remaining is validation +} \ No newline at end of file diff --git a/configs/eval_params.defaults.json b/configs/eval_params.defaults.json new file mode 100644 index 0000000000000000000000000000000000000000..babd892d4a56a86bc8825d6c04138a0be8066ef5 --- /dev/null +++ b/configs/eval_params.defaults.json @@ -0,0 +1,44 @@ +{ + "results_dirname": "eval_runs", + + "test_time_augmentation": false, + + "batch_size_mult": 64, // Inference has a certain batch_size, post-process (polygonization included) will use a batch of size batch_size*batch_size_mult + // ASM: + // eval_batch_size=32, batch_size_mult=1: finished 96 batches in 290s: 0,094401042s/sample + // eval_batch_size=32, batch_size_mult=16: finished 96 batches in 167s: 0,054361979s/sample + // eval_batch_size=32, batch_size_mult=8: finished 96 batches in 172s: 0,055989583s/sample + // eval_batch_size=12, batch_size_mult=96: finished 192 batches in 125s: 0,098958333s/sample + // eval_batch_size=8, batch_size_mult=128: finished 256 batches in 114s: 0,055664063s/sample + // eval_batch_size=16, batch_size_mult=64: finished 128 batches in 107s: 0,052246094s/sample + // (no saving of output) eval_batch_size=16, batch_size_mult=64: finished 128 batches in 89s: 0,043457031s/sample + + // ACM: + // (no saving of output) eval_batch_size=16, batch_size_mult=64: finished 128 batches in 85s: 0,041503906s/sample + // With TTA (x8) eval_batch_size=16, batch_size_mult=64: finished 128 batches in 366s: 0,178710938s/sample + + "patch_size": null, + "patch_overlap": 200, + + "seg_threshold": 0.5, + + "save_individual_outputs": { + "image": false, + "seg_gt": false, + "seg": false, + "seg_mask": false, + "seg_opencities_mask": false, + "seg_luxcarta": false, + "crossfield": false, + "uv_angles": false, + "poly_shapefile": false, + "poly_geojson": false, + "poly_viz": false + }, + + "save_aggregated_outputs": { + "stats": false, + "seg_coco": false, + "poly_coco": false + } +} diff --git a/configs/eval_params.inria_dataset.json b/configs/eval_params.inria_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..cb2323c104dba6553532210b469446621d58f437 --- /dev/null +++ b/configs/eval_params.inria_dataset.json @@ -0,0 +1,20 @@ +{ + "defaults_filepath": "configs/eval_params.defaults.json", + + "test_time_augmentation": true, + + "batch_size_mult": 1, + + "patch_size": 1024, + + "save_individual_outputs": { + "seg": true, + "seg_mask": true, + "poly_shapefile": true, + "poly_viz": false + }, + + "save_aggregated_outputs": { + "stats": true + } +} \ No newline at end of file diff --git a/configs/eval_params.luxcarta_dataset.json b/configs/eval_params.luxcarta_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..b17744046e2609bc581725e68937c66cbdecefb2 --- /dev/null +++ b/configs/eval_params.luxcarta_dataset.json @@ -0,0 +1,15 @@ +{ + "defaults_filepath": "configs/eval_params.defaults.json", + + "patch_size": 1024, + + "save_individual_outputs": { + "seg": true, + "poly_shapefile": true, + "poly_viz": true + }, + + "save_aggregated_outputs": { + "stats": true + } +} \ No newline at end of file diff --git a/configs/eval_params.mapping_dataset.json b/configs/eval_params.mapping_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..165015c18e50b9bf28f17e0b730ec474be79a369 --- /dev/null +++ b/configs/eval_params.mapping_dataset.json @@ -0,0 +1,14 @@ +{ + "defaults_filepath": "configs/eval_params.defaults.json", + + "save_individual_outputs": { + "seg": false, + "poly_viz": false + }, + + "save_aggregated_outputs": { + "stats": false, + "seg_coco": true, + "poly_coco": true + } +} \ No newline at end of file diff --git a/configs/eval_params.segbuildings_dataset.json b/configs/eval_params.segbuildings_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..154038641d4257a0feeffcd86ccf188d94de9c4c --- /dev/null +++ b/configs/eval_params.segbuildings_dataset.json @@ -0,0 +1,23 @@ +{ + "patch_size": null, + "patch_overlap": 200, + + "seg_threshold": 0.5, + + "save_outputs": { + "image": false, + "stats": true, + "seg_gt": false, + "seg": false, + "seg_mask": false, + "seg_opencities_mask": true, + "seg_luxcarta": false, + "seg_coco": false, + "crossfield": false, + "uv_angles": false, + "poly_shapefile": false, + "poly_geojson": false, + "poly_viz": false, + "poly_coco": false + } +} diff --git a/configs/eval_params.xview2_dataset.json b/configs/eval_params.xview2_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..451bffd7c248888e62a5627c91d649221ac099f6 --- /dev/null +++ b/configs/eval_params.xview2_dataset.json @@ -0,0 +1,20 @@ +{ + "defaults_filepath": "configs/eval_params.defaults.json", + + "test_time_augmentation": true, + + "batch_size_mult": 1, + + "patch_size": 1024, + + "save_individual_outputs": { + "seg": true, + "seg_mask": false, + "poly_shapefile": false, + "poly_viz": false + }, + + "save_aggregated_outputs": { + "stats": true + } +} \ No newline at end of file diff --git a/configs/loss_params.json b/configs/loss_params.json new file mode 100644 index 0000000000000000000000000000000000000000..804c51941ef618977ae5dcbdb4b168ac77a9df91 --- /dev/null +++ b/configs/loss_params.json @@ -0,0 +1,26 @@ +{ + "multiloss": { + "normalization_params": { + "min_samples": 10, // Per GPU + "max_samples": 1000 // Per GPU + }, + "coefs": { + "epoch_thresholds": [0, 5, 10], // From 0 to 5: gradually go from coefs[0] to coefs[1] for list coefs + "seg": 10, + "crossfield_align": 1, + "crossfield_align90": 0.2, + "crossfield_smooth": 0.005, + "seg_interior_crossfield": [0, 0, 0.2], + "seg_edge_crossfield": [0, 0, 0.2], + "seg_edge_interior": [0, 0, 0.2] + } + }, + "seg_loss_params": { // https://github.com/neptune-ai/open-solution-mapping-challenge/blob/master/neptune.yaml + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, // Dist weights as in the original U-Net paper + "use_size": true, // Size weights increasing importance of smaller buildings + "w0": 50, // From original U-Net paper: distance weight to increase loss between objects + "sigma": 10 // From original U-Net paper: distance weight to increase loss between objects + } +} \ No newline at end of file diff --git a/configs/optim_params.json b/configs/optim_params.json new file mode 100644 index 0000000000000000000000000000000000000000..e93bd03dab445178d7c9d7dc5f308151ddd8bc6f --- /dev/null +++ b/configs/optim_params.json @@ -0,0 +1,15 @@ +{ + "optimizer": "Adam", + "batch_size": 8, // Batch size per GPU. The effective batch size is effective_batch_size=world_size*batch_size + "base_lr": 1e-3, // Will be multiplied by the effective_batch_size=world_size*batch_size. + "max_lr": 1e-1, // Maximum resulting learning rate + "gamma": 0.95, // Gamma of exponential learning rate scheduler + "weight_decay": 0, // Not used + "dropout_keep_prob": 1.0, // Not used + "max_epoch": 1000, + "log_steps": 200, + "checkpoint_epoch": 1, + "checkpoints_to_keep": 5, // outputs + "logs_dirname": "logs", + "checkpoints_dirname": "checkpoints" +} \ No newline at end of file diff --git a/configs/polygonize_params.json b/configs/polygonize_params.json new file mode 100644 index 0000000000000000000000000000000000000000..d6de1dfacacb9fa0495f9e0af0eda10706391ec6 --- /dev/null +++ b/configs/polygonize_params.json @@ -0,0 +1,58 @@ +{ + + "method": ["simple", "acm"], // Possible options: simple, acm, asm, tracing or a list: ["simple", "asm", "asm"] to apply several + + "common_params": { + "init_data_level": 0.5 // Seg level to init contours for all methods + }, + + "simple_method": { + "data_level": 0.5, // Seg level to init contours + "tolerance": [0.125, 1], // Can be a list of values to apply + "seg_threshold": 0.5, + "min_area": 10 + }, + + "asm_method": { + "init_method": "skeleton", // Can be either "skeleton" or "marching_squares" + "data_level": 0.5, + "loss_params": { + "coefs": { + "step_thresholds": [0, 100, 200, 300], // From 0 to 500: gradually go from coefs[0] to coefs[1] + "data": [1.0, 0.1, 0.0, 0.0], + "crossfield": [0.0, 0.05, 0.0, 0.0], + "length": [0.1, 0.01, 0.0, 0.0], + "curvature": [0.0, 0.0, 1.0, 0.0], + "corner": [0.0, 0.0, 0.5, 0.0], + "junction": [0.0, 0.0, 0.5, 0.0] + }, + "curvature_dissimilarity_threshold": 2, // In pixels: for each sub-paths, if the dissimilarity (in the same sense as in the Ramer-Douglas-Peucker alg) is lower than straightness_threshold, then optimize the curve angles to be zero. + "corner_angles": [45, 90, 135], // In degrees: target angles for corners. + "corner_angle_threshold": 22.5, // If a corner angle is less than this threshold away from any angle in corner_angles, optimize it. + "junction_angles": [0, 45, 90, 135], // In degrees: target angles for junction corners. + "junction_angle_weights": [1, 0.01, 0.1, 0.01], // Order of decreassing importance: straight, right-angle, then 45° junction corners. + "junction_angle_threshold": 22.5 // If a junction corner angle is less than this threshold away from any angle in junction_angles, optimize it. + }, + "lr": 0.1, + "gamma": 0.995, + "device": "cuda", + "tolerance": [0.125, 1], + "seg_threshold": 0.5, + "min_area": 10 +}, + + "acm_method": { + "steps": 500, + "data_level": 0.5, // Seg level to optimize contours (better set it equal to common_params.init_data_level + "data_coef": 0.1, + "length_coef": 0.4, + "crossfield_coef": 0.5, + "poly_lr": 0.01, + "warmup_iters": 100, + "warmup_factor": 0.1, + "device": "cuda", + "tolerance": [0.125, 1], // Can be a list of values to apply + "seg_threshold": 0.5, // Remove polygons below that threshold + "min_area": 10 + } +} diff --git a/dataset_folds.py b/dataset_folds.py new file mode 100644 index 0000000000000000000000000000000000000000..23847459d1062177fe613c5d980d824bd58e77e6 --- /dev/null +++ b/dataset_folds.py @@ -0,0 +1,234 @@ +import functools + +import torch +import torch.utils.data + +from frame_field_learning import data_transforms +from lydorn_utils import print_utils + + +def inria_aerial_train_tile_filter(tile, train_val_split_point): + return tile["number"] <= train_val_split_point + + +def inria_aerial_val_tile_filter(tile, train_val_split_point): + return train_val_split_point < tile["number"] + + +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_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 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 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 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_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() diff --git a/eval_coco.py b/eval_coco.py new file mode 100644 index 0000000000000000000000000000000000000000..c1be9ae2371360c1c5300fb3adb282f3b395b295 --- /dev/null +++ b/eval_coco.py @@ -0,0 +1,712 @@ +import os +import fnmatch + +import shapely.geometry +from tqdm import tqdm +from multiprocess import Pool +import json + +# COCO: +from pycocotools.coco import COCO +from pycocotools import mask as maskUtils +from pycocotools.cocoeval import Params +import datetime +import time +from collections import defaultdict +import copy +from functools import partial +import numpy as np + +from lydorn_utils import python_utils, run_utils +from lydorn_utils import print_utils +from lydorn_utils import polygon_utils + + +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 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 compute_contour_metrics(gts_dts): + gts, dts = gts_dts + gt_polygons = [shapely.geometry.Polygon(np.array(coords).reshape(-1, 2)) for ann in gts + for coords in ann["segmentation"]] + dt_polygons = [shapely.geometry.Polygon(np.array(coords).reshape(-1, 2)) for ann in dts + for coords in ann["segmentation"]] + fixed_gt_polygons = polygon_utils.fix_polygons(gt_polygons, buffer=0.0001) # Buffer adds vertices but is needed to repair some geometries + fixed_dt_polygons = polygon_utils.fix_polygons(dt_polygons) + # cosine_similarities, edge_distances = \ + # polygon_utils.compute_polygon_contour_measures(dt_polygons, gt_polygons, sampling_spacing=2.0, min_precision=0.5, + # max_stretch=2) + max_angle_diffs = polygon_utils.compute_polygon_contour_measures(fixed_dt_polygons, fixed_gt_polygons, sampling_spacing=2.0, min_precision=0.5, max_stretch=2) + + return max_angle_diffs + + +class ContourEval: + def __init__(self, coco_gt, coco_dt): + """ + + @param coco_gt: coco object with ground truth annotations + @param coco_dt: coco object with detection results + """ + self.coco_gt = coco_gt # ground truth COCO API + self.coco_dt = coco_dt # detections COCO API + + self.img_ids = sorted(coco_gt.getImgIds()) + self.cat_ids = sorted(coco_dt.getCatIds()) + + def evaluate(self, pool=None): + gts = self.coco_gt.loadAnns(self.coco_gt.getAnnIds(imgIds=self.img_ids)) + dts = self.coco_dt.loadAnns(self.coco_dt.getAnnIds(imgIds=self.img_ids)) + + _gts = defaultdict(list) # gt for evaluation + _dts = defaultdict(list) # dt for evaluation + for gt in gts: + _gts[gt['image_id'], gt['category_id']].append(gt) + for dt in dts: + _dts[dt['image_id'], dt['category_id']].append(dt) + evalImgs = defaultdict(list) # per-image per-category evaluation results + + # Compute metric + args_list = [] + # i = 1000 + for img_id in self.img_ids: + for cat_id in self.cat_ids: + gts = _gts[img_id, cat_id] + dts = _dts[img_id, cat_id] + args_list.append((gts, dts)) + # i -= 1 + # if i <= 0: + # break + + if pool is None: + measures_list = [] + for args in tqdm(args_list, desc="Contour metrics"): + measures_list.append(compute_contour_metrics(args)) + else: + measures_list = list(tqdm(pool.imap(compute_contour_metrics, args_list), desc="Contour metrics", total=len(args_list))) + measures_list = [measure for measures in measures_list for measure in measures] # Flatten list + # half_tangent_cosine_similarities_list, edge_distances_list = zip(*measures_list) + # half_tangent_cosine_similarities_list = [item for item in half_tangent_cosine_similarities_list if item is not None] + measures_list = [value for value in measures_list if value is not None] + max_angle_diffs = np.array(measures_list) + max_angle_diffs = max_angle_diffs * 180 / np.pi # Convert to degrees + + return max_angle_diffs + + +class COCOeval: + # Interface for evaluating detection on the Microsoft COCO dataset. + # + # The usage for CocoEval is as follows: + # cocoGt=..., cocoDt=... # load dataset and results + # E = CocoEval(cocoGt,cocoDt); # initialize CocoEval object + # E.params.recThrs = ...; # set parameters as desired + # E.evaluate(); # run per image evaluation + # E.accumulate(); # accumulate per image results + # E.summarize(); # display summary metrics of results + # For example usage see evalDemo.m and http://mscoco.org/. + # + # The evaluation parameters are as follows (defaults in brackets): + # imgIds - [all] N img ids to use for evaluation + # catIds - [all] K cat ids to use for evaluation + # iouThrs - [.5:.05:.95] T=10 IoU thresholds for evaluation + # recThrs - [0:.01:1] R=101 recall thresholds for evaluation + # areaRng - [...] A=4 object area ranges for evaluation + # maxDets - [1 10 100] M=3 thresholds on max detections per image + # iouType - ['segm'] set iouType to 'segm', 'bbox' or 'keypoints' + # iouType replaced the now DEPRECATED useSegm parameter. + # useCats - [1] if true use category labels for evaluation + # Note: if useCats=0 category labels are ignored as in proposal scoring. + # Note: multiple areaRngs [Ax2] and maxDets [Mx1] can be specified. + # + # evaluate(): evaluates detections on every image and every category and + # concats the results into the "evalImgs" with fields: + # dtIds - [1xD] id for each of the D detections (dt) + # gtIds - [1xG] id for each of the G ground truths (gt) + # dtMatches - [TxD] matching gt id at each IoU or 0 + # gtMatches - [TxG] matching dt id at each IoU or 0 + # dtScores - [1xD] confidence of each dt + # gtIgnore - [1xG] ignore flag for each gt + # dtIgnore - [TxD] ignore flag for each dt at each IoU + # + # accumulate(): accumulates the per-image, per-category evaluation + # results in "evalImgs" into the dictionary "eval" with fields: + # params - parameters used for evaluation + # date - date evaluation was performed + # counts - [T,R,K,A,M] parameter dimensions (see above) + # precision - [TxRxKxAxM] precision for every evaluation setting + # recall - [TxKxAxM] max recall for every evaluation setting + # Note: precision and recall==-1 for settings with no gt objects. + # + # See also coco, mask, pycocoDemo, pycocoEvalDemo + # + # Microsoft COCO Toolbox. version 2.0 + # Data, paper, and tutorials available at: http://mscoco.org/ + # Code written by Piotr Dollar and Tsung-Yi Lin, 2015. + # Licensed under the Simplified BSD License [see coco/license.txt] + def __init__(self, cocoGt=None, cocoDt=None, iouType='segm'): + ''' + Initialize CocoEval using coco APIs for gt and dt + :param cocoGt: coco object with ground truth annotations + :param cocoDt: coco object with detection results + :return: None + ''' + if not iouType: + print('iouType not specified. use default iouType segm') + self.cocoGt = cocoGt # ground truth COCO API + self.cocoDt = cocoDt # detections COCO API + self.params = {} # evaluation parameters + self.evalImgs = defaultdict(list) # per-image per-category evaluation results [KxAxI] elements + self.eval = {} # accumulated evaluation results + self._gts = defaultdict(list) # gt for evaluation + self._dts = defaultdict(list) # dt for evaluation + self.params = Params(iouType=iouType) # parameters + self._paramsEval = {} # parameters for evaluation + self.stats = [] # result summarization + self.ious = {} # ious between all gts and dts + if cocoGt is not None: + self.params.imgIds = sorted(cocoGt.getImgIds()) + self.params.catIds = sorted(cocoGt.getCatIds()) + + def _prepare(self): + ''' + Prepare ._gts and ._dts for evaluation based on params + :return: None + ''' + + def _toMask(anns, coco): + # modify ann['segmentation'] by reference + for ann in anns: + rle = coco.annToRLE(ann) + ann['rle'] = rle + + p = self.params + if p.useCats: + gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)) + dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)) + else: + gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds)) + dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds)) + + # convert ground truth to mask if iouType == 'segm' + if p.iouType == 'segm': + _toMask(gts, self.cocoGt) + _toMask(dts, self.cocoDt) + # set ignore flag + for gt in gts: + gt['ignore'] = gt['ignore'] if 'ignore' in gt else 0 + gt['ignore'] = 'iscrowd' in gt and gt['iscrowd'] + if p.iouType == 'keypoints': + gt['ignore'] = (gt['num_keypoints'] == 0) or gt['ignore'] + self._gts = defaultdict(list) # gt for evaluation + self._dts = defaultdict(list) # dt for evaluation + for gt in gts: + self._gts[gt['image_id'], gt['category_id']].append(gt) + for dt in dts: + self._dts[dt['image_id'], dt['category_id']].append(dt) + self.evalImgs = defaultdict(list) # per-image per-category evaluation results + self.eval = {} # accumulated evaluation results + + def evaluate(self): + ''' + Run per image evaluation on given images and store results (a list of dict) in self.evalImgs + :return: None + ''' + tic = time.time() + print('Running per image evaluation...') + p = self.params + # add backward compatibility if useSegm is specified in params + if p.useSegm is not None: + p.iouType = 'segm' if p.useSegm == 1 else 'bbox' + print('useSegm (deprecated) is not None. Running {} evaluation'.format(p.iouType)) + print('Evaluate annotation type *{}*'.format(p.iouType)) + p.imgIds = list(np.unique(p.imgIds)) + if p.useCats: + p.catIds = list(np.unique(p.catIds)) + p.maxDets = sorted(p.maxDets) + self.params = p + + self._prepare() + # loop through images, area range, max detection number + catIds = p.catIds if p.useCats else [-1] + + if p.iouType == 'segm' or p.iouType == 'bbox': + computeIoU = self.computeIoU + elif p.iouType == 'keypoints': + computeIoU = self.computeOks + self.ious = {(imgId, catId): computeIoU(imgId, catId) \ + for imgId in p.imgIds + for catId in catIds} + + evaluateImg = self.evaluateImg + maxDet = p.maxDets[-1] + self.evalImgs = [evaluateImg(imgId, catId, areaRng, maxDet) + for catId in catIds + for areaRng in p.areaRng + for imgId in p.imgIds + ] + self._paramsEval = copy.deepcopy(self.params) + toc = time.time() + print('DONE (t={:0.2f}s).'.format(toc - tic)) + + def computeIoU(self, imgId, catId): + p = self.params + if p.useCats: + gt = self._gts[imgId, catId] + dt = self._dts[imgId, catId] + else: + gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]] + dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]] + if len(gt) == 0 and len(dt) == 0: + return [] + inds = np.argsort([-d['score'] for d in dt], kind='mergesort') + dt = [dt[i] for i in inds] + if len(dt) > p.maxDets[-1]: + dt = dt[0:p.maxDets[-1]] + + if p.iouType == 'segm': + g = [g['rle'] for g in gt] + d = [d['rle'] for d in dt] + elif p.iouType == 'bbox': + g = [g['bbox'] for g in gt] + d = [d['bbox'] for d in dt] + else: + raise Exception('unknown iouType for iou computation') + + # compute iou between each dt and gt region + iscrowd = [int(o['iscrowd']) for o in gt] + ious = maskUtils.iou(d, g, iscrowd) + return ious + + def computeOks(self, imgId, catId): + p = self.params + # dimention here should be Nxm + gts = self._gts[imgId, catId] + dts = self._dts[imgId, catId] + inds = np.argsort([-d['score'] for d in dts], kind='mergesort') + dts = [dts[i] for i in inds] + if len(dts) > p.maxDets[-1]: + dts = dts[0:p.maxDets[-1]] + # if len(gts) == 0 and len(dts) == 0: + if len(gts) == 0 or len(dts) == 0: + return [] + ious = np.zeros((len(dts), len(gts))) + sigmas = np.array( + [.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89]) / 10.0 + vars = (sigmas * 2) ** 2 + k = len(sigmas) + # compute oks between each detection and ground truth object + for j, gt in enumerate(gts): + # create bounds for ignore regions(double the gt bbox) + g = np.array(gt['keypoints']) + xg = g[0::3]; + yg = g[1::3]; + vg = g[2::3] + k1 = np.count_nonzero(vg > 0) + bb = gt['bbox'] + x0 = bb[0] - bb[2]; + x1 = bb[0] + bb[2] * 2 + y0 = bb[1] - bb[3]; + y1 = bb[1] + bb[3] * 2 + for i, dt in enumerate(dts): + d = np.array(dt['keypoints']) + xd = d[0::3]; + yd = d[1::3] + if k1 > 0: + # measure the per-keypoint distance if keypoints visible + dx = xd - xg + dy = yd - yg + else: + # measure minimum distance to keypoints in (x0,y0) & (x1,y1) + z = np.zeros((k)) + dx = np.max((z, x0 - xd), axis=0) + np.max((z, xd - x1), axis=0) + dy = np.max((z, y0 - yd), axis=0) + np.max((z, yd - y1), axis=0) + e = (dx ** 2 + dy ** 2) / vars / (gt['area'] + np.spacing(1)) / 2 + if k1 > 0: + e = e[vg > 0] + ious[i, j] = np.sum(np.exp(-e)) / e.shape[0] + return ious + + def evaluateImg(self, imgId, catId, aRng, maxDet): + ''' + perform evaluation for single category and image + :return: dict (single image results) + ''' + p = self.params + if p.useCats: + gt = self._gts[imgId, catId] + dt = self._dts[imgId, catId] + else: + gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]] + dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]] + if len(gt) == 0 and len(dt) == 0: + return None + + for g in gt: + if g['ignore'] or (g['area'] < aRng[0] or g['area'] > aRng[1]): + g['_ignore'] = 1 + else: + g['_ignore'] = 0 + + # sort dt highest score first, sort gt ignore last + gtind = np.argsort([g['_ignore'] for g in gt], kind='mergesort') + gt = [gt[i] for i in gtind] + dtind = np.argsort([-d['score'] for d in dt], kind='mergesort') + dt = [dt[i] for i in dtind[0:maxDet]] + iscrowd = [int(o['iscrowd']) for o in gt] + # load computed ious + ious = self.ious[imgId, catId][:, gtind] if len(self.ious[imgId, catId]) > 0 else self.ious[imgId, catId] + + T = len(p.iouThrs) + G = len(gt) + D = len(dt) + gtm = np.zeros((T, G)) + dtm = np.zeros((T, D)) + gtIg = np.array([g['_ignore'] for g in gt]) + dtIg = np.zeros((T, D)) + if len(ious): + for tind, t in enumerate(p.iouThrs): + for dind, d in enumerate(dt): + # information about best match so far (m=-1 -> unmatched) + iou = min([t, 1 - 1e-10]) + m = -1 + for gind, g in enumerate(gt): + # if this gt already matched, and not a crowd, continue + if gtm[tind, gind] > 0 and not iscrowd[gind]: + continue + # if dt matched to reg gt, and on ignore gt, stop + if m > -1 and gtIg[m] == 0 and gtIg[gind] == 1: + break + # continue to next gt unless better match made + if ious[dind, gind] < iou: + continue + # if match successful and best so far, store appropriately + iou = ious[dind, gind] + m = gind + # if match made store id of match for both dt and gt + if m == -1: + continue + dtIg[tind, dind] = gtIg[m] + dtm[tind, dind] = gt[m]['id'] + gtm[tind, m] = d['id'] + # set unmatched detections outside of area range to ignore + a = np.array([d['area'] < aRng[0] or d['area'] > aRng[1] for d in dt]).reshape((1, len(dt))) + dtIg = np.logical_or(dtIg, np.logical_and(dtm == 0, np.repeat(a, T, 0))) + # store results for given image and category + return { + 'image_id': imgId, + 'category_id': catId, + 'aRng': aRng, + 'maxDet': maxDet, + 'dtIds': [d['id'] for d in dt], + 'gtIds': [g['id'] for g in gt], + 'dtMatches': dtm, + 'gtMatches': gtm, + 'dtScores': [d['score'] for d in dt], + 'gtIgnore': gtIg, + 'dtIgnore': dtIg, + } + + def accumulate(self, p=None): + ''' + Accumulate per image evaluation results and store the result in self.eval + :param p: input params for evaluation + :return: None + ''' + print('Accumulating evaluation results...') + tic = time.time() + if not self.evalImgs: + print('Please run evaluate() first') + # allows input customized parameters + if p is None: + p = self.params + p.catIds = p.catIds if p.useCats == 1 else [-1] + T = len(p.iouThrs) + R = len(p.recThrs) + K = len(p.catIds) if p.useCats else 1 + A = len(p.areaRng) + M = len(p.maxDets) + precision = -np.ones((T, R, K, A, M)) # -1 for the precision of absent categories + recall = -np.ones((T, K, A, M)) + scores = -np.ones((T, R, K, A, M)) + + # create dictionary for future indexing + _pe = self._paramsEval + catIds = _pe.catIds if _pe.useCats else [-1] + setK = set(catIds) + setA = set(map(tuple, _pe.areaRng)) + setM = set(_pe.maxDets) + setI = set(_pe.imgIds) + # get inds to evaluate + k_list = [n for n, k in enumerate(p.catIds) if k in setK] + m_list = [m for n, m in enumerate(p.maxDets) if m in setM] + a_list = [n for n, a in enumerate(map(lambda x: tuple(x), p.areaRng)) if a in setA] + i_list = [n for n, i in enumerate(p.imgIds) if i in setI] + I0 = len(_pe.imgIds) + A0 = len(_pe.areaRng) + # retrieve E at each category, area range, and max number of detections + for k, k0 in enumerate(k_list): + Nk = k0 * A0 * I0 + for a, a0 in enumerate(a_list): + Na = a0 * I0 + for m, maxDet in enumerate(m_list): + E = [self.evalImgs[Nk + Na + i] for i in i_list] + E = [e for e in E if not e is None] + if len(E) == 0: + continue + dtScores = np.concatenate([e['dtScores'][0:maxDet] for e in E]) + + # different sorting method generates slightly different results. + # mergesort is used to be consistent as Matlab implementation. + inds = np.argsort(-dtScores, kind='mergesort') + dtScoresSorted = dtScores[inds] + + dtm = np.concatenate([e['dtMatches'][:, 0:maxDet] for e in E], axis=1)[:, inds] + dtIg = np.concatenate([e['dtIgnore'][:, 0:maxDet] for e in E], axis=1)[:, inds] + gtIg = np.concatenate([e['gtIgnore'] for e in E]) + npig = np.count_nonzero(gtIg == 0) + if npig == 0: + continue + tps = np.logical_and(dtm, np.logical_not(dtIg)) + fps = np.logical_and(np.logical_not(dtm), np.logical_not(dtIg)) + + tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float) + fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float) + for t, (tp, fp) in enumerate(zip(tp_sum, fp_sum)): + tp = np.array(tp) + fp = np.array(fp) + nd = len(tp) + rc = tp / npig + pr = tp / (fp + tp + np.spacing(1)) + q = np.zeros((R,)) + ss = np.zeros((R,)) + + if nd: + recall[t, k, a, m] = rc[-1] + else: + recall[t, k, a, m] = 0 + + # numpy is slow without cython optimization for accessing elements + # use python array gets significant speed improvement + pr = pr.tolist() + q = q.tolist() + + for i in range(nd - 1, 0, -1): + if pr[i] > pr[i - 1]: + pr[i - 1] = pr[i] + + inds = np.searchsorted(rc, p.recThrs, side='left') + try: + for ri, pi in enumerate(inds): + q[ri] = pr[pi] + ss[ri] = dtScoresSorted[pi] + except: + pass + precision[t, :, k, a, m] = np.array(q) + scores[t, :, k, a, m] = np.array(ss) + self.eval = { + 'params': p, + 'counts': [T, R, K, A, M], + 'date': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), + 'precision': precision, + 'recall': recall, + 'scores': scores, + } + toc = time.time() + print('DONE (t={:0.2f}s).'.format(toc - tic)) + + def summarize(self): + ''' + Compute and display summary metrics for evaluation results. + Note this function can *only* be applied on the default parameter setting + ''' + + def _summarize(ap=1, iouThr=None, areaRng='all', maxDets=100): + p = self.params + iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}' + titleStr = 'Average Precision' if ap == 1 else 'Average Recall' + typeStr = '(AP)' if ap == 1 else '(AR)' + iouStr = '{:0.2f}:{:0.2f}'.format(p.iouThrs[0], p.iouThrs[-1]) \ + if iouThr is None else '{:0.2f}'.format(iouThr) + + aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng] + mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets] + if ap == 1: + # dimension of precision: [TxRxKxAxM] + s = self.eval['precision'] + # IoU + if iouThr is not None: + t = np.where(iouThr == p.iouThrs)[0] + s = s[t] + s = s[:, :, :, aind, mind] + else: + # dimension of recall: [TxKxAxM] + s = self.eval['recall'] + if iouThr is not None: + t = np.where(iouThr == p.iouThrs)[0] + s = s[t] + s = s[:, :, aind, mind] + if len(s[s > -1]) == 0: + mean_s = -1 + else: + mean_s = np.mean(s[s > -1]) + print(iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s)) + return mean_s + + def _summarizeDets(): + stats = np.zeros((12,)) + stats[0] = _summarize(1) + stats[1] = _summarize(1, iouThr=.5, maxDets=self.params.maxDets[2]) + stats[2] = _summarize(1, iouThr=.75, maxDets=self.params.maxDets[2]) + stats[3] = _summarize(1, areaRng='small', maxDets=self.params.maxDets[2]) + stats[4] = _summarize(1, areaRng='medium', maxDets=self.params.maxDets[2]) + stats[5] = _summarize(1, areaRng='large', maxDets=self.params.maxDets[2]) + stats[6] = _summarize(0) + stats[7] = _summarize(0, iouThr=0.5, maxDets=self.params.maxDets[2]) + stats[8] = _summarize(0, iouThr=0.75, maxDets=self.params.maxDets[2]) + stats[9] = _summarize(0, areaRng='small', maxDets=self.params.maxDets[2]) + stats[10] = _summarize(0, areaRng='medium', maxDets=self.params.maxDets[2]) + stats[11] = _summarize(0, areaRng='large', maxDets=self.params.maxDets[2]) + return stats + + def _summarizeKps(): + stats = np.zeros((10,)) + stats[0] = _summarize(1, maxDets=20) + stats[1] = _summarize(1, maxDets=20, iouThr=.5) + stats[2] = _summarize(1, maxDets=20, iouThr=.75) + stats[3] = _summarize(1, maxDets=20, areaRng='medium') + stats[4] = _summarize(1, maxDets=20, areaRng='large') + stats[5] = _summarize(0, maxDets=20) + stats[6] = _summarize(0, maxDets=20, iouThr=.5) + stats[7] = _summarize(0, maxDets=20, iouThr=.75) + stats[8] = _summarize(0, maxDets=20, areaRng='medium') + stats[9] = _summarize(0, maxDets=20, areaRng='large') + return stats + + if not self.eval: + raise Exception('Please run accumulate() first') + iouType = self.params.iouType + if iouType == 'segm' or iouType == 'bbox': + summarize = _summarizeDets + elif iouType == 'keypoints': + summarize = _summarizeKps + self.stats = summarize() + + def __str__(self): + self.summarize() diff --git a/frame_field_learning/README.md b/frame_field_learning/README.md new file mode 100644 index 0000000000000000000000000000000000000000..4f7a1b42a8bca92fa9fa5c29fba6c82cabc558ac --- /dev/null +++ b/frame_field_learning/README.md @@ -0,0 +1,172 @@ +This folder contains all the sources files of the "frame_field_learning" Python package. + +We briefly introduce each script in the following. + + +## data_transform.py + +Contains functions that return transformations applied to input data before feeding it to the network. +It includes pre-processing (whose result is store on disk), +CPU transforms which are applied on the loaded pre-processed data before being transferred to the GPU, +and finally GPU transforms which are applied on the GPU for speed (such as data augmentation). + + +## evaluate.py + +Defines the "evaluate" function called by ```main.py``` when ```--mode=eval``` which setups and instantiates an Evaluator object whose evaluate method is then called. + + +## evaluator.py + +Defines the "Evaluator" class used to run inference on a trained model followed by computing some measures and saving all results for all samples of a evaluation fold of a dataset. + + +## frame_field_utils.py + +This script defines all "frame field"-related functions. +For example the ```framefield_align_error``` function is used to compute the "align", +"align90" losses as well as the "frame field align" energy for ASm optimization in our paper. + +The ```LaplacianPenalty``` class is used for the "smooth" loss to ensure a smooth frame field. + +Both ```compute_closest_in_uv``` and ```detect_corners``` are use to detect corners sing the frame field. +They are somewhat redundant but the first is applied on torch tensors while the second is applied on numpy arrays. + + +## ictnet.py + +This script implements ICTNet (https://theictlab.org/lp/2019ICTNet) in PyTorch to try and add frame field learning to ICTNet. +We can use it as a backbone model with the ```ICTNetBackbone``` class. + +## inference.py + +This script defines the ```inference``` function to run a trained model on one image tile. It is used by the evaluation code. + +If the "patch_size" parameter is set in the "eval_params", then the image tile is split in small patches of size "patch_size", +inference is run on bacth_size patches at a time, +finally the result tile is stoched together from the results of all patches. +Typically those patches overlap by a few pixels and the result is linearly interpolated in overlapping areas. + +If there is no "patch_size" parameter, inference is run directly on the whole image tile. + + +## inference_from_filepath.py + +This script defines the ```inference_from_filepath``` function which is called by ```main.py``` when the ```--in_filepath``` argument is set. +It runs inference + polygonization on the images specified by the ```--in_filepath``` argument and saves the result. + + +## local_utils.py + +This script holds several utility functions for the project that do not really belong elsewhere. + + +## losses.py + +The script defines a ```Loss``` class that is the base class to be used for all loss functions. +It makes it easy to compute normalized losses for easier balancing. +It thus also defines a ```MultiLoss``` class to combine all necessary losses and apply multiplicative coefficients to their normalized version. + +All individual losses (segmentation loss, frame field align loss, frame field smooth loss, etc.) are defined as classes inheriting from the ```Loss``` class. + +Lastly the ```build_combined_loss``` function instantiates a ```MultiLoss``` object with all required losses depending on the specified config file. + + +## measures.py + +This script defines the ```iou``` function to compute the Intersection over Union (Iou), also calles the Jaccard index, and +the ```dice``` function to compute the differentiable version of the IoU in the form of the Dice coefficient. + + +## model.py + +This script defines the ```FrameFieldModel``` class which implements the final network. +It needs a backbone to be specified model and adds the necessary convolutions for the segmentation and frame field outputs. +Its forward method also performs transforms to be done on the GPU before pushing the result into the backbone. + + +## plot_utils.py + +This script defines several functions to plot segmentations, frame fields, polygons and polylines for visualization. + + +## polygonize.py + +This script is the entrypoint for all polygonization algorithms. It can then call these polygonization methods: +- Active Contours Model (ACM) from polygonize_acm.py +- Active Skeleton Model (ASM) from polygonize_asm.py +- Simple polygonization (Marching Cubes contour detection + Ramer-Douglas-Peucker simplification) from polygonize_simple.py + + +## polygonize_acm.py + +This script implements the Active Contours Model (ACM) polygonization algorithm. +It starts with contours detected with Marching Squares which are then optimized on the GPU with PyTorch to align to the frame field (in addition to other objectives, see paper). +It does not handle common walls between adjoining buildings. +For that reason we have also developed the Active Skeleton Model (ASM), see next section. + + +## polygonize_asm.py + +This script implements the Active Skeleton Model (ASM) polygonization algorithm. +It starts with a skeleton graph detected from the wall segmentation map. +It is then optimized on the GPU with PyTorch to align to the frame field (in addition to other objectives, see paper). +It can also be initialized with Marching Squares and essentially implements the ACM as well. +Thus, the polygonize_acm.py script should be redundant +(we still keep it around as it follows a different approach to the data structure which can be interesting). + + +## polygonize_simple.py + +This script implements the simple polygonization (Marching Cubes contour detection + Ramer-Douglas-Peucker simplification) used as baseline in the paper. + + +## polygonize_utils.py + +This scripts implements a few utility functions used by several of the above polygonization methods. +For example, the marching squares contour detection and computation of polygon probability computed from the building interior probability map. + + +## save_utils.py + +This script implements functions for saving results produced by the network and polygonization in various formats. + + +## train.py + +This script defines the ```train``` function which sets up the traiing procedure. +It instantiates a ```Trainer``` object which will then run the optimization. + +## trainer.py + +This script implements the ```Trainer``` class which is responsible for training the given model. +It implements multi-GPU training, loss normalization, restarting traiing from a checkpoint, etc. + + +## tta_utils.py + +When performing inference with the model, +Test Time Augmentation (TTA) can be used to increase the quality of the result. +We augment the input 8 times with right-angled rotations and flips. +The outputs are then merged with a given aggregation function. +We implemented several but it seems averaging works best. +The aggregation function is used on the 8 segmentation maps. +However, for the frame field, neither the uv nor the c0, c2 representations can be averaged easily (because of non-linearity). +Thus, we first search for the one segmentation out of the 8 ones that agress the most with the aggregated one. +Then the corresponding frame field is selected. +TTA is performed in the ```forward``` method of ```FrameFieldModel```. + + +## unet.py + +This script implements the ```UNetBackbone``` class for the original U-Net network to be used as a backbone. + + +## unet_resnet.py + +This script is an adapted code from https://github.com/neptune-ai/open-solution-mapping-challenge/blob/master/src/unet_models.py +to use the ResNet network as an encoder for a U-Net. This is used to instantiate the Unet-Resnet101 backbone we use in the paper. + + +### :warning: TODO: complete this README + diff --git a/frame_field_learning/data_transforms.py b/frame_field_learning/data_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..d66b4c7b17bd597561fa756c58374b5165a268be --- /dev/null +++ b/frame_field_learning/data_transforms.py @@ -0,0 +1,474 @@ +from collections import OrderedDict + +import PIL +import numpy as np +import torch +import torchvision +import kornia + +import torch_lydorn.kornia +import torch_lydorn.torchvision + +from lydorn_utils import print_utils + + +class Print(object): + """Convert polygons to a single graph""" + + def __init__(self): + pass + + def __call__(self, sample): + print("\n") + print(sample.keys()) + for key, item in sample.items(): + if type(item) == np.ndarray or type(item) == torch.Tensor: + if len(item.shape): + print(key, type(item), item.shape, item.dtype, item.min(), item.max()) + else: + print(key, type(item), item, item.dtype, item.min(), item.max()) + elif type(item) == PIL.Image.Image: + print(key, type(item), item.size, item.mode, np.array(item).min(), np.array(item).max()) + elif type(item) == list: + print(key, type(item[0]), len(item)) + # exit() + # print(sample["image"].dtype) + # print(sample["image"].shape) + # print(sample["image"].min()) + # print(sample["image"].max()) + # for key, value in sample.items(): + # print(key + ":") + # if type(value) == np.ndarray: + # print(value.shape) + # elif type(value) == list: + # print(len(value)) + # else: + # print("a") + return sample + + +class CudaDataAugmentation(object): + def __init__(self, input_patch_size: int, vflip: bool, affine: bool, scaling: list, color_jitter: bool): + self.input_patch_size = input_patch_size + self.vflip = vflip + self.affine = affine + self.scaling = scaling + self.color_jitter = None + if color_jitter: + self.color_jitter = kornia.augmentation.ColorJitter(brightness=0.05, contrast=0.05, saturation=.5, hue=.1) + self.tensor_keys_bilinear = ["image", "gt_polygons_image", "distances", + "valid_mask"] # Affine transform applied with bilinear sampling + self.tensor_keys_nearest = ["sizes", "gt_crossfield_angle"] # Affine transform applied with nearest sampling + + @staticmethod + def get_slices(batch, keys, last_slice_stop=0): + slices = OrderedDict() + for key in keys: + s = slice(last_slice_stop, last_slice_stop + batch[key].shape[1]) + last_slice_stop += batch[key].shape[1] + slices[key] = s + return slices + + def __call__(self, batch): + with torch.no_grad(): + batch_size, im_channels, height, width = batch["image"].shape + device = batch["image"].device + batch["valid_mask"] = torch.ones((batch_size, 1, height, width), dtype=torch.float, + device=device) # Apply losses only when valid_mask is True (pad with 0 when rotating) + + # Combine all images into one for faster/easier processing (store slices to separate them later on): + tensor_keys_bilinear = [key for key in self.tensor_keys_bilinear if key in batch] + tensor_keys_nearest = [key for key in self.tensor_keys_nearest if key in batch] + tensor_keys = tensor_keys_bilinear + tensor_keys_nearest + combined = torch.cat([batch[tensor_key] for tensor_key in tensor_keys], dim=1) + slices_bilinear = self.get_slices(batch, tensor_keys_bilinear, last_slice_stop=0) + slices_nearest = self.get_slices(batch, tensor_keys_nearest, + last_slice_stop=slices_bilinear[tensor_keys_bilinear[-1]].stop) + bilinear_slice = slice(slices_bilinear[tensor_keys_bilinear[0]].start, + slices_bilinear[tensor_keys_bilinear[-1]].stop) + nearest_slice = slice(slices_nearest[tensor_keys_nearest[0]].start, + slices_nearest[tensor_keys_nearest[-1]].stop) + + # Rotation (and translation) + if self.affine: + angle: torch.Tensor = torch.empty(batch_size, device=device).uniform_(-180, 180) + # To include corner pixels if angle=45 (coords are between -1 and 1 for grid_sample): + max_offset = np.sqrt(2) - 1 + offset: torch.Tensor = torch.empty((batch_size, 2), device=device).uniform_(-max_offset, max_offset) + downscale_factor = None + if self.scaling is not None: + downscale_factor: torch.Tensor = torch.empty(batch_size, device=device).uniform_(*self.scaling) + affine_grid = torch_lydorn.kornia.geometry.transform.get_affine_grid(combined, angle, offset, + downscale_factor) + combined[:, bilinear_slice, ...] = \ + torch.nn.functional.grid_sample(combined[:, bilinear_slice, ...], + affine_grid, mode='bilinear') + + + + # Rotate sizes and anglefield with mode='nearest' + # because it makes no sense to interpolate size values and angle values: + combined[:, nearest_slice, ...] = torch.nn.functional.grid_sample(combined[:, nearest_slice, ...], + affine_grid, mode='nearest') + + # Additionally the angle field's values themselves have to be rotated: + combined[:, slices_nearest["gt_crossfield_angle"], + ...] = torch_lydorn.torchvision.transforms.functional.rotate_anglefield( + combined[:, slices_nearest["gt_crossfield_angle"], ...], angle) + + # The sizes and distances should be adjusted as well because of the scaling. + if downscale_factor is not None: + if "sizes" in slices_nearest: + size_equals_one = combined[:, slices_nearest["sizes"], ...] == 1 + combined[:, slices_nearest["sizes"], :, :] /= downscale_factor[:, None, None, None] + combined[:, slices_nearest["sizes"], ...][size_equals_one] = 1 + if "distances" in slices_bilinear: + distance_equals_one = combined[:, slices_bilinear["distances"], ...] == 1 + combined[:, slices_bilinear["distances"], :, :] /= downscale_factor[:, None, None, None] + combined[:, slices_bilinear["distances"], ...][distance_equals_one] = 1 + + # Center crop + if self.input_patch_size is not None: + prev_image_norm = combined.shape[2] + combined.shape[3] + combined = torch_lydorn.torchvision.transforms.functional.center_crop(combined, self.input_patch_size) + current_image_norm = combined.shape[2] + combined.shape[3] + # Sizes and distances are affected by this because they are relative to the image's size. + # All non-one pixels have to be renormalized: + size_ratio = prev_image_norm / current_image_norm + if "sizes" in slices_nearest: + combined[:, slices_nearest["sizes"], ...][ + combined[:, slices_nearest["sizes"], ...] != 1] *= size_ratio + if "distances" in slices_bilinear: + combined[:, slices_bilinear["distances"], ...][ + combined[:, slices_bilinear["distances"], ...] != 1] *= size_ratio + + # vflip + if self.vflip: + to_flip: torch.Tensor = torch.empty(batch_size, device=device).uniform_(0, 1) < 0.5 + combined[to_flip] = kornia.geometry.transform.vflip(combined[to_flip]) + combined[ + to_flip, slices_nearest[ + "gt_crossfield_angle"], ...] = torch_lydorn.torchvision.transforms.functional.vflip_anglefield( + combined[to_flip, slices_nearest["gt_crossfield_angle"], ...]) + + # Split data: + batch["image"] = combined[:, slices_bilinear["image"], ...] + if "gt_polygons_image" in slices_bilinear: + batch["gt_polygons_image"] = combined[:, slices_bilinear["gt_polygons_image"], ...] + if "distances" in slices_bilinear: + batch["distances"] = combined[:, slices_bilinear["distances"], ...] + batch["valid_mask"] = 0.99 < combined[:, slices_bilinear["valid_mask"], + ...] # Take a very high threshold to remove fuzzy pixels + + if "sizes" in slices_nearest: + batch["sizes"] = combined[:, slices_nearest["sizes"], ...] + batch["gt_crossfield_angle"] = combined[:, slices_nearest["gt_crossfield_angle"], ...] + + # Color jitter + if self.color_jitter is not None and batch["image"].shape[1] == 3: + batch["image"] = self.color_jitter(batch["image"]) + + # --- Zero padding of sizes and distances is not correct, they should be padded with ones: + if self.affine: + if "sizes" in slices_nearest: + batch["sizes"][~batch["valid_mask"]] = 1 + if "distances" in slices_bilinear: + batch["distances"][~batch["valid_mask"]] = 1 + + return batch + + +class CudaCrop(object): + def __init__(self, input_patch_size: int): + self.input_patch_size = input_patch_size + self.tensor_keys = ["image", "gt_polygons_image", "distances", "valid_mask", "sizes", "gt_crossfield_angle"] + + def __call__(self, batch): + for tensor_key in self.tensor_keys: + if tensor_key in batch: + batch[tensor_key] = torch_lydorn.torchvision.transforms.functional.center_crop(batch[tensor_key], + self.input_patch_size) + return batch + + +def get_offline_transform(config, augmentations=False, to_patches=True): + data_patch_size = config["dataset_params"]["data_patch_size"] if augmentations else config["dataset_params"][ + "input_patch_size"] + transform_list = [ + torch_lydorn.torchvision.transforms.Map( + transform=torch_lydorn.torchvision.transforms.TransformByKey( + transform=torchvision.transforms.Compose([ + torch_lydorn.torchvision.transforms.RemoveDoubles(epsilon=0.01), + torch_lydorn.torchvision.transforms.FilterPolyVertexCount(min=3), + torch_lydorn.torchvision.transforms.ApproximatePolygon(tolerance=0.01), + torch_lydorn.torchvision.transforms.FilterPolyVertexCount(min=3) + ]), key="gt_polygons")), + + torch_lydorn.torchvision.transforms.FilterEmptyPolygons(key="gt_polygons"), + ] + if to_patches: + transform_list.extend([ + torch_lydorn.torchvision.transforms.ToPatches(stride=config["dataset_params"]["input_patch_size"], + size=data_patch_size), + torch_lydorn.torchvision.transforms.FilterEmptyPolygons(key="gt_polygons"), + ]) + transform_list.extend([ + torch_lydorn.torchvision.transforms.Map( + transform=torchvision.transforms.Compose([ + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torch_lydorn.torchvision.transforms.Rasterize(fill=True, edges=True, vertices=True, + line_width=4, antialiasing=True), + key=["image", "gt_polygons"], outkey="gt_polygons_image"), + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torch_lydorn.torchvision.transforms.AngleFieldInit(line_width=6), + key=["image", "gt_polygons"], + outkey="gt_crossfield_angle") + ])), + ]) + offline_transform = torchvision.transforms.Compose(transform_list) + return offline_transform + + +def get_offline_transform_patch(raster: bool = True, fill: bool = True, edges: bool = True, vertices: bool = True, + distances: bool = True, sizes: bool = True, angle_field: bool = True): + transform_list = [] + if raster: + if not distances and not sizes: + rasterize_transform = torch_lydorn.torchvision.transforms.TransformByKey( + transform=torch_lydorn.torchvision.transforms.Rasterize(fill=fill, edges=edges, vertices=vertices, + line_width=4, antialiasing=True, + return_distances=False, + return_sizes=False), + key=["image", "gt_polygons"], outkey="gt_polygons_image") + elif distances and sizes: + rasterize_transform = torch_lydorn.torchvision.transforms.TransformByKey( + transform=torch_lydorn.torchvision.transforms.Rasterize(fill=fill, edges=edges, vertices=vertices, + line_width=4, antialiasing=True, + return_distances=True, + return_sizes=True), + key=["image", "gt_polygons"], outkey=["gt_polygons_image", "distances", "sizes"]) + else: + raise NotImplementedError + transform_list.append(rasterize_transform) + if angle_field: + transform_list.append( + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torch_lydorn.torchvision.transforms.AngleFieldInit(line_width=6), + key=["image", "gt_polygons"], + outkey="gt_crossfield_angle") + ) + + return torchvision.transforms.Compose(transform_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 get_eval_online_cpu_transform(): + online_transform = torchvision.transforms.Compose([ + # 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"), + torch_lydorn.torchvision.transforms.TransformByKey(torch_lydorn.torchvision.transforms.ToTensor(), + key="gt_crossfield_angle"), + torch_lydorn.torchvision.transforms.TransformByKey(transform=torch_lydorn.torchvision.transforms.ToTensor(), + key="distances"), + torch_lydorn.torchvision.transforms.TransformByKey(transform=torch_lydorn.torchvision.transforms.ToTensor(), + key="sizes"), + torch_lydorn.torchvision.transforms.RemoveKeys(keys=["gt_polygons"]) + ]) + return online_transform + + +def get_online_cuda_transform(config, augmentations=False): + device_transform_list = [ + torch_lydorn.torchvision.transforms.TransformByKey(transform=torchvision.transforms.Compose([ + torchvision.transforms.Lambda(lambda tensor: tensor.float().div(255)) + ]), key="image"), + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torchvision.transforms.Lambda(lambda tensor: tensor.float().div(255)), + key="gt_polygons_image"), + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torchvision.transforms.Lambda(lambda tensor: np.pi * tensor.float().div(255)), + key="gt_crossfield_angle"), + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torchvision.transforms.Lambda(lambda tensor: tensor.float()), + key="distances", ignore_key_error=True), + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torchvision.transforms.Lambda(lambda tensor: tensor.float()), + key="sizes", ignore_key_error=True), + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torchvision.transforms.Lambda(lambda tensor: tensor.float()), + key="class_freq", ignore_key_error=True), + ] + if augmentations and config["data_aug_params"]["device"] == "cpu": + if config["data_aug_params"]["affine"]: + # Rotate angle field + device_transform_list.append( + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torch_lydorn.torchvision.transforms.functional.rotate_anglefield, + key=["gt_crossfield_angle", "rand_angle"], + outkey="gt_crossfield_angle")) + if config["data_aug_params"]["vflip"]: + device_transform_list.append( + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torch_lydorn.torchvision.transforms.functional.vflip_anglefield, + key=["gt_crossfield_angle", "rand_flip"], + outkey="gt_crossfield_angle")) + if config["data_aug_params"]["device"] == "cuda": + input_patch_size = config["dataset_params"]["input_patch_size"] if "input_patch_size" in config[ + "dataset_params"] else None # No crop if None + if augmentations: + device_transform_list.append(CudaDataAugmentation(input_patch_size, + config["data_aug_params"]["vflip"], + config["data_aug_params"]["affine"], + config["data_aug_params"]["scaling"], + config["data_aug_params"]["color_jitter"])) + elif input_patch_size is not None: + device_transform_list.append(CudaCrop(input_patch_size)) + device_transform_list.append( + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torch_lydorn.torchvision.transforms.functional.batch_normalize, + key=["image", "image_mean", "image_std"], + outkey="image"), ) + device_transform = torchvision.transforms.Compose(device_transform_list) + return device_transform + + +def get_eval_online_cuda_transform(config): + device_transform_list = [ + torch_lydorn.torchvision.transforms.TransformByKey(transform=torchvision.transforms.Compose([ + torchvision.transforms.Lambda(lambda tensor: tensor.float().div(255)) + ]), key="image"), + torch_lydorn.torchvision.transforms.TransformByKey( + transform=torch_lydorn.torchvision.transforms.functional.batch_normalize, + key=["image", "image_mean", "image_std"], + outkey="image") + ] + device_transform = torchvision.transforms.Compose(device_transform_list) + return device_transform diff --git a/frame_field_learning/evaluate.py b/frame_field_learning/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..c9392421e56f08a240458831dbb3933883d1bc62 --- /dev/null +++ b/frame_field_learning/evaluate.py @@ -0,0 +1,62 @@ +import random +import torch +import torch.distributed +import torch.utils.data + +from . import data_transforms +from .model import FrameFieldModel +from .evaluator import Evaluator + +from lydorn_utils import print_utils + +try: + import apex + from apex import amp + APEX_AVAILABLE = True +except ModuleNotFoundError: + APEX_AVAILABLE = False + + +def evaluate(gpu: int, config: dict, shared_dict, barrier, eval_ds, backbone): + # --- Setup DistributedDataParallel --- # + rank = config["nr"] * config["gpus"] + gpu + torch.distributed.init_process_group( + backend='nccl', + init_method='env://', + world_size=config["world_size"], + rank=rank + ) + + if gpu == 0: + print("# --- Start evaluating --- #") + + # Choose device + torch.cuda.set_device(gpu) + + # --- Online transform performed on the device (GPU): + eval_online_cuda_transform = data_transforms.get_eval_online_cuda_transform(config) + + if "samples" in config: + rng_samples = random.Random(0) + eval_ds = torch.utils.data.Subset(eval_ds, rng_samples.sample(range(len(eval_ds)), config["samples"])) + # eval_ds = torch.utils.data.Subset(eval_ds, range(config["samples"])) + + eval_sampler = torch.utils.data.distributed.DistributedSampler(eval_ds, num_replicas=config["world_size"], rank=rank) + + eval_ds = torch.utils.data.DataLoader(eval_ds, batch_size=config["optim_params"]["eval_batch_size"], pin_memory=True, sampler=eval_sampler, num_workers=config["num_workers"]) + + model = FrameFieldModel(config, backbone=backbone, eval_transform=eval_online_cuda_transform) + model.cuda(gpu) + + if config["use_amp"] and APEX_AVAILABLE: + amp.register_float_function(torch, 'sigmoid') + model = amp.initialize(model, opt_level="O1") + elif config["use_amp"] and not APEX_AVAILABLE and gpu == 0: + print_utils.print_warning("WARNING: Cannot use amp because the apex library is not available!") + + # Wrap the model for distributed training + model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu]) + + evaluator = Evaluator(gpu, config, shared_dict, barrier, model, run_dirpath=config["eval_params"]["run_dirpath"]) + split_name = config["fold"][0] + evaluator.evaluate(split_name, eval_ds) diff --git a/frame_field_learning/evaluator.py b/frame_field_learning/evaluator.py new file mode 100644 index 0000000000000000000000000000000000000000..33684490aa790323b8822fc23ea71ce7d51f2092 --- /dev/null +++ b/frame_field_learning/evaluator.py @@ -0,0 +1,221 @@ +import os +import csv + +from tqdm import tqdm +from multiprocess import Pool, Process, Queue +from functools import partial +import time + +import torch +import torch.utils.data +# from pytorch_memlab import profile, profile_every + +from . import inference, save_utils, polygonize +from . import local_utils +from . import measures + +from lydorn_utils import run_utils +from lydorn_utils import python_utils +from lydorn_utils import print_utils +from lydorn_utils import async_utils + + +class Evaluator: + 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)) + + # @profile + def evaluate(self, split_name: str, ds: torch.utils.data.DataLoader): + + # Prepare data saving: + flag_filepath_format = os.path.join(self.eval_dirpath, split_name, "{}.flag") + + # Loading model + self.load_checkpoint() + self.model.eval() + + # Create pool for multiprocessing + pool = None + if not self.config["eval_params"]["patch_size"]: + # If single image is not being split up, then a pool to process each sample in the batch makes sense + pool = Pool(processes=self.config["num_workers"]) + + compute_polygonization = self.config["eval_params"]["save_individual_outputs"]["poly_shapefile"] or \ + self.config["eval_params"]["save_individual_outputs"]["poly_geojson"] or \ + self.config["eval_params"]["save_individual_outputs"]["poly_viz"] or \ + self.config["eval_params"]["save_aggregated_outputs"]["poly_coco"] + + # Saving individual outputs to disk: + save_individual_outputs = True in self.config["eval_params"]["save_individual_outputs"].values() + saver_async = None + if save_individual_outputs: + save_outputs_partial = partial(save_utils.save_outputs, config=self.config, eval_dirpath=self.eval_dirpath, + split_name=split_name, flag_filepath_format=flag_filepath_format) + saver_async = async_utils.Async(save_outputs_partial) + saver_async.start() + + # Saving aggregated outputs + save_aggregated_outputs = True in self.config["eval_params"]["save_aggregated_outputs"].values() + + tile_data_list = [] + + if self.gpu == 0: + tile_iterator = tqdm(ds, desc="Eval {}: ".format(split_name), leave=True) + else: + tile_iterator = ds + for tile_i, tile_data in enumerate(tile_iterator): + # --- Inference, add result to tile_data_list + if self.config["eval_params"]["patch_size"] is not None: + # Cut image into patches for inference + inference.inference_with_patching(self.config, self.model, tile_data) + else: + # Feed images as-is to the model + inference.inference_no_patching(self.config, self.model, tile_data) + + tile_data_list.append(tile_data) + + # --- Accumulate batches into tile_data_list until capacity is reached (or this is the last batch) + if self.config["eval_params"]["batch_size_mult"] <= len(tile_data_list)\ + or tile_i == len(tile_iterator) - 1: + # Concat tensors of tile_data_list + accumulated_tile_data = {} + for key in tile_data_list[0].keys(): + if isinstance(tile_data_list[0][key], list): + accumulated_tile_data[key] = [item for _tile_data in tile_data_list for item in _tile_data[key]] + elif isinstance(tile_data_list[0][key], torch.Tensor): + accumulated_tile_data[key] = torch.cat([_tile_data[key] for _tile_data in tile_data_list], dim=0) + else: + raise TypeError(f"Type {type(tile_data_list[0][key])} is not handled!") + tile_data_list = [] # Empty tile_data_list + else: + # tile_data_list is not full yet, continue running inference... + continue + + # --- Polygonize + if compute_polygonization: + crossfield = accumulated_tile_data["crossfield"] if "crossfield" in accumulated_tile_data else None + accumulated_tile_data["polygons"], accumulated_tile_data["polygon_probs"] = polygonize.polygonize( + self.config["polygonize_params"], accumulated_tile_data["seg"], + crossfield_batch=crossfield, + pool=pool) + + # --- Save output + if self.config["eval_params"]["save_individual_outputs"]["seg_mask"] or \ + self.config["eval_params"]["save_aggregated_outputs"]["seg_coco"]: + # Take seg_interior: + seg_pred_mask = self.config["eval_params"]["seg_threshold"] < accumulated_tile_data["seg"][:, 0, ...] + accumulated_tile_data["seg_mask"] = seg_pred_mask + + accumulated_tile_data = local_utils.batch_to_cpu(accumulated_tile_data) + sample_list = local_utils.split_batch(accumulated_tile_data) + + # Save individual outputs: + if save_individual_outputs: + for sample in sample_list: + saver_async.add_work(sample) + + # Store aggregated outputs: + if save_aggregated_outputs: + self.shared_dict["name_list"].extend(accumulated_tile_data["name"]) + if self.config["eval_params"]["save_aggregated_outputs"]["stats"]: + y_pred = accumulated_tile_data["seg"][:, 0, ...].cpu() + if "gt_mask" in accumulated_tile_data: + y_true = accumulated_tile_data["gt_mask"][:, 0, ...] + elif "gt_polygons_image" in accumulated_tile_data: + y_true = accumulated_tile_data["gt_polygons_image"][:, 0, ...] + else: + raise ValueError("Either gt_mask or gt_polygons_image should be in accumulated_tile_data") + iou = measures.iou(y_pred.reshape(y_pred.shape[0], -1), y_true.reshape(y_true.shape[0], -1), + threshold=self.config["eval_params"]["seg_threshold"]) + self.shared_dict["iou_list"].extend(iou.cpu().numpy()) + if self.config["eval_params"]["save_aggregated_outputs"]["seg_coco"]: + for sample in sample_list: + annotations = save_utils.seg_coco(sample) + self.shared_dict["seg_coco_list"].extend(annotations) + if self.config["eval_params"]["save_aggregated_outputs"]["poly_coco"]: + for sample in sample_list: + annotations = save_utils.poly_coco(sample["polygons"], sample["polygon_probs"], sample["image_id"].item()) + self.shared_dict["poly_coco_list"].append(annotations) # annotations could be a dict, or a list + # END of loop over samples + + # Save aggregated results + if save_aggregated_outputs: + self.barrier.wait() # Wait on all processes so that shared_dict is synchronized. + if self.gpu == 0: + if self.config["eval_params"]["save_aggregated_outputs"]["stats"]: + print("Start saving stats:") + # Save sample_stats in CSV: + t1 = time.time() + stats_filepath = os.path.join(self.eval_dirpath, "{}.stats.csv".format(split_name)) + stats_file = open(stats_filepath, "w") + fnames = ["name", "iou"] + writer = csv.DictWriter(stats_file, fieldnames=fnames) + writer.writeheader() + for name, iou in sorted(zip(self.shared_dict["name_list"], self.shared_dict["iou_list"]), key=lambda pair: pair[0]): + writer.writerow({ + "name": name, + "iou": iou + }) + stats_file.close() + print(f"Finished in {time.time() - t1:02}s") + + if self.config["eval_params"]["save_aggregated_outputs"]["seg_coco"]: + print("Start saving seg_coco:") + t1 = time.time() + seg_coco_filepath = os.path.join(self.eval_dirpath, "{}.annotation.seg.json".format(split_name)) + python_utils.save_json(seg_coco_filepath, list(self.shared_dict["seg_coco_list"])) + print(f"Finished in {time.time() - t1:02}s") + + if self.config["eval_params"]["save_aggregated_outputs"]["poly_coco"]: + print("Start saving poly_coco:") + poly_coco_base_filepath = os.path.join(self.eval_dirpath, f"{split_name}.annotation.poly") + t1 = time.time() + save_utils.save_poly_coco(self.shared_dict["poly_coco_list"], poly_coco_base_filepath) + print(f"Finished in {time.time() - t1:02}s") + + # Sync point of individual outputs + if save_individual_outputs: + print_utils.print_info(f"GPU {self.gpu} -> INFO: Finishing saving individual outputs.") + saver_async.join() + self.barrier.wait() # Wait on all processes so that all saver_asyncs are finished + + def load_checkpoint(self): + """ + Loads best val checkpoint in checkpoints_dirpath + """ + filepaths = python_utils.get_filepaths(self.checkpoints_dirpath, startswith_str="checkpoint.best_val.", + endswith_str=".tar") + if len(filepaths): + filepaths = sorted(filepaths) + filepath = filepaths[-1] # Last best val checkpoint filepath in case there is more than one + if self.gpu == 0: + print_utils.print_info("Loading best val checkpoint: {}".format(filepath)) + else: + # No best val checkpoint fount: find last checkpoint: + filepaths = python_utils.get_filepaths(self.checkpoints_dirpath, endswith_str=".tar", + startswith_str="checkpoint.") + if len(filepaths) == 0: + raise FileNotFoundError("No checkpoint could be found at that location.") + filepaths = sorted(filepaths) + filepath = filepaths[-1] # Last checkpoint + if self.gpu == 0: + print_utils.print_info("Loading last checkpoint: {}".format(filepath)) + # map_location is used to load on current device: + checkpoint = torch.load(filepath, map_location="cuda:{}".format(self.gpu)) + + self.model.module.load_state_dict(checkpoint['model_state_dict']) diff --git a/frame_field_learning/frame_field_utils.py b/frame_field_learning/frame_field_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3ae770d8f5e739730b699dc5c77911fa857eccf6 --- /dev/null +++ b/frame_field_learning/frame_field_utils.py @@ -0,0 +1,115 @@ +import numpy as np + +import torch +from torch.nn import functional as F + +from torch_lydorn.torch.utils.complex import complex_mul, complex_sqrt, complex_abs_squared + + +def framefield_align_error(c0, c2, z, complex_dim=-1): + assert c0.shape == c2.shape == z.shape, \ + "All inputs should have the same shape. Currently c0: {}, c2: {}, z: {}".format(c0.shape, c2.shape, z.shape) + assert c0.shape[complex_dim] == c2.shape[complex_dim] == z.shape[complex_dim] == 2, \ + "All inputs should have their complex_dim size equal 2 (real and imag parts)" + + z_squared = complex_mul(z, z, complex_dim=complex_dim) + z_pow_4 = complex_mul(z_squared, z_squared, complex_dim=complex_dim) + # All tensors are assimilated as being complex so adding that way works (adding a scalar wouldn't work): + f_z = z_pow_4 + complex_mul(c2, z_squared, complex_dim=complex_dim) + c0 + loss = complex_abs_squared(f_z, complex_dim) # Square of the absolute value of f_z + return loss + + +class LaplacianPenalty: + def __init__(self, channels: int): + self.channels = channels + self.filter = torch.tensor([[0.5, 1.0, 0.5], + [1.0, -6., 1.0], + [0.5, 1.0, 0.5]]) / 12 + self.filter = self.filter[None, None, ...].expand(self.channels, -1, -1, -1) + + def laplacian_filter(self, tensor): + # with torch.autograd.profiler.profile(use_cuda=True) as prof: + penalty_tensor = F.conv2d(tensor, self.filter.to(tensor.device), padding=1, + groups=self.channels) + # print("penalty_tensor min={}, max={}".format(penalty_tensor.min(), penalty_tensor.max())) + # print(prof.key_averages().table(sort_by="cuda_time_total")) + return torch.abs(penalty_tensor) + + def __call__(self, tensor: torch.Tensor) -> torch.Tensor: + return self.laplacian_filter(tensor) + + +def c0c2_to_uv(c0c2: torch.Tensor) -> torch.Tensor: + c0, c2 = torch.chunk(c0c2, 2, dim=1) + c2_squared = complex_mul(c2, c2, complex_dim=1) + c2_squared_minus_4c0 = c2_squared - 4 * c0 + sqrt_c2_squared_minus_4c0 = complex_sqrt(c2_squared_minus_4c0, complex_dim=1) + u_squared = (c2 + sqrt_c2_squared_minus_4c0) / 2 + v_squared = (c2 - sqrt_c2_squared_minus_4c0) / 2 + uv_squared = torch.stack([u_squared, v_squared], dim=1) # Shape (B, 'uv': 2, 'complex': 2, H, W) + uv = complex_sqrt(uv_squared, complex_dim=2) + return uv + + +def compute_closest_in_uv(directions: torch.Tensor, uv: torch.Tensor) -> torch.Tensor: + """ + For each direction, compute if it is more aligned with {u, -u} (output 0) or {v, -v} (output 1). + + @param directions: Tensor of shape (N, 2) + @param uv: Tensor of shape (N, 'uv': 2, 'complex': 2) + @return: closest_in_uv of shape (N,) with the index in the 'uv' dimension of the closest vector in uv to direction + """ + uv_dot_dir = torch.sum(uv * directions[:, None, :], dim=2) + abs_uv_dot_dir = torch.abs(uv_dot_dir) + + closest_in_uv = torch.argmin(abs_uv_dot_dir, dim=1) + + return closest_in_uv + + +def detect_corners(polylines, u, v): + def compute_direction_score(ij, edges, field_dir): + values = field_dir[ij[:, 0], ij[:, 1]] + edge_dot_dir = edges[:, 0] * values.real + edges[:, 1] * values.imag + abs_edge_dot_dir = np.abs(edge_dot_dir) + return abs_edge_dot_dir + + def compute_is_corner(points, left_edges, right_edges): + if points.shape[0] == 0: + return np.empty(0, dtype=np.bool) + + coords = np.round(points).astype(np.int) + coords[:, 0] = np.clip(coords[:, 0], 0, u.shape[0] - 1) + coords[:, 1] = np.clip(coords[:, 1], 0, u.shape[1] - 1) + left_u_score = compute_direction_score(coords, left_edges, u) + left_v_score = compute_direction_score(coords, left_edges, v) + right_u_score = compute_direction_score(coords, right_edges, u) + right_v_score = compute_direction_score(coords, right_edges, v) + + left_is_u_aligned = left_v_score < left_u_score + right_is_u_aligned = right_v_score < right_u_score + + return np.logical_xor(left_is_u_aligned, right_is_u_aligned) + + corner_masks = [] + for polyline in polylines: + corner_mask = np.zeros(polyline.shape[0], dtype=np.bool) + if np.max(np.abs(polyline[0] - polyline[-1])) < 1e-6: + # Closed polyline + left_edges = np.concatenate([polyline[-2:-1] - polyline[-1:], polyline[:-2] - polyline[1:-1]], axis=0) + right_edges = polyline[1:] - polyline[:-1] + corner_mask[:-1] = compute_is_corner(polyline[:-1, :], left_edges, right_edges) + # left_edges and right_edges do not include the redundant last vertex, thus we have to do this assignment: + corner_mask[-1] = corner_mask[0] + else: + # Open polyline + corner_mask[0] = True + corner_mask[-1] = True + left_edges = polyline[:-2] - polyline[1:-1] + right_edges = polyline[2:] - polyline[1:-1] + corner_mask[1:-1] = compute_is_corner(polyline[1:-1, :], left_edges, right_edges) + corner_masks.append(corner_mask) + + return corner_masks + diff --git a/frame_field_learning/ictnet.py b/frame_field_learning/ictnet.py new file mode 100644 index 0000000000000000000000000000000000000000..6a01471594dd835846673f3ef4a07fbf7d30e534 --- /dev/null +++ b/frame_field_learning/ictnet.py @@ -0,0 +1,357 @@ +from collections import OrderedDict +from torch import nn +from torch.nn import functional as F +import torch +import torch.utils.checkpoint + +# from pytorch_memlab import profile, profile_every + + +def humanbytes(B): + 'Return the given bytes as a human friendly KB, MB, GB, or TB string' + B = float(B) + KB = float(1024) + MB = float(KB ** 2) # 1,048,576 + GB = float(KB ** 3) # 1,073,741,824 + TB = float(KB ** 4) # 1,099,511,627,776 + + if B < KB: + return '{0} {1}'.format(B, 'Bytes' if 0 == B > 1 else 'Byte') + elif KB <= B < MB: + return '{0:.2f} KB'.format(B / KB) + elif MB <= B < GB: + return '{0:.2f} MB'.format(B / MB) + elif GB <= B < TB: + return '{0:.2f} GB'.format(B / GB) + elif TB <= B: + return '{0:.2f} TB'.format(B / TB) + + +def get_preact_conv(in_channels, out_channels, kernel_size=3, padding=1, dropout_2d=0.2): + block = nn.Sequential( + nn.BatchNorm2d(in_channels), + nn.ReLU(inplace=True), + nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding), + nn.Dropout2d(dropout_2d) + ) + return block + + +def _dense_layer_function_factory(norm, relu, conv): + def bn_function(*inputs): + concated_features = torch.cat(inputs, 1) + bottleneck_output = conv(relu(norm(concated_features))) + return bottleneck_output + + return bn_function + + +class DenseLayer(nn.Module): + def __init__(self, in_channels, out_channels, dropout_2d=0.2, efficient=False): + super(DenseLayer, self).__init__() + self.add_module('norm', nn.BatchNorm2d(in_channels)), + self.add_module('relu', nn.ReLU(inplace=True)), + self.add_module('conv', nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)), + self.dropout_2d = dropout_2d + self.efficient = efficient + + def forward(self, *prev_features): + dense_layer_function = _dense_layer_function_factory(self.norm, self.relu, self.conv) + if self.efficient and any(prev_feature.requires_grad for prev_feature in prev_features): + new_features = torch.utils.checkpoint.checkpoint(dense_layer_function, *prev_features) + else: + new_features = dense_layer_function(*prev_features) + if 0 < self.dropout_2d: + new_features = F.dropout2d(new_features, p=self.dropout_2d, training=self.training) + return new_features + + +class SELayer(nn.Module): + def __init__(self, in_channels, ratio): + super(SELayer, self).__init__() + self.block = nn.Sequential( + nn.AdaptiveAvgPool2d(1), # Global average pooling + nn.Flatten(), # Prepare for fully-connected layers + nn.Linear(in_channels, in_channels // ratio), + nn.ReLU(inplace=True), + nn.Linear(in_channels // ratio, in_channels), + nn.Sigmoid() + ) + + def forward(self, x): + excitation = self.block(x) + x *= excitation[:, :, None, None] + return x + + +class DenseBlock(nn.Module): + def __init__(self, in_channels, n_layers, growth_rate, dropout_2d, return_only_new=False, efficient=False): + super(DenseBlock, self).__init__() + assert 0 < n_layers, "n_layers should be at least 1" + self.in_channels = in_channels + self.return_only_new = return_only_new + + channels = in_channels + self.layers = torch.nn.ModuleList() + for j in range(n_layers): + # Compute new feature maps + layer = DenseLayer(channels, growth_rate, dropout_2d=dropout_2d, efficient=efficient) + self.layers.append(layer) + channels += growth_rate + + if return_only_new: + se_layer_in_channel = channels - in_channels # Remove input, only keep new features + else: + se_layer_in_channel = channels + self.se_layer = SELayer(se_layer_in_channel, ratio=1) + + # @profile_every(1) + def forward(self, x): + features = [x] + for layer in self.layers: + new_features = layer(*features) + features.append(new_features) + + if self.return_only_new: + features = features[1:] + + features = torch.cat(features, 1) + features = self.se_layer(features) + + return features + + +def get_transition_down(in_channels, out_channels, dropout_2d=0.2): + block = nn.Sequential( + get_preact_conv(in_channels, out_channels, kernel_size=1, padding=0, dropout_2d=dropout_2d), + nn.MaxPool2d(kernel_size=2, stride=2) + ) + return block + + +def cat_non_matching(x1, x2): + diffY = x1.size()[2] - x2.size()[2] + diffX = x1.size()[3] - x2.size()[3] + + x2 = F.pad(x2, (diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2)) + + # for padding issues, see + # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a + # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd + + x = torch.cat([x1, x2], dim=1) + return x + + +class TransitionUp(nn.Module): + def __init__(self, in_channels, n_filters_keep): + super(TransitionUp, self).__init__() + self.conv_transpose_2d = nn.ConvTranspose2d(in_channels, n_filters_keep, kernel_size=4, stride=2, padding=1) + + def forward(self, x, skip_connection): + x = self.conv_transpose_2d(x) + x = cat_non_matching(x, skip_connection) + return x + + +class ICTNetBackbone(nn.Module): + """ + ICTNet model: https://theictlab.org/lp/2019ICTNet. + """ + def __init__(self, preset_model='FC-DenseNet56', in_channels=3, out_channels=2, n_filters_first_conv=48, n_pool=5, growth_rate=12, n_layers_per_block=4, dropout_2d=0.2, efficient=False): + super().__init__() + + # --- Handle args + if preset_model == 'FC-DenseNet56': + n_pool = 5 + growth_rate = 12 + n_layers_per_block = 4 + elif preset_model == 'FC-DenseNet67': + n_pool = 5 + growth_rate = 16 + n_layers_per_block = 5 + elif preset_model == 'FC-DenseNet103': + n_pool = 5 + growth_rate = 16 + n_layers_per_block = [4, 5, 7, 10, 12, 15, 12, 10, 7, 5, 4] + else: + n_pool = n_pool + growth_rate = growth_rate + n_layers_per_block = n_layers_per_block + + if type(n_layers_per_block) == list: + assert (len(n_layers_per_block) == 2 * n_pool + 1) + elif type(n_layers_per_block) == int: + n_layers_per_block = [n_layers_per_block] * (2 * n_pool + 1) + else: + raise ValueError + + # --- Instantiate layers + self.first_conv = nn.Conv2d(in_channels, n_filters_first_conv, 3, padding=1) + + # Downsampling path + channels = n_filters_first_conv + self.down_dense_blocks = torch.nn.ModuleList() + self.transition_downs = torch.nn.ModuleList() + skip_connection_channels = [] + for i in range(n_pool): + # Dense Block + self.down_dense_blocks.append(DenseBlock(in_channels=channels, n_layers=n_layers_per_block[i], growth_rate=growth_rate, dropout_2d=dropout_2d, return_only_new=False, efficient=efficient)) + channels += growth_rate * n_layers_per_block[i] + skip_connection_channels.append(channels) + # Transition Down + self.transition_downs.append(get_transition_down(in_channels=channels, out_channels=channels, dropout_2d=dropout_2d)) + + # Bottleneck Dense Block + self.bottleneck_dense_block = DenseBlock(in_channels=channels, n_layers=n_layers_per_block[n_pool], growth_rate=growth_rate, dropout_2d=dropout_2d, return_only_new=True, efficient=efficient) + up_in_channels = n_layers_per_block[n_pool] * growth_rate # We will only upsample the new feature maps + + # Upsampling path + self.transition_ups = torch.nn.ModuleList() + self.up_dense_blocks = torch.nn.ModuleList() + for i in range(n_pool): + # Transition Up (Upsampling + concatenation with the skip connection) + n_filters_keep = growth_rate * n_layers_per_block[n_pool + i] + self.transition_ups.append(TransitionUp(in_channels=up_in_channels, n_filters_keep=n_filters_keep)) + up_out_channels = skip_connection_channels[n_pool - i - 1] + n_filters_keep # After concatenation + + # Dense Block + # We will only upsample the new feature maps + self.up_dense_blocks.append( + DenseBlock(in_channels=up_out_channels, n_layers=n_layers_per_block[n_pool + i + 1], growth_rate=growth_rate, + dropout_2d=dropout_2d, return_only_new=True, efficient=efficient)) + up_in_channels = growth_rate * n_layers_per_block[n_pool + i + 1] # We will only upsample the new feature maps + + # Last layer + self.final_conv = nn.Conv2d(up_in_channels, out_channels, 1, padding=0) + + # @profile + def forward(self, x): + stack = self.first_conv(x) + + skip_connection_list = [] + # print(humanbytes(torch.cuda.memory_allocated())) + for down_dense_block, transition_down in zip(self.down_dense_blocks, self.transition_downs): + # Dense Block + stack = down_dense_block(stack) + + # At the end of the dense block, the current stack is stored in the skip_connections list + skip_connection_list.append(stack) + + # Transition Down + stack = transition_down(stack) + # print(humanbytes(torch.cuda.memory_allocated())) + + skip_connection_list = skip_connection_list[::-1] + + # Bottleneck Dense Block + # We will only upsample the new feature maps + stack = self.bottleneck_dense_block(stack) + + # Upsampling path + # print(humanbytes(torch.cuda.memory_allocated())) + for transition_up, up_dense_block, skip_connection in zip(self.transition_ups, self.up_dense_blocks, skip_connection_list): + # Transition Up ( Upsampling + concatenation with the skip connection) + stack = transition_up(stack, skip_connection) + + # Dense Block + # We will only upsample the new feature maps + stack = up_dense_block(stack) + # print(humanbytes(torch.cuda.memory_allocated())) + + # Final conv + stack = self.final_conv(stack) + + result = OrderedDict() + result["out"] = stack + + # print(humanbytes(torch.cuda.memory_allocated())) + + return result + + +def count_trainable_params(model): + count = 0 + for param in model.parameters(): + if param.requires_grad: + count += param.numel() + return count + + +def main(): + device = "cuda" + b = 2 + c = 3 + h = 512 + w = 512 + features = 32 + + # Init input + x = torch.rand((b, c, h, w), device=device) + print("x: ", x.shape, x.min().item(), x.max().item()) + + # # Test SELayer + # print("--- Test SELayer:") + # se_layer = SELayer(in_channels=c, ratio=1) + # y = se_layer(x) + # print("y: ", y.shape) + # print("------") + + # # Test DenseBlock + # print("--- Test DenseBlock:") + # dense_block = DenseBlock(in_channels=c, n_layers=5, growth_rate=16, dropout_2d=0.2, path_type="down") + # y, new_y = dense_block(x) + # print("y: ", y.shape) + # print("new_y: ", new_y.shape) + # print("------") + + # # Test transition_down + # print("--- Test transition_down:") + # transition_down = get_transition_down(in_channels=c, out_channels=features, dropout_2d=0.2) + # x_down = transition_down(x) + # print("x_down: ", x_down.shape) + # print("------") + # + # # Test TransitionUp + # print("--- Test TransitionUp:") + # transition_up = TransitionUp(in_channels=features, n_filters_keep=features//2) + # y = transition_up(x_down, x) + # print("y: ", y.shape) + # print("------") + + # Test ICTNetBackboneICTNetBackboneTest SELayer:") + backbone = ICTNetBackbone(out_channels=features, preset_model="FC-DenseNet103", dropout_2d=0.0, efficient=True) + print("ICTNetBackbone has {} trainable params".format(count_trainable_params(backbone))) + # print(backbone) + backbone.to(device) + result = backbone(x) + y = result["out"] + print("y: ", y.shape) + print("------") + + print("Back-prop:") + loss = torch.sum(y) + loss.backward() + + +if __name__ == "__main__": + main() + + + + + + + + + + + + + + + + + + + diff --git a/frame_field_learning/inference.py b/frame_field_learning/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..f861fe675bd5c479f8663f7fd1f807f009e03ac1 --- /dev/null +++ b/frame_field_learning/inference.py @@ -0,0 +1,167 @@ +import sys + +from tqdm import tqdm +import scipy + +import numpy as np +import torch + +from . import local_utils +from . import polygonize + +from lydorn_utils import image_utils +from lydorn_utils import print_utils +from lydorn_utils import python_utils + + +def network_inference(config, model, batch): + if config['device'] == 'cuda': + batch = local_utils.batch_to_cuda(batch) + pred, batch = model(batch, tta=config["eval_params"]["test_time_augmentation"]) + return pred, batch + + +def inference(config, model, tile_data, compute_polygonization=False, pool=None): + if config["eval_params"]["patch_size"] is not None: + # Cut image into patches for inference + inference_with_patching(config, model, tile_data) + single_sample = True + else: + # Feed images as-is to the model + inference_no_patching(config, model, tile_data) + single_sample = False + + # Polygonize: + if compute_polygonization: + pool = None if single_sample else pool # A single big image is being processed + crossfield = tile_data["crossfield"] if "crossfield" in tile_data else None + polygons_batch, probs_batch = polygonize.polygonize(config["polygonize_params"], tile_data["seg"], crossfield_batch=crossfield, + pool=pool) + tile_data["polygons"] = polygons_batch + tile_data["polygon_probs"] = probs_batch + + return tile_data + + +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 command-line argument." + .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"])) + sys.exit() + + tile_data["seg"] = pred["seg"] + if "crossfield" in pred: + tile_data["crossfield"] = pred["crossfield"] + + return tile_data + + +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 load_checkpoint(model, checkpoints_dirpath, device): + """ + Loads best val checkpoint in checkpoints_dirpath + """ + filepaths = python_utils.get_filepaths(checkpoints_dirpath, startswith_str="checkpoint.best_val.", + endswith_str=".tar") + if len(filepaths): + filepaths = sorted(filepaths) + filepath = filepaths[-1] # Last best val checkpoint filepath in case there is more than one + print_utils.print_info("Loading best val checkpoint: {}".format(filepath)) + else: + # No best val checkpoint fount: find last checkpoint: + filepaths = python_utils.get_filepaths(checkpoints_dirpath, endswith_str=".tar", + startswith_str="checkpoint.") + filepaths = sorted(filepaths) + filepath = filepaths[-1] # Last checkpoint + print_utils.print_info("Loading last checkpoint: {}".format(filepath)) + + device = torch.device(device) + checkpoint = torch.load(filepath, map_location=device) # map_location is used to load on current device + + model.load_state_dict(checkpoint['model_state_dict']) + + return model diff --git a/frame_field_learning/inference_from_filepath.py b/frame_field_learning/inference_from_filepath.py new file mode 100644 index 0000000000000000000000000000000000000000..3d308a22503e780b6809d0ee62b86c956cfc2664 --- /dev/null +++ b/frame_field_learning/inference_from_filepath.py @@ -0,0 +1,107 @@ +import os +import numpy as np +import skimage.io +import torch + +from tqdm import tqdm + +from . import data_transforms, save_utils +from .model import FrameFieldModel +from . import inference +from . import local_utils + +from torch_lydorn import torchvision + +from lydorn_utils import print_utils +from lydorn_utils import run_utils + + +def inference_from_filepath(config, in_filepaths, backbone, out_dirpath=None): + # --- 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 + in_filepath_list = [os.path.join(in_filepaths, in_filepath) for in_filepath in os.listdir(in_filepaths) if in_filepath.endswith(('.JPG','.PNG','.png','.jpg','.jepg','bmp'))] + pbar = tqdm(in_filepath_list, desc="Infer images") + for in_filepath in pbar: + print(in_filepath) + 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] + + # --- Saving outputs --- # + + pbar.set_postfix(status="Saving output") + + # Figuring out_base_filepath out: + if out_dirpath is None: + out_dirpath = os.path.dirname(in_filepath) + base_filename = os.path.splitext(os.path.basename(in_filepath))[0] + out_base_filepath = (out_dirpath, base_filename) + + if config["compute_seg"]: + if config["eval_params"]["save_individual_outputs"]["seg_mask"]: + seg_mask = 0.5 < tile_data["seg"][0] + save_utils.save_seg_mask(seg_mask, out_base_filepath, "mask", tile_data["image_filepath"]) + if config["eval_params"]["save_individual_outputs"]["seg"]: + save_utils.save_seg(tile_data["seg"], out_base_filepath, "seg", tile_data["image_filepath"]) + if config["eval_params"]["save_individual_outputs"]["seg_luxcarta"]: + save_utils.save_seg_luxcarta_format(tile_data["seg"], out_base_filepath, "seg_luxcarta_format", tile_data["image_filepath"]) + + if config["compute_crossfield"] and config["eval_params"]["save_individual_outputs"]["crossfield"]: + save_utils.save_crossfield(tile_data["crossfield"], out_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"], out_base_filepath, "poly_viz") + if config["eval_params"]["save_individual_outputs"]["poly_shapefile"]: + save_utils.save_shapefile(tile_data["polygons"], out_base_filepath, "poly_shapefile", tile_data["image_filepath"]) + + # if config["eval_params"]["save_individual_outputs"]["seg_gt"]: + # save_utils.save_seg(tile_data["gt_polygons_image"], base_filepath, "seg.gt", tile_data["image_filepath"]) + # if config["eval_params"]["save_individual_outputs"]["seg"]: + # save_utils.save_seg(tile_data["seg"], base_filepath, "seg", tile_data["image_filepath"]) + # if config["eval_params"]["save_individual_outputs"]["seg_mask"]: + # save_utils.save_seg_mask(tile_data["seg_mask"], base_filepath, "seg_mask", tile_data["image_filepath"]) + # if config["eval_params"]["save_individual_outputs"]["seg_opencities_mask"]: + # save_utils.save_opencities_mask(tile_data["seg_mask"], base_filepath, "drivendata", + # tile_data["image_filepath"]) + # if config["eval_params"]["save_individual_outputs"]["seg_luxcarta"]: + # save_utils.save_seg_luxcarta_format(tile_data["seg"], base_filepath, "seg_luxcarta_format", + # tile_data["image_filepath"]) + # if config["eval_params"]["save_individual_outputs"]["crossfield"]: + # save_utils.save_crossfield(tile_data["crossfield"], base_filepath, "crossfield") + # if config["eval_params"]["save_individual_outputs"]["uv_angles"]: + # save_utils.save_uv_angles(tile_data["crossfield"], base_filepath, "uv_angles", tile_data["image_filepath"]) + # + # if "polygons" in tile_data: + # save_utils.save_polygons(tile_data["polygons"], base_filepath, "polygons", tile_data["image_filepath"]) diff --git a/frame_field_learning/local_utils.py b/frame_field_learning/local_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f0ee02609477714a9d801a646e6171e29a679913 --- /dev/null +++ b/frame_field_learning/local_utils.py @@ -0,0 +1,217 @@ +import functools +import os +import subprocess +import sys +import time + +from lydorn_utils import run_utils, print_utils +from lydorn_utils import python_utils + + +def compute_max_disp(disp_params): + m_g_t = disp_params["max_global_translation"] + m_g_h = disp_params["max_global_homography"] + m_p_t = disp_params["max_poly_translation"] + m_p_h = disp_params["max_poly_homography"] + m_h_c = disp_params["max_homography_coef"] + return (m_g_t + m_h_c*m_g_h) + (m_p_t + m_h_c*m_p_h) + + +def get_git_revision_hash(): + try: + hash = subprocess.check_output(['git', 'rev-parse', 'HEAD'], stderr=subprocess.STDOUT).decode("utf-8")[:-1] + except subprocess.CalledProcessError: + hash = None + return hash + + +def setup_run(config): + run_name = config["run_name"] + new_run = config["new_run"] + init_run_name = config["init_run_name"] + + working_dir = os.path.dirname(os.path.abspath(__file__)) + runs_dir = os.path.join(working_dir, config["runs_dirpath"]) + + # setup init checkpoints directory path if one is specified: + if init_run_name is not None: + init_run_dirpath = run_utils.setup_run_dir(runs_dir, init_run_name) + _, init_checkpoints_dirpath = run_utils.setup_run_subdirs(init_run_dirpath) + else: + init_checkpoints_dirpath = None + + # setup run directory: + run_dirpath = run_utils.setup_run_dir(runs_dir, run_name, new_run) + + # save config in logs directory + run_utils.save_config(config, run_dirpath) + + # save args + args_filepath = os.path.join(run_dirpath, "args.json") + args_to_save = { + "run_name": run_name, + "new_run": new_run, + "init_run_name": init_run_name, + "batch_size": config["optim_params"]["batch_size"], + } + if "samples" in config: + args_to_save["samples"] = config["samples"] + python_utils.save_json(args_filepath, args_to_save) + + # save current commit hash + commit_hash = get_git_revision_hash() + if commit_hash is not None: + commit_hash_filepath = os.path.join(run_dirpath, "commit_history.json") + if os.path.exists(commit_hash_filepath): + commit_hashes = python_utils.load_json(commit_hash_filepath) + if commit_hashes[-1] != commit_hash: + commit_hashes.append(commit_hash) + python_utils.save_json(commit_hash_filepath, commit_hashes) + else: + commit_hashes = [commit_hash] + python_utils.save_json(commit_hash_filepath, commit_hashes) + + return run_dirpath, init_checkpoints_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 batch_to_cuda(batch): + # Send data to computing device: + for key, item in batch.items(): + if hasattr(item, "cuda"): + batch[key] = item.cuda(non_blocking=True) + return batch + + +def batch_to_cpu(batch): + # Send data to computing device: + for key, item in batch.items(): + if hasattr(item, "cuda"): + batch[key] = item.cpu() + return batch + + +def split_batch(tile_data): + assert len(tile_data["image"].shape) == 4, "tile_data[\"image\"] should be (N, C, H, W)" + tile_data_list = [] + for i in range(tile_data["image"].shape[0]): + individual_tile_data = {} + for key, item in tile_data.items(): + if not i < len(item): + print(key, len(item)) + individual_tile_data[key] = item[i] + tile_data_list.append(individual_tile_data) + return tile_data_list + + +def _concat_dictionaries(dict1, dict2): + """ + Recursive concat dictionaries. Dict 1 and Dict 2 must have the same key hierarchy (this is not checked). + + :param dict1: Dictionary to add to. + :param dict2: Dictionary to add from + :return: Merged dictionary dict1 + """ + for key in dict1.keys(): + item1 = dict1[key] + item2 = dict2[key] + if isinstance(item1, dict): # And item2 is dict too. + dict1[key] = _concat_dictionaries(item1, item2) + else: + dict1[key].extend(item2) + return dict1 + + +def _root_concat_dictionaries(dict1, dict2): + + t0 = time.time() + dict1 = _concat_dictionaries(dict1, dict2) + print(f"_root_concat_dictionaries: {time.time() - t0:02}s") + return dict1 + + +def list_of_dicts_to_dict_of_lists(list_of_dicts): + """ + Works recursively by using _concat_dictionaries which is recursive + + @param list_of_dicts: + @return: dict_of_lists + """ + return functools.reduce(_concat_dictionaries, list_of_dicts) + + +def flatten_dict(_dict): + """ + Makes a hierarchy of dicts flat + + @param _dict: + @return: + """ + new_dict = {} + for key, item in _dict.items(): + if isinstance(item, dict): + item = flatten_dict(item) + for k in item.keys(): + new_dict[key + "." + k] = item[k] + else: + new_dict[key] = item + return new_dict + + +def _generate_list_of_dicts(list_length, methods_count, submethods_count, annotation_count, segmentation_length): + list_of_dicts = [] + for i in range(list_length): + d = {} + for method_i in range(methods_count): + d[f"method_{method_i}"] = {} + for submethod_i in range(submethods_count): + d[f"method_{method_i}"][f"submethod_{submethod_i}"] = [] + for annotation_i in range(annotation_count): + annotation = { + "image_id": 0, + "segmentation": [list(range(segmentation_length))], + "category_id": 100, # Building + "bbox": [0, 1, 0, 1], + "score": 1.0 + } + d[f"method_{method_i}"][f"submethod_{submethod_i}"].append(annotation) + list_of_dicts.append(d) + return list_of_dicts + + +def main(): + # list_of_dicts = [ + # { + # "method1": { + # "submethod1": [[0, 1, 2, 3], [4, 5, 6]] + # } + # }, + # { + # "method1": { + # "submethod1": [[7, 8, 9], [10, 11, 12, 13, 14, 15]] + # } + # }, + # ] + t0 = time.time() + list_of_dicts = _generate_list_of_dicts(list_length=2000, methods_count=2, submethods_count=2, annotation_count=100, segmentation_length=200) + print(f"_generate_list_of_dicts: {time.time() - t0:02}s") + + t0 = time.time() + dict_of_lists = list_of_dicts_to_dict_of_lists(list_of_dicts) + print(f"list_of_dicts_to_dict_of_lists: {time.time() - t0:02}s") + + flat_dict_of_lists = flatten_dict(dict_of_lists) + + +if __name__ == "__main__": + main() diff --git a/frame_field_learning/losses.py b/frame_field_learning/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..89507a9c9263af9e3c1f42fccccadf421da4626c --- /dev/null +++ b/frame_field_learning/losses.py @@ -0,0 +1,441 @@ +import math +from functools import partial + +import scipy.interpolate +import numpy as np +import torch +import torch.distributed +from torch.nn import functional as F + +from . import measures +from . import frame_field_utils + +import torch_lydorn.kornia + +from lydorn_utils import math_utils, print_utils + + +# --- Base classes --- # + + +class Loss(torch.nn.Module): + def __init__(self, name): + """ + Attribute extra_info can be used in self.compute() to add intermediary results of loss computation for + visualization for example. + It is the second output of self.__call__() + + :param name: + """ + super(Loss, self).__init__() + self.name = name + self.norm_meter = None + self.norm = torch.nn.parameter.Parameter(torch.Tensor(1), requires_grad=False) + self.reset_norm() + self.extra_info = {} # + + def reset_norm(self): + self.norm_meter = math_utils.AverageMeter("{}_norm".format(self.name), init_val=1) + self.norm[0] = self.norm_meter.val + + def update_norm(self, pred_batch, gt_batch, nums): + loss = self.compute(pred_batch, gt_batch) + self.norm_meter.update(loss, nums) + self.norm[0] = self.norm_meter.val + + def sync(self, world_size): + """ + This method should be used to synchronize loss norms across GPUs when using distributed training + :return: + """ + torch.distributed.all_reduce(self.norm) + self.norm /= world_size + + def compute(self, pred_batch, gt_batch): + raise NotImplementedError + + def forward(self, pred_batch, gt_batch, normalize=True): + loss = self.compute(pred_batch, gt_batch) + if normalize: + assert 1e-9 < self.norm[0], "self.norm[0] <= 1e-9 -> this might lead to numerical instabilities." + loss = loss / self.norm[0] + extra_info = self.extra_info + self.extra_info = {} # Re-init extra_info + # contains_nan = bool(torch.sum(torch.isnan(loss)).item()) + # assert not contains_nan, f"loss {str(self)} is Nan!" + return loss, extra_info + + def __repr__(self): + return "{} (name={}, norm={:0.06})".format(self.__class__.__name__, self.name, self.norm[0]) + + +class MultiLoss(torch.nn.Module): + def __init__(self, loss_funcs, weights, epoch_thresholds=None, pre_processes=None): + """ + + @param loss_funcs: + @param weights: + @param pre_processes: List of functions to call with 2 arguments (which are updated): pred_batch, gt_batch to compute only one values used by several losses. + """ + super(MultiLoss, self).__init__() + assert len(loss_funcs) == len(weights), \ + "Should have the same amount of loss_funcs ({}) and weights ({})".format(len(loss_funcs), len(weights)) + self.loss_funcs = torch.nn.ModuleList(loss_funcs) + + self.weights = [] + for weight in weights: + if isinstance(weight, list): + # Weight is a list of coefs corresponding to epoch_thresholds, they will be interpolated in-between + self.weights.append(scipy.interpolate.interp1d(epoch_thresholds, weight, bounds_error=False, fill_value=(weight[0], weight[-1]))) + elif isinstance(weight, float) or isinstance(weight, int): + self.weights.append(float(weight)) + else: + raise TypeError(f"Type {type(weight)} not supported as a loss coef weight.") + + self.pre_processes = pre_processes + + for loss_func, weight in zip(self.loss_funcs, self.weights): + if weight == 0: + print_utils.print_info(f"INFO: loss '{loss_func.name}' has a weight of zero and thus won't affect grad update.") + + def reset_norm(self): + for loss_func in self.loss_funcs: + loss_func.reset_norm() + + def update_norm(self, pred_batch, gt_batch, nums): + if self.pre_processes is not None: + for pre_process in self.pre_processes: + pred_batch, gt_batch = pre_process(pred_batch, gt_batch) + for loss_func in self.loss_funcs: + loss_func.update_norm(pred_batch, gt_batch, nums) + + def sync(self, world_size): + """ + This method should be used to synchronize loss norms across GPUs when using distributed training + :return: + """ + for loss_func in self.loss_funcs: + loss_func.sync(world_size) + + def forward(self, pred_batch, gt_batch, normalize=True, epoch=None): + if self.pre_processes is not None: + for pre_process in self.pre_processes: + pred_batch, gt_batch = pre_process(pred_batch, gt_batch) + total_loss = 0 + # total_weight = 0 + individual_losses_dict = {} + extra_dict = {} + for loss_func_i, weight_i in zip(self.loss_funcs, self.weights): + loss_i, extra_dict_i = loss_func_i(pred_batch, gt_batch, normalize=normalize) + if isinstance(weight_i, scipy.interpolate.interpolate.interp1d) and epoch is not None: + current_weight = float(weight_i(epoch)) + else: + current_weight = weight_i + total_loss += current_weight * loss_i + # total_weight += weight_i + individual_losses_dict[loss_func_i.name] = loss_i + extra_dict[loss_func_i.name] = extra_dict_i + # total_loss /= total_weight + return total_loss, individual_losses_dict, extra_dict + + def __repr__(self): + ret = "\n\t".join([str(loss_func) for loss_func in self.loss_funcs]) + return "{}:\n\t{}".format(self.__class__.__name__, ret) + + +# --- Build combined loss: --- # +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 compute_gt_field(pred_batch, gt_batch): + gt_crossfield_angle = gt_batch["gt_crossfield_angle"] + gt_field = torch.cat([torch.cos(gt_crossfield_angle), + torch.sin(gt_crossfield_angle)], dim=1) + gt_batch["gt_field"] = gt_field + return pred_batch, gt_batch + + +class ComputeSegGrads: + def __init__(self, device): + self.spatial_gradient = torch_lydorn.kornia.filters.SpatialGradient(mode="scharr", coord="ij", normalized=True, device=device) + + def __call__(self, pred_batch, gt_batch): + seg = pred_batch["seg"] # (b, c, h, w) + seg_grads = 2 * self.spatial_gradient(seg) # (b, c, 2, h, w), Normalize (kornia normalizes to -0.5, 0.5 for input in [0, 1]) + seg_grad_norm = seg_grads.norm(dim=2) # (b, c, h, w) + seg_grads_normed = seg_grads / (seg_grad_norm[:, :, None, ...] + 1e-6) # (b, c, 2, h, w) + pred_batch["seg_grads"] = seg_grads + pred_batch["seg_grad_norm"] = seg_grad_norm + pred_batch["seg_grads_normed"] = seg_grads_normed + return pred_batch, gt_batch + + +def build_combined_loss(config): + pre_processes = [] + loss_funcs = [] + weights = [] + if config["compute_seg"]: + partial_compute_seg_loss_weigths = partial(compute_seg_loss_weigths, config=config) + pre_processes.append(partial_compute_seg_loss_weigths) + gt_channel_selector = [config["seg_params"]["compute_interior"], config["seg_params"]["compute_edge"], config["seg_params"]["compute_vertex"]] + loss_funcs.append(SegLoss(name="seg", + gt_channel_selector=gt_channel_selector, + bce_coef=config["loss_params"]["seg_loss_params"]["bce_coef"], + dice_coef=config["loss_params"]["seg_loss_params"]["dice_coef"])) + weights.append(config["loss_params"]["multiloss"]["coefs"]["seg"]) + + if config["compute_crossfield"]: + pre_processes.append(compute_gt_field) + loss_funcs.append(CrossfieldAlignLoss(name="crossfield_align")) + weights.append(config["loss_params"]["multiloss"]["coefs"]["crossfield_align"]) + loss_funcs.append(CrossfieldAlign90Loss(name="crossfield_align90")) + weights.append(config["loss_params"]["multiloss"]["coefs"]["crossfield_align90"]) + loss_funcs.append(CrossfieldSmoothLoss(name="crossfield_smooth")) + weights.append(config["loss_params"]["multiloss"]["coefs"]["crossfield_smooth"]) + + # --- Coupling losses: + if config["compute_seg"]: + need_seg_grads = False + pred_channel = -1 + # Seg interior <-> Crossfield coupling: + if config["seg_params"]["compute_interior"] and config["compute_crossfield"]: + need_seg_grads = True + pred_channel += 1 + loss_funcs.append(SegCrossfieldLoss(name="seg_interior_crossfield", pred_channel=pred_channel)) + weights.append(config["loss_params"]["multiloss"]["coefs"]["seg_interior_crossfield"]) + # Seg edge <-> Crossfield coupling: + if config["seg_params"]["compute_edge"] and config["compute_crossfield"]: + need_seg_grads = True + pred_channel += 1 + loss_funcs.append(SegCrossfieldLoss(name="seg_edge_crossfield", pred_channel=pred_channel)) + weights.append(config["loss_params"]["multiloss"]["coefs"]["seg_edge_crossfield"]) + + # Seg edge <-> seg interior coupling: + if config["seg_params"]["compute_interior"] and config["seg_params"]["compute_edge"]: + need_seg_grads = True + loss_funcs.append(SegEdgeInteriorLoss(name="seg_edge_interior")) + weights.append(config["loss_params"]["multiloss"]["coefs"]["seg_edge_interior"]) + + if need_seg_grads: + pre_processes.append(ComputeSegGrads(config["device"])) + + combined_loss = MultiLoss(loss_funcs, weights, epoch_thresholds=config["loss_params"]["multiloss"]["coefs"]["epoch_thresholds"], pre_processes=pre_processes) + return combined_loss + + +# --- Specific losses --- # +class SegLoss(Loss): + def __init__(self, name, gt_channel_selector, bce_coef=0.5, dice_coef=0.5): + """ + :param name: + :param gt_channel_selector: used to select which channels gt_polygons_image to use to compare to predicted seg + (see docstring of method compute() for more details). + """ + super(SegLoss, self).__init__(name) + self.gt_channel_selector = gt_channel_selector + self.bce_coef = bce_coef + self.dice_coef = dice_coef + + def compute(self, pred_batch, gt_batch): + """ + seg and gt_polygons_image do not necessarily have the same channel count. + gt_selector is used to select which channels of gt_polygons_image to use. + For example, if seg has C_pred=2 (interior and edge) and + gt_polygons_image has C_gt=3 (interior, edge and vertex), use gt_channel_selector=slice(0, 2) + + @param pred_batch: key "seg" is shape (N, C_pred, H, W) + @param gt_batch: key "gt_polygons_image" is shape (N, C_gt, H, W) + @return: + """ + # print(self.name) + pred_seg = pred_batch["seg"] + gt_seg = gt_batch["gt_polygons_image"][:, self.gt_channel_selector, ...] + weights = gt_batch["seg_loss_weights"][:, self.gt_channel_selector, ...] + dice = measures.dice_loss(pred_seg, gt_seg) + mean_dice = torch.mean(dice) + mean_cross_entropy = F.binary_cross_entropy(pred_seg, gt_seg, weight=weights, reduction="mean") + + # Display: + # dispaly_pred_seg = pred_seg[0, 0].cpu().detach().numpy() + # print(f'{self.name}_pred:', dispaly_pred_seg.shape, dispaly_pred_seg.min(), dispaly_pred_seg.max()) + # skimage.io.imsave(f'{self.name}_pred.png', dispaly_pred_seg) + # dispaly_gt_seg = gt_seg[0].cpu().detach().numpy() + # skimage.io.imsave(f'{self.name}_gt.png', dispaly_gt_seg) + + return self.bce_coef * mean_cross_entropy + self.dice_coef * mean_dice + + +class CrossfieldAlignLoss(Loss): + def __init__(self, name): + super(CrossfieldAlignLoss, self).__init__(name) + + def compute(self, pred_batch, gt_batch): + c0 = pred_batch["crossfield"][:, :2] + c2 = pred_batch["crossfield"][:, 2:] + z = gt_batch["gt_field"] + gt_polygons_image = gt_batch["gt_polygons_image"] + assert 2 <= gt_polygons_image.shape[1], \ + "gt_polygons_image should have at least 2 channels for interior and edges" + gt_edges = gt_polygons_image[:, 1, ...] + align_loss = frame_field_utils.framefield_align_error(c0, c2, z, complex_dim=1) + avg_align_loss = torch.mean(align_loss * gt_edges) + + self.extra_info["gt_field"] = gt_batch["gt_field"] + return avg_align_loss + + +class CrossfieldAlign90Loss(Loss): + def __init__(self, name): + super(CrossfieldAlign90Loss, self).__init__(name) + + def compute(self, pred_batch, gt_batch): + c0 = pred_batch["crossfield"][:, :2] + c2 = pred_batch["crossfield"][:, 2:] + z = gt_batch["gt_field"] + z_90deg = torch.cat((- z[:, 1:2, ...], z[:, 0:1, ...]), dim=1) + gt_polygons_image = gt_batch["gt_polygons_image"] + assert gt_polygons_image.shape[1] == 3, \ + "gt_polygons_image should have 3 channels for interior, edges and vertices" + gt_edges = gt_polygons_image[:, 1, ...] + gt_vertices = gt_polygons_image[:, 2, ...] + gt_edges_minus_vertices = gt_edges - gt_vertices + gt_edges_minus_vertices = gt_edges_minus_vertices.clamp(0, 1) + align90_loss = frame_field_utils.framefield_align_error(c0, c2, z_90deg, complex_dim=1) + avg_align90_loss = torch.mean(align90_loss * gt_edges_minus_vertices) + return avg_align90_loss + + +class CrossfieldSmoothLoss(Loss): + def __init__(self, name): + super(CrossfieldSmoothLoss, self).__init__(name) + self.laplacian_penalty = frame_field_utils.LaplacianPenalty(channels=4) + + def compute(self, pred_batch, gt_batch): + c0c2 = pred_batch["crossfield"] + gt_polygons_image = gt_batch["gt_polygons_image"] + gt_edges_inv = 1 - gt_polygons_image[:, 1, ...] + penalty = self.laplacian_penalty(c0c2) + avg_penalty = torch.mean(penalty * gt_edges_inv[:, None, ...]) + return avg_penalty + + +class SegCrossfieldLoss(Loss): + def __init__(self, name, pred_channel): + super(SegCrossfieldLoss, self).__init__(name) + self.pred_channel = pred_channel + + def compute(self, pred_batch, gt_batch): + # TODO: don't apply on corners: corner_map = gt_batch["gt_polygons_image"][:, 2, :, :] + # TODO: apply on all seg at once? Like seg is now? + c0 = pred_batch["crossfield"][:, :2] + c2 = pred_batch["crossfield"][:, 2:] + seg_slice_grads_normed = pred_batch["seg_grads_normed"][:, self.pred_channel, ...] + seg_slice_grad_norm = pred_batch["seg_grad_norm"][:, self.pred_channel, ...] + align_loss = frame_field_utils.framefield_align_error(c0, c2, seg_slice_grads_normed, complex_dim=1) + # normed_align_loss = align_loss * seg_slice_grad_norm + # avg_align_loss = torch.sum(normed_align_loss) / (torch.sum(seg_slice_grad_norm) + 1e-6) + avg_align_loss = torch.mean(align_loss * seg_slice_grad_norm.detach()) + # (prev line) Don't back-propagate to seg_slice_grad_norm so that seg smoothness is not encouraged + + # Save extra info for viz: + self.extra_info["seg_slice_grads"] = pred_batch["seg_grads"][:, self.pred_channel, ...] + return avg_align_loss + + +class SegEdgeInteriorLoss(Loss): + """ + Enforce seg edge to be equal to interior grad norm except inside buildings + """ + + def __init__(self, name): + super(SegEdgeInteriorLoss, self).__init__(name) + + def compute(self, pred_batch, batch): + seg_interior = pred_batch["seg"][:, 0, ...] + seg_edge = pred_batch["seg"][:, 1, ...] + seg_interior_grad_norm = pred_batch["seg_grad_norm"][:, 0, ...] + raw_loss = torch.abs(seg_edge - seg_interior_grad_norm) + # Apply the loss only on interior boundaries and outside of objects + outside_mask = (torch.cos(np.pi * seg_interior) + 1) / 2 + boundary_mask = (1 - torch.cos(np.pi * seg_interior_grad_norm)) / 2 + mask = torch.max(outside_mask, boundary_mask).float() + avg_loss = torch.mean(raw_loss * mask) + return avg_loss diff --git a/frame_field_learning/measures.py b/frame_field_learning/measures.py new file mode 100644 index 0000000000000000000000000000000000000000..f81ca531a4a8c66f1118b8e32b41ac42cdf431d5 --- /dev/null +++ b/frame_field_learning/measures.py @@ -0,0 +1,58 @@ +import torch + + +def iou(y_pred, y_true, threshold): + assert len(y_pred.shape) == len(y_true.shape) == 2, "Input tensor shapes should be (N, .)" + mask_pred = threshold < y_pred + mask_true = threshold < y_true + intersection = torch.sum(mask_pred * mask_true, dim=-1) + union = torch.sum(mask_pred + mask_true, dim=-1) + r = intersection.float() / union.float() + r[union == 0] = 1 + return r + + +def dice_loss(y_pred, y_true, smooth=1, eps=1e-7): + """ + + @param y_pred: (N, C, H, W) + @param y_true: (N, C, H, W) + @param smooth: + @param eps: + @return: (N, C) + """ + numerator = 2 * torch.sum(y_true * y_pred, dim=(-1, -2)) + denominator = torch.sum(y_true, dim=(-1, -2)) + torch.sum(y_pred, dim=(-1, -2)) + return 1 - (numerator + smooth) / (denominator + smooth + eps) + + +def main(): + # import kornia + # spatial_gradient_function = kornia.filters.SpatialGradient() + # + # image = torch.zeros((7, 7)) + # image[2:5, 2:5] = 1 + # print(image) + # + # grads = spatial_gradient_function(image[None, None, ...])[0, 0, ...] / 4 + # print(grads[0]) + # print(grads[1]) + + y_true = torch.tensor([ + [0, 0, 0, 0, 1, 1, 1], + [0, 0, 0, 0, 1, 1, 1], + [0, 0, 0, 0, 0, 0, 0] + ]) + y_pred = torch.tensor([ + [0, 0, 0, 0, 1, 1, 1], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + ]) + print(y_true.shape) + print(y_pred.shape) + r = iou(y_pred, y_true, threshold=0.5) + print(r) + + +if __name__ == "__main__": + main() diff --git a/frame_field_learning/model.py b/frame_field_learning/model.py new file mode 100644 index 0000000000000000000000000000000000000000..753a3ddecc43176810dc017e17f4baf0b56ac035 --- /dev/null +++ b/frame_field_learning/model.py @@ -0,0 +1,108 @@ +import torch +from torchvision.models.segmentation._utils import _SimpleSegmentationModel +# from pytorch_memlab import profile, profile_every +from frame_field_learning import tta_utils + + +def get_out_channels(module): + if hasattr(module, "out_channels"): + return module.out_channels + children = list(module.children()) + i = 1 + out_channels = None + while out_channels is None and i <= len(children): + last_child = children[-i] + out_channels = get_out_channels(last_child) + i += 1 + # If we get out of the loop but out_channels is None, then the prev child of the parent module will be checked, etc. + return out_channels + + +class FrameFieldModel(torch.nn.Module): + def __init__(self, config: dict, backbone, train_transform=None, eval_transform=None): + """ + + :param config: + :param backbone: A _SimpleSegmentationModel network, its output features will be used to compute seg and framefield. + :param train_transform: transform applied to the inputs when self.training is True + :param eval_transform: transform applied to the inputs when self.training is False + """ + super(FrameFieldModel, self).__init__() + assert config["compute_seg"] or config["compute_crossfield"], \ + "Model has to compute at least one of those:\n" \ + "\t- segmentation\n" \ + "\t- cross-field" + assert isinstance(backbone, _SimpleSegmentationModel), \ + "backbone should be an instance of _SimpleSegmentationModel" + self.config = config + self.backbone = backbone + self.train_transform = train_transform + self.eval_transform = eval_transform + + backbone_out_features = get_out_channels(self.backbone) + + # --- Add other modules if activated in config: + seg_channels = 0 + if self.config["compute_seg"]: + seg_channels = self.config["seg_params"]["compute_vertex"]\ + + self.config["seg_params"]["compute_edge"]\ + + self.config["seg_params"]["compute_interior"] + self.seg_module = torch.nn.Sequential( + torch.nn.Conv2d(backbone_out_features, backbone_out_features, 3, padding=1), + torch.nn.BatchNorm2d(backbone_out_features), + torch.nn.ELU(), + torch.nn.Conv2d(backbone_out_features, seg_channels, 1), + torch.nn.Sigmoid(),) + + if self.config["compute_crossfield"]: + crossfield_channels = 4 + self.crossfield_module = torch.nn.Sequential( + torch.nn.Conv2d(backbone_out_features + seg_channels, backbone_out_features, 3, padding=1), + torch.nn.BatchNorm2d(backbone_out_features), + torch.nn.ELU(), + torch.nn.Conv2d(backbone_out_features, crossfield_channels, 1), + torch.nn.Tanh(), + ) + + def inference(self, image): + outputs = {} + + # --- Extract features for every pixel of the image with a U-Net --- # + backbone_features = self.backbone(image)["out"] + + if self.config["compute_seg"]: + # --- Output a segmentation of the image --- # + seg = self.seg_module(backbone_features) + seg_to_cat = seg.clone().detach() + backbone_features = torch.cat([backbone_features, seg_to_cat], dim=1) # Add seg to image features + outputs["seg"] = seg + + if self.config["compute_crossfield"]: + # --- Output a cross-field of the image --- # + crossfield = 2 * self.crossfield_module(backbone_features) # Outputs c_0, c_2 values in [-2, 2] + outputs["crossfield"] = crossfield + + return outputs + + # @profile + def forward(self, xb, tta=False): + # print("\n### --- PolyRefine.forward(xb) --- ####") + if self.training: + if self.train_transform is not None: + xb = self.train_transform(xb) + else: + if self.eval_transform is not None: + xb = self.eval_transform(xb) + + if not tta: + final_outputs = self.inference(xb["image"]) + else: + final_outputs = tta_utils.tta_inference(self, xb, self.config["eval_params"]["seg_threshold"]) + + # # Save image + # image_seg_display = plot_utils.get_tensorboard_image_seg_display(image_display, final_outputs["seg"], + # crossfield=final_outputs["crossfield"]) + # image_seg_display = image_seg_display[1].cpu().detach().numpy().transpose(1, 2, 0) + # skimage.io.imsave(f"out_final.png", image_seg_display) + + return final_outputs, xb diff --git a/frame_field_learning/plot_utils.py b/frame_field_learning/plot_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..658fdd1f4c5f7df0e91f6fd6a129d8a85a425361 --- /dev/null +++ b/frame_field_learning/plot_utils.py @@ -0,0 +1,265 @@ +import random + +import skimage.io +from descartes import PolygonPatch +from matplotlib.collections import PatchCollection +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.figure import Figure +import matplotlib.pyplot as plt + +import numpy as np +import torch +import shapely.geometry + +from lydorn_utils import math_utils +from torch_lydorn import torchvision + + +def get_seg_display(seg): + dtype = seg.dtype + seg_display = np.zeros([seg.shape[0], seg.shape[1], 4], dtype=dtype) + if len(seg.shape) == 2: + seg_display[..., 0] = seg + seg_display[..., 3] = seg + else: + for i in range(seg.shape[-1]): + seg_display[..., i] = seg[..., i] + clip_max = 255 if dtype == np.uint8 else 1 + seg_display[..., 3] = np.clip(np.sum(seg, axis=-1), 0, clip_max) + return seg_display + + +def get_tensorboard_image_seg_display(image, seg, crossfield=None): + assert len(image.shape) == 4 and image.shape[1] == 3, f"image should be (N, 3, H, W), not {image.shape}." + assert len(seg.shape) == 4 and seg.shape[1] <= 3, f"image should be (N, C<=3, H, W), not {seg.shape}." + assert image.shape[0] == seg.shape[0], "image and seg should have the same batch size." + assert image.shape[2] == seg.shape[2], "image and seg should have the same image height." + assert image.shape[3] == seg.shape[3], "image and seg should have the same image width." + if crossfield is not None: + assert len(crossfield.shape) == 4 and crossfield.shape[ + 1] == 4, f"crossfield should be (N, 4, H, W), not {crossfield.shape}." + assert image.shape[0] == crossfield.shape[0], "image and crossfield should have the same batch size." + assert image.shape[2] == crossfield.shape[2], "image and crossfield should have the same image height." + assert image.shape[3] == crossfield.shape[3], "image and crossfield should have the same image width." + + alpha = torch.clamp(torch.sum(seg, dim=1, keepdim=True), 0, 1) + + # Add missing seg channels + seg_display = torch.zeros_like(image) + seg_display[:, :seg.shape[1], ...] = seg + + image_seg_display = (1 - alpha) * image + alpha * seg_display + image_seg_display = image_seg_display.cpu() + + if crossfield is not None: + np_crossfield = crossfield.cpu().detach().numpy().transpose(0, 2, 3, 1) + image_plot_crossfield_list = [get_image_plot_crossfield(_crossfield, crossfield_stride=10) for _crossfield in + np_crossfield] + image_plot_crossfield_list = [torchvision.transforms.functional.to_tensor(image_plot_crossfield).float() / 255 + for image_plot_crossfield in image_plot_crossfield_list] + image_plot_crossfield = torch.stack(image_plot_crossfield_list, dim=0) + alpha = image_plot_crossfield[:, 3:4, :, :] + image_seg_display = (1 - alpha) * image_seg_display + alpha * image_plot_crossfield[:, :3, :, :] + # image_seg_display = image_plot_crossfield[:, :3, :, :] + + return image_seg_display + + +def plot_crossfield(axis, crossfield, crossfield_stride, alpha=0.5, width=0.5, add_scale=1, invert_y=True): + x = np.arange(0, crossfield.shape[1], crossfield_stride) + y = np.arange(0, crossfield.shape[0], crossfield_stride) + x, y = np.meshgrid(x, y) + i = y + if invert_y: + i = crossfield.shape[0] - 1 - y + j = x + scale = add_scale * 1 / crossfield_stride + + c0c2 = crossfield[i, j, :] + u, v = math_utils.compute_crossfield_uv(c0c2) + + # u_angle = 0.5 + # u.real = np.cos(u_angle) + # u.imag = np.sin(u_angle) + # v *= 0 + + quiveropts = dict(color=(0, 0, 1, alpha), headaxislength=0, headlength=0, pivot='middle', angles="xy", units='xy', + scale=scale, width=width, headwidth=1) + axis.quiver(x, y, u.imag, -u.real, **quiveropts) + axis.quiver(x, y, v.imag, -v.real, **quiveropts) + + +def get_image_plot_crossfield(crossfield, crossfield_stride): + fig = Figure(figsize=(crossfield.shape[1] / 100, crossfield.shape[0] / 100), dpi=100) + canvas = FigureCanvasAgg(fig) + ax = fig.gca() + + plot_crossfield(ax, crossfield, crossfield_stride, alpha=1.0, width=2.0, add_scale=1) + + ax.axis('off') + fig.tight_layout(pad=0) + # To remove the huge white borders + ax.margins(0) + + canvas.draw() + image_from_plot = np.frombuffer(canvas.tostring_argb(), dtype=np.uint8) + image_from_plot = image_from_plot.reshape(canvas.get_width_height()[::-1] + (4,)) + image_from_plot = np.roll(image_from_plot, -1, axis=-1) # Convert ARGB to RGBA + + # Fix alpha (white to alpha) + # mask = np.sum(image_from_plot[:, :, :3], axis=2) == 3*255 + # image_from_plot[mask, 3] = 0 + mini = image_from_plot.min() + image_from_plot[:, :, 3] = np.max(255 - image_from_plot[:, :, :3] + mini, axis=2) + + return image_from_plot + + +def plot_polygons(axis, polygons, polygon_probs=None, draw_vertices=True, linewidths=2, markersize=10, alpha=0.2, + color_choices=None): + if len(polygons) == 0: + return + patches = [] + for i, geometry in enumerate(polygons): + polygon = shapely.geometry.Polygon(geometry) + if not polygon.is_empty: + patch = PolygonPatch(polygon) + patches.append(patch) + random.seed(1) + if color_choices is None: + color_choices = [ + [0, 0, 1, 1], + [0, 1, 0, 1], + [1, 0, 0, 1], + [1, 1, 0, 1], + [1, 0, 1, 1], + [0, 1, 1, 1], + [0.5, 1, 0, 1], + [1, 0.5, 0, 1], + [0.5, 0, 1, 1], + [1, 0, 0.5, 1], + [0, 0.5, 1, 1], + [0, 1, 0.5, 1], + ] + colors = random.choices(color_choices, k=len(patches)) + edgecolors = np.array(colors, dtype=np.float) + facecolors = edgecolors.copy() + if polygon_probs is not None: + facecolors[:, -1] = alpha * np.array(polygon_probs) + 0.1 + else: + facecolors[:, -1] = alpha + p = PatchCollection(patches, facecolors=facecolors, edgecolors=edgecolors, linewidths=linewidths) + axis.add_collection(p) + + if draw_vertices: + for i, polygon in enumerate(polygons): + axis.plot(*polygon.exterior.xy, marker="o", color=edgecolors[i], markersize=markersize) + for interior in polygon.interiors: + axis.plot(*interior.xy, marker="o", color=edgecolors[i], markersize=markersize) + + +def plot_line_strings(axis, line_strings, draw_vertices=True, linewidths=2, markersize=5): + artists = [] + marker = "o" if draw_vertices else None + for line_string in line_strings: + artist, = axis.plot(*line_string.xy, marker=marker, markersize=markersize) + artists.append(artist) + return artists + + +def plot_geometries(axis, geometries, draw_vertices=True, linewidths=2, markersize=3): + polygons = [] + line_strings = [] + for geometry in geometries: + if isinstance(geometry, shapely.geometry.Polygon): + polygons.append(geometry) + elif isinstance(geometry, shapely.geometry.LineString): + line_strings.append(geometry) + elif isinstance(geometry, shapely.geometry.MultiLineString): + for line_string in geometry: + line_strings.append(line_string) + else: + raise NotImplementedError(f"Geometry type {type(geometry)} not implemented") + + if len(polygons): + plot_polygons(axis, polygons, draw_vertices=draw_vertices, linewidths=linewidths, markersize=markersize) + + if len(line_strings): + artists = plot_line_strings(axis, line_strings, draw_vertices=draw_vertices, linewidths=linewidths, markersize=markersize) + return artists + + +def save_poly_viz(image, polygons, out_filepath, linewidths=2, markersize=20, alpha=0.2, draw_vertices=True, + corners=None, crossfield=None, polygon_probs=None, seg=None, color_choices=None, dpi=10): + assert isinstance(polygons, list), f"polygons should be of type list, not {type(polygons)}" + if len(polygons): + assert (type(polygons[0]) == np.ndarray or type(polygons[0]) == shapely.geometry.Polygon), \ + f"Item of the polygons list should be of type ndarray or shapely Polygon, not {type(polygons[0])}" + if polygon_probs is not None: + assert type(polygon_probs) == list + assert len(polygons) == len(polygon_probs), \ + "len(polygons)={} should be equal to len(polygon_probs)={}".format(len(polygons), len(polygon_probs)) + # Setup plot + height = image.shape[0] + width = image.shape[1] + f, axis = plt.subplots(1, 1, figsize=(width / 10, height / 10), dpi=10) + + axis.imshow(image) + + if seg is not None: + seg *= 0.9 + axis.imshow(seg) + + if crossfield is not None: + plot_crossfield(axis, crossfield, crossfield_stride=1, alpha=0.5, width=0.1, add_scale=1.1, invert_y=False) + + plot_polygons(axis, polygons, polygon_probs=polygon_probs, draw_vertices=draw_vertices, linewidths=linewidths, + markersize=markersize, alpha=alpha, color_choices=color_choices) + + if corners is not None and len(corners): + assert len(corners[0].shape) == 2 + for corner_array in corners: + plt.plot(corner_array[:, 0], corner_array[:, 1], marker="o", linewidth=0, markersize=20, color="red") + + axis.autoscale(False) + axis.axis('equal') + axis.axis('off') + plt.subplots_adjust(left=0, right=1, top=1, bottom=0) # Plot without margins + plt.savefig(out_filepath, transparent=True, dpi=dpi) + plt.close() + + +def main(): + image = torch.zeros((2, 3, 512, 512)) + 0.5 + seg = torch.zeros((2, 2, 512, 512)) + seg[:, 0, 100:200, 100:200] = 1 + crossfield = torch.zeros((2, 4, 512, 512)) + # u_angle = np.random.random(10000) * np.pi + # v_angle = np.random.random(10000) * np.pi + u_angle = 0.25 + v_angle = u_angle + np.pi / 2 + u = np.cos(u_angle) + 1j * np.sin(u_angle) + v = np.cos(v_angle) + 1j * np.sin(v_angle) + c0 = np.power(u, 2) * np.power(v, 2) + c2 = - (np.power(u, 2) + np.power(v, 2)) + # print("c0:") + # print(np.abs(c0).min(), np.abs(c0).mean(), np.abs(c0).max()) + # print(c0.real.min(), c0.real.mean(), c0.real.mean()) + # print(c0.imag.min(), c0.imag.mean(), c0.imag.max()) + # print("c2:") + # print(np.abs(c2).min(), np.abs(c2).mean(), np.abs(c2).max()) + # print(c2.real.min(), c2.real.mean(), c2.real.max()) + # print(c2.real.min(), c2.imag.mean(), c2.imag.max()) + + crossfield[:, 0, :, :] = c0.real + crossfield[:, 1, :, :] = c0.imag + crossfield[:, 2, :, :] = c2.real + crossfield[:, 3, :, :] = c2.imag + + image_seg_display = get_tensorboard_image_seg_display(image, seg, crossfield=crossfield) + image_seg_display = image_seg_display.cpu().numpy().transpose(0, 2, 3, 1) + skimage.io.imsave("image_seg_display.png", image_seg_display[0]) + + +if __name__ == "__main__": + main() diff --git a/frame_field_learning/polygonize.py b/frame_field_learning/polygonize.py new file mode 100644 index 0000000000000000000000000000000000000000..de4b1e34ae8c1d383e05d7ec61df4a870b71832a --- /dev/null +++ b/frame_field_learning/polygonize.py @@ -0,0 +1,82 @@ +from . import polygonize_utils +from . import polygonize_acm +from . import polygonize_simple + +from lydorn_utils import print_utils + + +class Polygonizer(): + def __init__(self, polygonize_params, pool=None): + self.pool = pool + self.polygonizer_asm = None + + 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": + from . import polygonize_asm + if crossfield_batch is None: + # Cannot run the ASM method + return None, None + if self.polygonizer_asm is None: + self.polygonizer_asm = polygonize_asm.PolygonizerASM(polygonize_params["asm_method"], pool=self.pool) + 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 polygonize(polygonize_params, seg_batch, crossfield_batch=None, pool=None, pre_computed=None): + polygonizer = Polygonizer(polygonize_params, pool=pool) + return polygonizer(polygonize_params, seg_batch, crossfield_batch=crossfield_batch, pre_computed=pre_computed) diff --git a/frame_field_learning/polygonize_acm.py b/frame_field_learning/polygonize_acm.py new file mode 100644 index 0000000000000000000000000000000000000000..9ee71be8cc6e8a8e9d0136aa5bd426d2134199bf --- /dev/null +++ b/frame_field_learning/polygonize_acm.py @@ -0,0 +1,596 @@ +import argparse +import fnmatch +import time + +import numpy as np +import skimage +import skimage.measure +import skimage.io +from tqdm import tqdm +import shapely.geometry +import shapely.ops +import shapely.prepared +import cv2 + +from functools import partial + +import torch + +from frame_field_learning import polygonize_utils +from frame_field_learning import frame_field_utils + +from torch_lydorn.torch.nn.functionnal import bilinear_interpolate +from torch_lydorn.torchvision.transforms import polygons_to_tensorpoly, tensorpoly_pad + +from lydorn_utils import math_utils +from lydorn_utils import python_utils +from lydorn_utils import print_utils + + +DEBUG = False + + +def debug_print(s: str): + if DEBUG: + print_utils.print_debug(s) + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--raw_pred', + nargs='*', + type=str, + help='Filepath to the raw pred file(s)') + argparser.add_argument( + '--im_filepath', + type=str, + help='Filepath to input image. Will retrieve seg and crossfield in the same directory') + argparser.add_argument( + '--dirpath', + type=str, + help='Path to directory containing seg and crossfield files. Will perform polygonization on all.') + argparser.add_argument( + '--bbox', + nargs='*', + type=int, + help='Selects area in bbox for computation: [min_row, min_col, max_row, max_col]') + argparser.add_argument( + '--steps', + type=int, + help='Optim steps') + + args = argparser.parse_args() + return args + + +class PolygonAlignLoss: + def __init__(self, indicator, level, c0c2, data_coef, length_coef, crossfield_coef, dist=None, dist_coef=None): + self.indicator = indicator + self.level = level + self.c0c2 = c0c2 + self.dist = dist + + self.data_coef = data_coef + self.length_coef = length_coef + self.crossfield_coef = crossfield_coef + self.dist_coef = dist_coef + + def __call__(self, tensorpoly): + """ + + :param tensorpoly: closed polygon + :return: + """ + polygon = tensorpoly.pos[tensorpoly.to_padded_index] + polygon_batch = tensorpoly.batch[tensorpoly.to_padded_index] + + # Compute edges: + edges = polygon[1:] - polygon[:-1] + # Compute edge mask to remove edges that connect two different polygons from loss + # Also note the last poly_slice is not used, because the last edge of the last polygon is not connected to a non-existant next polygon: + edge_mask = torch.ones((edges.shape[0]), device=edges.device) + edge_mask[tensorpoly.to_unpadded_poly_slice[:-1, 1]] = 0 + + midpoints = (polygon[1:] + polygon[:-1]) / 2 + midpoints_batch = polygon_batch[1:] + + midpoints_int = midpoints.round().long() + midpoints_int[:, 0] = torch.clamp(midpoints_int[:, 0], 0, self.c0c2.shape[2] - 1) + midpoints_int[:, 1] = torch.clamp(midpoints_int[:, 1], 0, self.c0c2.shape[3] - 1) + midpoints_c0 = self.c0c2[midpoints_batch, :2, midpoints_int[:, 0], midpoints_int[:, 1]] + midpoints_c2 = self.c0c2[midpoints_batch, 2:, midpoints_int[:, 0], midpoints_int[:, 1]] + + norms = torch.norm(edges, dim=-1) + # Add edges with small norms to the edge mask so that losses are not computed on them + edge_mask[norms < 0.1] = 0 # Less than 10% of a pixel + z = edges / (norms[:, None] + 1e-3) + + # Align to crossfield + align_loss = frame_field_utils.framefield_align_error(midpoints_c0, midpoints_c2, z, complex_dim=1) + align_loss = align_loss * edge_mask + total_align_loss = torch.sum(align_loss) + + # Align to level set of indicator: + pos_indicator_value = bilinear_interpolate(self.indicator[:, None, ...], tensorpoly.pos, batch=tensorpoly.batch) + # TODO: Try to use grid_sample with batch for speed: put batch dim to height dim and make a single big image. + # TODO: Convert pos accordingly and take care of borders + # height = self.indicator.shape[1] + # width = self.indicator.shape[2] + # normed_xy = tensorpoly.pos.roll(shifts=1, dims=-1) + # normed_xy[: 0] /= (width-1) + # normed_xy[: 1] /= (height-1) + # centered_xy = 2*normed_xy - 1 + # pos_value = torch.nn.functional.grid_sample(self.indicator[None, None, ...], centered_batch_xy[None, None, ...], align_corners=True).squeeze() + level_loss = torch.sum(torch.pow(pos_indicator_value - self.level, 2)) + + # Align to minimum distance from the boundary + dist_loss = None + if self.dist is not None: + pos_dist_value = bilinear_interpolate(self.dist[:, None, ...], tensorpoly.pos, batch=tensorpoly.batch) + dist_loss = torch.sum(torch.pow(pos_dist_value, 2)) + + length_penalty = torch.sum( + torch.pow(norms * edge_mask, 2)) # Sum of squared norm to penalise uneven edge lengths + # length_penalty = torch.sum(norms) + + losses_dict = { + "align": total_align_loss.item(), + "level": level_loss.item(), + "length": length_penalty.item(), + } + coef_sum = self.data_coef + self.length_coef + self.crossfield_coef + total_loss = (self.data_coef * level_loss + self.length_coef * length_penalty + self.crossfield_coef * total_align_loss) + if dist_loss is not None: + losses_dict["dist"] = dist_loss.item() + total_loss += self.dist_coef * dist_loss + coef_sum += self.dist_coef + total_loss /= coef_sum + return total_loss, losses_dict + + +class TensorPolyOptimizer: + def __init__(self, config, tensorpoly, indicator, c0c2, data_coef, length_coef, crossfield_coef, dist=None, dist_coef=None): + assert len(indicator.shape) == 3, "indicator: (N, H, W)" + assert len(c0c2.shape) == 4 and c0c2.shape[1] == 4, "c0c2: (N, 4, H, W)" + if dist is not None: + assert len(dist.shape) == 3, "dist: (N, H, W)" + + + self.config = config + self.tensorpoly = tensorpoly + + # Require grads for graph.pos: this is what is optimized + self.tensorpoly.pos.requires_grad = True + + # Save pos of endpoints so that they can be reset after each step (endpoints are not meant to be moved) + self.endpoint_pos = self.tensorpoly.pos[self.tensorpoly.is_endpoint].clone() + + self.criterion = PolygonAlignLoss(indicator, config["data_level"], c0c2, data_coef, length_coef, + crossfield_coef, dist=dist, dist_coef=dist_coef) + self.optimizer = torch.optim.SGD([tensorpoly.pos], lr=config["poly_lr"]) + + def lr_warmup_func(iter): + if iter < config["warmup_iters"]: + coef = 1 + (config["warmup_factor"] - 1) * (config["warmup_iters"] - iter) / config["warmup_iters"] + else: + coef = 1 + return coef + + self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lr_warmup_func) + + def step(self, iter_num): + self.optimizer.zero_grad() + loss, losses_dict = self.criterion(self.tensorpoly) + # print("loss:", loss.item()) + loss.backward() + # print(polygon_tensor.grad[0]) + self.optimizer.step() + self.lr_scheduler.step(iter_num) + + # Move endpoints back: + with torch.no_grad(): + self.tensorpoly.pos[self.tensorpoly.is_endpoint] = self.endpoint_pos + return loss.item(), losses_dict + + def optimize(self): + # if DEBUG: + # optim_iter = tqdm(range(self.config["steps"]), desc="Gradient descent", leave=True) + # else: + # optim_iter = range(self.config["steps"]) + # # print("---------------------------------------------") + # for iter_num in optim_iter: + # loss, losses_dict = self.step(iter_num) + # if DEBUG: + # optim_iter.set_postfix(loss=loss, **losses_dict) + optim_iter = range(self.config["steps"]) + for iter_num in optim_iter: + loss, losses_dict = self.step(iter_num) + return self.tensorpoly + + +def contours_batch_to_tensorpoly(contours_batch): + # Convert a batch of contours to a TensorPoly representation with PyTorch tensors + tensorpoly = polygons_to_tensorpoly(contours_batch) + # Pad contours so that we can treat them as closed: + tensorpoly = tensorpoly_pad(tensorpoly, padding=(0, 1)) + return tensorpoly + + +def tensorpoly_to_contours_batch(tensorpoly): + # Convert back to contours + contours_batch = [[] for _ in range(tensorpoly.batch_size)] + for poly_i in range(tensorpoly.poly_slice.shape[0]): + s = tensorpoly.poly_slice[poly_i, :] + contour = np.array(tensorpoly.pos[s[0]:s[1], :].detach().cpu()) + is_open = tensorpoly.is_endpoint[s[0]] # Is open = if first vertex is an endpoint + if not is_open: + # Close contour + contour = np.concatenate([contour, contour[:1, :]], axis=0) + batch_i = tensorpoly.batch[s[0]] # Batch of polygon = batch of first vertex + contours_batch[batch_i].append(contour) + return contours_batch + + +def print_contours_stats(contours): + min_length = contours[0].shape[0] + max_length = contours[0].shape[0] + nb_vertices = 0 + for contour in contours: + nb_vertices += contour.shape[0] + if contour.shape[0] < min_length: + min_length = contour.shape[0] + if max_length < contour.shape[0]: + max_length = contour.shape[0] + print("Nb polygon:", len(contours), "Nb vertices:", nb_vertices, "Min lengh:", min_length, "Max lengh:", max_length) + + +def shapely_postprocess(contours, u, v, np_indicator, tolerance, config): + if type(tolerance) == list: + # Use several tolerance values for simplification. return a dict with all results + out_polygons_dict = {} + out_probs_dict = {} + for tol in tolerance: + out_polygons, out_probs = shapely_postprocess(contours, u, v, np_indicator, tol, config) + out_polygons_dict["tol_{}".format(tol)] = out_polygons + out_probs_dict["tol_{}".format(tol)] = out_probs + return out_polygons_dict, out_probs_dict + else: + height = np_indicator.shape[0] + width = np_indicator.shape[1] + + # debug_print("Corner-aware simplification") + # Simplify contours a little to avoid some close-together corner-detection: + # TODO: handle close-together corners better + contours = [skimage.measure.approximate_polygon(contour, tolerance=min(1, tolerance)) for contour in contours] + corner_masks = frame_field_utils.detect_corners(contours, u, v) + contours = polygonize_utils.split_polylines_corner(contours, corner_masks) + + # Convert to Shapely: + line_string_list = [shapely.geometry.LineString(out_contour[:, ::-1]) for out_contour in contours] + + line_string_list = [line_string.simplify(tolerance, preserve_topology=True) for line_string in line_string_list] + + # Add image boundary line_strings for border polygons + line_string_list.append( + shapely.geometry.LinearRing([ + (0, 0), + (0, height - 1), + (width - 1, height - 1), + (width - 1, 0), + ])) + + # debug_print("Merge polylines") + + # Merge polylines (for border polygons): + multi_line_string = shapely.ops.unary_union(line_string_list) + + # debug_print("polygonize_full") + + # Find polygons: + polygons, dangles, cuts, invalids = shapely.ops.polygonize_full(multi_line_string) + polygons = list(polygons) + + # debug_print("Remove small polygons") + + # Remove small polygons + polygons = [polygon for polygon in polygons if + config["min_area"] < polygon.area] + + # debug_print("Remove low prob polygons") + + # Remove low prob polygons + filtered_polygons = [] + filtered_polygon_probs = [] + for polygon in polygons: + prob = polygonize_utils.compute_geom_prob(polygon, np_indicator) + # print("acm:", np_indicator.min(), np_indicator.mean(), np_indicator.max(), prob) + if config["seg_threshold"] < prob: + filtered_polygons.append(polygon) + filtered_polygon_probs.append(prob) + + return filtered_polygons, filtered_polygon_probs + + +def post_process(contours, np_seg, np_crossfield, config): + u, v = math_utils.compute_crossfield_uv(np_crossfield) # u, v are complex arrays + + np_indicator = np_seg[:, :, 0] + polygons, probs = shapely_postprocess(contours, u, v, np_indicator, config["tolerance"], config) + + return polygons, probs + + +def polygonize(seg_batch, crossfield_batch, config, pool=None, pre_computed=None): + tic_start = time.time() + + assert len(seg_batch.shape) == 4 and seg_batch.shape[ + 1] <= 3, "seg_batch should be (N, C, H, W) with C <= 3, not {}".format(seg_batch.shape) + assert len(crossfield_batch.shape) == 4 and crossfield_batch.shape[ + 1] == 4, "crossfield_batch should be (N, 4, H, W)" + assert seg_batch.shape[0] == crossfield_batch.shape[0], "Batch size for seg and crossfield should match" + + + # Indicator + # tic = time.time() + indicator_batch = seg_batch[:, 0, :, :] + np_indicator_batch = indicator_batch.cpu().numpy() + indicator_batch = indicator_batch.to(config["device"]) + # toc = time.time() + # debug_print(f"Indicator to cpu: {toc - tic}s") + + # Distance image + dist_batch = None + if "dist_coef" in config: + # tic = time.time() + np_dist_batch = np.empty(np_indicator_batch.shape) + for batch_i in range(np_indicator_batch.shape[0]): + dist_1 = cv2.distanceTransform(np_indicator_batch[batch_i].astype(np.uint8), distanceType=cv2.DIST_L2, maskSize=cv2.DIST_MASK_5, dstType=cv2.CV_64F) + dist_2 = cv2.distanceTransform(1 - np_indicator_batch[batch_i].astype(np.uint8), distanceType=cv2.DIST_L2, maskSize=cv2.DIST_MASK_5, dstType=cv2.CV_64F) + np_dist_batch[0] = dist_1 + dist_2 - 1 + dist_batch = torch.from_numpy(np_dist_batch) + dist_batch = dist_batch.to(config["device"]) + # skimage.io.imsave("dist.png", np_dist_batch[0]) + # toc = time.time() + # debug_print(f"Distance image: {toc - tic}s") + + # debug_print("Init contours") + if pre_computed is None or "init_contours_batch" not in pre_computed: + # tic = time.time() + init_contours_batch = polygonize_utils.compute_init_contours_batch(np_indicator_batch, config["data_level"], pool=pool) + # toc = time.time() + # debug_print(f"Init contours: {toc - tic}s") + else: + init_contours_batch = pre_computed["init_contours_batch"] + + # debug_print("Convert contours to tensorpoly") + tensorpoly = contours_batch_to_tensorpoly(init_contours_batch) + + # debug_print("Optimize") + + # --- Optimize + # tic = time.time() + + tensorpoly.to(config["device"]) + crossfield_batch = crossfield_batch.to(config["device"]) + dist_coef = config["dist_coef"] if "dist_coef" in config else None + tensorpoly_optimizer = TensorPolyOptimizer(config, tensorpoly, indicator_batch, crossfield_batch, + config["data_coef"], + config["length_coef"], config["crossfield_coef"], dist=dist_batch, dist_coef=dist_coef) + tensorpoly = tensorpoly_optimizer.optimize() + + out_contours_batch = tensorpoly_to_contours_batch(tensorpoly) + + # toc = time.time() + # debug_print(f"Optimize contours: {toc - tic}s") + + # --- Post-process: + # debug_print("Post-process") + # tic = time.time() + + np_seg_batch = np.transpose(seg_batch.cpu().numpy(), (0, 2, 3, 1)) + np_crossfield_batch = np.transpose(crossfield_batch.cpu().numpy(), (0, 2, 3, 1)) + if pool is not None: + post_process_partial = partial(post_process, config=config) + polygons_probs_batch = pool.starmap(post_process_partial, zip(out_contours_batch, np_seg_batch, np_crossfield_batch)) + polygons_batch, probs_batch = zip(*polygons_probs_batch) + else: + polygons_batch = [] + probs_batch = [] + for i, out_contours in enumerate(out_contours_batch): + polygons, probs = post_process(out_contours, np_seg_batch[i], np_crossfield_batch[i], config) + polygons_batch.append(polygons) + probs_batch.append(probs) + + # toc = time.time() + # debug_print(f"Shapely post-process: {toc - tic}s") + + # toc = time.time() + # print(f"Post-process: {toc - tic}s") + # --- + + toc_end = time.time() + # debug_print(f"Total: {toc_end - tic_start}s") + + return polygons_batch, probs_batch + + +def main(): + from frame_field_learning import framefield, inference + import os + + def save_gt_poly(raw_pred_filepath, name): + filapth_format = "/data/mapping_challenge_dataset/processed/val/data_{}.pt" + sample = torch.load(filapth_format.format(name)) + polygon_arrays = sample["gt_polygons"] + polygons = [shapely.geometry.Polygon(polygon[:, ::-1]) for polygon in polygon_arrays] + base_filepath = os.path.join(os.path.dirname(raw_pred_filepath), name) + filepath = base_filepath + "." + name + ".pdf" + plot_utils.save_poly_viz(image, polygons, filepath) + + config = { + "indicator_add_edge": False, + "steps": 500, + "data_level": 0.5, + "data_coef": 0.1, + "length_coef": 0.4, + "crossfield_coef": 0.5, + "poly_lr": 0.01, + "warmup_iters": 100, + "warmup_factor": 0.1, + "device": "cuda", + "tolerance": 0.5, + "seg_threshold": 0.5, + "min_area": 1, + + "inner_polylines_params": { + "enable": False, + "max_traces": 1000, + "seed_threshold": 0.5, + "low_threshold": 0.1, + "min_width": 2, # Minimum width of trace to take into account + "max_width": 8, + "step_size": 1, + } + } + # --- Process args --- # + args = get_args() + if args.steps is not None: + config["steps"] = args.steps + + if args.raw_pred is not None: + # Load raw_pred(s) + image_list = [] + name_list = [] + seg_list = [] + crossfield_list = [] + for raw_pred_filepath in args.raw_pred: + raw_pred = torch.load(raw_pred_filepath) + image_list.append(raw_pred["image"]) + name_list.append(raw_pred["name"]) + seg_list.append(raw_pred["seg"]) + crossfield_list.append(raw_pred["crossfield"]) + seg_batch = torch.stack(seg_list, dim=0) + crossfield_batch = torch.stack(crossfield_list, dim=0) + + out_contours_batch, out_probs_batch = polygonize(seg_batch, crossfield_batch, config) + + for i, raw_pred_filepath in enumerate(args.raw_pred): + image = image_list[i] + name = name_list[i] + polygons = out_contours_batch[i] + base_filepath = os.path.join(os.path.dirname(raw_pred_filepath), name) + filepath = base_filepath + ".poly_acm.pdf" + plot_utils.save_poly_viz(image, polygons, filepath) + + # Load gt polygons + save_gt_poly(raw_pred_filepath, name) + elif args.im_filepath: + # Load from filepath, look for seg and crossfield next to the image + # Load data + image = skimage.io.imread(args.im_filepath) + base_filepath = os.path.splitext(args.im_filepath)[0] + seg = skimage.io.imread(base_filepath + ".seg.tif") / 255 + crossfield = np.load(base_filepath + ".crossfield.npy", allow_pickle=True) + + # Select bbox for dev + if args.bbox is not None: + assert len(args.bbox) == 4, "bbox should have 4 values" + bbox = args.bbox + # bbox = [1440, 210, 1800, 650] # vienna12 + # bbox = [2808, 2393, 3124, 2772] # innsbruck19 + image = image[bbox[0]:bbox[2], bbox[1]:bbox[3]] + seg = seg[bbox[0]:bbox[2], bbox[1]:bbox[3]] + crossfield = crossfield[bbox[0]:bbox[2], bbox[1]:bbox[3]] + extra_name = ".bbox_{}_{}_{}_{}".format(*bbox) + else: + extra_name = "" + + # Convert to torch and add batch dim + seg_batch = torch.tensor(np.transpose(seg[:, :, :2], (2, 0, 1)), dtype=torch.float)[None, ...] + crossfield_batch = torch.tensor(np.transpose(crossfield, (2, 0, 1)), dtype=torch.float)[None, ...] + + out_contours_batch, out_probs_batch = polygonize(seg_batch, crossfield_batch, config) + + polygons = out_contours_batch[0] + # Save shapefile + # save_utils.save_shapefile(polygons, base_filepath + extra_name, "poly_acm", args.im_filepath) + + # Save pdf viz + filepath = base_filepath + extra_name + ".poly_acm.pdf" + plot_utils.save_poly_viz(image, polygons, filepath, linewidths=1, draw_vertices=True, color_choices=[[0, 1, 0, 1]]) + elif args.dirpath: + seg_filename_list = fnmatch.filter(os.listdir(args.dirpath), "*.seg.tif") + sorted(seg_filename_list) + pbar = tqdm(seg_filename_list, desc="Poly files") + for id, seg_filename in enumerate(pbar): + basename = seg_filename[:-len(".seg.tif")] + # shp_filepath = os.path.join(args.dirpath, basename + ".poly_acm.shp") + # Verify if image has already been polygonized + # if os.path.exists(shp_filepath): + # continue + + pbar.set_postfix(name=basename, status="Loading data...") + crossfield_filename = basename + ".crossfield.npy" + metadata_filename = basename + ".metadata.json" + seg = skimage.io.imread(os.path.join(args.dirpath, seg_filename)) / 255 + crossfield = np.load(os.path.join(args.dirpath, crossfield_filename), allow_pickle=True) + metadata = python_utils.load_json(os.path.join(args.dirpath, metadata_filename)) + # image_filepath = metadata["image_filepath"] + # as_shp_filename = os.path.splitext(os.path.basename(image_filepath))[0] + # as_shp_filepath = os.path.join(os.path.dirname(os.path.dirname(image_filepath)), "gt_polygons", as_shp_filename + ".shp") + + # Convert to torch and add batch dim + seg_batch = torch.tensor(np.transpose(seg[:, :, :2], (2, 0, 1)), dtype=torch.float)[None, ...] + crossfield_batch = torch.tensor(np.transpose(crossfield, (2, 0, 1)), dtype=torch.float)[None, ...] + + pbar.set_postfix(name=basename, status="Polygonazing...") + out_contours_batch, out_probs_batch = polygonize(seg_batch, crossfield_batch, config) + + polygons = out_contours_batch[0] + + # Save as shp + # pbar.set_postfix(name=basename, status="Saving .shp...") + # geo_utils.save_shapefile_from_shapely_polygons(polygons, shp_filepath, as_shp_filepath) + + # Save as COCO annotation + base_filepath = os.path.join(args.dirpath, basename) + inference.save_poly_coco(polygons, id, base_filepath, "annotation.poly") + else: + print("Showcase on a very simple example:") + seg = np.zeros((6, 8, 3)) + # Triangle: + seg[1, 4] = 1 + seg[2, 3:5] = 1 + seg[3, 2:5] = 1 + seg[4, 1:5] = 1 + # L extension: + seg[3:5, 5:7] = 1 + + u = np.zeros((6, 8), dtype=np.complex) + v = np.zeros((6, 8), dtype=np.complex) + # Init with grid + u.real = 1 + v.imag = 1 + # Add slope + u[:4, :4] *= np.exp(1j * np.pi/4) + v[:4, :4] *= np.exp(1j * np.pi/4) + # Add slope corners + # u[:2, 4:6] *= np.exp(1j * np.pi / 4) + # v[4:, :2] *= np.exp(- 1j * np.pi / 4) + + crossfield = math_utils.compute_crossfield_c0c2(u, v) + + seg_batch = torch.tensor(np.transpose(seg[:, :, :2], (2, 0, 1)), dtype=torch.float)[None, ...] + crossfield_batch = torch.tensor(np.transpose(crossfield, (2, 0, 1)), dtype=torch.float)[None, ...] + + out_contours_batch, out_probs_batch = polygonize(seg_batch, crossfield_batch, config) + + polygons = out_contours_batch[0] + + filepath = "demo_poly_acm.pdf" + plot_utils.save_poly_viz(seg, polygons, filepath, linewidths=0.5, draw_vertices=True, crossfield=crossfield) + + +if __name__ == '__main__': + main() diff --git a/frame_field_learning/polygonize_asm.py b/frame_field_learning/polygonize_asm.py new file mode 100644 index 0000000000000000000000000000000000000000..7e833863a20ecf34544fe7c1e979ca0c9a5fff1c --- /dev/null +++ b/frame_field_learning/polygonize_asm.py @@ -0,0 +1,1181 @@ +import argparse +import fnmatch +import functools +import glob +import time +from typing import List + +import numpy as np +import skan +import skimage +import skimage.measure +import skimage.morphology +import skimage.io +from tqdm import tqdm +import shapely.geometry +import shapely.ops +import shapely.prepared +import scipy.interpolate + +from functools import partial + +import torch +import torch_scatter + +from frame_field_learning import polygonize_utils, plot_utils, frame_field_utils, save_utils + +from torch_lydorn.torch.nn.functionnal import bilinear_interpolate +from torch_lydorn.torchvision.transforms import Paths, Skeleton, TensorSkeleton, skeletons_to_tensorskeleton, tensorskeleton_to_skeletons +import torch_lydorn.kornia + +from lydorn_utils import math_utils +from lydorn_utils import python_utils +from lydorn_utils import print_utils + +DEBUG = False + + +def debug_print(s: str): + if DEBUG: + print_utils.print_debug(s) + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--raw_pred', + nargs='*', + type=str, + help='Filepath to the raw pred file(s)') + argparser.add_argument( + '--im_filepath', + type=str, + help='Filepath to input image. Will retrieve seg and crossfield in the same directory') + argparser.add_argument( + '--seg_filepath', + type=str, + help='Filepath to input segmentation image.') + argparser.add_argument( + '--angles_map_filepath', + type=str, + help='Filepath to frame field angles map.') + argparser.add_argument( + '--dirpath', + type=str, + help='Path to directory containing seg and crossfield files. Will perform polygonization on all.') + argparser.add_argument( + '--bbox', + nargs='*', + type=int, + help='Selects area in bbox for computation: [min_row, min_col, max_row, max_col]') + argparser.add_argument( + '--steps', + type=int, + help='Optim steps') + + args = argparser.parse_args() + return args + + +def get_junction_corner_index(tensorskeleton): + """ + Returns as a tensor the list of 3-tuples each representing a corner of a junction. + The 3-tuple contains the indices of the 3 vertices making up the corner. + + In the text below, we use the following notation: + - J: the number of junction nodes + - Sd: the sum of the degrees of all the junction nodes + - T: number of tip nodes + @return: junction_corner_index of shape (Sd*J - T, 3) which is a list of 3-tuples (for each junction corner) + """ + # --- Compute all junction edges: + junction_edge_index = torch.empty((2 * tensorskeleton.num_paths, 2), dtype=torch.long, device=tensorskeleton.path_index.device) + junction_edge_index[:tensorskeleton.num_paths, 0] = tensorskeleton.path_index[tensorskeleton.path_delim[:-1]] + junction_edge_index[:tensorskeleton.num_paths, 1] = tensorskeleton.path_index[tensorskeleton.path_delim[:-1] + 1] + junction_edge_index[tensorskeleton.num_paths:, 0] = tensorskeleton.path_index[tensorskeleton.path_delim[1:] - 1] + junction_edge_index[tensorskeleton.num_paths:, 1] = tensorskeleton.path_index[tensorskeleton.path_delim[1:] - 2] + # --- Remove tip junctions + degrees = tensorskeleton.degrees[junction_edge_index[:, 0]] + junction_edge_index = junction_edge_index[1 < degrees, :] + # --- Group by junction by sorting + group_indices = torch.argsort(junction_edge_index[:, 0], dim=0) + grouped_junction_edge_index = junction_edge_index[group_indices, :] + # --- Compute angle to vertical axis of each junction edge + junction_edge = tensorskeleton.pos.detach()[grouped_junction_edge_index, :] + junction_tangent = junction_edge[:, 1, :] - junction_edge[:, 0, :] + junction_angle_to_axis = torch.atan2(junction_tangent[:, 1], junction_tangent[:, 0]) + # --- Sort by angle for each junction separately and build junction_corner_index + unique = torch.unique_consecutive(grouped_junction_edge_index[:, 0]) + count = tensorskeleton.degrees[unique] + junction_end_index = torch.cumsum(count, dim=0) + slice_start = 0 + junction_corner_index = torch.empty((grouped_junction_edge_index.shape[0], 3), dtype=torch.long, device=tensorskeleton.path_index.device) + for slice_end in junction_end_index: + slice_angle_to_axis = junction_angle_to_axis[slice_start:slice_end] + slice_junction_edge_index = grouped_junction_edge_index[slice_start:slice_end] + sort_indices = torch.argsort(slice_angle_to_axis, dim=0) + slice_junction_edge_index = slice_junction_edge_index[sort_indices] + junction_corner_index[slice_start:slice_end, 0] = slice_junction_edge_index[:, 1] + junction_corner_index[slice_start:slice_end, 1] = slice_junction_edge_index[:, 0] + junction_corner_index[slice_start:slice_end, 2] = slice_junction_edge_index[:, 1].roll(-1, dims=0) + slice_start = slice_end + return junction_corner_index + + +class AlignLoss: + def __init__(self, tensorskeleton: TensorSkeleton, indicator: torch.Tensor, level: float, c0c2: torch.Tensor, loss_params): + """ + :param tensorskeleton: skeleton graph in tensor format + :return: + """ + self.tensorskeleton = tensorskeleton + self.indicator = indicator + self.level = level + self.c0c2 = c0c2 + # self.uv = frame_field_utils.c0c2_to_uv(c0c2) + + # Prepare junction_corner_index: + + # TODO: junction_corner_index: list + self.junction_corner_index = get_junction_corner_index(tensorskeleton) + + # Loss coefs + self.data_coef_interp = scipy.interpolate.interp1d(loss_params["coefs"]["step_thresholds"], + loss_params["coefs"]["data"]) + self.length_coef_interp = scipy.interpolate.interp1d(loss_params["coefs"]["step_thresholds"], + loss_params["coefs"]["length"]) + self.crossfield_coef_interp = scipy.interpolate.interp1d(loss_params["coefs"]["step_thresholds"], + loss_params["coefs"]["crossfield"]) + self.curvature_coef_interp = scipy.interpolate.interp1d(loss_params["coefs"]["step_thresholds"], + loss_params["coefs"]["curvature"]) + self.corner_coef_interp = scipy.interpolate.interp1d(loss_params["coefs"]["step_thresholds"], + loss_params["coefs"]["corner"]) + self.junction_coef_interp = scipy.interpolate.interp1d(loss_params["coefs"]["step_thresholds"], + loss_params["coefs"]["junction"]) + + self.curvature_dissimilarity_threshold = loss_params["curvature_dissimilarity_threshold"] + self.corner_angles = np.pi * torch.tensor(loss_params["corner_angles"]) / 180 # Convert to radians + self.corner_angle_threshold = np.pi * loss_params["corner_angle_threshold"] / 180 # Convert to radians + self.junction_angles = np.pi * torch.tensor(loss_params["junction_angles"]) / 180 # Convert to radians + self.junction_angle_weights = torch.tensor(loss_params["junction_angle_weights"]) + self.junction_angle_threshold = np.pi * loss_params["junction_angle_threshold"] / 180 # Convert to radians + + # Pre-compute useful pointers + # edge_index_start = tensorskeleton.path_index[:-1] + # edge_index_end = tensorskeleton.path_index[1:] + # + # self.tensorskeleton.edge_index = edge_index + + def __call__(self, pos: torch.Tensor, iter_num: int): + # --- Align to frame field loss + path_pos = pos[self.tensorskeleton.path_index] + detached_path_pos = path_pos.detach() + path_batch = self.tensorskeleton.batch[self.tensorskeleton.path_index] + tangents = path_pos[1:] - path_pos[:-1] + # Compute edge mask to remove edges that connect two different paths from loss + edge_mask = torch.ones((tangents.shape[0]), device=tangents.device) + edge_mask[self.tensorskeleton.path_delim[1:-1] - 1] = 0 # Zero out edges between paths + + midpoints = (path_pos[1:] + path_pos[:-1]) / 2 + midpoints_batch = self.tensorskeleton.batch[self.tensorskeleton.path_index[:-1]] # Same as start point of edge + + midpoints_int = midpoints.round().long() + midpoints_int[:, 0] = torch.clamp(midpoints_int[:, 0], 0, self.c0c2.shape[2] - 1) + midpoints_int[:, 1] = torch.clamp(midpoints_int[:, 1], 0, self.c0c2.shape[3] - 1) + midpoints_c0 = self.c0c2[midpoints_batch, :2, midpoints_int[:, 0], midpoints_int[:, 1]] + midpoints_c2 = self.c0c2[midpoints_batch, 2:, midpoints_int[:, 0], midpoints_int[:, 1]] + + norms = torch.norm(tangents, dim=-1) + edge_mask[norms < 0.1] = 0 # Zero out very small edges + normed_tangents = tangents / (norms[:, None] + 1e-6) + + align_loss = frame_field_utils.framefield_align_error(midpoints_c0, midpoints_c2, normed_tangents, complex_dim=1) + align_loss = align_loss * edge_mask + total_align_loss = torch.sum(align_loss) + + # --- Align to level set of indicator: + pos_value = bilinear_interpolate(self.indicator[:, None, ...], pos, batch=self.tensorskeleton.batch) + # TODO: use grid_sample with batch: put batch dim to height dim and make a single big image. + # TODO: Convert pos accordingly and take care of borders + # height = self.indicator.shape[1] + # width = self.indicator.shape[2] + # normed_xy = tensorskeleton.pos.roll(shifts=1, dims=-1) + # normed_xy[: 0] /= (width-1) + # normed_xy[: 1] /= (height-1) + # centered_xy = 2*normed_xy - 1 + # pos_value = torch.nn.functional.grid_sample(self.indicator[None, None, ...], + # centered_batch_xy[None, None, ...], align_corners=True).squeeze() + level_loss = torch.sum(torch.pow(pos_value - self.level, 2)) + + # --- Prepare useful tensors for curvature loss: + prev_pos = detached_path_pos[:-2] + middle_pos = path_pos[1:-1] + next_pos = detached_path_pos[2:] + prev_tangent = middle_pos - prev_pos + next_tangent = next_pos - middle_pos + prev_norm = torch.norm(prev_tangent, dim=-1) + next_norm = torch.norm(next_tangent, dim=-1) + + # --- Apply length penalty with sum of squared norm to penalize uneven edge lengths on selected edges + prev_length_loss = torch.pow(prev_norm, 2) + next_length_loss = torch.pow(next_norm, 2) + prev_length_loss[self.tensorskeleton.path_delim[1:-1] - 1] = 0 # Zero out invalid norms between paths + prev_length_loss[self.tensorskeleton.path_delim[1:-1] - 2] = 0 # Zero out unwanted contribution to loss + next_length_loss[self.tensorskeleton.path_delim[1:-1] - 1] = 0 # Zero out unwanted contribution to loss + next_length_loss[self.tensorskeleton.path_delim[1:-1] - 2] = 0 # Zero out invalid norms between paths + length_loss = prev_length_loss + next_length_loss + total_length_loss = torch.sum(length_loss) + + # --- Detect corners: + with torch.no_grad(): + middle_pos_int = middle_pos.round().long() + middle_pos_int[:, 0] = torch.clamp(middle_pos_int[:, 0], 0, self.c0c2.shape[2] - 1) + middle_pos_int[:, 1] = torch.clamp(middle_pos_int[:, 1], 0, self.c0c2.shape[3] - 1) + middle_batch = path_batch[1:-1] + middle_c0c2 = self.c0c2[middle_batch, :, middle_pos_int[:, 0], middle_pos_int[:, 1]] + middle_uv = frame_field_utils.c0c2_to_uv(middle_c0c2) + prev_tangent_closest_in_uv = frame_field_utils.compute_closest_in_uv(prev_tangent, middle_uv) + next_tangent_closest_in_uv = frame_field_utils.compute_closest_in_uv(next_tangent, middle_uv) + is_corner = prev_tangent_closest_in_uv != next_tangent_closest_in_uv + is_corner[self.tensorskeleton.path_delim[1:-1] - 2] = 0 # Zero out invalid corners between sub-paths + is_corner[self.tensorskeleton.path_delim[1:-1] - 1] = 0 # Zero out invalid corners between sub-paths + is_corner_index = torch.nonzero(is_corner)[:, 0] + 1 # Shift due to first vertex not being represented in is_corner + # TODO: evaluate running time of torch.sort: does it slow down the optimization much? + sub_path_delim, sub_path_sort_indices = torch.sort(torch.cat([self.tensorskeleton.path_delim, is_corner_index])) + sub_path_delim_is_corner = self.tensorskeleton.path_delim.shape[0] <= sub_path_sort_indices # If condition is true, then the delimiter is from is_corner_index + + # --- Compute sub-path dissimilarity in the sense of the Ramer-Douglas-Peucker alg + # dissimilarity is equal to the max distance of vertices to the straight line connecting the start and end points of the sub-path. + with torch.no_grad(): + sub_path_start_index = sub_path_delim[:-1] + sub_path_end_index = sub_path_delim[1:].clone() + sub_path_end_index[~sub_path_delim_is_corner[1:]] -= 1 # For non-corner delimitators, have to shift + sub_path_start_pos = path_pos[sub_path_start_index] + sub_path_end_pos = path_pos[sub_path_end_index] + sub_path_normal = sub_path_end_pos - sub_path_start_pos + sub_path_normal = sub_path_normal / (torch.norm(sub_path_normal, dim=1)[:, None] + 1e-6) + expanded_sub_path_start_pos = torch_scatter.gather_csr(sub_path_start_pos, + sub_path_delim) + expanded_sub_path_normal = torch_scatter.gather_csr(sub_path_normal, + sub_path_delim) + relative_path_pos = path_pos - expanded_sub_path_start_pos + relative_path_pos_projected_lengh = torch.sum(relative_path_pos * expanded_sub_path_normal, dim=1) + relative_path_pos_projected = relative_path_pos_projected_lengh[:, None] * expanded_sub_path_normal + path_pos_distance = torch.norm(relative_path_pos - relative_path_pos_projected, dim=1) + sub_path_max_distance = torch_scatter.segment_max_csr(path_pos_distance, sub_path_delim)[0] + sub_path_small_dissimilarity_mask = sub_path_max_distance < self.curvature_dissimilarity_threshold + + # --- Compute curvature loss: + # print("prev_norm:", prev_norm.min().item(), prev_norm.max().item()) + prev_dir = prev_tangent / (prev_norm[:, None] + 1e-6) + next_dir = next_tangent / (next_norm[:, None] + 1e-6) + dot = prev_dir[:, 0] * next_dir[:, 0] + \ + prev_dir[:, 1] * next_dir[:, 1] # dot product + det = prev_dir[:, 0] * next_dir[:, 1] - \ + prev_dir[:, 1] * next_dir[:, 0] # determinant + vertex_angles = torch.acos(dot) * torch.sign(det) # TODO: remove acos for speed? Switch everything to signed dot product? + # Save angles of detected corners: + corner_angles = vertex_angles[is_corner_index - 1] # -1 because of the shift of vertex_angles relative to path_pos + # Compute the mean vertex angle for each sub-path separately: + vertex_angles[sub_path_delim[1:-1] - 1] = 0 # Zero out invalid angles between paths as well as corner angles + vertex_angles[self.tensorskeleton.path_delim[1:-1] - 2] = 0 # Zero out invalid angles between paths (caused by the junction points being in all paths of the junction) + sub_path_vertex_angle_delim = sub_path_delim.clone() + sub_path_vertex_angle_delim[-1] -= 2 + sub_path_sum_vertex_angle = torch_scatter.segment_sum_csr(vertex_angles, sub_path_vertex_angle_delim) + sub_path_lengths = sub_path_delim[1:] - sub_path_delim[:-1] + sub_path_lengths[sub_path_delim_is_corner[1:]] += 1 # Fix length of paths split by corners + sub_path_valid_angle_count = sub_path_lengths - 2 + # print("sub_path_valid_angle_count:", sub_path_valid_angle_count.min().item(), sub_path_valid_angle_count.max().item()) + sub_path_mean_vertex_angles = sub_path_sum_vertex_angle / sub_path_valid_angle_count + sub_path_mean_vertex_angles[sub_path_small_dissimilarity_mask] = 0 # Optimize sub-path with a small dissimilarity to have straight edges + expanded_sub_path_mean_vertex_angles = torch_scatter.gather_csr(sub_path_mean_vertex_angles, + sub_path_vertex_angle_delim) + curvature_loss = torch.pow(vertex_angles - expanded_sub_path_mean_vertex_angles, 2) + curvature_loss[sub_path_delim[1:-1] - 1] = 0 # Zero out loss for start vertex of inner sub-paths + curvature_loss[self.tensorskeleton.path_delim[1:-1] - 2] = 0 # Zero out loss for end vertex of inner paths (caused by the junction points being in all paths of the junction) + total_curvature_loss = torch.sum(curvature_loss) + + # --- Computer corner loss: + corner_abs_angles = torch.abs(corner_angles) + self.corner_angles = self.corner_angles.to(corner_abs_angles.device) + corner_snap_dist = torch.abs(corner_abs_angles[:, None] - self.corner_angles) + corner_snap_dist_optim_mask = corner_snap_dist < self.corner_angle_threshold + corner_snap_dist_optim = corner_snap_dist[corner_snap_dist_optim_mask] + corner_loss = torch.pow(corner_snap_dist_optim, 2) + total_corner_loss = torch.sum(corner_loss) + + # --- Compute junction corner loss + junction_corner = pos[self.junction_corner_index, :] + junction_prev_tangent = junction_corner[:, 1, :] - junction_corner[:, 0, :] + junction_next_tangent = junction_corner[:, 2, :] - junction_corner[:, 1, :] + junction_prev_dir = junction_prev_tangent / (torch.norm(junction_prev_tangent, dim=-1)[:, None] + 1e-6) + junction_next_dir = junction_next_tangent / (torch.norm(junction_next_tangent, dim=-1)[:, None] + 1e-6) + junction_dot = junction_prev_dir[:, 0] * junction_next_dir[:, 0] + \ + junction_prev_dir[:, 1] * junction_next_dir[:, 1] # dot product + junction_abs_angles = torch.acos(junction_dot) + self.junction_angles = self.junction_angles.to(junction_abs_angles.device) + self.junction_angle_weights = self.junction_angle_weights.to(junction_abs_angles.device) + junction_snap_dist = torch.abs(junction_abs_angles[:, None] - self.junction_angles) + junction_snap_dist_optim_mask = junction_snap_dist < self.junction_angle_threshold + junction_snap_dist *= self.junction_angle_weights[None, :] # Apply weights per target angle (as we use the L1 norm, it works applying before the norm) + junction_snap_dist_optim = junction_snap_dist[junction_snap_dist_optim_mask] + junction_loss = torch.abs(junction_snap_dist_optim) + total_junction_loss = torch.sum(junction_loss) + + losses_dict = { + "align": total_align_loss.item(), + "level": level_loss.item(), + "length": total_length_loss.item(), + "curvature": total_curvature_loss.item(), + "corner": total_corner_loss.item(), + "junction": total_junction_loss.item(), + } + # Get the loss coefs depending on the current step: + data_coef = float(self.data_coef_interp(iter_num)) + length_coef = float(self.length_coef_interp(iter_num)) + crossfield_coef = float(self.crossfield_coef_interp(iter_num)) + curvature_coef = float(self.curvature_coef_interp(iter_num)) + corner_coef = float(self.corner_coef_interp(iter_num)) + junction_coef = float(self.junction_coef_interp(iter_num)) + # total_loss = data_coef * level_loss + length_coef * total_length_loss + crossfield_coef * total_align_loss + \ + # curvature_coef * total_curvature_loss + corner_coef * total_corner_loss + junction_coef * total_junction_loss + # total_loss = data_coef * level_loss + length_coef * total_length_loss + crossfield_coef * total_align_loss + \ + # curvature_coef * total_curvature_loss + corner_coef * total_corner_loss + junction_coef * total_junction_loss + # TODO: Debug adding curvature_coef * total_curvature_loss + corner_coef * total_corner_loss + junction_coef * total_junction_loss + total_loss = data_coef * level_loss + length_coef * total_length_loss + crossfield_coef * total_align_loss + + # print(iter_num) + # input("...") + + return total_loss, losses_dict + + +class TensorSkeletonOptimizer: + def __init__(self, config: dict, tensorskeleton: TensorSkeleton, indicator: torch.Tensor, c0c2: torch.Tensor): + assert len(indicator.shape) == 3, f"indicator should be of shape (N, H, W), not {indicator.shape}" + assert len(c0c2.shape) == 4 and c0c2.shape[1] == 4, f"c0c2 should be of shape (N, 4, H, W), not {c0c2.shape}" + + self.config = config + self.tensorskeleton = tensorskeleton + + # Save endpoints that are tips so that they can be reset after each step (tips are not meant to be moved) + self.is_tip = self.tensorskeleton.degrees == 1 + self.tip_pos = self.tensorskeleton.pos[self.is_tip] + + # Require grads for graph.pos: this is what is optimized + self.tensorskeleton.pos.requires_grad = True + + level = config["data_level"] + self.criterion = AlignLoss(self.tensorskeleton, indicator, level, c0c2, config["loss_params"]) + self.optimizer = torch.optim.RMSprop([tensorskeleton.pos], lr=config["lr"], alpha=0.9) + self.lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(self.optimizer, config["gamma"]) + + def step(self, iter_num): + self.optimizer.zero_grad() + + # tic = time.time() + loss, losses_dict = self.criterion(self.tensorskeleton.pos, iter_num) + + # toc = time.time() + # print(f"Forward: {toc - tic}s") + + # print("loss:", loss.item()) + # tic = time.time() + loss.backward() + + pos_gard_is_nan = torch.isnan(self.tensorskeleton.pos.grad).any().item() + if pos_gard_is_nan: + print(f"{iter_num} pos.grad is nan") + + # print(self.tensorskeleton.pos.grad) + # print(torch.norm(self.tensorskeleton.pos.grad, dim=1).max().item()) + # toc = time.time() + # print(f"Backward: {toc - tic}s") + self.optimizer.step() + + # Move tips back: + with torch.no_grad(): + # TODO: snap to nearest image border + self.tensorskeleton.pos[self.is_tip] = self.tip_pos + + if self.lr_scheduler is not None: + self.lr_scheduler.step() + + return loss.item(), losses_dict + + def optimize(self) -> TensorSkeleton: + if DEBUG: + optim_iter = tqdm(range(self.config["loss_params"]["coefs"]["step_thresholds"][-1]), desc="Gradient descent", leave=True) + for iter_num in optim_iter: + loss, losses_dict = self.step(iter_num) + optim_iter.set_postfix(loss=loss, **losses_dict) + else: + for iter_num in range(self.config["loss_params"]["coefs"]["step_thresholds"][-1]): + loss, losses_dict = self.step(iter_num) + # for iter_num in range(self.config["loss_params"]["coefs"]["step_thresholds"][-1]): + # loss, losses_dict = self.step(iter_num) + return self.tensorskeleton + + +def shapely_postprocess(polylines, np_indicator, tolerance, config): + if type(tolerance) == list: + # Use several tolerance values for simplification. return a dict with all results + out_polygons_dict = {} + out_probs_dict = {} + for tol in tolerance: + out_polygons, out_probs = shapely_postprocess(polylines, np_indicator, tol, config) + out_polygons_dict["tol_{}".format(tol)] = out_polygons + out_probs_dict["tol_{}".format(tol)] = out_probs + return out_polygons_dict, out_probs_dict + else: + height = np_indicator.shape[0] + width = np_indicator.shape[1] + + # Convert to Shapely: + # tic = time.time() + line_string_list = [shapely.geometry.LineString(polyline[:, ::-1]) for polyline in polylines] + line_string_list = [line_string.simplify(tolerance, preserve_topology=True) for line_string in line_string_list] + # toc = time.time() + # print(f"simplify: {toc - tic}s") + + # Add image boundary line_strings for border polygons + line_string_list.append( + shapely.geometry.LinearRing([ + (0, 0), + (0, height - 1), + (width - 1, height - 1), + (width - 1, 0), + ])) + + # debug_print("Merge polylines") + + # Merge polylines (for border polygons): + + # tic = time.time() + multi_line_string = shapely.ops.unary_union(line_string_list) + # toc = time.time() + # print(f"shapely.ops.unary_union: {toc - tic}s") + + # debug_print("polygonize_full") + + # Find polygons: + polygons = shapely.ops.polygonize(multi_line_string) + polygons = list(polygons) + + # debug_print("Remove small polygons") + + # Remove small polygons + # tic = time.time() + polygons = [polygon for polygon in polygons if + config["min_area"] < polygon.area] + # toc = time.time() + # print(f"Remove small polygons: {toc - tic}s") + + # debug_print("Remove low prob polygons") + + # Remove low prob polygons + # tic = time.time() + + filtered_polygons = [] + filtered_polygon_probs = [] + for polygon in polygons: + prob = polygonize_utils.compute_geom_prob(polygon, np_indicator) + # print("acm:", np_indicator.min(), np_indicator.mean(), np_indicator.max(), prob) + if config["seg_threshold"] < prob: + filtered_polygons.append(polygon) + filtered_polygon_probs.append(prob) + + # toc = time.time() + # print(f"Remove low prob polygons: {toc - tic}s") + + return filtered_polygons, filtered_polygon_probs + + +def post_process(polylines, np_indicator, np_crossfield, config): + + # debug_print("Corner-aware simplification") + # Simplify contours a little to avoid some close-together corner-detection: + # tic = time.time() + u, v = math_utils.compute_crossfield_uv(np_crossfield) # u, v are complex arrays + corner_masks = frame_field_utils.detect_corners(polylines, u, v) + polylines = polygonize_utils.split_polylines_corner(polylines, corner_masks) + # toc = time.time() + # print(f"Corner detect: {toc - tic}s") + + polygons, probs = shapely_postprocess(polylines, np_indicator, config["tolerance"], config) + return polygons, probs + + +def get_skeleton(np_edge_mask, config): + """ + + @param np_edge_mask: + @param config: + @return: + """ + # --- Skeletonize + # tic = time.time() + # Pad np_edge_mask first otherwise pixels on the bottom and right are lost after skeletonize: + pad_width = 2 + np_edge_mask_padded = np.pad(np_edge_mask, pad_width=pad_width, mode="edge") + skeleton_image = skimage.morphology.skeletonize(np_edge_mask_padded) + skeleton_image = skeleton_image[pad_width:-pad_width, pad_width:-pad_width] + + # toc = time.time() + # debug_print(f"skimage.morphology.skeletonize: {toc - tic}s") + + # tic = time.time() + + # if skeleton_image.max() == False: + # # There is no polylines to be detected + # return [], np.empty((0, 2), dtype=np.bool) + + skeleton = Skeleton() + if 0 < skeleton_image.sum(): + # skan does not work in some cases (paths of 2 pixels or less, etc) which raises a ValueError, in witch case we continue with an empty skeleton. + try: + skeleton = skan.Skeleton(skeleton_image, keep_images=False) + # skan.skeleton sometimes returns skeleton.coordinates.shape[0] != skeleton.degrees.shape[0] or + # skeleton.coordinates.shape[0] != skeleton.paths.indices.max() + 1 + # Slice coordinates accordingly + skeleton.coordinates = skeleton.coordinates[:skeleton.paths.indices.max() + 1] + if skeleton.coordinates.shape[0] != skeleton.degrees.shape[0]: + raise ValueError(f"skeleton.coordinates.shape[0] = {skeleton.coordinates.shape[0]} while skeleton.degrees.shape[0] = {skeleton.degrees.shape[0]}. They should be of same size.") + except ValueError as e: + if DEBUG: + print_utils.print_warning( + f"WARNING: skan.Skeleton raised a ValueError({e}). skeleton_image has {skeleton_image.sum()} true values. Continuing without detecting skeleton in this image...") + skimage.io.imsave("np_edge_mask.png", np_edge_mask.astype(np.uint8) * 255) + skimage.io.imsave("skeleton_image.png", skeleton_image.astype(np.uint8) * 255) + + # toc = time.time() + #debug_print(f"skan.Skeleton: {toc - tic}s") + + # tic = time.time() + + # # --- For each endpoint, see if it's a tip or not + # endpoints_src = skeleton.paths.indices[skeleton.paths.indptr[:-1]] + # endpoints_dst = skeleton.paths.indices[skeleton.paths.indptr[1:] - 1] + # deg_src = skeleton.degrees[endpoints_src] + # deg_dst = skeleton.degrees[endpoints_dst] + # is_tip_array = np.stack([deg_src == 1, deg_dst == 1], axis=1) + + # toc = time.time() + # debug_print(f"Convert to polylines: {toc - tic}s") + + return skeleton + + +def get_marching_squares_skeleton(np_int_prob, config): + """ + + @param np_int_prob: + @param config: + @return: + """ + # tic = time.time() + contours = skimage.measure.find_contours(np_int_prob, config["data_level"], fully_connected='low', positive_orientation='high') + # Keep contours with more than 3 vertices and large enough area + contours = [contour for contour in contours if 3 <= contour.shape[0] and + config["min_area"] < shapely.geometry.Polygon(contour).area] + + # If there are no contours, return empty skeleton + if len(contours) == 0: + return Skeleton() + + toc = time.time() + #debug_print(f"get_skeleton_polylines: {toc - tic}s") + # Simplify contours a tiny bit: + # contours = [skimage.measure.approximate_polygon(contour, tolerance=0.001) for contour in contours] + + # Convert into skeleton representation + coordinates = [] + indices_offset = 0 + indices = [] + indptr = [0] + degrees = [] + + for i, contour in enumerate(contours): + # Check if it is a closed contour + is_closed = np.max(np.abs(contour[0] - contour[-1])) < 1e-6 + if is_closed: + _coordinates = contour[:-1, :] # Don't include redundant vertex in coordinates + else: + _coordinates = contour + _degrees = 2 * np.ones(_coordinates.shape[0], dtype=np.long) + if not is_closed: + _degrees[0] = 1 + _degrees[-1] = 1 + _indices = list(range(indices_offset, indices_offset + _coordinates.shape[0])) + if is_closed: + _indices.append(_indices[0]) # Close contour with indices + coordinates.append(_coordinates) + degrees.append(_degrees) + indices.extend(_indices) + indptr.append(indptr[-1] + len(_indices)) + indices_offset += _coordinates.shape[0] + + coordinates = np.concatenate(coordinates, axis=0) + degrees = np.concatenate(degrees, axis=0) + indices = np.array(indices) + indptr = np.array(indptr) + + paths = Paths(indices, indptr) + skeleton = Skeleton(coordinates, paths, degrees) + + return skeleton + + +# @profile +def compute_skeletons(seg_batch, config, spatial_gradient, pool=None) -> List[Skeleton]: + assert len(seg_batch.shape) == 4 and seg_batch.shape[ + 1] <= 3, "seg_batch should be (N, C, H, W) with C <= 3, not {}".format(seg_batch.shape) + + int_prob_batch = seg_batch[:, 0, :, :] + if config["init_method"] == "marching_squares": + # Only interior segmentation is available, initialize with marching squares + np_int_prob_batch = int_prob_batch.cpu().numpy() + get_marching_squares_skeleton_partial = functools.partial(get_marching_squares_skeleton, config=config) + if pool is not None: + skeletons_batch = pool.map(get_marching_squares_skeleton_partial, np_int_prob_batch) + else: + skeletons_batch = list(map(get_marching_squares_skeleton_partial, np_int_prob_batch)) + elif config["init_method"] == "skeleton": + tic_correct = time.time() + # Edge segmentation is also available, initialize with skan.Squeleton + corrected_edge_prob_batch = config["data_level"] < int_prob_batch # Convet to mask + corrected_edge_prob_batch = corrected_edge_prob_batch[:, None, :, :].float() # Convet to float for spatial grads + corrected_edge_prob_batch = 2 * spatial_gradient(corrected_edge_prob_batch)[:, 0, :, :] # (b, 2, h, w), Normalize (kornia normalizes to -0.5, 0.5 for input in [0, 1]) + corrected_edge_prob_batch = corrected_edge_prob_batch.norm(dim=1) # (b, h, w), take the gradient norm + # int_contours_mask_batch = compute_contours_mask(int_mask_batch[:, None, :, :])[:, 0, :, :] + # corrected_edge_prob_batch = int_contours_mask_batch.float() + if 2 <= seg_batch.shape[1]: + corrected_edge_prob_batch = torch.clamp(seg_batch[:, 1, :, :] + corrected_edge_prob_batch, 0, 1) + # Save for viz + # save_edge_prob_map = (corrected_edge_prob_batch[0].cpu().numpy() * 255).astype(np.uint8)[:, :, None] + # skimage.io.imsave("corrected_edge_prob_batch.png", save_edge_prob_map) + + toc_correct = time.time() + #debug_print(f"Correct edge prob map: {toc_correct - tic_correct}s") + + # --- Init skeleton + corrected_edge_mask_batch = config["data_level"] < corrected_edge_prob_batch + np_corrected_edge_mask_batch = corrected_edge_mask_batch.cpu().numpy() + + get_skeleton_partial = functools.partial(get_skeleton, config=config) + # polylines_batch = [] + # is_tip_batch = [] + # for np_corrected_edge_mask in np_corrected_edge_mask_batch: + # polylines, is_tip_array = get_skeleton_polylines_partial(np_corrected_edge_mask) + # polylines_batch.append(polylines) + # is_tip_batch.append(is_tip_array) + # tic = time.time() + if pool is not None: + skeletons_batch = pool.map(get_skeleton_partial, np_corrected_edge_mask_batch) + else: + skeletons_batch = list(map(get_skeleton_partial, np_corrected_edge_mask_batch)) + # toc = time.time() + #debug_print(f"get_skeleton_polylines: {toc - tic}s") + else: + raise NotImplementedError(f"init_method '{config['init_method']}' not recognized. Valid init methods are 'skeleton' and 'marching_squares'") + + return skeletons_batch + + +def skeleton_to_polylines(skeleton: Skeleton) -> List[np.ndarray]: + polylines = [] + for path_i in range(skeleton.paths.indptr.shape[0] - 1): + start, stop = skeleton.paths.indptr[path_i:path_i + 2] + path_indices = skeleton.paths.indices[start:stop] + path_coordinates = skeleton.coordinates[path_indices] + polylines.append(path_coordinates) + return polylines + + +class PolygonizerASM: + def __init__(self, config, pool=None): + self.config = config + self.pool = pool + self.spatial_gradient = torch_lydorn.kornia.filters.SpatialGradient(mode="scharr", coord="ij", normalized=True, + device=self.config["device"], dtype=torch.float) + + # @profile + def __call__(self, seg_batch, crossfield_batch, pre_computed=None): + tic_start = time.time() + + assert len(seg_batch.shape) == 4 and seg_batch.shape[ + 1] <= 3, "seg_batch should be (N, C, H, W) with C <= 3, not {}".format(seg_batch.shape) + assert len(crossfield_batch.shape) == 4 and crossfield_batch.shape[ + 1] == 4, "crossfield_batch should be (N, 4, H, W)" + assert seg_batch.shape[0] == crossfield_batch.shape[0], "Batch size for seg and crossfield should match" + + + seg_batch = seg_batch.to(self.config["device"]) + crossfield_batch = crossfield_batch.to(self.config["device"]) + + # --- Get initial polylines + # tic = time.time() + skeletons_batch = compute_skeletons(seg_batch, self.config, self.spatial_gradient, pool=self.pool) + # toc = time.time() + # debug_print(f"Init polylines: {toc - tic}s") + + # # --- Compute distance transform + # tic = time.time() + # + # np_int_mask_batch = int_mask_batch.cpu().numpy() + # np_dist_batch = np.empty(np_int_mask_batch.shape) + # for batch_i in range(np_int_mask_batch.shape[0]): + # dist_1 = cv.distanceTransform(np_int_mask_batch[batch_i].astype(np.uint8), distanceType=cv.DIST_L2, maskSize=cv.DIST_MASK_5, dstType=cv.CV_64F) + # dist_2 = cv.distanceTransform(1 - np_int_mask_batch[batch_i].astype(np.uint8), distanceType=cv.DIST_L2, maskSize=cv.DIST_MASK_5, dstType=cv.CV_64F) + # np_dist_batch[0] = dist_1 + dist_2 + # dist_batch = torch.from_numpy(np_dist_batch) + # + # toc = time.time() + # print(f"Distance transform: {toc - tic}s") + + # --- Optimize skeleton: + tensorskeleton = skeletons_to_tensorskeleton(skeletons_batch, device=self.config["device"]) + + # --- Check if tensorskeleton is empty + if tensorskeleton.num_paths == 0: + batch_size = seg_batch.shape[0] + polygons_batch = [[]]*batch_size + probs_batch = [[]]*batch_size + return polygons_batch, probs_batch + + int_prob_batch = seg_batch[:, 0, :, :] + # dist_batch = dist_batch.to(config["device"]) + tensorskeleton_optimizer = TensorSkeletonOptimizer(self.config, tensorskeleton, int_prob_batch, + crossfield_batch) + + if DEBUG: + # Animation of optimization + import matplotlib.pyplot as plt + import matplotlib.animation as animation + + fig, ax = plt.subplots(figsize=(10, 10)) + ax.autoscale(False) + ax.axis('equal') + ax.axis('off') + plt.subplots_adjust(left=0, right=1, top=1, bottom=0) # Plot without margins + + image = int_prob_batch.cpu().numpy()[0] + ax.imshow(image, cmap=plt.cm.gray) + + out_skeletons_batch = tensorskeleton_to_skeletons(tensorskeleton) + polylines_batch = [skeleton_to_polylines(skeleton) for skeleton in out_skeletons_batch] + out_polylines = [shapely.geometry.LineString(polyline[:, ::-1]) for polyline in polylines_batch[0]] + artists = plot_utils.plot_geometries(ax, out_polylines, draw_vertices=True, linewidths=1) + + optim_pbar = tqdm(desc="Gradient descent", leave=True, total=self.config["loss_params"]["coefs"]["step_thresholds"][-1]) + + def init(): # only required for blitting to give a clean slate. + for artist, polyline in zip(artists, polylines_batch[0]): + artist.set_xdata([np.nan] * polyline.shape[0]) + artist.set_ydata([np.nan] * polyline.shape[0]) + return artists + + def animate(i): + loss, losses_dict = tensorskeleton_optimizer.step(i) + optim_pbar.update(int(2 * i / self.config["loss_params"]["coefs"]["step_thresholds"][-1])) + optim_pbar.set_postfix(loss=loss, **losses_dict) + out_skeletons_batch = tensorskeleton_to_skeletons(tensorskeleton) + polylines_batch = [skeleton_to_polylines(skeleton) for skeleton in out_skeletons_batch] + for artist, polyline in zip(artists, polylines_batch[0]): + artist.set_xdata(polyline[:, 1]) + artist.set_ydata(polyline[:, 0]) + return artists + + ani = animation.FuncAnimation( + fig, animate, init_func=init, interval=0, blit=True, frames=self.config["loss_params"]["coefs"]["step_thresholds"][-1], repeat=False) + + # To save the animation, use e.g. + # + # ani.save("movie.mp4") + # + # or + # + # writer = animation.FFMpegWriter( + # fps=15, metadata=dict(artist='Me'), bitrate=1800) + # ani.save("movie.mp4", writer=writer) + + plt.show() + else: + tensorskeleton = tensorskeleton_optimizer.optimize() + + out_skeletons_batch = tensorskeleton_to_skeletons(tensorskeleton) + + # --- Convert the skeleton representation into polylines + polylines_batch = [skeleton_to_polylines(skeleton) for skeleton in out_skeletons_batch] + + # toc = time.time() + #debug_print(f"Optimize skeleton: {toc - tic}s") + + # --- Post-process: + # debug_print("Post-process") + # tic = time.time() + + np_crossfield_batch = np.transpose(crossfield_batch.cpu().numpy(), (0, 2, 3, 1)) + np_int_prob_batch = int_prob_batch.cpu().numpy() + post_process_partial = partial(post_process, config=self.config) + if self.pool is not None: + polygons_probs_batch = self.pool.starmap(post_process_partial, + zip(polylines_batch, np_int_prob_batch, np_crossfield_batch)) + else: + polygons_probs_batch = map(post_process_partial, polylines_batch, np_int_prob_batch, + np_crossfield_batch) + polygons_batch, probs_batch = zip(*polygons_probs_batch) + + # toc = time.time() + #debug_print(f"Post-process: {toc - tic}s") + + toc_end = time.time() + #debug_print(f"Total: {toc_end - tic_start}s") + + if DEBUG: + # --- display results + import matplotlib.pyplot as plt + image = np_int_prob_batch[0] + polygons = polygons_batch[0] + out_polylines = [shapely.geometry.LineString(polyline[:, ::-1]) for polyline in polylines_batch[0]] + + fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(16, 16), sharex=True, sharey=True) + ax = axes.ravel() + + ax[0].imshow(image, cmap=plt.cm.gray) + plot_utils.plot_geometries(ax[0], out_polylines, draw_vertices=True, linewidths=1) + ax[0].axis('off') + ax[0].set_title('original', fontsize=20) + + # ax[1].imshow(skeleton, cmap=plt.cm.gray) + # ax[1].axis('off') + # ax[1].set_title('skeleton', fontsize=20) + + fig.tight_layout() + plt.show() + + return polygons_batch, probs_batch + + +def polygonize(seg_batch, crossfield_batch, config, pool=None, pre_computed=None): + polygonizer_asm = PolygonizerASM(config, pool=pool) + return polygonizer_asm(seg_batch, crossfield_batch, pre_computed=pre_computed) + + +def main(): + from frame_field_learning import inference + import os + + def save_gt_poly(raw_pred_filepath, name): + filapth_format = "/data/mapping_challenge_dataset/processed/val/data_{}.pt" + sample = torch.load(filapth_format.format(name)) + polygon_arrays = sample["gt_polygons"] + polygons = [shapely.geometry.Polygon(polygon[:, ::-1]) for polygon in polygon_arrays] + base_filepath = os.path.join(os.path.dirname(raw_pred_filepath), name) + filepath = base_filepath + "." + name + ".pdf" + plot_utils.save_poly_viz(image, polygons, filepath) + + config = { + "init_method": "skeleton", # Can be either skeleton or marching_squares + "data_level": 0.5, + "loss_params": { + "coefs": { + "step_thresholds": [0, 100, 200, 300], # From 0 to 500: gradually go from coefs[0] to coefs[1] + "data": [1.0, 0.1, 0.0, 0], + "crossfield": [0.0, 0.05, 0.0, 0], + "length": [0.1, 0.01, 0.0, 0], + "curvature": [0.0, 0.0, 1.0, 1e-6], + "corner": [0.0, 0.0, 0.5, 1e-6], + "junction": [0.0, 0.0, 0.5, 1e-6], + }, + "curvature_dissimilarity_threshold": 2, # In pixels: for each sub-paths, if the dissimilarity (in the same sense as in the Ramer-Douglas-Peucker alg) is lower than curvature_dissimilarity_threshold, then optimize the curve angles to be zero. + "corner_angles": [45, 90, 135], # In degrees: target angles for corners. + "corner_angle_threshold": 22.5, # If a corner angle is less than this threshold away from any angle in corner_angles, optimize it. + "junction_angles": [0, 45, 90, 135], # In degrees: target angles for junction corners. + "junction_angle_weights": [1, 0.01, 0.1, 0.01], # Order of decreassing importance: straight, right-angle, then 45° junction corners. + "junction_angle_threshold": 22.5, # If a junction corner angle is less than this threshold away from any angle in junction_angles, optimize it. + }, + "lr": 0.1, + "gamma": 0.995, + "device": "cuda", + "tolerance": 1.0, + "seg_threshold": 0.5, + "min_area": 10, + } + # --- Process args --- # + args = get_args() + if args.steps is not None: + config["steps"] = args.steps + + if args.raw_pred is not None: + # Load raw_pred(s) + image_list = [] + name_list = [] + seg_list = [] + crossfield_list = [] + for raw_pred_filepath in args.raw_pred: + raw_pred = torch.load(raw_pred_filepath) + image_list.append(raw_pred["image"]) + name_list.append(raw_pred["name"]) + seg_list.append(raw_pred["seg"]) + crossfield_list.append(raw_pred["crossfield"]) + seg_batch = torch.stack(seg_list, dim=0) + crossfield_batch = torch.stack(crossfield_list, dim=0) + + out_contours_batch, out_probs_batch = polygonize(seg_batch, crossfield_batch, config) + + for i, raw_pred_filepath in enumerate(args.raw_pred): + image = image_list[i] + name = name_list[i] + polygons = out_contours_batch[i] + base_filepath = os.path.join(os.path.dirname(raw_pred_filepath), name) + filepath = base_filepath + ".poly_acm.pdf" + plot_utils.save_poly_viz(image, polygons, filepath) + + # Load gt polygons + save_gt_poly(raw_pred_filepath, name) + elif args.im_filepath: + # Load from filepath, look for seg and crossfield next to the image + # Load data + image = skimage.io.imread(args.im_filepath) + base_filepath = os.path.splitext(args.im_filepath)[0] + if args.seg_filepath is not None: + seg = skimage.io.imread(args.seg_filepath) / 255 + else: + seg = skimage.io.imread(base_filepath + ".seg.tif") / 255 + crossfield = np.load(base_filepath + ".crossfield.npy", allow_pickle=True) + + # Select bbox for dev + if args.bbox is not None: + assert len(args.bbox) == 4, "bbox should have 4 values" + bbox = args.bbox + # bbox = [1440, 210, 1800, 650] # vienna12 + # bbox = [2808, 2393, 3124, 2772] # innsbruck19 + image = image[bbox[0]:bbox[2], bbox[1]:bbox[3]] + seg = seg[bbox[0]:bbox[2], bbox[1]:bbox[3]] + crossfield = crossfield[bbox[0]:bbox[2], bbox[1]:bbox[3]] + extra_name = ".bbox_{}_{}_{}_{}".format(*bbox) + else: + extra_name = "" + + # Convert to torch and add batch dim + seg_batch = torch.tensor(np.transpose(seg[:, :, :2], (2, 0, 1)), dtype=torch.float)[None, ...] + crossfield_batch = torch.tensor(np.transpose(crossfield, (2, 0, 1)), dtype=torch.float)[None, ...] + + # # Add samples to batch to increase batch size for testing + # batch_size = 4 + # seg_batch = seg_batch.repeat((batch_size, 1, 1, 1)) + # crossfield_batch = crossfield_batch.repeat((batch_size, 1, 1, 1)) + + out_contours_batch, out_probs_batch = polygonize(seg_batch, crossfield_batch, config) + + polygons = out_contours_batch[0] + + # Save geojson + # save_utils.save_geojson(polygons, base_filepath + extra_name, name="poly_asm", image_filepath=args.im_filepath) + + # Save shapefile + save_utils.save_shapefile(polygons, base_filepath + extra_name, "poly_asm", args.im_filepath) + + # Save pdf viz + filepath = base_filepath + extra_name + ".poly_asm.pdf" + # plot_utils.save_poly_viz(image, polygons, filepath, linewidths=1, draw_vertices=True, color_choices=[[0, 1, 0, 1]]) + plot_utils.save_poly_viz(image, polygons, filepath, markersize=30, linewidths=1, draw_vertices=True) + elif args.seg_filepath is not None and args.angles_map_filepath is not None: + total_t1 = time.time() + print("Loading data in image format") + seg_filepaths = sorted(glob.glob(args.seg_filepath)) + angles_map_filepaths = sorted(glob.glob(args.angles_map_filepath)) + assert len(seg_filepaths) == len(angles_map_filepaths) + + for seg_filepath, angles_map_filepath in zip(seg_filepaths, angles_map_filepaths): + print("Running on:", seg_filepath, angles_map_filepath) + base_filepath = os.path.splitext(seg_filepath)[0] + # --- Load seg (or prob map) and frame field angles from file + config = { + "init_method": "skeleton", # Can be either skeleton or marching_squares + "data_level": 0.5, + "loss_params": { + "coefs": { + "step_thresholds": [0, 100, 200], # From 0 to 500: gradually go from coefs[0] to coefs[1] + "data": [1.0, 0.1, 0.0], + "crossfield": [0.0, 0.05, 0.0], + "length": [0.1, 0.01, 0.0], + "curvature": [0.0, 0.0, 0.0], + "corner": [0.0, 0.0, 0.0], + "junction": [0.0, 0.0, 0.0], + }, + "curvature_dissimilarity_threshold": 2, + # In pixels: for each sub-paths, if the dissimilarity (in the same sense as in the Ramer-Douglas-Peucker alg) is lower than curvature_dissimilarity_threshold, then optimize the curve angles to be zero. + "corner_angles": [45, 90, 135], # In degrees: target angles for corners. + "corner_angle_threshold": 22.5, + # If a corner angle is less than this threshold away from any angle in corner_angles, optimize it. + "junction_angles": [0, 45, 90, 135], # In degrees: target angles for junction corners. + "junction_angle_weights": [1, 0.01, 0.1, 0.01], + # Order of decreassing importance: straight, right-angle, then 45° junction corners. + "junction_angle_threshold": 22.5, + # If a junction corner angle is less than this threshold away from any angle in junction_angles, optimize it. + }, + "lr": 0.1, + "gamma": 0.995, + "device": "cuda", + "tolerance": 1.0, + "seg_threshold": 0.5, + "min_area": 10, + } + input_seg = skimage.io.imread(seg_filepath) / 255 + seg = input_seg[:, :, [1, 2]] # LuxCarta channels are (background, interior, wall), re-arrange to be (interior, wall) + angles_map = np.pi * skimage.io.imread(angles_map_filepath) / 255 + + t1 = time.time() + + u_angle = angles_map[:, :, 0] + v_angle = angles_map[:, :, 1] + u = np.cos(u_angle) - 1j * np.sin(u_angle) # y-axis inverted + v = np.cos(v_angle) - 1j * np.sin(v_angle) # y-axis inverted + crossfield = math_utils.compute_crossfield_c0c2(u, v) + + # Convert to torch and add batch dim + seg_batch = torch.tensor(np.transpose(seg[:, :, :2], (2, 0, 1)), dtype=torch.float)[None, ...] + crossfield_batch = torch.tensor(np.transpose(crossfield, (2, 0, 1)), dtype=torch.float)[None, ...] + + # # Add samples to batch to increase batch size for testing + # batch_size = 4 + # seg_batch = seg_batch.repeat((batch_size, 1, 1, 1)) + # crossfield_batch = crossfield_batch.repeat((batch_size, 1, 1, 1)) + + try: + out_contours_batch, out_probs_batch = polygonize(seg_batch, crossfield_batch, config) + + t2 = time.time() + + print(f"Time: {t2 - t1:02f}s") + + polygons = out_contours_batch[0] + + # Save geojson + # save_utils.save_geojson(polygons, base_filepath + extra_name, name="poly_asm", image_filepath=args.im_filepath) + + # Save shapefile + save_utils.save_shapefile(polygons, base_filepath, "poly_asm", seg_filepath) + + # # Save pdf viz + # filepath = base_filepath + ".poly_asm.pdf" + # # plot_utils.save_poly_viz(image, polygons, filepath, linewidths=1, draw_vertices=True, color_choices=[[0, 1, 0, 1]]) + # plot_utils.save_poly_viz(input_seg, polygons, filepath, markersize=30, linewidths=1, draw_vertices=True) + except ValueError as e: + print("ERROR:", e) + total_t2 = time.time() + print(f"Total time: {total_t2 - total_t1:02f}s") + elif args.dirpath: + seg_filename_list = fnmatch.filter(os.listdir(args.dirpath), "*.seg.tif") + sorted(seg_filename_list) + pbar = tqdm(seg_filename_list, desc="Poly files") + for id, seg_filename in enumerate(pbar): + basename = seg_filename[:-len(".seg.tif")] + # shp_filepath = os.path.join(args.dirpath, basename + ".poly_acm.shp") + # Verify if image has already been polygonized + # if os.path.exists(shp_filepath): + # continue + + pbar.set_postfix(name=basename, status="Loading data...") + crossfield_filename = basename + ".crossfield.npy" + metadata_filename = basename + ".metadata.json" + seg = skimage.io.imread(os.path.join(args.dirpath, seg_filename)) / 255 + crossfield = np.load(os.path.join(args.dirpath, crossfield_filename), allow_pickle=True) + metadata = python_utils.load_json(os.path.join(args.dirpath, metadata_filename)) + # image_filepath = metadata["image_filepath"] + # as_shp_filename = os.path.splitext(os.path.basename(image_filepath))[0] + # as_shp_filepath = os.path.join(os.path.dirname(os.path.dirname(image_filepath)), "gt_polygons", as_shp_filename + ".shp") + + # Convert to torch and add batch dim + seg_batch = torch.tensor(np.transpose(seg[:, :, :2], (2, 0, 1)), dtype=torch.float)[None, ...] + crossfield_batch = torch.tensor(np.transpose(crossfield, (2, 0, 1)), dtype=torch.float)[None, ...] + + pbar.set_postfix(name=basename, status="Polygonazing...") + out_contours_batch, out_probs_batch = polygonize(seg_batch, crossfield_batch, config) + + polygons = out_contours_batch[0] + + # Save as shp + # pbar.set_postfix(name=basename, status="Saving .shp...") + # geo_utils.save_shapefile_from_shapely_polygons(polygons, shp_filepath, as_shp_filepath) + + # Save as COCO annotation + base_filepath = os.path.join(args.dirpath, basename) + inference.save_poly_coco(polygons, id, base_filepath, "annotation.poly") + else: + print("Showcase on a very simple example:") + config = { + "init_method": "marching_squares", # Can be either skeleton or marching_squares + "data_level": 0.5, + "loss_params": { + "coefs": { + "step_thresholds": [0, 100, 200, 300], # From 0 to 500: gradually go from coefs[0] to coefs[1] + "data": [1.0, 0.1, 0.0, 0.0], + "crossfield": [0.0, 0.05, 0.0, 0.0], + "length": [0.1, 0.01, 0.0, 0.0], + "curvature": [0.0, 0.0, 0.0, 0.0], + "corner": [0.0, 0.0, 0.0, 0.0], + "junction": [0.0, 0.0, 0.0, 0.0], + }, + "curvature_dissimilarity_threshold": 2, + # In pixels: for each sub-paths, if the dissimilarity (in the same sense as in the Ramer-Douglas-Peucker alg) is lower than straightness_threshold, then optimize the curve angles to be zero. + "corner_angles": [45, 90, 135], # In degrees: target angles for corners. + "corner_angle_threshold": 22.5, + # If a corner angle is less than this threshold away from any angle in corner_angles, optimize it. + "junction_angles": [0, 45, 90, 135], # In degrees: target angles for junction corners. + "junction_angle_weights": [1, 0.01, 0.1, 0.01], + # Order of decreassing importance: straight, right-angle, then 45° junction corners. + "junction_angle_threshold": 22.5, + # If a junction corner angle is less than this threshold away from any angle in junction_angles, optimize it. + }, + "lr": 0.01, + "gamma": 0.995, + "device": "cuda", + "tolerance": 0.5, + "seg_threshold": 0.5, + "min_area": 10, + } + + seg = np.zeros((6, 8, 1)) + # Triangle: + seg[1, 4] = 1 + seg[2, 3:5] = 1 + seg[3, 2:5] = 1 + seg[4, 1:5] = 1 + # L extension: + seg[3:5, 5:7] = 1 + + u = np.zeros((6, 8), dtype=np.complex) + v = np.zeros((6, 8), dtype=np.complex) + # Init with grid + u.real = 1 + v.imag = 1 + # Add slope + u[:4, :4] *= np.exp(1j * np.pi / 4) + v[:4, :4] *= np.exp(1j * np.pi / 4) + # Add slope corners + # u[:2, 4:6] *= np.exp(1j * np.pi / 4) + # v[4:, :2] *= np.exp(- 1j * np.pi / 4) + + crossfield = math_utils.compute_crossfield_c0c2(u, v) + + seg_batch = torch.tensor(np.transpose(seg[:, :, :2], (2, 0, 1)), dtype=torch.float)[None, ...] + crossfield_batch = torch.tensor(np.transpose(crossfield, (2, 0, 1)), dtype=torch.float)[None, ...] + + # Add samples to batch to increase batch size + batch_size = 16 + seg_batch = seg_batch.repeat((batch_size, 1, 1, 1)) + crossfield_batch = crossfield_batch.repeat((batch_size, 1, 1, 1)) + + out_contours_batch, out_probs_batch = polygonize(seg_batch, crossfield_batch, config) + + polygons = out_contours_batch[0] + + filepath = "demo_poly_asm.pdf" + plot_utils.save_poly_viz(seg[:, :, 0], polygons, filepath, linewidths=0.5, draw_vertices=True, crossfield=crossfield) + + +if __name__ == '__main__': + main() diff --git a/frame_field_learning/polygonize_simple.py b/frame_field_learning/polygonize_simple.py new file mode 100644 index 0000000000000000000000000000000000000000..6579c0099796b83186289efa2115cd7640fad13c --- /dev/null +++ b/frame_field_learning/polygonize_simple.py @@ -0,0 +1,234 @@ +import argparse + +import os +import torch + +import numpy as np +import skimage +import skimage.measure +import skimage.io +import shapely.geometry +import shapely.ops +from PIL import Image +from multiprocess import Pool +from tqdm import tqdm + +from functools import partial + +from lydorn_utils import print_utils, geo_utils + +from frame_field_learning import polygonize_utils, plot_utils + +DEBUG = False + + +def debug_print(s: str): + if DEBUG: + print_utils.print_debug(s) + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--seg_filepath', + required=True, + nargs='*', + type=str, + help='Filepath(s) to input segmentation/mask image.') + argparser.add_argument( + '--im_dirpath', + required=True, + type=str, + help='Path to the directory containing the corresponding images os the segmentation/mask. ' + 'Files must have the same filename as --seg_filepath.' + 'Used for vizualization or saving the shapefile with the same coordinate system as that image.') + argparser.add_argument( + '--out_dirpath', + required=True, + type=str, + help='Path to the output directory.') + argparser.add_argument( + '--out_ext', + type=str, + default="shp", + choices=['pdf', 'shp'], + help="File extension of the output. " + "'pdf': pdf visualization (requires --im_dirpath for the image), 'shp': shapefile") + argparser.add_argument( + '--bbox', + nargs='*', + type=int, + help='Selects area in bbox for computation.') + + args = argparser.parse_args() + return args + + +def simplify(polygons, probs, tolerance): + if type(tolerance) == list: + out_polygons_dict = {} + out_probs_dict = {} + for tol in tolerance: + out_polygons, out_probs = simplify(polygons, probs, tol) + out_polygons_dict["tol_{}".format(tol)] = out_polygons + out_probs_dict["tol_{}".format(tol)] = out_probs + return out_polygons_dict, out_probs_dict + else: + out_polygons = [polygon.simplify(tolerance, preserve_topology=True) for polygon in polygons] + return out_polygons, probs + + +def shapely_postprocess(out_contours, np_indicator, config): + height = np_indicator.shape[0] + width = np_indicator.shape[1] + + # Handle holes: + line_string_list = [shapely.geometry.LineString(out_contour[:, ::-1]) for out_contour in out_contours] + + # Add image boundary line_strings for border polygons + line_string_list.append( + shapely.geometry.LinearRing([ + (0, 0), + (0, height - 1), + (width - 1, height - 1), + (width - 1, 0), + ])) + + # Merge polylines (for border polygons): + multi_line_string = shapely.ops.unary_union(line_string_list) + + # Find polygons: + polygons, dangles, cuts, invalids = shapely.ops.polygonize_full(multi_line_string) + + polygons = list(polygons) + + # Remove small polygons + polygons = [polygon for polygon in polygons if + config["min_area"] < polygon.area] + + # Remove low prob polygons + filtered_polygons = [] + filtered_polygon_probs = [] + for polygon in polygons: + prob = polygonize_utils.compute_geom_prob(polygon, np_indicator) + # print("simple:", np_indicator.min(), np_indicator.mean(), np_indicator.max(), prob) + if config["seg_threshold"] < prob: + filtered_polygons.append(polygon) + filtered_polygon_probs.append(prob) + + polygons, probs = simplify(filtered_polygons, filtered_polygon_probs, config["tolerance"]) + + return polygons, probs + + +def polygonize(seg_batch, config, pool=None, pre_computed=None): + # tic_total = time.time() + + assert len(seg_batch.shape) == 4 and seg_batch.shape[ + 1] <= 3, "seg_batch should be (N, C, H, W) with C <= 3, not {}".format(seg_batch.shape) + + # Indicator + # tic = time.time() + indicator_batch = seg_batch[:, 0, :, :] + np_indicator_batch = indicator_batch.cpu().numpy() + # toc = time.time() + # debug_print(f"Indicator to cpu: {toc - tic}s") + + if pre_computed is None or "init_contours_batch" not in pre_computed: + # tic = time.time() + init_contours_batch = polygonize_utils.compute_init_contours_batch(np_indicator_batch, config["data_level"], pool=pool) + # toc = time.time() + # debug_print(f"Init contours: {toc - tic}s") + else: + init_contours_batch = pre_computed["init_contours_batch"] + + # tic = time.time() + # Convert contours to shapely polygons to handle holes: + if pool is not None: + shapely_postprocess_partial = partial(shapely_postprocess, config=config) + polygons_probs_batch = pool.starmap(shapely_postprocess_partial, zip(init_contours_batch, np_indicator_batch)) + polygons_batch, probs_batch = zip(*polygons_probs_batch) + else: + polygons_batch = [] + probs_batch = [] + for i, out_contours in enumerate(init_contours_batch): + polygons, probs = shapely_postprocess(out_contours, np_indicator_batch[i], config) + polygons_batch.append(polygons) + probs_batch.append(probs) + + # toc = time.time() + # debug_print(f"Shapely post-process: {toc - tic}s") + + # toc_total = time.time() + # debug_print(f"Total: {toc_total - tic_total}s") + + return polygons_batch, probs_batch + + +def run_one(seg_filepath, out_dirpath, config, im_dirpath, out_ext=None, bbox=None): + filename = os.path.basename(seg_filepath) + name = os.path.splitext(filename)[0] + + # Load image + image = None + im_filepath = os.path.join(im_dirpath, name + ".tif") + if out_ext == "pdf": + image = skimage.io.imread(im_filepath) + + # seg = skimage.io.imread(seg_filepath) / 255 + seg_img = Image.open(seg_filepath) + seg = np.array(seg_img) + if seg.dtype == np.uint8: + seg = seg / 255 + elif seg.dtype == np.bool: + seg = seg.astype(np.float) + + # Select bbox for dev + if bbox is not None: + assert len(bbox) == 4, "bbox should have 4 values" + # bbox = [1440, 210, 1800, 650] # vienna12 + # bbox = [2808, 2393, 3124, 2772] # innsbruck19 + if image is not None: + image = image[bbox[0]:bbox[2], bbox[1]:bbox[3]] + seg = seg[bbox[0]:bbox[2], bbox[1]:bbox[3]] + extra_name = ".bbox_{}_{}_{}_{}".format(*bbox) + else: + extra_name = "" + + # Convert to torch and add batch dim + if len(seg.shape) < 3: + seg = seg[:, :, None] + seg_batch = torch.tensor(np.transpose(seg[:, :, :2], (2, 0, 1)), dtype=torch.float)[None, ...] + + out_contours_batch, out_probs_batch = polygonize(seg_batch, config) + + polygons = out_contours_batch[0] + + if out_ext == "shp": + out_filepath = os.path.join(out_dirpath, name + ".shp") + geo_utils.save_shapefile_from_shapely_polygons(polygons, im_filepath, out_filepath) + elif out_ext == "pdf": + base_filepath = os.path.splitext(seg_filepath)[0] + filepath = base_filepath + extra_name + ".poly_simple.pdf" + # plot_utils.save_poly_viz(image, polygons, filepath, linewidths=1, draw_vertices=True, color_choices=[[0, 1, 0, 1]]) + plot_utils.save_poly_viz(image, polygons, filepath, markersize=30, linewidths=1, draw_vertices=True) + else: + raise ValueError(f"out_ext should be shp or pdf, not {out_ext}") + + +def main(): + config = { + "data_level": 0.5, + "tolerance": 1.0, + "seg_threshold": 0.5, + "min_area": 10 + } + # --- Process args --- # + args = get_args() + + pool = Pool() + list(tqdm(pool.imap(partial(run_one, out_dirpath=args.out_dirpath, config=config, im_dirpath=args.im_dirpath, out_ext=args.out_ext, bbox=args.bbox), args.seg_filepath), desc="Simple poly.", total=len(args.seg_filepath))) + + +if __name__ == '__main__': + main() diff --git a/frame_field_learning/polygonize_utils.py b/frame_field_learning/polygonize_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6955ba6d88cbe6bfc882d9eae8ed37d64a5a37c6 --- /dev/null +++ b/frame_field_learning/polygonize_utils.py @@ -0,0 +1,97 @@ +from functools import partial +from collections import Iterable + +import shapely.geometry +import shapely.ops +import shapely.affinity +import numpy as np +import cv2 + +import skimage.measure + +from lydorn_utils import print_utils + + +def compute_init_contours(np_indicator, level): + assert isinstance(np_indicator, np.ndarray) and len(np_indicator.shape) == 2, "indicator should have shape (H, W)" + # Using marching squares + contours = skimage.measure.find_contours(np_indicator, level, fully_connected='low', positive_orientation='high') + + # Using OpenCV: + # u, contours, _ = cv2.findContours((level < np_indicator).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) + # contours = [contour[:, 0, ::-1] for contour in contours] + # contours = [np.concatenate((contour, contour[0:1, :]), axis=0) for contour in contours] + + # Using rasterio + # import rasterio.features + # shapes = rasterio.features.shapes((level < np_indicator).astype(np.uint8)) + # contours = [] + # for shape in shapes: + # for coords in shape[0]["coordinates"]: + # contours.append(np.array(coords)[:, ::-1] - 0.5) + + # Simplify contours a tiny bit: + # contours = [skimage.measure.approximate_polygon(contour, tolerance=0.01) for contour in contours] + return contours + + +def compute_init_contours_batch(np_indicator_batch, level, pool=None): + post_process_partial = partial(compute_init_contours, level=level) + if pool is not None: + init_contours_batch = pool.map(post_process_partial, np_indicator_batch) + else: + init_contours_batch = list(map(post_process_partial, np_indicator_batch)) + return init_contours_batch + + +def split_polylines_corner(polylines, corner_masks): + new_polylines = [] + for polyline, corner_mask in zip(polylines, corner_masks): + splits, = np.where(corner_mask) + if len(splits) == 0: + new_polylines.append(polyline) + continue + slice_list = [(splits[i], splits[i+1] + 1) for i in range(len(splits) - 1)] + for s in slice_list: + new_polylines.append(polyline[s[0]:s[1]]) + # Possibly add a merged polyline if start and end vertices are not corners (or endpoints of open polylines) + if ~corner_mask[0] and ~corner_mask[-1]: # In fact any of those conditon should be enough + new_polyline = np.concatenate([polyline[splits[-1]:], polyline[:splits[0] + 1]], axis=0) + new_polylines.append(new_polyline) + return new_polylines + + +def compute_geom_prob(geom, prob_map, output_debug=False): + assert len(prob_map.shape) == 2, "prob_map should have size (H, W), not {}".format(prob_map.shape) + + if isinstance(geom, Iterable): + return [compute_geom_prob(_geom, prob_map, output_debug=output_debug) for _geom in geom] + elif isinstance(geom, shapely.geometry.Polygon): + # --- Cut with geom bounds: + minx, miny, maxx, maxy = geom.bounds + minx = int(minx) + miny = int(miny) + maxx = int(maxx) + 1 + maxy = int(maxy) + 1 + geom = shapely.affinity.translate(geom, xoff=-minx, yoff=-miny) + prob_map = prob_map[miny:maxy, minx:maxx] + + # --- Rasterize TODO: better rasterization (or sampling) of polygon ? + raster = np.zeros(prob_map.shape, dtype=np.uint8) + exterior_array = np.round(np.array(geom.exterior.coords)).astype(np.int32) + interior_array_list = [np.round(np.array(interior.coords)).astype(np.int32) for interior in geom.interiors] + cv2.fillPoly(raster, [exterior_array], color=1) + cv2.fillPoly(raster, interior_array_list, color=0) + + raster_sum = np.sum(raster) + if 0 < raster_sum: + polygon_prob = np.sum(raster * prob_map) / raster_sum + else: + polygon_prob = 0 + if output_debug: + print_utils.print_warning("WARNING: empty polygon raster in polygonize_tracing.compute_polygon_prob().") + + return polygon_prob + else: + raise NotImplementedError(f"Geometry of type {type(geom)} not implemented!") + diff --git a/frame_field_learning/runs/mapping_dataset.unet_resnet101_pretrained.train_val _ 2020-09-07 11_28_51/args.json b/frame_field_learning/runs/mapping_dataset.unet_resnet101_pretrained.train_val _ 2020-09-07 11_28_51/args.json new file mode 100644 index 0000000000000000000000000000000000000000..fcfa20b5ba8fb8cb827523854cbc39cecbbc341a --- /dev/null +++ b/frame_field_learning/runs/mapping_dataset.unet_resnet101_pretrained.train_val _ 2020-09-07 11_28_51/args.json @@ -0,0 +1 @@ +{"run_name": "mapping_dataset.unet_resnet101_pretrained.train_val", "new_run": false, "init_run_name": null, "batch_size": 2} \ No newline at end of file diff --git a/frame_field_learning/runs/mapping_dataset.unet_resnet101_pretrained.train_val _ 2020-09-07 11_28_51/config.json b/frame_field_learning/runs/mapping_dataset.unet_resnet101_pretrained.train_val _ 2020-09-07 11_28_51/config.json new file mode 100644 index 0000000000000000000000000000000000000000..d1065a5d2822c72ad920c7ced8affc88b31b500d --- /dev/null +++ b/frame_field_learning/runs/mapping_dataset.unet_resnet101_pretrained.train_val _ 2020-09-07 11_28_51/config.json @@ -0,0 +1,249 @@ +{ + "run_name": "mapping_dataset.unet_resnet101_pretrained.train_val", + "master_addr": "localhost", + "master_port": "6666", + "fold": [ + "train", + "val" + ], + "seg_params": { + "compute_interior": true, + "compute_edge": true, + "compute_vertex": false + }, + "backbone_params": { + "pretrained": true, + "name": "unet_resnet", + "encoder_depth": 101, + "input_features": 3, + "num_filters": 32, + "dropout_2d": 0.2, + "is_deconv": false + }, + "optim_params": { + "batch_size": 10, + "max_epoch": 47, + "optimizer": "Adam", + "base_lr": 0.001, + "max_lr": 0.1, + "gamma": 0.95, + "weight_decay": 0, + "dropout_keep_prob": 1.0, + "log_steps": 200, + "checkpoint_epoch": 1, + "checkpoints_to_keep": 5, + "logs_dirname": "logs", + "checkpoints_dirname": "checkpoints" + }, + "runs_dirpath": "runs", + "new_run": false, + "init_run_name": null, + "nodes": 1, + "gpus": 0, + "nr": 0, + "world_size": 4, + "dataset_params": { + "root_dirname": "mapping_challenge_dataset", + "small": false, + "seed": 0, + "train_fraction": 0.75 + }, + "eval_params": { + "save_individual_outputs": { + "seg": true, + "poly_viz": true, + "image": false, + "seg_gt": false, + "seg_mask": true, + "seg_opencities_mask": true, + "seg_luxcarta": true, + "crossfield": false, + "uv_angles": true, + "poly_shapefile": true, + "poly_geojson": false + }, + "save_aggregated_outputs": { + "stats": false, + "seg_coco": true, + "poly_coco": true + }, + "results_dirname": "eval_runs", + "test_time_augmentation": false, + "batch_size_mult": 64, + "patch_size": 1024, + "patch_overlap": 200, + "seg_threshold": 0.5 + }, + "data_dir_candidates": [ + "/data/titane/user/nigirard/data", + "~/data", + "/data" + ], + "num_workers": 12, + "data_aug_params": { + "enable": true, + "vflip": true, + "affine": true, + "scaling": [ + 0.75, + 1.5 + ], + "color_jitter": true, + "device": "cpu" + }, + "device": "cpu", + "use_amp": false, + "compute_seg": true, + "compute_crossfield": true, + "loss_params": { + "multiloss": { + "normalization_params": { + "min_samples": 10, + "max_samples": 1000 + }, + "coefs": { + "epoch_thresholds": [ + 0, + 5 + ], + "seg": 10, + "crossfield_align": 1, + "crossfield_align90": 0.2, + "crossfield_smooth": 0.005, + "seg_interior_crossfield": [ + 0.0, + 0.2 + ], + "seg_edge_crossfield": [ + 0.0, + 0.2 + ], + "seg_edge_interior": [ + 0.0, + 0.2 + ] + } + }, + "seg_loss_params": { + "bce_coef": 1.0, + "dice_coef": 0.2, + "use_dist": true, + "use_size": true, + "w0": 50, + "sigma": 10 + } + }, + "polygonize_params": { + "method": [ + "simple", + "acm" + ], + "common_params": { + "init_data_level": 0.5 + }, + "simple_method": { + "data_level": 0.5, + "tolerance": [ + 0.125 + ], + "seg_threshold": 0.5, + "min_area": 10 + }, + "asm_method": { + "init_method": "marching_squares", + "data_level": 0.5, + "loss_params": { + "coefs": { + "step_thresholds": [ + 0, + 100, + 200, + 300 + ], + "data": [ + 1.0, + 0.1, + 0.0, + 0.0 + ], + "crossfield": [ + 0.0, + 0.05, + 0.0, + 0.0 + ], + "length": [ + 0.1, + 0.01, + 0.0, + 0.0 + ], + "curvature": [ + 0.0, + 0.0, + 1.0, + 0.0 + ], + "corner": [ + 0.0, + 0.0, + 0.5, + 0.0 + ], + "junction": [ + 0.0, + 0.0, + 0.5, + 0.0 + ] + }, + "curvature_dissimilarity_threshold": 2, + "corner_angles": [ + 45, + 90, + 135 + ], + "corner_angle_threshold": 22.5, + "junction_angles": [ + 0, + 45, + 90, + 135 + ], + "junction_angle_weights": [ + 1, + 0.01, + 0.1, + 0.01 + ], + "junction_angle_threshold": 22.5 + }, + "lr": 0.1, + "gamma": 0.995, + "device": "cpu", + "tolerance": [ + 0.125, + 1 + ], + "seg_threshold": 0.5, + "min_area": 10 + }, + "acm_method": { + "steps": 500, + "data_level": 0.5, + "data_coef": 0.1, + "length_coef": 0.4, + "crossfield_coef": 0.5, + "poly_lr": 0.01, + "warmup_iters": 100, + "warmup_factor": 0.1, + "device": "cpu", + "tolerance": [ + 0.125, + 1 + ], + "seg_threshold": 0.5, + "min_area": 10 + } + } +} diff --git a/frame_field_learning/save_utils.py b/frame_field_learning/save_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..13846f3e1a05e7157f1c04df110ab540da65c39d --- /dev/null +++ b/frame_field_learning/save_utils.py @@ -0,0 +1,283 @@ +import itertools +import os + +import numpy as np +import pycocotools.mask +import shapely.geometry +import torch +import skimage.morphology +import skimage.measure + +from lydorn_utils import python_utils, geo_utils, math_utils +import cv2 +from frame_field_learning import plot_utils, local_utils +import tifffile + + +def get_save_filepath(base_filepath, name=None, ext=""): + if type(base_filepath) is tuple: + if name is not None: + save_filepath = os.path.join(base_filepath[0], name, base_filepath[1] + ext) + else: + save_filepath = os.path.join(base_filepath[0], base_filepath[1] + ext) + elif type(base_filepath) is str: + if name is not None: + save_filepath = base_filepath + "." + name + ext + else: + save_filepath = base_filepath + ext + else: + raise TypeError(f"base_filepath should be either of tuple or str, not {type(base_filepath)}") + os.makedirs(os.path.dirname(save_filepath), exist_ok=True) + return save_filepath + + +def save_outputs(tile_data, config, eval_dirpath, split_name, flag_filepath_format): + # print("--- save_outputs() ---") + split_eval_dirpath = os.path.join(eval_dirpath, split_name) + if not os.path.exists(split_eval_dirpath): + os.makedirs(split_eval_dirpath) + + base_filepath = os.path.join(split_eval_dirpath, tile_data["name"]) + if "image_relative_filepath" in tile_data: + image_filepath = os.path.join(config["data_root_dir"], tile_data["image_relative_filepath"]) + elif "image_filepath" in tile_data: + image_filepath = tile_data["image_filepath"] + else: + raise ValueError("Could not get image_filepath from tile_data") + + if config["eval_params"]["save_individual_outputs"]["image"]: + src_filepath = "/data" + image_filepath # Because we are executing in Docker, this is a hack! + filepath = base_filepath + ".image" + if os.path.islink(filepath): + os.remove(filepath) + os.symlink(src_filepath, filepath) + if config["eval_params"]["save_individual_outputs"]["seg_gt"]: + save_seg(tile_data["gt_polygons_image"], base_filepath, "seg.gt", image_filepath) + if config["eval_params"]["save_individual_outputs"]["seg"]: + save_seg(tile_data["seg"], base_filepath, "seg", image_filepath) + if config["eval_params"]["save_individual_outputs"]["seg_mask"]: + save_seg_mask(tile_data["seg_mask"], (split_eval_dirpath, tile_data["name"]), "seg_mask", image_filepath) + if config["eval_params"]["save_individual_outputs"]["seg_opencities_mask"]: + save_opencities_mask(tile_data["seg_mask"], base_filepath, "drivendata", + image_filepath) + if config["eval_params"]["save_individual_outputs"]["seg_luxcarta"]: + save_seg_luxcarta_format(tile_data["seg"], base_filepath, "seg_luxcarta_format", + image_filepath) + + if config["eval_params"]["save_individual_outputs"]["crossfield"]: + save_crossfield(tile_data["crossfield"], base_filepath, "crossfield") + if config["eval_params"]["save_individual_outputs"]["uv_angles"]: + save_uv_angles(tile_data["crossfield"], base_filepath, "uv_angles", + image_filepath) + + if config["eval_params"]["save_individual_outputs"]["poly_shapefile"]: + save_shapefile(tile_data["polygons"], base_filepath, "poly_shapefile", image_filepath) + if config["eval_params"]["save_individual_outputs"]["poly_geojson"]: + save_geojson(tile_data["polygons"], base_filepath, "poly_geojson", image_filepath) + if "poly_viz" in config["eval_params"]["save_individual_outputs"] and config["eval_params"]["save_individual_outputs"]["poly_viz"]: + save_poly_viz(tile_data["image"], tile_data["polygons"], tile_data["polygon_probs"], base_filepath, "poly_viz") + + if "raw_pred" in config["eval_params"]["save_individual_outputs"] and config["eval_params"]["save_individual_outputs"]["raw_pred"]: + save_raw_pred(tile_data, base_filepath, "raw_pred") + + # Save a flag file to mark this sample as evaluated + # pathlib.Path(flag_filepath_format.format(tile_data["name"])).touch() + + # print("Finished saving") + + +def save_seg(seg, base_filepath, name, image_filepath): + seg = np.transpose(seg.numpy(), (1, 2, 0)) + # seg = torch_utils.to_numpy_image(seg) + seg_display = plot_utils.get_seg_display(seg) + if seg_display.dtype != np.uint8: + seg_display = (255 * seg_display).astype(np.uint8) + seg_display_filepath = get_save_filepath(base_filepath, name, ".tif") + # with warnings.catch_warnings(): + # warnings.simplefilter("ignore") + # skimage.io.imsave(seg_display_filepath_jpg, seg_display) + geo_utils.save_image_as_geotiff(seg_display_filepath, seg_display, image_filepath) + return seg_display_filepath + + +def save_seg_mask(seg_mask, base_filepath, name, image_filepath): + seg_mask = seg_mask.numpy() + out = (255 * seg_mask).astype(np.uint8)[:, :, None] + out_filepath = get_save_filepath(base_filepath, name, ".tif") + geo_utils.save_image_as_geotiff(out_filepath, out, image_filepath) + return out_filepath + + +def save_seg_luxcarta_format(seg, base_filepath, name, image_filepath): + seg = np.transpose(seg.numpy(), (1, 2, 0)) + seg_luxcarta = np.zeros((seg.shape[0], seg.shape[1], 1), dtype=np.uint8) + seg_interior = 0.5 < seg[..., 0] + seg_luxcarta[seg_interior] = 1 + if 2 <= seg.shape[2]: + seg_edge = 0.5 < seg[..., 1] + seg_luxcarta[seg_edge] = 2 + seg_luxcarta_filepath = get_save_filepath(base_filepath, name, ".tif") + geo_utils.save_image_as_geotiff(seg_luxcarta_filepath, seg_luxcarta, image_filepath) + + +def save_poly_viz(image, polygons, polygon_probs, base_filepath, name): + if type(image) == torch.Tensor: + image = np.transpose(image.numpy(), (1, 2, 0)) + + if type(polygons) == dict: + # Means several methods/settings were used + for key in polygons.keys(): + save_poly_viz(image, polygons[key], polygon_probs[key], base_filepath, name + "." + key) + elif type(polygons) == list: + filepath = get_save_filepath(base_filepath, name, ".pdf") + plot_utils.save_poly_viz(image, polygons, filepath, polygon_probs=polygon_probs) + else: + raise TypeError("polygons has unrecognized type {}".format(type(polygons))) + + + +def save_shapefile(polygons, base_filepath, name, image_filepath): + if type(polygons) == dict: + # Means several methods/settings were used + for key, item in polygons.items(): + save_shapefile(item, base_filepath, name + "." + key, image_filepath) + elif type(polygons) == list: + filepath = get_save_filepath(base_filepath, name, ".shp") + if type(polygons[0]) == np.array: + geo_utils.save_shapefile_from_polygons(polygons, image_filepath, filepath) + elif type(polygons[0]) == shapely.geometry.polygon.Polygon: + geo_utils.save_shapefile_from_shapely_polygons(polygons, image_filepath, filepath) + else: + raise TypeError("polygons has unrecognized type {}".format(type(polygons))) + + +def save_geojson(polygons, base_filepath, name=None, image_filepath=None): + # TODO: add georef and features + filepath = get_save_filepath(base_filepath, name, ".geojson") + polygons_geometry_collection = shapely.geometry.collection.GeometryCollection(polygons) + geojson = shapely.geometry.mapping(polygons_geometry_collection) + python_utils.save_json(filepath, geojson) + + +def poly_coco(polygons: list, polygon_probs: list, image_id: int): + if type(polygons) == dict: + # Means several methods/settings were used + annotations_dict = {} + for key in polygons.keys(): + _polygons = polygons[key] + _polygon_probs = polygon_probs[key] + annotations_dict[key] = poly_coco(_polygons, _polygon_probs, image_id) + return annotations_dict + elif type(polygons) == list: + annotations = [] + for polygon, prob in zip(polygons, polygon_probs): + bbox = np.round([polygon.bounds[0], polygon.bounds[1], + polygon.bounds[2] - polygon.bounds[0], polygon.bounds[3] - polygon.bounds[1]], 2) + exterior = list(np.round(np.array(polygon.exterior.coords).reshape(-1), 2)) + # interiors = [list(np.round(np.array(interior.coords).reshape(-1), 2)) for interior in polygon.interiors] + # segmentation = [exterior, *interiors] + segmentation = [exterior] + score = prob + annotation = { + "category_id": 100, # Building + "bbox": list(bbox), + "segmentation": segmentation, + "score": score, + "image_id": image_id} + annotations.append(annotation) + return annotations + else: + raise TypeError("polygons has unrecognized type {}".format(type(polygons))) + + +def save_poly_coco(annotations: list, base_filepath: str): + """ + + @param annotations: Either [[annotation1 of im1, annotation2 of im1, ...], ...] or [{"method1": [annotation1 of im1, ...]}, ...] + @param base_filepath: + @return: + """ + # seg_coco_list is either a list of annotations or a list of dictionaries for each method (and sub-methods) used + if type(annotations[0]) == dict: + # Means several methods/settings were used + # Transform list of dicts to a dict of lists: + dictionary = local_utils.list_of_dicts_to_dict_of_lists(annotations) + + dictionary = local_utils.flatten_dict(dictionary) + + for key, _annotations in dictionary.items(): + out_filepath = base_filepath + "." + key + ".json" + python_utils.save_json(out_filepath, _annotations) + elif type(annotations[0]) == list: + # Concatenate all lists + flattened_annotations = list(itertools.chain(*annotations)) + out_filepath = get_save_filepath(base_filepath, None, ".json") + python_utils.save_json(out_filepath, flattened_annotations) + else: + raise TypeError("annotations has unrecognized type {}".format(type(annotations))) + + +def seg_coco(sample): + annotations = [] + # Have to convert binary mask to a list of annotations + seg_interior = sample["seg"][0, :, :].numpy() + seg_mask = sample["seg_mask"].numpy() + labels = skimage.morphology.label(seg_mask) + properties = skimage.measure.regionprops(labels, cache=True) + for i, contour_props in enumerate(properties): + skimage_bbox = contour_props["bbox"] + obj_mask = seg_mask[skimage_bbox[0]:skimage_bbox[2], skimage_bbox[1]:skimage_bbox[3]] + obj_seg_interior = seg_interior[skimage_bbox[0]:skimage_bbox[2], skimage_bbox[1]:skimage_bbox[3]] + score = float(np.mean(obj_seg_interior * obj_mask)) + + coco_bbox = [skimage_bbox[1], skimage_bbox[0], + skimage_bbox[3] - skimage_bbox[1], skimage_bbox[2] - skimage_bbox[0]] + + image_mask = labels == (i + 1) # The mask has to span the whole image + rle = pycocotools.mask.encode(np.asfortranarray(image_mask)) + rle["counts"] = rle["counts"].decode("utf-8") + image_id = sample["image_id"].item() + annotation = { + "category_id": 100, # Building + "bbox": coco_bbox, + "segmentation": rle, + "score": score, + "image_id": image_id} + annotations.append(annotation) + return annotations + + +def save_seg_coco(sample, base_filepath, name): + filepath = get_save_filepath(base_filepath, name, ".json") + annotations = seg_coco(sample) + python_utils.save_json(filepath, annotations) + + +def save_crossfield(crossfield, base_filepath, name): + # TODO: optimize crossfield disk space + # Save raw crossfield + crossfield = np.transpose(crossfield.numpy(), (1, 2, 0)) + crossfield_filepath =get_save_filepath(base_filepath, name, ".npy") + np.save(crossfield_filepath, crossfield) + + +def save_uv_angles(crossfield, base_filepath, name, image_filepath): + crossfield = np.transpose(crossfield.numpy(), (1, 2, 0)) + u, v = math_utils.compute_crossfield_uv(crossfield) # u, v are complex arrays + u_angle, v_angle = np.angle(u), np.angle(v) + u_angle, v_angle = np.mod(u_angle, np.pi), np.mod(v_angle, np.pi) + uv_angles = np.stack([u_angle, v_angle], axis=-1) + uv_angles_as_image = np.round(uv_angles * 255 / np.pi).astype(np.uint8) + save_filepath = get_save_filepath(base_filepath, name, ".tif") + geo_utils.save_image_as_geotiff(save_filepath, uv_angles_as_image, image_filepath) + + +def save_raw_pred(sample, base_filepath, name): + save_filepath =get_save_filepath(base_filepath, name, ".pt") + torch.save(sample, save_filepath) + + +def save_opencities_mask(seg_mask, base_filepath, name, image_filepath): + seg_mask = seg_mask.numpy() + out_filepath = get_save_filepath(base_filepath, name, ".tif") + tifffile.imwrite(out_filepath, np.logical_not((np.array(seg_mask).astype(np.bool)))) diff --git a/frame_field_learning/train.py b/frame_field_learning/train.py new file mode 100644 index 0000000000000000000000000000000000000000..2a161e042fe6ed9750ccf1aa6b63d7ea782ee015 --- /dev/null +++ b/frame_field_learning/train.py @@ -0,0 +1,146 @@ +import random +import torch +import torch.utils.data +import torch.distributed + +from . import data_transforms +from .model import FrameFieldModel +from .trainer import Trainer +from . import losses +from . import local_utils + +from lydorn_utils import print_utils + +try: + import apex + from apex import amp + + APEX_AVAILABLE = True +except ModuleNotFoundError: + APEX_AVAILABLE = False + + +def count_trainable_params(model): + count = 0 + for param in model.parameters(): + if param.requires_grad: + count += param.numel() + return count + + +def train(gpu, config, shared_dict, barrier, train_ds, val_ds, backbone): + # --- Set seeds --- # + torch.manual_seed(2) # For DistributedDataParallel: make sure all models are initialized identically + # torch.backends.cudnn.deterministic = True + # torch.backends.cudnn.benchmark = False + # os.environ['CUDA_LAUNCH_BLOCKING'] = 1 + torch.autograd.set_detect_anomaly(True) + + # --- Setup DistributedDataParallel --- # + rank = config["nr"] * config["gpus"] + gpu + torch.distributed.init_process_group( + backend='nccl', + init_method='env://', + world_size=config["world_size"], + rank=rank + ) + + if gpu == 0: + print("# --- Start training --- #") + + # --- Setup run --- # + # Setup run on process 0: + if gpu == 0: + shared_dict["run_dirpath"], shared_dict["init_checkpoints_dirpath"] = local_utils.setup_run(config) + barrier.wait() # Wait on all processes so that shared_dict is synchronized. + + # Choose device + torch.cuda.set_device(gpu) + + # --- Online transform performed on the device (GPU): + train_online_cuda_transform = data_transforms.get_online_cuda_transform(config, + augmentations=config["data_aug_params"][ + "enable"]) + if val_ds is not None: + eval_online_cuda_transform = data_transforms.get_online_cuda_transform(config, augmentations=False) + else: + eval_online_cuda_transform = None + + if "samples" in config: + rng_samples = random.Random(0) + train_ds = torch.utils.data.Subset(train_ds, rng_samples.sample(range(len(train_ds)), config["samples"])) + if val_ds is not None: + val_ds = torch.utils.data.Subset(val_ds, rng_samples.sample(range(len(val_ds)), config["samples"])) + # test_ds = torch.utils.data.Subset(test_ds, list(range(config["samples"]))) + + if gpu == 0: + print(f"Train dataset has {len(train_ds)} samples.") + + train_sampler = torch.utils.data.distributed.DistributedSampler(train_ds, + num_replicas=config["world_size"], rank=rank) + val_sampler = None + if val_ds is not None: + val_sampler = torch.utils.data.distributed.DistributedSampler(val_ds, + num_replicas=config["world_size"], rank=rank) + if "samples" in config: + eval_batch_size = min(2 * config["optim_params"]["batch_size"], config["samples"]) + else: + eval_batch_size = 2 * config["optim_params"]["batch_size"] + + init_dl = torch.utils.data.DataLoader(train_ds, batch_size=eval_batch_size, pin_memory=True, + sampler=train_sampler, num_workers=config["num_workers"], drop_last=True) + train_dl = torch.utils.data.DataLoader(train_ds, batch_size=config["optim_params"]["batch_size"], shuffle=False, + pin_memory=True, sampler=train_sampler, num_workers=config["num_workers"], + drop_last=True) + if val_ds is not None: + val_dl = torch.utils.data.DataLoader(val_ds, batch_size=eval_batch_size, pin_memory=True, + sampler=val_sampler, num_workers=config["num_workers"], drop_last=True) + else: + val_dl = None + + model = FrameFieldModel(config, backbone=backbone, train_transform=train_online_cuda_transform, + eval_transform=eval_online_cuda_transform) + model.cuda(gpu) + if gpu == 0: + print("Model has {} trainable params".format(count_trainable_params(model))) + + loss_func = losses.build_combined_loss(config).cuda(gpu) + # Compute learning rate + lr = min(config["optim_params"]["base_lr"] * config["optim_params"]["batch_size"] * config["world_size"], config["optim_params"]["max_lr"]) + + if config["optim_params"]["optimizer"] == "Adam": + optimizer = torch.optim.Adam(model.parameters(), + lr=lr, + # weight_decay=config["optim_params"]["weight_decay"], + eps=1e-8 # Increase if instability is detected + ) + elif config["optim_params"]["optimizer"] == "RMSProp": + optimizer = torch.optim.RMSprop(model.parameters(), lr=lr) + else: + raise NotImplementedError(f"Optimizer {config['optim_params']['optimizer']} not recognized") + # optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9) + + if config["use_amp"] and APEX_AVAILABLE: + amp.register_float_function(torch, 'sigmoid') + model, optimizer = amp.initialize(model, optimizer, opt_level="O1") + elif config["use_amp"] and not APEX_AVAILABLE and gpu == 0: + print_utils.print_warning("WARNING: Cannot use amp because the apex library is not available!") + + # Wrap the model for distributed training + model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu], find_unused_parameters=True) + + # def lr_warmup_func(epoch): + # if epoch < config["warmup_epochs"]: + # coef = 1 + (config["warmup_factor"] - 1) * (config["warmup_epochs"] - epoch) / config["warmup_epochs"] + # else: + # coef = 1 + # return coef + # lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_warmup_func) + # lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', verbose=True) + lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, config["optim_params"]["gamma"]) + + trainer = Trainer(rank, gpu, config, model, optimizer, loss_func, + run_dirpath=shared_dict["run_dirpath"], + init_checkpoints_dirpath=shared_dict["init_checkpoints_dirpath"], + lr_scheduler=lr_scheduler) + trainer.fit(train_dl, val_dl=val_dl, init_dl=init_dl) diff --git a/frame_field_learning/trainer.py b/frame_field_learning/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..121cff1392c49cd08f5ca60b4ac6eca5578e5203 --- /dev/null +++ b/frame_field_learning/trainer.py @@ -0,0 +1,449 @@ +import os + +import torch_lydorn.torchvision +from tqdm import tqdm + +import torch +import torch.distributed + +import warnings + +with warnings.catch_warnings(): + warnings.simplefilter("ignore") + from torch.utils.tensorboard import SummaryWriter + +# from pytorch_memlab import profile, profile_every + +from . import measures, plot_utils +from . import local_utils + +from lydorn_utils import run_utils +from lydorn_utils import python_utils +from lydorn_utils import math_utils + +try: + from apex import amp + + APEX_AVAILABLE = True +except ModuleNotFoundError: + APEX_AVAILABLE = False + + +def humanbytes(B): + 'Return the given bytes as a human friendly KB, MB, GB, or TB string' + B = float(B) + KB = float(1024) + MB = float(KB ** 2) # 1,048,576 + GB = float(KB ** 3) # 1,073,741,824 + TB = float(KB ** 4) # 1,099,511,627,776 + + if B < KB: + return '{0} {1}'.format(B, 'Bytes' if 0 == B > 1 else 'Byte') + elif KB <= B < MB: + return '{0:.2f} KB'.format(B / KB) + elif MB <= B < GB: + return '{0:.2f} MB'.format(B / MB) + elif GB <= B < TB: + return '{0:.2f} GB'.format(B / GB) + elif TB <= B: + return '{0:.2f} TB'.format(B / TB) + + +class Trainer: + 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 log_weights(self, module, module_name, step): + weight_list = module.parameters() + for i, weight in enumerate(weight_list): + if len(weight.shape) == 4: + weight_type = "4d" + elif len(weight.shape) == 1: + weight_type = "1d" + elif len(weight.shape) == 2: + weight_type = "2d" + else: + weight_type = "" + self.train_writer.add_histogram('{}/{}/{}/hist'.format(module_name, i, weight_type), weight, step) + # self.writer.add_scalar('{}/{}/mean'.format(module_name, i), mean, step) + # self.writer.add_scalar('{}/{}/max'.format(module_name, i), maxi, step) + + # def log_pr_curve(self, name, pred, batch, iter_step): + # num_thresholds = 100 + # thresholds = torch.linspace(0, 2 * self.config["max_disp_global"] + self.config["max_disp_poly"], steps=num_thresholds) + # dists = measures.pos_dists(pred, batch).cpu() + # tiled_dists = dists.repeat(num_thresholds, 1) + # tiled_thresholds = thresholds.repeat(dists.shape[0], 1).t() + # true_positives = tiled_dists < tiled_thresholds + # true_positive_counts = torch.sum(true_positives, dim=1) + # recall = true_positive_counts.float() / true_positives.shape[1] + # + # precision = 1 - thresholds / (2 * self.config["max_disp_global"] + self.config["max_disp_poly"]) + # + # false_positive_counts = true_positives.shape[1] - true_positive_counts + # true_negative_counts = torch.zeros(num_thresholds) + # false_negative_counts = torch.zeros(num_thresholds) + # self.writer.add_pr_curve_raw(name, true_positive_counts, + # false_positive_counts, + # true_negative_counts, + # false_negative_counts, + # precision, + # recall, + # global_step=iter_step, + # num_thresholds=num_thresholds) + + def sync_outputs(self, loss, individual_metrics_dict): + # Reduce to rank 0: + torch.distributed.reduce(loss, dst=0) + for key in individual_metrics_dict.keys(): + torch.distributed.reduce(individual_metrics_dict[key], dst=0) + # Average on rank 0: + if self.rank == 0: + loss /= self.config["world_size"] + for key in individual_metrics_dict.keys(): + individual_metrics_dict[key] /= self.config["world_size"] + + # from pytorch_memlab import profile + # @profile + def loss_batch(self, batch, opt=None, epoch=None): + # print("Forward pass:") + # t0 = time.time() + pred, batch = self.model(batch) + # print(f"{time.time() - t0}s") + + # print("Loss computation:") + # t0 = time.time() + loss, individual_metrics_dict, extra_dict = self.loss_func(pred, batch, epoch=epoch) + # print(f"{time.time() - t0}s") + + # Compute IoUs at different thresholds + if "seg" in pred: + y_pred = pred["seg"][:, 0, ...] + y_true = batch["gt_polygons_image"][:, 0, ...] + iou_thresholds = [0.1, 0.25, 0.5, 0.75, 0.9] + for iou_threshold in iou_thresholds: + iou = measures.iou(y_pred.reshape(y_pred.shape[0], -1), y_true.reshape(y_true.shape[0], -1), threshold=iou_threshold) + mean_iou = torch.mean(iou) + individual_metrics_dict[f"IoU_{iou_threshold}"] = mean_iou + + # print("Backward pass:") + # t0 = time.time() + if opt is not None: + # Detect if loss is nan + # contains_nan = bool(torch.sum(torch.isnan(loss)).item()) + # if contains_nan: + # raise ValueError("NaN values detected, aborting...") + if self.config["use_amp"] and APEX_AVAILABLE: + with amp.scale_loss(loss, self.optimizer) as scaled_loss: + scaled_loss.backward() + else: + loss.backward() + + # torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) + + # all_grads = [] + # for param in self.model.parameters(): + # # print("shape: {}".format(param.shape)) + # if param.grad is not None: + # all_grads.append(param.grad.view(-1)) + # all_grads = torch.cat(all_grads) + # all_grads_abs = torch.abs(all_grads) + + opt.step() + opt.zero_grad() + # print(f"{time.time() - t0}s") + + # Synchronize losses/accuracies to GPU 0 so that they can be logged + self.sync_outputs(loss, individual_metrics_dict) + + for key in individual_metrics_dict: + individual_metrics_dict[key] = individual_metrics_dict[key].item() + + # Log IoU if exists + log_iou = None + iou_name = f"IoU_{0.5}" # Progress bars will show that IoU and it will be saved in checkpoints + if iou_name in individual_metrics_dict: + log_iou = individual_metrics_dict[iou_name] + + return pred, batch, loss.item(), individual_metrics_dict, extra_dict, log_iou, batch["image"].shape[0] + + def run_epoch(self, split_name, dl, epoch, log_steps=None, opt=None, iter_step=None): + assert split_name in ["train", "val"] + if split_name == "train": + writer = self.train_writer + elif split_name == "val": + writer = self.val_writer + assert iter_step is not None + else: + writer = None + + running_loss_meter = math_utils.AverageMeter("running_loss") + running_losses_meter_dict = {loss_func.name: math_utils.AverageMeter(loss_func.name) for loss_func in + self.loss_func.loss_funcs} + total_running_loss_meter = math_utils.AverageMeter("total_running_loss") + running_iou_meter = math_utils.AverageMeter("running_iou") + total_running_iou_meter = math_utils.AverageMeter("total_running_iou") + + # batch_index_offset = 0 + epoch_iterator = dl + if self.gpu == 0: + epoch_iterator = tqdm(epoch_iterator, desc="{}: ".format(split_name), leave=False) + for i, batch in enumerate(epoch_iterator): + # Send batch to device + batch = local_utils.batch_to_cuda(batch) + + # with torch.autograd.detect_anomaly(): # TODO: comment when not debugging + pred, batch, total_loss, metrics_dict, loss_extra_dict, log_iou, nums = self.loss_batch(batch, opt=opt, epoch=epoch) + # with torch.autograd.profiler.profile(use_cuda=True) as prof: + # loss, nums = self.loss_batch(batch, opt=opt) + # print(prof.key_averages().table(sort_by="cuda_time_total")) + + running_loss_meter.update(total_loss, nums) + for name, loss in metrics_dict.items(): + if name not in running_losses_meter_dict: # Init + running_losses_meter_dict[name] = math_utils.AverageMeter(name) + running_losses_meter_dict[name].update(loss, nums) + total_running_loss_meter.update(total_loss, nums) + if log_iou is not None: + running_iou_meter.update(log_iou, nums) + total_running_iou_meter.update(log_iou, nums) + + # Log values + # batch_index = i + batch_index_offset + if split_name == "train": + iter_step = epoch * len(epoch_iterator) + i + if split_name == "train" and (iter_step % log_steps == 0) or \ + split_name == "val" and i == (len(epoch_iterator) - 1): + # if iter_step % log_steps == 0: + if self.gpu == 0: + epoch_iterator.set_postfix(loss="{:.4f}".format(running_loss_meter.get_avg()), + iou="{:.4f}".format(running_iou_meter.get_avg())) + + # Logs + if self.rank == 0: + writer.add_scalar("Metrics/Loss", running_loss_meter.get_avg(), iter_step) + for key, meter in running_losses_meter_dict.items(): + writer.add_scalar(f"Metrics/{key}", meter.get_avg(), iter_step) + + image_display = torch_lydorn.torchvision.transforms.functional.batch_denormalize(batch["image"], + batch[ + "image_mean"], + batch["image_std"]) + # # Save image overlaid with gt_seg to tensorboard: + # image_gt_seg_display = plot_utils.get_tensorboard_image_seg_display(image_display, batch["gt_polygons_image"]) + # writer.add_images('gt_seg', image_gt_seg_display, iter_step) + + # Save image overlaid with seg to tensorboard: + if "seg" in pred: + crossfield = pred["crossfield"] if "crossfield" in pred else None + image_seg_display = plot_utils.get_tensorboard_image_seg_display(image_display, pred["seg"], crossfield=crossfield) + writer.add_images('seg', image_seg_display, iter_step) + + # self.log_pr_curve("PR curve/{}".format(name), pred, batch, iter_step) + + # self.log_weights(self.model.module.backbone, "backbone", iter_step) + # if hasattr(self.model.module, "seg_module"): + # self.log_weights(self.model.module.seg_module, "seg_module", iter_step) + # if hasattr(self.model.module, "crossfield_module"): + # self.log_weights(self.model.module.crossfield_module, "crossfield_module", iter_step) + + # self.writer.flush() + # im = batch["image"][0] + # self.writer.add_image('image', im) + running_loss_meter.reset() + for key, meter in running_losses_meter_dict.items(): + meter.reset() + running_iou_meter.reset() + + return total_running_loss_meter.get_avg(), total_running_iou_meter.get_avg(), iter_step + + def compute_loss_norms(self, dl, total_batches): + self.loss_func.reset_norm() + + t = None + if self.gpu == 0: + t = tqdm(total=total_batches, desc="Init loss norms", leave=True) # Initialise + + batch_i = 0 + while batch_i < total_batches: + for batch in dl: + # Update loss norms + batch = local_utils.batch_to_cuda(batch) + pred, batch = self.model(batch) + self.loss_func.update_norm(pred, batch, batch["image"].shape[0]) + if t is not None: + t.update(1) + batch_i += 1 + if not batch_i < total_batches: + break + + # Now sync loss norms across GPUs: + self.loss_func.sync(self.config["world_size"]) + + def fit(self, train_dl, val_dl=None, init_dl=None): + # Try loading previous model + checkpoint = self.load_checkpoint(self.checkpoints_dirpath) # Try last checkpoint + if checkpoint is None and self.init_checkpoints_dirpath is not None: + # Try with init_checkpoints_dirpath: + checkpoint = self.load_checkpoint(self.init_checkpoints_dirpath) + checkpoint["epoch"] = 0 # Re-start from 0 + if checkpoint is None: + checkpoint = { + "epoch": 0, + } + if init_dl is not None: + # --- Compute norms of losses on several epochs: + self.model.train() # Important for batchnorm and dropout, even in computing loss norms + with torch.no_grad(): + loss_norm_batches_min = self.config["loss_params"]["multiloss"]["normalization_params"]["min_samples"] // (2 * self.config["optim_params"]["batch_size"]) + 1 + loss_norm_batches_max = self.config["loss_params"]["multiloss"]["normalization_params"]["max_samples"] // (2 * self.config["optim_params"]["batch_size"]) + 1 + loss_norm_batches = max(loss_norm_batches_min, min(loss_norm_batches_max, len(init_dl))) + self.compute_loss_norms(init_dl, loss_norm_batches) + + if self.gpu == 0: + # Prints loss norms: + print(self.loss_func) + + start_epoch = checkpoint["epoch"] # Start at next epoch + + fit_iterator = range(start_epoch, self.config["optim_params"]["max_epoch"]) + if self.gpu == 0: + fit_iterator = tqdm(fit_iterator, desc="Fitting: ", initial=start_epoch, + total=self.config["optim_params"]["max_epoch"]) + + train_loss = None + val_loss = None + train_iou = None + epoch = None + for epoch in fit_iterator: + + self.model.train() + train_loss, train_iou, iter_step = self.run_epoch("train", train_dl, epoch, self.config["optim_params"]["log_steps"], + opt=self.optimizer) + + if val_dl is not None: + self.model.eval() + with torch.no_grad(): + val_loss, val_iou, _ = self.run_epoch("val", val_dl, epoch, self.config["optim_params"]["log_steps"], iter_step=iter_step) + else: + val_loss = None + val_iou = None + + if val_loss is not None: + self.lr_scheduler.step() + else: + self.lr_scheduler.step() + + if self.gpu == 0: + postfix_args = {"t_loss": "{:.4f}".format(train_loss), "t_iou": "{:.4f}".format(train_iou)} + if val_loss is not None: + postfix_args["v_loss"] = "{:.4f}".format(val_loss) + if val_loss is not None: + postfix_args["v_iou"] = "{:.4f}".format(val_iou) + fit_iterator.set_postfix(**postfix_args) + if self.rank == 0: + if (epoch + 1) % self.config["optim_params"]["checkpoint_epoch"] == 0: + self.save_last_checkpoint(epoch + 1, train_loss, val_loss, train_iou, + val_iou) # Save the last completed epoch, hence the "+1" + self.delete_old_checkpoint(epoch + 1) + if val_loss is not None: + self.save_best_val_checkpoint(epoch + 1, train_loss, val_loss, train_iou, val_iou) + if self.rank == 0 and epoch is not None: + self.save_last_checkpoint(epoch + 1, train_loss, val_loss, train_iou, + val_iou) # Save the last completed epoch, hence the "+1" + + def load_checkpoint(self, checkpoints_dirpath): + """ + Loads last checkpoint in checkpoints_dirpath + :param checkpoints_dirpath: + :return: + """ + try: + filepaths = python_utils.get_filepaths(checkpoints_dirpath, endswith_str=".tar", + startswith_str="checkpoint.") + if len(filepaths) == 0: + return None + + filepaths = sorted(filepaths) + filepath = filepaths[-1] # Last checkpoint + + checkpoint = torch.load(filepath, map_location="cuda:{}".format( + self.gpu)) # map_location is used to load on current device + + self.model.module.load_state_dict(checkpoint['model_state_dict']) + + self.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + self.lr_scheduler.load_state_dict(checkpoint['lr_scheduler_state_dict']) + self.loss_func.load_state_dict(checkpoint['loss_func_state_dict']) + epoch = checkpoint['epoch'] + + return { + "epoch": epoch, + } + except NotADirectoryError: + return None + + def save_checkpoint(self, filepath, epoch, train_loss, val_loss, train_acc, val_acc): + torch.save({ + 'epoch': epoch, + 'model_state_dict': self.model.module.state_dict(), # model is a DistributedDataParallel module + 'optimizer_state_dict': self.optimizer.state_dict(), + 'lr_scheduler_state_dict': self.lr_scheduler.state_dict(), + 'loss_func_state_dict': self.loss_func.state_dict(), + 'train_loss': train_loss, + 'val_loss': val_loss, + 'train_acc': train_acc, + 'val_acc': val_acc, + }, filepath) + + def save_last_checkpoint(self, epoch, train_loss, val_loss, train_acc, val_acc): + filename_format = "checkpoint.epoch_{:06d}.tar" + filepath = os.path.join(self.checkpoints_dirpath, filename_format.format(epoch)) + self.save_checkpoint(filepath, epoch, train_loss, val_loss, train_acc, val_acc) + + def delete_old_checkpoint(self, current_epoch): + filename_format = "checkpoint.epoch_{:06d}.tar" + to_delete_epoch = current_epoch - self.config["optim_params"]["checkpoints_to_keep"] * self.config["optim_params"]["checkpoint_epoch"] + filepath = os.path.join(self.checkpoints_dirpath, filename_format.format(to_delete_epoch)) + if os.path.exists(filepath): + os.remove(filepath) + + def save_best_val_checkpoint(self, epoch, train_loss, val_loss, train_acc, val_acc): + filepath = os.path.join(self.checkpoints_dirpath, "checkpoint.best_val.epoch_{:06d}.tar".format(epoch)) + + # Search for a prev best val checkpoint: + prev_filepaths = python_utils.get_filepaths(self.checkpoints_dirpath, startswith_str="checkpoint.best_val.", + endswith_str=".tar") + + if len(prev_filepaths): + prev_filepaths = sorted(prev_filepaths) + prev_filepath = prev_filepaths[-1] # Last best val checkpoint filepath in case there is more than one + + prev_best_val_checkpoint = torch.load(prev_filepath) + prev_best_loss = prev_best_val_checkpoint["val_loss"] + if val_loss < prev_best_loss: + self.save_checkpoint(filepath, epoch, train_loss, val_loss, train_acc, val_acc) + # Delete prev best val + [os.remove(prev_filepath) for prev_filepath in prev_filepaths] + else: + self.save_checkpoint(filepath, epoch, train_loss, val_loss, train_acc, val_acc) diff --git a/frame_field_learning/tta_utils.py b/frame_field_learning/tta_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b3c4bab6f879b412a65d4a3eef693b7de52516b1 --- /dev/null +++ b/frame_field_learning/tta_utils.py @@ -0,0 +1,229 @@ +import kornia +import torch +import torch_lydorn.torchvision +# from pytorch_memlab import profile, profile_every +from frame_field_learning import measures +import cv2 as cv +import numpy as np + + +def compute_distance_transform(tensor: torch.Tensor) -> torch.Tensor: + device = tensor.device + array = tensor.cpu().numpy() + shape = array.shape + array = array.reshape(-1, *shape[-2:]).astype(np.uint8) + dist_trans = np.empty(array.shape, dtype=np.float32) + for i in range(array.shape[0]): + dist_trans[i] = cv.distanceTransform(array[i], distanceType=cv.DIST_L2, maskSize=cv.DIST_MASK_5, dstType=cv.CV_64F) + dist_trans = dist_trans.reshape(shape) + dist_trans = torch.tensor(dist_trans, device=device) + return dist_trans + + +def select_crossfield(all_outputs, final_seg): + # Choose frame field from the replicate that best matches the final seg interior + dice_loss = measures.dice_loss(all_outputs["seg"][:, :, 0, :, :], final_seg[None, :, 0, :, :]) + # Get index of the replicate that achieves the min dice_loss (as it's a loss, lower is better) + indices_best = torch.argmin(dice_loss, dim=0) + batch_range = torch.arange(all_outputs["seg"].shape[1]) # batch size + # For each batch select frame field from the replicate in indices_best + final_crossfield = all_outputs["crossfield"][indices_best, batch_range] + return final_crossfield + + +def aggr_mean(all_outputs): + final_outputs = {} + if "seg" in all_outputs: + final_seg = torch.mean(all_outputs["seg"], dim=0) + final_outputs["seg"] = final_seg # Final seg is between min and max: positive pixels are closer to min + if "crossfield" in all_outputs: + final_outputs["crossfield"] = select_crossfield(all_outputs, final_seg) + else: + raise NotImplementedError("Test Time Augmentation requires segmentation to be computed.") + return final_outputs + + +def aggr_median(all_outputs): + final_outputs = {} + if "seg" in all_outputs: + final_seg, _ = torch.median(all_outputs["seg"], dim=0) + final_outputs["seg"] = final_seg # Final seg is between min and max: positive pixels are closer to min + if "crossfield" in all_outputs: + final_outputs["crossfield"] = select_crossfield(all_outputs, final_seg) + else: + raise NotImplementedError("Test Time Augmentation requires segmentation to be computed.") + return final_outputs + + +def aggr_dist_trans(all_outputs, seg_threshold): + final_outputs = {} + if "seg" in all_outputs: + min_seg, _ = torch.min(all_outputs["seg"], dim=0) + max_seg, _ = torch.max(all_outputs["seg"], dim=0) + # Final seg will be between min and max seg. The idea is that we don't loose the sharp corners (which taking the mean does) + dist_ext_to_min_seg = compute_distance_transform(min_seg < seg_threshold) + dist_int_to_max_seg = compute_distance_transform(seg_threshold < max_seg) + final_seg = dist_ext_to_min_seg < dist_int_to_max_seg + final_outputs["seg"] = final_seg # Final seg is between min and max: positive pixels are closer to min + if "crossfield" in all_outputs: + final_outputs["crossfield"] = select_crossfield(all_outputs, final_seg) + else: + raise NotImplementedError("Test Time Augmentation requires segmentation to be computed.") + return final_outputs + + +def aggr_translated(all_outputs, seg_threshold, image_display=None): + final_outputs = {} + if "seg" in all_outputs: + # Cleanup all_seg by multiplying with the mean seg + all_seg = all_outputs["seg"] + all_seg_mask: torch.Tensor = seg_threshold < all_seg + mean_seg = torch.mean(all_seg_mask.float(), dim=0) + mean_seg_mask = seg_threshold < mean_seg + all_cleaned_seg = all_seg_mask * mean_seg[None, ...] + # all_cleaned_seg_mask = seg_threshold < all_cleaned_seg + # all_cleaned_seg[~all_cleaned_seg_mask] = 0 # Put 0 where seg is below threshold + + # # --- DEBUG SAVE + # image_seg_display = plot_utils.get_tensorboard_image_seg_display(image_display, mean_seg) + # image_seg_display = image_seg_display[0].cpu().detach().numpy().transpose(1, 2, 0) + # skimage.io.imsave(f"image_seg_display_mean_seg.png", image_seg_display) + # for i, cleaned_seg in enumerate(all_cleaned_seg): + # image_seg_display = plot_utils.get_tensorboard_image_seg_display(image_display, cleaned_seg) + # image_seg_display = image_seg_display[0].cpu().detach().numpy().transpose(1, 2, 0) + # skimage.io.imsave(f"image_seg_display_replicate_cleaned_{i}.png", image_seg_display) + # # --- + + # Compute barycenter of all cleaned segs + range_x = torch.arange(all_cleaned_seg.shape[4], device=all_cleaned_seg.device) + range_y = torch.arange(all_cleaned_seg.shape[3], device=all_cleaned_seg.device) + grid_y, grid_x = torch.meshgrid([range_x, range_y]) + grid_xy = torch.stack([grid_x, grid_y], dim=-1) + + # Average of coordinates, weighted by segmentation confidence + spatial_mean_xy = torch.sum(grid_xy[None, None, None, :, :, :] * all_cleaned_seg[:, :, :, :, :, None], dim=(3, 4)) / torch.sum(all_cleaned_seg[:, :, :, :, :, None], dim=(3, 4)) + # Median of all replicate's means + median_spatial_mean_xy, _ = torch.median(spatial_mean_xy, dim=0) + # Compute shift between each replicates and the average + shift_xy = median_spatial_mean_xy[None, :, :, :] - spatial_mean_xy + shift_xy *= 2 # The shift for the original segs is twice the shift between cleaned segs (assuming homogenous shifts and enough segs) + shift_xy = shift_xy.view(-1, shift_xy.shape[-1]) + shape = all_outputs["seg"].shape + shifted_seg = kornia.geometry.translate(all_outputs["seg"].view(-1, *shape[-3:]), shift_xy).view(shape) + + # # --- DEBUG SAVE + # for i, replicate_shifted_seg in enumerate(shifted_seg): + # image_seg_display = plot_utils.get_tensorboard_image_seg_display(image_display, replicate_shifted_seg) + # image_seg_display = image_seg_display[0].cpu().detach().numpy().transpose(1, 2, 0) + # skimage.io.imsave(f"image_seg_display_replicate_shifted_{i}.png", image_seg_display) + # # --- + + # Compute mean shifted seg + mean_shifted_seg = torch.mean(shifted_seg, dim=0) + # Select replicate seg (and crossfield) that best matches mean_shifted_seg + dice_loss = measures.dice_loss(shifted_seg[:, :, 0, :, :], mean_shifted_seg[None, :, 0, :, :]) + # Get index of the replicate that achieves the min dice_loss (as it's a loss, lower is better) + indices_best = torch.argmin(dice_loss, dim=0) + batch_range = torch.arange(all_outputs["seg"].shape[1]) # batch size + # For each batch select seg and frame field from the replicate in indices_best + final_outputs["seg"] = shifted_seg[indices_best, batch_range] + if "crossfield" in all_outputs: + final_outputs["crossfield"] = all_outputs["crossfield"][indices_best, batch_range] + + # if "crossfield" in all_outputs: + # final_outputs["crossfield"] = select_crossfield(all_outputs, final_seg) + else: + raise NotImplementedError("Test Time Augmentation requires segmentation to be computed.") + return final_outputs + + +def tta_inference(model, xb, seg_threshold): + # Perform inference several times with transformed input image and aggregate results + replicates = 4 * 2 # 4 rotations, each with vflip/no vflip + + # Init results tensors + notrans_outputs = model.inference(xb["image"]) + output_keys = notrans_outputs.keys() + all_outputs = {} + for key in output_keys: + all_outputs[key] = torch.empty((replicates, *notrans_outputs[key].shape), dtype=notrans_outputs[key].dtype, + device=notrans_outputs[key].device) + all_outputs[key][0] = notrans_outputs[key] + # Flip image + flipped_image = kornia.geometry.transform.vflip(xb["image"]) + flipped_outputs = model.inference(flipped_image) + for key in output_keys: + reversed_output = kornia.geometry.transform.vflip(flipped_outputs[key]) + all_outputs[key][1] = reversed_output + + # --- Apply transforms one by one and add results to all_outputs + for k in range(1, 4): + rotated_image = torch.rot90(xb["image"], k=k, dims=(-2, -1)) + rotated_outputs = model.inference(rotated_image) + for key in output_keys: + reversed_output = torch.rot90(rotated_outputs[key], k=-k, dims=(-2, -1)) + if key == "crossfield": + angle = -k * 90 + # TODO: use a faster implementation of rotate_framefield that only handles angles [0, 90, 180, 270] + reversed_output = torch_lydorn.torchvision.transforms.functional.rotate_framefield(reversed_output, + angle) + all_outputs[key][2 * k] = reversed_output + + # Flip rotated image + flipped_rotated_image = kornia.geometry.transform.vflip(rotated_image) + flipped_rotated_outputs = model.inference(flipped_rotated_image) + for key in output_keys: + reversed_output = torch.rot90(kornia.geometry.transform.vflip(flipped_rotated_outputs[key]), k=-k, + dims=(-2, -1)) + if key == "crossfield": + angle = -k * 90 + reversed_output = torch_lydorn.torchvision.transforms.functional.vflip_framefield(reversed_output) + reversed_output = torch_lydorn.torchvision.transforms.functional.rotate_framefield(reversed_output, + angle) + all_outputs[key][2 * k + 1] = reversed_output + + # --- DEBUG + # all_outputs["seg"] *= 0 + # for i in range(all_outputs["seg"].shape[0]): + # center = 512 + # size = 100 + # shift_x = random.randint(-20, 20) + # shift_y = random.randint(-20, 20) + # all_outputs["seg"][i][..., center + shift_y - size:center + shift_y + size, center + shift_x - size:center + shift_x + size] = 1 + # # Add noise + # noise_center_x = random.randint(100, 1024-100) + # noise_center_y = random.randint(100, 1024-100) + # noise_size = 10 + # all_outputs["seg"][i][..., noise_center_y - noise_size:noise_center_y + noise_size, noise_center_x - noise_size:noise_center_x + noise_size] = 1 + # # Add more noise + # all_outputs["seg"][i] += 0.25 * torch.rand(all_outputs["seg"][i].shape, device=all_outputs["seg"][i].device) + # all_outputs["seg"][i] = torch.clamp(all_outputs["seg"][i], 0, 1) + + + # # --- DEBUG SAVE + # image_display = torch_lydorn.torchvision.transforms.functional.batch_denormalize(xb["image"], + # xb["image_mean"], + # xb["image_std"]) + # for i, replicate_seg in enumerate(all_outputs["seg"]): + # image_seg_display = plot_utils.get_tensorboard_image_seg_display(image_display, replicate_seg) + # image_seg_display = image_seg_display[0].cpu().detach().numpy().transpose(1, 2, 0) + # skimage.io.imsave(f"image_seg_display_replicate_{i}.png", image_seg_display) + # # --- + + + # --- Aggregate results + # final_outputs = aggr_dist_trans(all_outputs, seg_threshold) + # final_outputs = aggr_translated(all_outputs, seg_threshold, image_display=image_display) + # final_outputs = aggr_translated(all_outputs, seg_threshold) + final_outputs = aggr_mean(all_outputs) + # final_outputs = aggr_median(all_outputs) + + # # --- DEBUG SAVE + # image_seg_display = plot_utils.get_tensorboard_image_seg_display(image_display, final_outputs["seg"]) + # image_seg_display = image_seg_display[0].cpu().detach().numpy().transpose(1, 2, 0) + # skimage.io.imsave("image_seg_display_final.png", image_seg_display) + # # --- + + # input("Press ...") + + return final_outputs \ No newline at end of file diff --git a/frame_field_learning/unet.py b/frame_field_learning/unet.py new file mode 100644 index 0000000000000000000000000000000000000000..49c295a309e181e4af4718009bfc75cbddf27101 --- /dev/null +++ b/frame_field_learning/unet.py @@ -0,0 +1,99 @@ +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class UNetBackbone(nn.Module): + def __init__(self, n_channels, n_hidden_base, no_padding=False): + super(UNetBackbone, self).__init__() + self.no_padding = no_padding + self.inc = InConv(n_channels, n_hidden_base, no_padding) + self.down1 = Down(n_hidden_base, n_hidden_base*2, no_padding) + self.down2 = Down(n_hidden_base*2, n_hidden_base*4, no_padding) + self.down3 = Down(n_hidden_base*4, n_hidden_base*8, no_padding) + self.down4 = Down(n_hidden_base*8, n_hidden_base*16, no_padding) + + self.up1 = Up(n_hidden_base*16, n_hidden_base*8, n_hidden_base*8, no_padding) + self.up2 = Up(n_hidden_base*8, n_hidden_base*4, n_hidden_base*4, no_padding) + self.up3 = Up(n_hidden_base*4, n_hidden_base*2, n_hidden_base*2, no_padding) + self.up4 = Up(n_hidden_base*2, n_hidden_base, n_hidden_base, no_padding) + + def forward(self, x): + x0 = self.inc.forward(x) + x1 = self.down1.forward(x0) + x2 = self.down2.forward(x1) + x3 = self.down3.forward(x2) + y4 = self.down4.forward(x3) + y3 = self.up1.forward(y4, x3) + y2 = self.up2.forward(y3, x2) + y1 = self.up3.forward(y2, x1) + y0 = self.up4.forward(y1, x0) + + result = OrderedDict() + result["out"] = y0 + + return result + + +class DoubleConv(nn.Module): + """(conv => BN => ReLU) * 2""" + + def __init__(self, in_ch, out_ch, no_padding): + super(DoubleConv, self).__init__() + self.conv = nn.Sequential( + nn.Conv2d(in_ch, out_ch, 3, padding=0 if no_padding else 1, bias=True), + nn.BatchNorm2d(out_ch), + nn.ELU(), + nn.Conv2d(out_ch, out_ch, 3, padding=0 if no_padding else 1, bias=True), + nn.BatchNorm2d(out_ch), + nn.ELU() + ) + + def forward(self, x): + x = self.conv(x) + return x + + +class InConv(nn.Module): + def __init__(self, in_ch, out_ch, no_padding): + super(InConv, self).__init__() + self.conv = DoubleConv(in_ch, out_ch, no_padding) + + def forward(self, x): + x = self.conv.forward(x) + return x + + +class Down(nn.Module): + def __init__(self, in_ch, out_ch, no_padding): + super(Down, self).__init__() + self.mpconv = nn.Sequential( + nn.MaxPool2d(2), + DoubleConv(in_ch, out_ch, no_padding) + ) + + def forward(self, x): + x = self.mpconv(x) + return x + + +class Up(nn.Module): + def __init__(self, in_ch_1, in_ch_2, out_ch, no_padding): + super(Up, self).__init__() + self.conv = DoubleConv(in_ch_1 + in_ch_2, out_ch, no_padding) + + def forward(self, x1, x2): + x1 = F.interpolate(x1, scale_factor=2, mode='bilinear', align_corners=False) + + # input is CHW + diffY = x2.size()[2] - x1.size()[2] + diffX = x2.size()[3] - x1.size()[3] + + x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, + diffY // 2, diffY - diffY // 2]) + + x = torch.cat([x2, x1], dim=1) + x = self.conv.forward(x) + return x diff --git a/frame_field_learning/unet_resnet.py b/frame_field_learning/unet_resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..2fea757c83aa1c2ba1b72981babb26206025a219 --- /dev/null +++ b/frame_field_learning/unet_resnet.py @@ -0,0 +1,155 @@ +from collections import OrderedDict +from torch import nn +from torch.nn import functional as F +import torch +import torchvision + + +def conv3x3(in_, out): + return nn.Conv2d(in_, out, 3, padding=1) + + +class ConvRelu(nn.Module): + def __init__(self, in_, out): + super().__init__() + self.conv = conv3x3(in_, out) + self.activation = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.conv(x) + x = self.activation(x) + return x + + +class DecoderBlockV2(nn.Module): + def __init__(self, in_channels, middle_channels, out_channels, is_deconv=True): + super(DecoderBlockV2, self).__init__() + self.in_channels = in_channels + + if is_deconv: + """ + Parameters for Deconvolution were chosen to avoid artifacts, following + link https://distill.pub/2016/deconv-checkerboard/ + """ + + self.block = nn.Sequential( + ConvRelu(in_channels, middle_channels), + nn.ConvTranspose2d(middle_channels, out_channels, kernel_size=4, stride=2, + padding=1), + nn.ReLU(inplace=True) + ) + else: + self.block = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), + nn.Conv2d(in_channels, middle_channels, 3, padding=1, bias=True), + nn.BatchNorm2d(middle_channels), + nn.ELU(), + nn.Conv2d(middle_channels, out_channels, 3, padding=1, bias=True), + nn.BatchNorm2d(out_channels), + nn.ELU() + ) + + def forward(self, x): + return self.block(x) + + +def cat_non_matching(x1, x2): + diffY = x1.size()[2] - x2.size()[2] + diffX = x1.size()[3] - x2.size()[3] + + x2 = F.pad(x2, (diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2)) + + # for padding issues, see + # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a + # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd + + x = torch.cat([x1, x2], dim=1) + return x + + +class UNetResNetBackbone(nn.Module): + """PyTorch U-Net model using ResNet(34, 101 or 152) encoder. + UNet: https://arxiv.org/abs/1505.04597 + ResNet: https://arxiv.org/abs/1512.03385 + Proposed by Alexander Buslaev: https://www.linkedin.com/in/al-buslaev/ + Args: + encoder_depth (int): Depth of a ResNet encoder (34, 101 or 152). + num_filters (int, optional): Number of filters in the last layer of decoder. Defaults to 32. + dropout_2d (float, optional): Probability factor of dropout layer before output layer. Defaults to 0.2. + pretrained (bool, optional): + False - no pre-trained weights are being used. + True - ResNet encoder is pre-trained on ImageNet. + Defaults to False. + is_deconv (bool, optional): + False: bilinear interpolation is used in decoder. + True: deconvolution is used in decoder. + Defaults to False. + """ + + def __init__(self, encoder_depth, num_filters=32, dropout_2d=0.2, + pretrained=False, is_deconv=False): + super().__init__() + self.dropout_2d = dropout_2d + + if encoder_depth == 34: + self.encoder = torchvision.models.resnet34(pretrained=pretrained) + bottom_channel_nr = 512 + elif encoder_depth == 101: + self.encoder = torchvision.models.resnet101(pretrained=pretrained) + bottom_channel_nr = 2048 + elif encoder_depth == 152: + self.encoder = torchvision.models.resnet152(pretrained=pretrained) + bottom_channel_nr = 2048 + else: + raise NotImplementedError('only 34, 101, 152 version of ResNet are implemented') + + self.pool = nn.MaxPool2d(2, 2) + + self.relu = nn.ReLU(inplace=True) + + self.conv1 = nn.Sequential(self.encoder.conv1, + self.encoder.bn1, + self.encoder.relu, + self.pool) + + self.conv2 = self.encoder.layer1 + + self.conv3 = self.encoder.layer2 + + self.conv4 = self.encoder.layer3 + + self.conv5 = self.encoder.layer4 + + self.center = DecoderBlockV2(bottom_channel_nr, num_filters * 8 * 2, num_filters * 8, is_deconv) + self.dec5 = DecoderBlockV2(bottom_channel_nr + num_filters * 8, num_filters * 8 * 2, num_filters * 8, is_deconv) + self.dec4 = DecoderBlockV2(bottom_channel_nr // 2 + num_filters * 8, num_filters * 8 * 2, num_filters * 8, + is_deconv) + self.dec3 = DecoderBlockV2(bottom_channel_nr // 4 + num_filters * 8, num_filters * 4 * 2, num_filters * 2, + is_deconv) + self.dec2 = DecoderBlockV2(bottom_channel_nr // 8 + num_filters * 2, num_filters * 2 * 2, num_filters * 2 * 2, + is_deconv) + self.dec1 = DecoderBlockV2(num_filters * 2 * 2, num_filters * 2 * 2, num_filters, is_deconv) + + def forward(self, x): + conv1 = self.conv1(x) + conv2 = self.conv2(conv1) + conv3 = self.conv3(conv2) + conv4 = self.conv4(conv3) + conv5 = self.conv5(conv4) + + pool = self.pool(conv5) + center = self.center(pool) + + dec5 = self.dec5(cat_non_matching(conv5, center)) + + dec4 = self.dec4(cat_non_matching(conv4, dec5)) + dec3 = self.dec3(cat_non_matching(conv3, dec4)) + dec2 = self.dec2(cat_non_matching(conv2, dec3)) + dec1 = self.dec1(dec2) + + y = F.dropout2d(dec1, p=self.dropout_2d) + + result = OrderedDict() + result["out"] = y + + return result diff --git a/lydorn_utils/README.md b/lydorn_utils/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b53bd8bc93527fe1db9803b6e0340250a704b756 --- /dev/null +++ b/lydorn_utils/README.md @@ -0,0 +1,15 @@ +# lydorn_utils +Various utilities for deep learning projects. + +## print_utils + +Print in color: warning, error, info, etc. messages. + +## python_utils + +Some path utilities and more... + +## run_utils + +Manage training runs. + diff --git a/lydorn_utils/__init__.py b/lydorn_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lydorn_utils/async_utils.py b/lydorn_utils/async_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..83acfc072a1e436375e42246ff29cb0e84a12cb7 --- /dev/null +++ b/lydorn_utils/async_utils.py @@ -0,0 +1,106 @@ +import multiprocessing +import json +import time + + +def async_func_wrapper(async_func, out_queue): + while True: + if not out_queue.empty(): + data = out_queue.get() + if data is not None: + async_func(data) + else: + break + + +class AsyncMultiprocess(object): + def __init__(self, async_func, num_workers=1): + """ + + :param async_func: Takes a queue as input. Listens to the queue and perform operations + on queue elements as long as elements are not None. Stops at the first None element encountered. + :param num_workers: + """ + assert 0 < num_workers, "num_workers should be at least 1." + self.num_workers = num_workers + self.queues = [multiprocessing.Queue() for _ in range(self.num_workers)] + self.subprocesses = [ + multiprocessing.Process(target=async_func_wrapper, args=(async_func, self.queues[i])) for i in range(self.num_workers) + ] + + self.current_process = 0 + + def start(self): + for subprocess in self.subprocesses: + subprocess.start() + + def add_work(self, data): + # TODO: add work to the least busy process (shortest queue) + self.queues[self.current_process].put(data) + self.current_process = (self.current_process + 1) % self.num_workers + + def join(self): + for subprocess, queue in zip(self.subprocesses, self.queues): + queue.put(None) + subprocess.join() + + +class Async(object): + def __init__(self, async_func): + """ + + :param async_func: Takes a queue as input. Listens to the queue and perform operations + on queue elements as long as elements are not None. Stops at the first None element encountered. + """ + self.queue = multiprocessing.Queue() + self.subprocess = multiprocessing.Process(target=async_func_wrapper, args=(async_func, self.queue)) + + def start(self): + self.subprocess.start() + + def add_work(self, data): + self.queue.put(data) + + def join(self): + self.queue.put(None) + self.subprocess.join() + + +def main(): + def process(data): + print("--- process() ---") + data["out_numbers"] = [number * number for number in data["in_numbers"]] + time.sleep(0.5) + return data + + def save(data): + print("--- save() ---") + time.sleep(1) + print("Finished saving") + + num_workers = 8 + data_list = [ + { + "filepath": "out/data.{}.json".format(i), + "in_numbers": list(range(1000)) + + } for i in range(5) + ] + + # AsyncMultiprocess + print("AsyncMultiprocess") + saver_async_multiprocess = AsyncMultiprocess(save, num_workers) + saver_async_multiprocess.start() + + t0 = time.time() + for data in data_list: + data = process(data) + saver_async_multiprocess.add_work(data) + saver_async_multiprocess.join() + print("Done in {}s".format(time.time() - t0)) + + print("Finished all!") + + +if __name__ == "__main__": + main() diff --git a/lydorn_utils/geo_utils.py b/lydorn_utils/geo_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6d16cba2cf5be5c0f8678383c0c8207b2d3cc8c0 --- /dev/null +++ b/lydorn_utils/geo_utils.py @@ -0,0 +1,486 @@ +import numpy as np +import time +import json +import os.path +from tqdm import tqdm +import functools + +import rasterio +from osgeo import gdal, ogr +from osgeo import osr +import overpy +from pyproj import Proj, transform, Transformer +import fiona +import fiona.crs +import shapely.geometry +import shapely.ops + +from . import polygon_utils +from . import math_utils +from . import print_utils + +# --- Params --- # + +QUERY_BASE = \ + """ + + + + + + + + + + + """ + +WGS84_WKT = """ + GEOGCS["GCS_WGS_1984", + DATUM["WGS_1984", + SPHEROID["WGS_84",6378137,298.257223563]], + PRIMEM["Greenwich",0], + UNIT["Degree",0.017453292519943295]] + """ + +CRS = {'no_defs': True, 'ellps': 'WGS84', 'datum': 'WGS84', 'proj': 'longlat'} + + +# --- --- # + + +def get_coor_in_space(image_filepath): + """ + + :param image_filepath: Path to geo-referenced tif image + :return: coor in original space and in wsg84 spatial reference and original geotransform + :return: geo transform (x_min, res, 0, y_max, 0, -res) + :return: [[OR_x_min,OR_y_min,OR_x_max,OR_y_max],[TR_x_min,TR_y_min,TR_x_max,TR_y_max]] + """ + # print(" get_coor_in_space(image_filepath)") + ds = gdal.Open(image_filepath) + width = ds.RasterXSize + height = ds.RasterYSize + gt = ds.GetGeoTransform() + + x_min = gt[0] + y_min = gt[3] + width * gt[4] + height * gt[5] + x_max = gt[0] + width * gt[1] + height * gt[2] + y_max = gt[3] + + prj = ds.GetProjection() + srs = osr.SpatialReference(wkt=prj) + + coor_sys = srs.GetAttrValue("PROJCS|AUTHORITY", 1) + + if coor_sys is None: + coor_sys = srs.GetAttrValue("GEOGCS|AUTHORITY", 1) + + new_cs = osr.SpatialReference() + new_cs.ImportFromWkt(WGS84_WKT) + + # print(srs, new_cs) + transform = osr.CoordinateTransformation(srs, new_cs) + + lat_long_min = transform.TransformPoint(x_min, y_min) + lat_long_max = transform.TransformPoint(x_max, y_max) + + coor = [[x_min, y_min, x_max, y_max], [lat_long_min[0], lat_long_min[1], lat_long_max[0], lat_long_max[1]]] + return coor, gt, coor_sys + + +def get_osm_data(coor_query): + """ + + :param coor_query: [x_min, min_z, x_max, y_max] + :return: OSM query result + """ + api = overpy.Overpass() + query_buildings = QUERY_BASE.format("building", coor_query[1], coor_query[0], coor_query[3], coor_query[2]) + query_successful = False + wait_duration = 60 + result = None + while not query_successful: + try: + result = api.query(query_buildings) + query_successful = True + except overpy.exception.OverpassGatewayTimeout or overpy.exception.OverpassTooManyRequests or ConnectionResetError: + print("OSM server overload. Waiting for {} seconds before querying again...".format(wait_duration)) + time.sleep(wait_duration) + wait_duration *= 2 # Multiply wait time by 2 for the next time + return result + + +def proj_to_epsg_space(nodes, coor_sys): + original = Proj(CRS) + destination = Proj(init='EPSG:{}'.format(coor_sys)) + polygon = [] + for node in nodes: + polygon.append(transform(original, destination, node.lon, node.lat)) + return np.array(polygon) + + +def compute_epsg_to_image_mat(coor, gt): + x_min = coor[0][0] + y_max = coor[0][3] + + transform_mat = np.array([ + [gt[1], 0, 0], + [0, gt[5], 0], + [x_min, y_max, 1], + ]) + return np.linalg.inv(transform_mat) + + +def compute_image_to_epsg_mat(coor, gt): + x_min = coor[0][0] + y_max = coor[0][3] + + transform_mat = np.array([ + [gt[1], 0, 0], + [0, gt[5], 0], + [x_min, y_max, 1], + ]) + return transform_mat + + +def apply_transform_mat(polygon_epsg_space, transform_mat): + polygon_epsg_space_homogeneous = math_utils.to_homogeneous(polygon_epsg_space) + polygon_image_space_homogeneous = np.matmul(polygon_epsg_space_homogeneous, transform_mat) + polygon_image_space = math_utils.to_euclidian(polygon_image_space_homogeneous) + return polygon_image_space + + +def get_polygons_from_osm(image_filepath, tag="", ij_coords=True): + coor, gt, coor_system = get_coor_in_space(image_filepath) + transform_mat = compute_epsg_to_image_mat(coor, gt) + osm_data = get_osm_data(coor[1]) + + polygons = [] + for way in osm_data.ways: + if way.tags.get(tag, "n/a") != 'n/a': + polygon = way.nodes + polygon_epsg_space = proj_to_epsg_space(polygon, coor_system) + polygon_image_space = apply_transform_mat(polygon_epsg_space, transform_mat) + if ij_coords: + polygon_image_space = polygon_utils.swap_coords(polygon_image_space) + polygons.append(polygon_image_space) + + return polygons + + +def get_polygons_from_shapefile(image_filepath, input_shapefile_filepath, progressbar=True): + def process_one_polygon(polygon): + assert len(polygon.shape) == 2, "polygon should have shape (n, d), not {}".format(polygon.shape) + if 2 < polygon.shape[1]: + print_utils.print_warning( + "WARNING: polygon from shapefile has shape {}. Will discard extra values to have polygon with shape ({}, 2)".format( + polygon.shape, polygon.shape[0])) + polygon = polygon[:, :2] + polygon_epsg_space = polygon + polygon_image_space = apply_transform_mat(polygon_epsg_space, transform_mat) + polygon_image_space = polygon_utils.swap_coords(polygon_image_space) + polygons.append(polygon_image_space) + + # Extract properties: + if "properties" in parsed_json: + properties = parsed_json["properties"] + properties_list.append(properties) + + coor, gt, coor_system = get_coor_in_space(image_filepath) + transform_mat = compute_epsg_to_image_mat(coor, gt) + + file = ogr.Open(input_shapefile_filepath) + assert file is not None, "File {} does not exist!".format(input_shapefile_filepath) + shape = file.GetLayer(0) + feature_count = shape.GetFeatureCount() + polygons = [] + properties_list = [] + if progressbar: + iterator = tqdm(range(feature_count), desc="Reading features", leave=False) + else: + iterator = range(feature_count) + for feature_index in iterator: + feature = shape.GetFeature(feature_index) + raw_json = feature.ExportToJson() + parsed_json = json.loads(raw_json) + + # Extract polygon: + geometry = parsed_json["geometry"] + if geometry["type"] == "Polygon": + polygon = np.array(geometry["coordinates"][0]) # TODO: handle polygons with holes (remove [0]) + process_one_polygon(polygon) + if geometry["type"] == "MultiPolygon": + for individual_coordinates in geometry["coordinates"]: + process_one_polygon(np.array(individual_coordinates[0])) # TODO: handle polygons with holes (remove [0]) + + if properties_list: + return polygons, properties_list + else: + return polygons + + +def create_ogr_polygon(polygon, transform_mat): + polygon_swapped_coords = polygon_utils.swap_coords(polygon) + polygon_epsg = apply_transform_mat(polygon_swapped_coords, transform_mat) + + ring = ogr.Geometry(ogr.wkbLinearRing) + for coord in polygon_epsg: + ring.AddPoint(coord[0], coord[1]) + + # Create polygon + poly = ogr.Geometry(ogr.wkbPolygon) + poly.AddGeometry(ring) + return poly.ExportToWkt() + + +def create_ogr_polygons(polygons, transform_mat): + ogr_polygons = [] + for polygon in polygons: + ogr_polygons.append(create_ogr_polygon(polygon, transform_mat)) + return ogr_polygons + + +def save_image_as_geotiff(save_filepath, image, source_geotiff_filepath): + # Get geo info from source image: + source_ds = gdal.Open(source_geotiff_filepath) + if source_ds is None: + raise FileNotFoundError(f"Could not load source file {source_geotiff_filepath}") + source_gt = source_ds.GetGeoTransform() + source_prj = source_ds.GetProjection() + + driver = gdal.GetDriverByName("GTiff") + outdata = driver.Create(save_filepath, image.shape[1], image.shape[0], image.shape[2]) + outdata.SetGeoTransform(source_gt) ##sets same geotransform as input + outdata.SetProjection(source_prj) ##sets same projection as input + for i in range(image.shape[2]): + outdata.GetRasterBand(i + 1).WriteArray(image[..., i]) + outdata.FlushCache() ##saves to disk!! + outdata = None + band = None + ds = None + + +def save_shapefile_from_polygons(polygons, image_filepath, output_shapefile_filepath, properties_list=None): + """ + https://gis.stackexchange.com/a/52708/8104 + """ + assert type(polygons) == list and type(polygons[0]) == np.ndarray and \ + len(polygons[0].shape) == 2 and polygons[0].shape[1] == 2, \ + "polygons should be a list of numpy arrays with shape (N, 2)" + if properties_list is not None: + assert len(polygons) == len(properties_list), "polygons and properties_list should have the same length" + + coor, gt, coor_system = get_coor_in_space(image_filepath) + transform_mat = compute_image_to_epsg_mat(coor, gt) + # Convert polygons to ogr_polygons + ogr_polygons = create_ogr_polygons(polygons, transform_mat) + + driver = ogr.GetDriverByName('Esri Shapefile') + ds = driver.CreateDataSource(output_shapefile_filepath) + + # create the spatial reference, WGS84 + srs = osr.SpatialReference() + srs.ImportFromEPSG(4326) + + layer = ds.CreateLayer('', None, ogr.wkbPolygon) + # Add one attribute + field_name_list = [] + field_type_list = [] + if properties_list is not None: + for properties in properties_list: + for (key, value) in properties.items(): + if key not in field_name_list: + field_name_list.append(key) + field_type_list.append(type(value)) + for (name, py_type) in zip(field_name_list, field_type_list): + if py_type == int: + ogr_type = ogr.OFTInteger + elif py_type == float: + print("is float") + ogr_type = ogr.OFTReal + elif py_type == str: + ogr_type = ogr.OFTString + else: + ogr_type = ogr.OFTInteger + layer.CreateField(ogr.FieldDefn(name, ogr_type)) + + defn = layer.GetLayerDefn() + + for index in range(len(ogr_polygons)): + ogr_polygon = ogr_polygons[index] + if properties_list is not None: + properties = properties_list[index] + else: + properties = {} + + # Create a new feature (attribute and geometry) + feat = ogr.Feature(defn) + for (key, value) in properties.items(): + feat.SetField(key, value) + + # Make a geometry, from Shapely object + geom = ogr.CreateGeometryFromWkt(ogr_polygon) + feat.SetGeometry(geom) + + layer.CreateFeature(feat) + feat = geom = None # destroy these + + # Save and close everything + ds = layer = feat = geom = None + + +def save_shapefile_from_shapely_polygons(polygons, image_filepath, output_shapefile_filepath): + # Define a polygon feature geometry with one attribute + schema = { + 'geometry': 'Polygon', + 'properties': {'id': 'int'}, + } + shp_crs = "EPSG:4326" + shp_srs = Proj(shp_crs) + raster = rasterio.open(image_filepath) + # raster_srs = Proj(raster.crs) + raster_proj = lambda x, y: raster.transform * (x, y) + # shp_proj = functools.partial(transform, raster_srs, shp_srs) + # shp_proj = Transformer.from_proj(raster_srs, shp_srs).transform + + # Write a new Shapefile + os.makedirs(os.path.dirname(output_shapefile_filepath), exist_ok=True) + with fiona.open(output_shapefile_filepath, 'w', driver='ESRI Shapefile', schema=schema, crs=fiona.crs.from_epsg(4326)) as c: + for id, polygon in enumerate(polygons): + # print("---") + # print(polygon) + raster_polygon = shapely.ops.transform(raster_proj, polygon) + # print(raster_polygon) + # shp_polygon = shapely.ops.transform(shp_proj, raster_polygon) + # print(shp_polygon) + + wkt_polygon = shapely.geometry.mapping(raster_polygon) + + c.write({ + 'geometry': wkt_polygon, + 'properties': {'id': id}, + }) + + +def indices_of_biggest_intersecting_polygon(polygon_list): + """ + Assumes polygons which intersect follow each other on the order given by polygon_list. + This avoids the huge complexity of looking for an intersection between every polygon. + + :param ori_gt_polygons: + :return: + """ + keep_index_list = [] + + current_cluster = [] # Indices of the polygons belonging to the current cluster (their union has one component) + + for index, polygon in enumerate(polygon_list): + # First, check if polygon intersects with current_cluster: + current_cluster_polygons = [polygon_list[index] for index in current_cluster] + is_intersection = polygon_utils.check_intersection_with_polygons(polygon, current_cluster_polygons) + if is_intersection: + # Just add polygon to the cluster, nothing else to do + current_cluster.append(index) + else: + # This mean the current polygon is part of the next cluster. + # First, find the biggest polygon in the current cluster + cluster_max_index = 0 + cluster_max_area = 0 + for cluster_polygon_index in current_cluster: + cluster_polygon = polygon_list[cluster_polygon_index] + area = polygon_utils.polygon_area(cluster_polygon) + if cluster_max_area < area: + cluster_max_area = area + cluster_max_index = cluster_polygon_index + # Add index of the biggest polygon to the keep_index_list: + keep_index_list.append(cluster_max_index) + + # Second, create a new cluster with the current polygon index + current_cluster = [index] + + return keep_index_list + + +def get_pixelsize(filepath): + raster = gdal.Open(filepath) + gt = raster.GetGeoTransform() + pixelsize_x = gt[1] + pixelsize_y = -gt[5] + pixelsize = (pixelsize_x + pixelsize_y) / 2 + return pixelsize + + +def crop_shapefile(input_filepath, mask_filepath, output_filepath): + shp_mask_filepath = os.path.join(os.path.dirname(input_filepath), "mask.shp") + # ogr2ogr.main(["", "-f", "ESRI Shapefile", shp_mask_filepath, mask_filepath]) + # # ogr2ogr.main(["", "-f", "KML", "-clipsrc", mask_filepath, output_filepath, input_filepath]) + # # script_filepath = os.path.join(os.path.dirname(__file__), "crop_shp_with_shp.sh") + # # subprocess.Popen(["ogr2ogr", "-clipsrc", mask_filepath, output_filepath, input_filepath]) + # + # print(input_filepath) + # print(mask_filepath) + # print(output_filepath) + # callstr = ['ogr2ogr', + # "-overwrite", + # "-t_srs", + # "EPSG:27700", + # '-clipsrc', + # shp_mask_filepath, + # output_filepath, + # input_filepath, + # "-skipfailures"] + # proc = subprocess.Popen(callstr, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + # stdout, stderr = proc.communicate() + # print(stdout) + # print(stderr) + + input_file = ogr.Open(input_filepath) + assert input_file is not None, "File {} does not exist!".format(input_filepath) + input_layer = input_file.GetLayer(0) + # for i in range(input_layer.GetFeatureCount()): + # feature = input_layer.GetFeature(i) + # raw_json = feature.ExportToJson() + # parsed_json = json.loads(raw_json) + # print(parsed_json) + # break + + mask_file = ogr.Open(shp_mask_filepath) + assert mask_file is not None, "File {} does not exist!".format(shp_mask_filepath) + mask_layer = mask_file.GetLayer(0) + print(mask_layer.GetFeatureCount()) + feature = mask_layer.GetFeature(0) + raw_json = feature.ExportToJson() + parsed_json = json.loads(raw_json) + print(parsed_json) + + # create empty result layer + ogrGeometryType = ogr.Geometry(ogr.wkbPolygon) + outDriver = ogr.GetDriverByName("ESRI Shapefile") + outDs = outDriver.CreateDataSource(output_filepath) + outLayer = outDs.CreateLayer('', None, ogr.wkbPolygon) + + input_layer.Intersection(mask_layer, outLayer, options=["SKIP_FAILURES=YES"]) + + +def main(): + main_dirpath = "/workspace/data/stereo_dataset/raw/leibnitz" + image_filepath = os.path.join(main_dirpath, "leibnitz_ortho_ref_RGB.tif") + input_shapefile_filepath = os.path.join(main_dirpath, "Leibnitz_buildings_ref.shp") + output_shapefile_filepath = os.path.join(main_dirpath, "Leibnitz_buildings_ref.shifted.shp") + + polygons, properties_list = get_polygons_from_shapefile(image_filepath, input_shapefile_filepath) + print(polygons[0]) + print(properties_list[0]) + + # Add shift + shift = np.array([0, 0]) + shifted_polygons = [polygon + shift for polygon in polygons] + print(shifted_polygons[0]) + + # Save shapefile + save_shapefile_from_polygons(shifted_polygons, image_filepath, output_shapefile_filepath, properties_list=properties_list) + + +if __name__ == "__main__": + main() diff --git a/lydorn_utils/image_utils.py b/lydorn_utils/image_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7aa6a1cafeddb863fadedf35fa8370be64f45c5e --- /dev/null +++ b/lydorn_utils/image_utils.py @@ -0,0 +1,252 @@ +from io import BytesIO +import math +import numpy as np +from PIL import Image +import skimage.draw + +from . import python_utils + +CV2 = False +if python_utils.module_exists("cv2"): + import cv2 + CV2 = True + +if python_utils.module_exists("matplotlib.pyplot"): + import matplotlib.pyplot as plt + + +def get_image_size(filepath): + im = Image.open(filepath) + return im.size + + +def load_image(image_filepath): + image = Image.open(image_filepath) + image.load() + image_array = np.array(image, dtype=np.uint8) + image.close() + return image_array + + +def padded_boundingbox(boundingbox, padding): + boundingbox_new = np.empty_like(boundingbox) + boundingbox_new[0:2] = boundingbox[0:2] + padding + boundingbox_new[2:4] = boundingbox[2:4] - padding + return boundingbox_new + + +def center_bbox(spatial_shape, output_shape): + """ + Return a bbox centered in spatial_shape with size output_shape + + :param spatial_shape: + :param output_shape: + :return: + """ + center = (spatial_shape[0] / 2, spatial_shape[1] / 2) + half_output_shape = (output_shape[0] / 2, output_shape[1] / 2) + bbox = [center[0] - half_output_shape[0], center[1] - half_output_shape[1], center[0] + half_output_shape[0], center[1] + half_output_shape[1]] + bbox = bbox_to_int(bbox) + return bbox + + +def bbox_add_margin(bbox, margin): + bbox_new = bbox.copy() + bbox_new[0:2] -= margin + bbox_new[2:4] += margin + return bbox_new + + +def bbox_to_int(bbox): + bbox_new = [ + int(np.floor(bbox[0])), + int(np.floor(bbox[1])), + int(np.ceil(bbox[2])), + int(np.ceil(bbox[3])), + ] + return bbox_new + + +def draw_line_aa_in_patch(edge, patch_bounds): + rr, cc, prob = skimage.draw.line_aa(edge[0][0], edge[0][1], edge[1][0], edge[1][1]) + keep_mask = (patch_bounds[0] <= rr) & (rr < patch_bounds[2]) \ + & (patch_bounds[1] <= cc) & (cc < patch_bounds[3]) + rr = rr[keep_mask] + cc = cc[keep_mask] + prob = prob[keep_mask] + return rr, cc, prob + + +def convert_array_to_jpg_bytes(image_array, mode=None): + img = Image.fromarray(image_array, mode=mode) + output = BytesIO() + img.save(output, format="JPEG", quality=90) + contents = output.getvalue() + output.close() + return contents + + +def displacement_map_to_transformation_maps(disp_field_map): + disp_field_map = disp_field_map.astype(np.float32) + i = np.arange(disp_field_map.shape[0], dtype=np.float32) + j = np.arange(disp_field_map.shape[1], dtype=np.float32) + iv, jv = np.meshgrid(i, j, indexing="ij") + reverse_map_i = iv + disp_field_map[:, :, 1] + reverse_map_j = jv + disp_field_map[:, :, 0] + return reverse_map_i, reverse_map_j + +if CV2: + def apply_displacement_field_to_image(image, disp_field_map): + trans_map_i, trans_map_j = displacement_map_to_transformation_maps(disp_field_map) + misaligned_image = cv2.remap(image, trans_map_j, trans_map_i, cv2.INTER_CUBIC) + return misaligned_image + + + def apply_displacement_fields_to_image(image, disp_field_maps): + disp_field_map_count = disp_field_maps.shape[0] + misaligned_image_list = [] + for i in range(disp_field_map_count): + misaligned_image = apply_displacement_field_to_image(image, disp_field_maps[i, :, :, :]) + misaligned_image_list.append(misaligned_image) + return misaligned_image_list +else: + def apply_displacement_fields_to_image(image, disp_field_map): + print("cv2 is not available, the apply_displacement_fields_to_image(image, disp_field_map) function cannot work!") + + def apply_displacement_fields_to_image(image, disp_field_maps): + print("cv2 is not available, the apply_displacement_fields_to_image(image, disp_field_maps) function cannot work!") + + +def get_axis_patch_count(length, stride, patch_res): + total_double_padding = patch_res - stride + patch_count = max(1, int(math.ceil((length - total_double_padding) / stride))) + return patch_count + + +def compute_patch_boundingboxes(image_size, stride, patch_res): + """ + + @param image_size: + @param stride: + @param patch_res: + @return: [[row_start, col_start, row_end, col_end], ...] + """ + im_rows = image_size[0] + im_cols = image_size[1] + + row_patch_count = get_axis_patch_count(im_rows, stride, patch_res) + col_patch_count = get_axis_patch_count(im_cols, stride, patch_res) + + patch_boundingboxes = [] + for i in range(0, row_patch_count): + if i < row_patch_count - 1: + row_slice_begin = i * stride + row_slice_end = row_slice_begin + patch_res + else: + row_slice_end = im_rows + row_slice_begin = row_slice_end - patch_res + for j in range(0, col_patch_count): + if j < col_patch_count - 1: + col_slice_begin = j*stride + col_slice_end = col_slice_begin + patch_res + else: + col_slice_end = im_cols + col_slice_begin = col_slice_end - patch_res + + patch_boundingbox = np.array([row_slice_begin, col_slice_begin, row_slice_end, col_slice_end], dtype=np.int) + assert row_slice_end - row_slice_begin == col_slice_end - col_slice_begin == patch_res, "ERROR: patch does not have the requested shape" + patch_boundingboxes.append(patch_boundingbox) + + return patch_boundingboxes + + +def clip_boundingbox(boundingbox, clip_list): + assert len(boundingbox) == len(clip_list), "len(boundingbox) should be equal to len(clip_values)" + clipped_boundingbox = [] + for bb_value, clip in zip(boundingbox[:2], clip_list[:2]): + clipped_value = max(clip, bb_value) + clipped_boundingbox.append(clipped_value) + for bb_value, clip in zip(boundingbox[2:], clip_list[2:]): + clipped_value = min(clip, bb_value) + clipped_boundingbox.append(clipped_value) + return clipped_boundingbox + + +def crop_or_pad_image_with_boundingbox(image, patch_boundingbox): + im_rows = image.shape[0] + im_cols = image.shape[1] + + row_padding_before = max(0, - patch_boundingbox[0]) + col_padding_before = max(0, - patch_boundingbox[1]) + row_padding_after = max(0, patch_boundingbox[2] - im_rows) + col_padding_after = max(0, patch_boundingbox[3] - im_cols) + + # Center padding: + row_padding = row_padding_before + row_padding_after + col_padding = col_padding_before + col_padding_after + row_padding_before = row_padding // 2 + col_padding_before = col_padding // 2 + row_padding_after = row_padding - row_padding // 2 + col_padding_after = col_padding - col_padding // 2 + + clipped_patch_boundingbox = clip_boundingbox(patch_boundingbox, [0, 0, im_rows, im_cols]) + + if len(image.shape) == 2: + patch = image[clipped_patch_boundingbox[0]:clipped_patch_boundingbox[2], clipped_patch_boundingbox[1]:clipped_patch_boundingbox[3]] + patch = np.pad(patch, [(row_padding_before, row_padding_after), (col_padding_before, col_padding_after)], mode="constant") + elif len(image.shape) == 3: + patch = image[clipped_patch_boundingbox[0]:clipped_patch_boundingbox[2], clipped_patch_boundingbox[1]:clipped_patch_boundingbox[3], :] + patch = np.pad(patch, [(row_padding_before, row_padding_after), (col_padding_before, col_padding_after), (0, 0)], mode="constant") + else: + print("Image input does not have the right shape/") + patch = None + return patch + + +def make_grid(images, padding=2, pad_value=0, return_offsets=False): + nmaps = images.shape[0] + ymaps = int(math.floor(math.sqrt(nmaps))) + xmaps = nmaps // ymaps + height, width = int(images.shape[1] + padding), int(images.shape[2] + padding) + grid = np.zeros((height * ymaps + padding, width * xmaps + padding, images.shape[3])) + pad_value + k = 0 + offsets = [] + for y in range(ymaps): + for x in range(xmaps): + if k >= nmaps: + break + x_offset = x * width + padding + y_offset = y * height + padding + grid[y * height + padding:(y+1) * height, x * width + padding:(x+1) * width, :] = images[k] + offsets.append((x_offset, y_offset)) + k = k + 1 + if return_offsets: + return grid, offsets + else: + return grid + + +if __name__ == "__main__": + im_rows = 5 + im_cols = 10 + stride = 1 + patch_res = 15 + + image = np.random.randint(0, 256, size=(im_rows, im_cols, 3), dtype=np.uint8) + image = Image.fromarray(image) + image = np.array(image) + plt.ion() + plt.figure(1) + plt.imshow(image) + plt.show() + + # Cut patches + patch_boundingboxes = compute_patch_boundingboxes(image.shape[0:2], stride, patch_res) + + plt.figure(2) + + for patch_boundingbox in patch_boundingboxes: + patch = crop_or_pad_image_with_boundingbox(image, patch_boundingbox) + plt.imshow(patch) + plt.show() + input("Press to finish...") diff --git a/lydorn_utils/math_utils.py b/lydorn_utils/math_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fa8ea7dbb0cb0bb4bcc4fbc4d97a1b0b2c7a477c --- /dev/null +++ b/lydorn_utils/math_utils.py @@ -0,0 +1,521 @@ +import numpy as np +import time +import sklearn.datasets +import skimage.transform +import scipy.stats + +from . import python_utils +from . import image_utils + +# if python_utils.module_exists("matplotlib.pyplot"): +# import matplotlib.pyplot as plt + +CV2 = False +if python_utils.module_exists("cv2"): + import cv2 + + CV2 = True + + +# import multiprocessing +# +# import python_utils +# +# if python_utils.module_exists("joblib"): +# from joblib import Parallel, delayed +# JOBLIB = True +# else: +# JOBLIB = False + + +# def plot_field_map(field_map): +# from mpl_toolkits.mplot3d import Axes3D +# +# row = np.linspace(0, 1, field_map.shape[0]) +# col = np.linspace(0, 1, field_map.shape[1]) +# rr, cc = np.meshgrid(row, col, indexing='ij') +# +# fig = plt.figure(figsize=(18, 9)) +# ax = fig.add_subplot(121, projection='3d') +# ax.plot_surface(rr, cc, field_map[:, :, 0], rstride=3, cstride=3, linewidth=1, antialiased=True) +# +# ax = fig.add_subplot(122, projection='3d') +# ax.plot_surface(rr, cc, field_map[:, :, 1], rstride=3, cstride=3, linewidth=1, antialiased=True) +# +# plt.show() + +# --- Classes --- # + +class AverageMeter(object): + """Computes and stores the average and current value""" + + def __init__(self, name="", init_val=0, fmt=':f'): + self.name = name + self.init_val = init_val + self.fmt = fmt + self.val = self.avg = self.init_val + self.sum = self.count = 0 + + def reset(self): + self.val = self.avg = self.init_val + self.sum = self.count = 0 + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / self.count + + def get_avg(self): + return self.avg + + def __str__(self): + fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})' + return fmtstr.format(**self.__dict__) + + +class RunningDecayingAverage(object): + """ + Updates average with val*(1 - decay) + avg*decay + """ + def __init__(self, decay, init_val=0): + assert 0 < decay < 1 + self.decay = decay + self.init_val = init_val + self.val = self.avg = self.init_val + + def reset(self): + self.val = self.avg = self.init_val + + def update(self, val): + self.val = val + self.avg = (1 - self.decay)*val + self.decay*self.avg + + def get_avg(self): + return self.avg + + +class DispFieldMapsPatchCreator: + def __init__(self, global_shape, patch_res, map_count, modes, gauss_mu_range, gauss_sig_scaling): + self.global_shape = global_shape + self.patch_res = patch_res + self.map_count = map_count + self.modes = modes + self.gauss_mu_range = gauss_mu_range + self.gauss_sig_scaling = gauss_sig_scaling + + self.current_patch_index = -1 + self.patch_boundingboxes = image_utils.compute_patch_boundingboxes(self.global_shape, stride=self.patch_res, + patch_res=self.patch_res) + self.disp_maps = None + self.create_new_disp_maps() + + def create_new_disp_maps(self): + print("DispFieldMapsPatchCreator.create_new_disp_maps()") + self.disp_maps = create_displacement_field_maps(self.global_shape, self.map_count, self.modes, + self.gauss_mu_range, self.gauss_sig_scaling) + + def get_patch(self): + self.current_patch_index += 1 + + if len(self.patch_boundingboxes) <= self.current_patch_index: + self.current_patch_index = 0 + self.create_new_disp_maps() + + patch_boundingbox = self.patch_boundingboxes[self.current_patch_index] + patch_disp_maps = self.disp_maps[:, patch_boundingbox[0]:patch_boundingbox[2], + patch_boundingbox[1]:patch_boundingbox[3], :] + return patch_disp_maps + + +# --- --- # + +def compute_crossfield_c0c2(u, v): + c0 = np.power(u, 2) * np.power(v, 2) + c2 = - (np.power(u, 2) + np.power(v, 2)) + crossfield = np.stack([c0.real, c0.imag, c2.real, c2.imag], axis=-1) + return crossfield + + +def compute_crossfield_uv(c0c2): + c0 = c0c2[..., 0] + 1j * c0c2[..., 1] + c2 = c0c2[..., 2] + 1j * c0c2[..., 3] + sqrt_c2_squared_minus_4c0 = np.sqrt(np.power(c2, 2) - 4 * c0) + u_squared = (c2 + sqrt_c2_squared_minus_4c0) / 2 + v_squared = (c2 - sqrt_c2_squared_minus_4c0) / 2 + u = np.sqrt(u_squared) + v = np.sqrt(v_squared) + return u, v + + +def to_homogeneous(array): + new_array = np.ones((array.shape[0], array.shape[1] + 1), dtype=array.dtype) + new_array[..., :-1] = array + return new_array + + +def to_euclidian(array_homogeneous): + array = array_homogeneous[:, 0:2] / array_homogeneous[:, 2:3] + return array + + +def stretch(array): + mini = np.min(array) + maxi = np.max(array) + if maxi - mini: + array -= mini + array *= 2 / (maxi - mini) + array -= 1 + return array + + +def crop_center(array, out_shape): + assert len(out_shape) == 2, "out_shape should be of length 2" + in_shape = np.array(array.shape[:2]) + start = in_shape // 2 - (out_shape // 2) + out_array = array[start[0]:start[0] + out_shape[0], start[1]:start[1] + out_shape[1], ...] + return out_array + + +def multivariate_gaussian(pos, mu, sigma): + """Return the multivariate Gaussian distribution on array pos. + + pos is an array constructed by packing the meshed arrays of variables + x_1, x_2, x_3, ..., x_k into its _last_ dimension. + + """ + + n = mu.shape[0] + sigma_det = np.linalg.det(sigma) + sigma_inv = np.linalg.inv(sigma) + N = np.sqrt((2 * np.pi) ** n * sigma_det) + # This einsum call calculates (x-mu)T.sigma-1.(x-mu) in a vectorized + # way across all the input variables. + + # print("\tStarting to create multivariate Gaussian") + # start = time.time() + + # print((pos - mu).shape) + # print(sigma_inv.shape) + try: + fac = np.einsum('...k,kl,...l->...', pos - mu, sigma_inv, pos - mu, optimize=True) + except: + fac = np.einsum('...k,kl,...l->...', pos - mu, sigma_inv, pos - mu) + # print(fac.shape) + + # end = time.time() + # print("\tFinished Gaussian in {}s".format(end - start)) + + return np.exp(-fac / 2) / N + + +def create_multivariate_gaussian_mixture_map(shape, mode_count, mu_range, sig_scaling): + shape = np.array(shape) + # print("Starting to create multivariate Gaussian mixture") + # main_start = time.time() + + dim_count = 2 + downsample_factor = 4 + dtype = np.float32 + + mu_scale = mu_range[1] - mu_range[0] + row = np.linspace(mu_range[0], mu_range[1], mu_scale * shape[0] / downsample_factor, dtype=dtype) + col = np.linspace(mu_range[0], mu_range[1], mu_scale * shape[1] / downsample_factor, dtype=dtype) + rr, cc = np.meshgrid(row, col, indexing='ij') + grid = np.stack([rr, cc], axis=2) + + mus = np.random.uniform(mu_range[0], mu_range[1], (mode_count, dim_count, 2)).astype(dtype) + # gams = np.random.rand(mode_count, dim_count, 2, 2).astype(dtype) + signs = np.random.choice([1, -1], size=(mode_count, dim_count)) + + # print("\tAdding gaussian mixtures one by one") + # start = time.time() + + # if JOBLIB: + # # Parallel computing of multivariate gaussians + # inputs = range(8) + # + # def processInput(i): + # size = 10 * i + 2000 + # a = np.random.random_sample((size, size)) + # b = np.random.random_sample((size, size)) + # n = np.dot(a, b) + # return n + # + # num_cores = multiprocessing.cpu_count() + # print("num_cores: {}".format(num_cores)) + # # num_cores = 1 + # + # results = Parallel(n_jobs=num_cores)(delayed(processInput)(i) for i in inputs) + # for result in results: + # print(result.shape) + # + # gaussian_mixture = np.zeros_like(grid) + # else: + gaussian_mixture = np.zeros_like(grid) + for mode_index in range(mode_count): + for dim in range(dim_count): + sig = (sig_scaling[1] - sig_scaling[0]) * sklearn.datasets.make_spd_matrix(2) + sig_scaling[0] + # sig = (sig_scaling[1] - sig_scaling[0]) * np.dot(gams[mode_index, dim], np.transpose(gams[mode_index, dim])) + sig_scaling[0] + sig = sig.astype(dtype) + multivariate_gaussian_grid = signs[mode_index, dim] * multivariate_gaussian(grid, mus[mode_index, dim], sig) + gaussian_mixture[:, :, dim] += multivariate_gaussian_grid + + # end = time.time() + # print("\tFinished adding gaussian mixtures in {}s".format(end - start)) + + # squared_gaussian_mixture = np.square(gaussian_mixture) + # magnitude_disp_field_map = np.sqrt(squared_gaussian_mixture[:, :, 0] + squared_gaussian_mixture[:, :, 1]) + # max_magnitude = magnitude_disp_field_map.max() + + gaussian_mixture[:, :, 0] = stretch(gaussian_mixture[:, :, 0]) + gaussian_mixture[:, :, 1] = stretch(gaussian_mixture[:, :, 1]) + + # Crop + gaussian_mixture = crop_center(gaussian_mixture, shape // downsample_factor) + + # plot_field_map(gaussian_mixture) + + # Upsample mixture + # gaussian_mixture = skimage.transform.rescale(gaussian_mixture, downsample_factor) + gaussian_mixture = skimage.transform.resize(gaussian_mixture, shape) + + main_end = time.time() + # print("Finished multivariate Gaussian mixture in {}s".format(main_end - main_start)) + + return gaussian_mixture + + +def create_displacement_field_maps(shape, map_count, modes, gauss_mu_range, gauss_sig_scaling, seed=None): + if seed is not None: + np.random.seed(seed) + disp_field_maps_list = [] + for disp_field_map_index in range(map_count): + disp_field_map_normed = create_multivariate_gaussian_mixture_map(shape, + modes, + gauss_mu_range, + gauss_sig_scaling) + disp_field_maps_list.append(disp_field_map_normed) + disp_field_maps = np.stack(disp_field_maps_list, axis=0) + + return disp_field_maps + + +def get_h_mat(t, theta, scale_offset, shear, p): + """ + Computes the homography matrix given the parameters + See https://medium.com/uruvideo/dataset-augmentation-with-random-homographies-a8f4b44830d4 + (fixed mistake in H_a) + + :param t: 2D translation vector + :param theta: Scalar angle + :param scale_offset: 2D scaling vector + :param shear: 2D shearing vector + :param p: 2D projection vector + :return: h_mat: shape (3, 3) + """ + cos_theta = np.cos(theta) + sin_theta = np.sin(theta) + h_e = np.array([ + [cos_theta, -sin_theta, t[0]], + [sin_theta, cos_theta, t[1]], + [0, 0, 1], + ]) + h_a = np.array([ + [1 + scale_offset[0], shear[1], 0], + [shear[0], 1 + scale_offset[1], 0], + [0, 0, 1], + ]) + h_p = np.array([ + [1, 0, 0], + [0, 1, 0], + [p[0], p[1], 1], + ]) + h_mat = h_e @ h_a @ h_p + return h_mat + + +if CV2: + def find_homography_4pt(src, dst): + """ + Estimates the homography that transforms src points into dst points. + Then converts the matrix representation into the 4 points representation. + + :param src: + :param dst: + :return: + """ + h_mat, _ = cv2.findHomography(src, dst) + h_4pt = convert_h_mat_to_4pt(h_mat) + return h_4pt + + + def convert_h_mat_to_4pt(h_mat): + src_4pt = np.array([[ + [-1, -1], + [1, -1], + [1, 1], + [-1, 1], + ]], dtype=np.float64) + h_4pt = cv2.perspectiveTransform(src_4pt, h_mat) + return h_4pt + + + def convert_h_4pt_to_mat(h_4pt): + src_4pt = np.array([ + [-1, -1], + [1, -1], + [1, 1], + [-1, 1], + ], dtype=np.float32) + h_4pt = h_4pt.astype(np.float32) + h_mat = cv2.getPerspectiveTransform(src_4pt, h_4pt) + return h_mat + + + def field_map_to_image(field_map): + mag, ang = cv2.cartToPolar(field_map[..., 0], field_map[..., 1]) + hsv = np.zeros((field_map.shape[0], field_map.shape[1], 3)) + hsv[..., 0] = ang * 180 / np.pi / 2 + hsv[..., 1] = 255 + hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX) + hsv = hsv.astype(np.uint8) + rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) + return rgb +else: + def find_homography_4pt(src, dst): + print("cv2 is not available, the find_homography_4pt(src, dst) function cannot work!") + + + def convert_h_mat_to_4pt(h_mat): + print("cv2 is not available, the convert_h_mat_to_4pt(h_mat) function cannot work!") + + + def convert_h_4pt_to_mat(h_4pt): + print("cv2 is not available, the convert_h_4pt_to_mat(h_4pt) function cannot work!") + + + def field_map_to_image(field_map): + print("cv2 is not available, the field_map_to_image(field_map) function cannot work!") + + +def circular_diff(a1, a2, range_max): + """ + Compute difference between a1 and a2 belonging to the circular interval [0, range_max). + For example to compute angle difference, use range_max=2*PI. + a1 and a2 must be between range_min and range_max! + Thus difference between 0 and range_max is 0. + :param a1: numpy array + :param a2: numpy array + :param range_max: + :return: + """ + d = range_max / 2 - np.abs(np.abs(a1 - a2) - range_max / 2) + return d + + +def invert_permutation(p): + '''The argument p is assumed to be some permutation of 0, 1, ..., len(p)-1. + Returns an array s, where s[i] gives the index of i in p. + ''' + s = np.empty(p.size, p.dtype) + s[p] = np.arange(p.size) + return s + + +def region_growing_1d(array, max_range, max_skew): + """ + :param array: + :param max_var: + :param max_mean_median_diff: + :return: + """ + + def verify_predicate(region): + """ + Region is sorted + :param region: + :return: + """ + skew = scipy.stats.skew(region) + return region[-1] - region[0] < max_range and abs(skew) < max_skew + + assert len(array.shape) == 1, "array should be 1d, not {}".format(array.shape) + p = np.argsort(array) + sorted_array = array[p] + + labels = np.zeros(len(sorted_array), dtype=np.long) + region_start = 0 + region_label = 1 + labels[region_start] = region_label + centers = [] + for i in range(1, len(sorted_array)): + region = sorted_array[region_start:i + 1] + if not verify_predicate(region): + # End current region + median = region[len(region) // 2] # region is sorted + centers.append(median) + # Begin a new region + region_start = i + region_label += 1 + labels[i] = region_label + centers.append(median) + + return labels[invert_permutation(p)], centers + + +def bilinear_interpolate(im, pos): + # From https://gist.github.com/peteflorence/a1da2c759ca1ac2b74af9a83f69ce20e + x = pos[..., 1] + y = pos[..., 0] + + x0 = np.floor(x).astype(int) + x1 = x0 + 1 + y0 = np.floor(y).astype(int) + y1 = y0 + 1 + + x0_clipped = np.clip(x0, 0, im.shape[1] - 1) + x1_clipped = np.clip(x1, 0, im.shape[1] - 1) + y0_clipped = np.clip(y0, 0, im.shape[0] - 1) + y1_clipped = np.clip(y1, 0, im.shape[0] - 1) + + Ia = im[y0_clipped, x0_clipped] + Ib = im[y1_clipped, x0_clipped] + Ic = im[y0_clipped, x1_clipped] + Id = im[y1_clipped, x1_clipped] + + wa = (x1 - x) * (y1 - y) + wb = (x1 - x) * (y - y0) + wc = (x - x0) * (y1 - y) + wd = (x - x0) * (y - y0) + + value = (Ia.T * wa).T + (Ib.T * wb).T + (Ic.T * wc).T + (Id.T * wd).T + + return value + + +def main(): + import matplotlib.pyplot as plt + # shape = (220, 220) + # mode_count = 30 + # mu_range = [0, 1] + # sig_scaling = [0.0, 0.002] + # create_multivariate_gaussian_mixture_map(shape, mode_count, mu_range, sig_scaling) + + # a1 = np.array([0.0]) + # a2 = np.array([3*np.pi/4]) + # range_max = np.pi + # d = circular_diff(a1, a2, range_max) + # print(d) + + array = np.concatenate([np.arange(1, 1.01, 0.001), np.arange(0, np.pi / 2, np.pi / 100)]) + print(array) + labels = region_growing_1d(array, max_range=np.pi / 10, max_skew=1) + print(labels) + + plt.plot(array, labels, ".") + plt.show() + + +if __name__ == "__main__": + main() diff --git a/lydorn_utils/ogr2ogr.py b/lydorn_utils/ogr2ogr.py new file mode 100644 index 0000000000000000000000000000000000000000..f3211ce987b9718b8fd4156b8c70cf43c605e3d7 --- /dev/null +++ b/lydorn_utils/ogr2ogr.py @@ -0,0 +1,1696 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +#***************************************************************************** +# $Id$ +# +# Project: OpenGIS Simple Features Reference Implementation +# Purpose: Python port of a simple client for translating between formats. +# Author: Even Rouault, +# +# Port from ogr2ogr.cpp whose author is Frank Warmerdam +# +#***************************************************************************** +# Copyright (c) 2010-2013, Even Rouault +# Copyright (c) 1999, Frank Warmerdam +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included +# in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +#************************************************************************** + +# Note : this is the most direct port of ogr2ogr.cpp possible +# It could be made much more Python'ish ! + +import sys +import os +import stat + +from osgeo import gdal +from osgeo import ogr +from osgeo import osr + +############################################################################### + +class ScaledProgressObject: + def __init__(self, min, max, cbk, cbk_data = None): + self.min = min + self.max = max + self.cbk = cbk + self.cbk_data = cbk_data + +############################################################################### + +def ScaledProgressFunc(pct, msg, data): + if data.cbk is None: + return True + return data.cbk(data.min + pct * (data.max - data.min), msg, data.cbk_data) + +############################################################################### + +def EQUAL(a, b): + return a.lower() == b.lower() + +############################################################################### +# Redefinition of GDALTermProgress, so that autotest/pyscripts/test_ogr2ogr_py.py +# can check that the progress bar is displayed + +nLastTick = -1 + +def TermProgress( dfComplete, pszMessage, pProgressArg ): + + global nLastTick + nThisTick = (int) (dfComplete * 40.0) + + if nThisTick < 0: + nThisTick = 0 + if nThisTick > 40: + nThisTick = 40 + + # Have we started a new progress run? + if nThisTick < nLastTick and nLastTick >= 39: + nLastTick = -1 + + if nThisTick <= nLastTick: + return True + + while nThisTick > nLastTick: + nLastTick = nLastTick + 1 + if (nLastTick % 4) == 0: + sys.stdout.write('%d' % ((nLastTick / 4) * 10)) + else: + sys.stdout.write('.') + + if nThisTick == 40: + print(" - done." ) + else: + sys.stdout.flush() + + return True + +class TargetLayerInfo: + def __init__(self): + self.poDstLayer = None + self.poCT = None + #self.papszTransformOptions = None + self.panMap = None + self.iSrcZField = None + +class AssociatedLayers: + def __init__(self): + self.poSrcLayer = None + self.psInfo = None + +#********************************************************************** +# main() +#********************************************************************** + +bSkipFailures = False +nGroupTransactions = 200 +bPreserveFID = False +nFIDToFetch = ogr.NullFID + +class Enum(set): + def __getattr__(self, name): + if name in self: + return name + raise AttributeError + +GeomOperation = Enum(["NONE", "SEGMENTIZE", "SIMPLIFY_PRESERVE_TOPOLOGY"]) + +def main(args = None, progress_func = TermProgress, progress_data = None): + + global bSkipFailures + global nGroupTransactions + global bPreserveFID + global nFIDToFetch + + pszFormat = "ESRI Shapefile" + pszDataSource = None + pszDestDataSource = None + papszLayers = [] + papszDSCO = [] + papszLCO = [] + bTransform = False + bAppend = False + bUpdate = False + bOverwrite = False + pszOutputSRSDef = None + pszSourceSRSDef = None + poOutputSRS = None + bNullifyOutputSRS = False + poSourceSRS = None + pszNewLayerName = None + pszWHERE = None + poSpatialFilter = None + pszSelect = None + papszSelFields = None + pszSQLStatement = None + eGType = -2 + bPromoteToMulti = False + eGeomOp = GeomOperation.NONE + dfGeomOpParam = 0 + papszFieldTypesToString = [] + bDisplayProgress = False + pfnProgress = None + pProgressArg = None + bClipSrc = False + bWrapDateline = False + poClipSrc = None + pszClipSrcDS = None + pszClipSrcSQL = None + pszClipSrcLayer = None + pszClipSrcWhere = None + poClipDst = None + pszClipDstDS = None + pszClipDstSQL = None + pszClipDstLayer = None + pszClipDstWhere = None + #pszSrcEncoding = None + #pszDstEncoding = None + bWrapDateline = False + bExplodeCollections = False + pszZField = None + nCoordDim = -1 + + if args is None: + args = sys.argv + + args = ogr.GeneralCmdLineProcessor( args ) + +# -------------------------------------------------------------------- +# Processing command line arguments. +# -------------------------------------------------------------------- + if args is None: + return False + + nArgc = len(args) + + iArg = 1 + while iArg < nArgc: + if EQUAL(args[iArg],"-f") and iArg < nArgc-1: + iArg = iArg + 1 + pszFormat = args[iArg] + + elif EQUAL(args[iArg],"-dsco") and iArg < nArgc-1: + iArg = iArg + 1 + papszDSCO.append(args[iArg] ) + + elif EQUAL(args[iArg],"-lco") and iArg < nArgc-1: + iArg = iArg + 1 + papszLCO.append(args[iArg] ) + + elif EQUAL(args[iArg],"-preserve_fid"): + bPreserveFID = True + + elif len(args[iArg]) >= 5 and EQUAL(args[iArg][0:5], "-skip"): + bSkipFailures = True + nGroupTransactions = 1 # #2409 + + elif EQUAL(args[iArg],"-append"): + bAppend = True + bUpdate = True + + elif EQUAL(args[iArg],"-overwrite"): + bOverwrite = True + bUpdate = True + + elif EQUAL(args[iArg],"-update"): + bUpdate = True + + elif EQUAL(args[iArg],"-fid") and iArg < nArgc-1: + iArg = iArg + 1 + nFIDToFetch = int(args[iArg]) + + elif EQUAL(args[iArg],"-sql") and iArg < nArgc-1: + iArg = iArg + 1 + pszSQLStatement = args[iArg] + + elif EQUAL(args[iArg],"-nln") and iArg < nArgc-1: + iArg = iArg + 1 + pszNewLayerName = args[iArg] + + elif EQUAL(args[iArg],"-nlt") and iArg < nArgc-1: + + if EQUAL(args[iArg+1],"NONE"): + eGType = ogr.wkbNone + elif EQUAL(args[iArg+1],"GEOMETRY"): + eGType = ogr.wkbUnknown + elif EQUAL(args[iArg+1],"PROMOTE_TO_MULTI"): + bPromoteToMulti = True + elif EQUAL(args[iArg+1],"POINT"): + eGType = ogr.wkbPoint + elif EQUAL(args[iArg+1],"LINESTRING"): + eGType = ogr.wkbLineString + elif EQUAL(args[iArg+1],"POLYGON"): + eGType = ogr.wkbPolygon + elif EQUAL(args[iArg+1],"GEOMETRYCOLLECTION"): + eGType = ogr.wkbGeometryCollection + elif EQUAL(args[iArg+1],"MULTIPOINT"): + eGType = ogr.wkbMultiPoint + elif EQUAL(args[iArg+1],"MULTILINESTRING"): + eGType = ogr.wkbMultiLineString + elif EQUAL(args[iArg+1],"MULTIPOLYGON"): + eGType = ogr.wkbMultiPolygon + elif EQUAL(args[iArg+1],"GEOMETRY25D"): + eGType = ogr.wkbUnknown | ogr.wkb25DBit + elif EQUAL(args[iArg+1],"POINT25D"): + eGType = ogr.wkbPoint25D + elif EQUAL(args[iArg+1],"LINESTRING25D"): + eGType = ogr.wkbLineString25D + elif EQUAL(args[iArg+1],"POLYGON25D"): + eGType = ogr.wkbPolygon25D + elif EQUAL(args[iArg+1],"GEOMETRYCOLLECTION25D"): + eGType = ogr.wkbGeometryCollection25D + elif EQUAL(args[iArg+1],"MULTIPOINT25D"): + eGType = ogr.wkbMultiPoint25D + elif EQUAL(args[iArg+1],"MULTILINESTRING25D"): + eGType = ogr.wkbMultiLineString25D + elif EQUAL(args[iArg+1],"MULTIPOLYGON25D"): + eGType = ogr.wkbMultiPolygon25D + else: + print("-nlt %s: type not recognised." % args[iArg+1]) + return False + + iArg = iArg + 1 + + elif EQUAL(args[iArg],"-dim") and iArg < nArgc-1: + + nCoordDim = int(args[iArg+1]) + if nCoordDim != 2 and nCoordDim != 3: + print("-dim %s: value not handled." % args[iArg+1]) + return False + iArg = iArg + 1 + + elif (EQUAL(args[iArg],"-tg") or \ + EQUAL(args[iArg],"-gt")) and iArg < nArgc-1: + iArg = iArg + 1 + nGroupTransactions = int(args[iArg]) + + elif EQUAL(args[iArg],"-s_srs") and iArg < nArgc-1: + iArg = iArg + 1 + pszSourceSRSDef = args[iArg] + + elif EQUAL(args[iArg],"-a_srs") and iArg < nArgc-1: + iArg = iArg + 1 + pszOutputSRSDef = args[iArg] + if EQUAL(pszOutputSRSDef, "NULL") or \ + EQUAL(pszOutputSRSDef, "NONE"): + pszOutputSRSDef = None + bNullifyOutputSRS = True + + elif EQUAL(args[iArg],"-t_srs") and iArg < nArgc-1: + iArg = iArg + 1 + pszOutputSRSDef = args[iArg] + bTransform = True + + elif EQUAL(args[iArg],"-spat") and iArg + 4 < nArgc: + oRing = ogr.Geometry(ogr.wkbLinearRing) + + oRing.AddPoint_2D( float(args[iArg+1]), float(args[iArg+2]) ) + oRing.AddPoint_2D( float(args[iArg+1]), float(args[iArg+4]) ) + oRing.AddPoint_2D( float(args[iArg+3]), float(args[iArg+4]) ) + oRing.AddPoint_2D( float(args[iArg+3]), float(args[iArg+2]) ) + oRing.AddPoint_2D( float(args[iArg+1]), float(args[iArg+2]) ) + + poSpatialFilter = ogr.Geometry(ogr.wkbPolygon) + poSpatialFilter.AddGeometry(oRing) + iArg = iArg + 4 + + elif EQUAL(args[iArg],"-where") and iArg < nArgc-1: + iArg = iArg + 1 + pszWHERE = args[iArg] + + elif EQUAL(args[iArg],"-select") and iArg < nArgc-1: + iArg = iArg + 1 + pszSelect = args[iArg] + if pszSelect.find(',') != -1: + papszSelFields = pszSelect.split(',') + else: + papszSelFields = pszSelect.split(' ') + if papszSelFields[0] == '': + papszSelFields = [] + + elif EQUAL(args[iArg],"-simplify") and iArg < nArgc-1: + iArg = iArg + 1 + eGeomOp = GeomOperation.SIMPLIFY_PRESERVE_TOPOLOGY + dfGeomOpParam = float(args[iArg]) + + elif EQUAL(args[iArg],"-segmentize") and iArg < nArgc-1: + iArg = iArg + 1 + eGeomOp = GeomOperation.SEGMENTIZE + dfGeomOpParam = float(args[iArg]) + + elif EQUAL(args[iArg],"-fieldTypeToString") and iArg < nArgc-1: + iArg = iArg + 1 + pszFieldTypeToString = args[iArg] + if pszFieldTypeToString.find(',') != -1: + tokens = pszFieldTypeToString.split(',') + else: + tokens = pszFieldTypeToString.split(' ') + + for token in tokens: + if EQUAL(token,"Integer") or \ + EQUAL(token,"Real") or \ + EQUAL(token,"String") or \ + EQUAL(token,"Date") or \ + EQUAL(token,"Time") or \ + EQUAL(token,"DateTime") or \ + EQUAL(token,"Binary") or \ + EQUAL(token,"IntegerList") or \ + EQUAL(token,"RealList") or \ + EQUAL(token,"StringList"): + + papszFieldTypesToString.append(token) + + elif EQUAL(token,"All"): + papszFieldTypesToString = [ 'All' ] + break + + else: + print("Unhandled type for fieldtypeasstring option : %s " % token) + return Usage() + + elif EQUAL(args[iArg],"-progress"): + bDisplayProgress = True + + #elif EQUAL(args[iArg],"-wrapdateline") ) + #{ + # bWrapDateline = True; + #} + # + elif EQUAL(args[iArg],"-clipsrc") and iArg < nArgc-1: + + bClipSrc = True + if IsNumber(args[iArg+1]) and iArg < nArgc - 4: + oRing = ogr.Geometry(ogr.wkbLinearRing) + + oRing.AddPoint_2D( float(args[iArg+1]), float(args[iArg+2]) ) + oRing.AddPoint_2D( float(args[iArg+1]), float(args[iArg+4]) ) + oRing.AddPoint_2D( float(args[iArg+3]), float(args[iArg+4]) ) + oRing.AddPoint_2D( float(args[iArg+3]), float(args[iArg+2]) ) + oRing.AddPoint_2D( float(args[iArg+1]), float(args[iArg+2]) ) + + poClipSrc = ogr.Geometry(ogr.wkbPolygon) + poClipSrc.AddGeometry(oRing) + iArg = iArg + 4 + + elif (len(args[iArg+1]) >= 7 and EQUAL(args[iArg+1][0:7],"POLYGON") ) or \ + (len(args[iArg+1]) >= 12 and EQUAL(args[iArg+1][0:12],"MULTIPOLYGON") ) : + poClipSrc = ogr.CreateGeometryFromWkt(args[iArg+1]) + if poClipSrc is None: + print("FAILURE: Invalid geometry. Must be a valid POLYGON or MULTIPOLYGON WKT\n") + return Usage() + + iArg = iArg + 1 + + elif EQUAL(args[iArg+1],"spat_extent"): + iArg = iArg + 1 + + else: + pszClipSrcDS = args[iArg+1] + iArg = iArg + 1 + + elif EQUAL(args[iArg],"-clipsrcsql") and iArg < nArgc-1: + pszClipSrcSQL = args[iArg+1] + iArg = iArg + 1 + + elif EQUAL(args[iArg],"-clipsrclayer") and iArg < nArgc-1: + pszClipSrcLayer = args[iArg+1] + iArg = iArg + 1 + + elif EQUAL(args[iArg],"-clipsrcwhere") and iArg < nArgc-1: + pszClipSrcWhere = args[iArg+1] + iArg = iArg + 1 + + elif EQUAL(args[iArg],"-clipdst") and iArg < nArgc-1: + + if IsNumber(args[iArg+1]) and iArg < nArgc - 4: + oRing = ogr.Geometry(ogr.wkbLinearRing) + + oRing.AddPoint_2D( float(args[iArg+1]), float(args[iArg+2]) ) + oRing.AddPoint_2D( float(args[iArg+1]), float(args[iArg+4]) ) + oRing.AddPoint_2D( float(args[iArg+3]), float(args[iArg+4]) ) + oRing.AddPoint_2D( float(args[iArg+3]), float(args[iArg+2]) ) + oRing.AddPoint_2D( float(args[iArg+1]), float(args[iArg+2]) ) + + poClipDst = ogr.Geometry(ogr.wkbPolygon) + poClipDst.AddGeometry(oRing) + iArg = iArg + 4 + + elif (len(args[iArg+1]) >= 7 and EQUAL(args[iArg+1][0:7],"POLYGON") ) or \ + (len(args[iArg+1]) >= 12 and EQUAL(args[iArg+1][0:12],"MULTIPOLYGON") ) : + poClipDst = ogr.CreateGeometryFromWkt(args[iArg+1]) + if poClipDst is None: + print("FAILURE: Invalid geometry. Must be a valid POLYGON or MULTIPOLYGON WKT\n") + return Usage() + + iArg = iArg + 1 + + elif EQUAL(args[iArg+1],"spat_extent"): + iArg = iArg + 1 + + else: + pszClipDstDS = args[iArg+1] + iArg = iArg + 1 + + elif EQUAL(args[iArg],"-clipdstsql") and iArg < nArgc-1: + pszClipDstSQL = args[iArg+1] + iArg = iArg + 1 + + elif EQUAL(args[iArg],"-clipdstlayer") and iArg < nArgc-1: + pszClipDstLayer = args[iArg+1] + iArg = iArg + 1 + + elif EQUAL(args[iArg],"-clipdstwhere") and iArg < nArgc-1: + pszClipDstWhere = args[iArg+1] + iArg = iArg + 1 + + elif EQUAL(args[iArg],"-explodecollections"): + bExplodeCollections = True + + elif EQUAL(args[iArg],"-zfield") and iArg < nArgc-1: + pszZField = args[iArg+1] + iArg = iArg + 1 + + elif args[iArg][0] == '-': + return Usage() + + elif pszDestDataSource is None: + pszDestDataSource = args[iArg] + elif pszDataSource is None: + pszDataSource = args[iArg] + else: + papszLayers.append (args[iArg] ) + + iArg = iArg + 1 + + if pszDataSource is None: + return Usage() + + if bPreserveFID and bExplodeCollections: + print("FAILURE: cannot use -preserve_fid and -explodecollections at the same time\n\n") + return Usage() + + if bClipSrc and pszClipSrcDS is not None: + poClipSrc = LoadGeometry(pszClipSrcDS, pszClipSrcSQL, pszClipSrcLayer, pszClipSrcWhere) + if poClipSrc is None: + print("FAILURE: cannot load source clip geometry\n" ) + return Usage() + + elif bClipSrc and poClipSrc is None: + if poSpatialFilter is not None: + poClipSrc = poSpatialFilter.Clone() + if poClipSrc is None: + print("FAILURE: -clipsrc must be used with -spat option or a\n" + \ + "bounding box, WKT string or datasource must be specified\n") + return Usage() + + if pszClipDstDS is not None: + poClipDst = LoadGeometry(pszClipDstDS, pszClipDstSQL, pszClipDstLayer, pszClipDstWhere) + if poClipDst is None: + print("FAILURE: cannot load dest clip geometry\n" ) + return Usage() + +# -------------------------------------------------------------------- +# Open data source. +# -------------------------------------------------------------------- + poDS = ogr.Open( pszDataSource, False ) + +# -------------------------------------------------------------------- +# Report failure +# -------------------------------------------------------------------- + if poDS is None: + print("FAILURE:\n" + \ + "Unable to open datasource `%s' with the following drivers." % pszDataSource) + + for iDriver in range(ogr.GetDriverCount()): + print(" -> " + ogr.GetDriver(iDriver).GetName() ) + + return False + +# -------------------------------------------------------------------- +# Try opening the output datasource as an existing, writable +# -------------------------------------------------------------------- + poODS = None + poDriver = None + + if bUpdate: + poODS = ogr.Open( pszDestDataSource, True ) + if poODS is None: + + if bOverwrite or bAppend: + poODS = ogr.Open( pszDestDataSource, False ) + if poODS is None: + # the datasource doesn't exist at all + bUpdate = False + else: + poODS.delete() + poODS = None + + if bUpdate: + print("FAILURE:\n" + + "Unable to open existing output datasource `%s'." % pszDestDataSource) + return False + + elif len(papszDSCO) > 0: + print("WARNING: Datasource creation options ignored since an existing datasource\n" + \ + " being updated." ) + + if poODS is not None: + poDriver = poODS.GetDriver() + +# -------------------------------------------------------------------- +# Find the output driver. +# -------------------------------------------------------------------- + if not bUpdate: + poDriver = ogr.GetDriverByName(pszFormat) + if poDriver is None: + print("Unable to find driver `%s'." % pszFormat) + print( "The following drivers are available:" ) + + for iDriver in range(ogr.GetDriverCount()): + print(" -> %s" % ogr.GetDriver(iDriver).GetName() ) + + return False + + if poDriver.TestCapability( ogr.ODrCCreateDataSource ) == False: + print( "%s driver does not support data source creation." % pszFormat) + return False + +# -------------------------------------------------------------------- +# Special case to improve user experience when translating +# a datasource with multiple layers into a shapefile. If the +# user gives a target datasource with .shp and it does not exist, +# the shapefile driver will try to create a file, but this is not +# appropriate because here we have several layers, so create +# a directory instead. +# -------------------------------------------------------------------- + if EQUAL(poDriver.GetName(), "ESRI Shapefile") and \ + pszSQLStatement is None and \ + (len(papszLayers) > 1 or \ + (len(papszLayers) == 0 and poDS.GetLayerCount() > 1)) and \ + pszNewLayerName is None and \ + EQUAL(os.path.splitext(pszDestDataSource)[1], ".SHP") : + + try: + os.stat(pszDestDataSource) + except: + try: + # decimal 493 = octal 0755. Python 3 needs 0o755, but + # this syntax is only supported by Python >= 2.6 + os.mkdir(pszDestDataSource, 493) + except: + print("Failed to create directory %s\n" + "for shapefile datastore.\n" % pszDestDataSource ) + return False + +# -------------------------------------------------------------------- +# Create the output data source. +# -------------------------------------------------------------------- + poODS = poDriver.CreateDataSource( pszDestDataSource, options = papszDSCO ) + if poODS is None: + print( "%s driver failed to create %s" % (pszFormat, pszDestDataSource )) + return False + +# -------------------------------------------------------------------- +# Parse the output SRS definition if possible. +# -------------------------------------------------------------------- + if pszOutputSRSDef is not None: + poOutputSRS = osr.SpatialReference() + if poOutputSRS.SetFromUserInput( pszOutputSRSDef ) != 0: + print( "Failed to process SRS definition: %s" % pszOutputSRSDef ) + return False + +# -------------------------------------------------------------------- +# Parse the source SRS definition if possible. +# -------------------------------------------------------------------- + if pszSourceSRSDef is not None: + poSourceSRS = osr.SpatialReference() + if poSourceSRS.SetFromUserInput( pszSourceSRSDef ) != 0: + print( "Failed to process SRS definition: %s" % pszSourceSRSDef ) + return False + +# -------------------------------------------------------------------- +# For OSM file. +# -------------------------------------------------------------------- + bSrcIsOSM = poDS.GetDriver() is not None and \ + poDS.GetDriver().GetName() == "OSM" + nSrcFileSize = 0 + if bSrcIsOSM and poDS.GetName() != "/vsistdin/": + sStat = gdal.VSIStatL(poDS.GetName()) + if sStat is not None: + nSrcFileSize = sStat.size + +# -------------------------------------------------------------------- +# Special case for -sql clause. No source layers required. +# -------------------------------------------------------------------- + if pszSQLStatement is not None: + if pszWHERE is not None: + print( "-where clause ignored in combination with -sql." ) + if len(papszLayers) > 0: + print( "layer names ignored in combination with -sql." ) + + poResultSet = poDS.ExecuteSQL( pszSQLStatement, poSpatialFilter, \ + None ) + + if poResultSet is not None: + nCountLayerFeatures = 0 + if bDisplayProgress: + if bSrcIsOSM: + pfnProgress = progress_func + pProgressArg = progress_data + + elif not poResultSet.TestCapability(ogr.OLCFastFeatureCount): + print( "Progress turned off as fast feature count is not available.") + bDisplayProgress = False + + else: + nCountLayerFeatures = poResultSet.GetFeatureCount() + pfnProgress = progress_func + pProgressArg = progress_data + +# -------------------------------------------------------------------- +# Special case to improve user experience when translating into +# single file shapefile and source has only one layer, and that +# the layer name isn't specified +# -------------------------------------------------------------------- + if EQUAL(poDriver.GetName(), "ESRI Shapefile") and \ + pszNewLayerName is None: + try: + mode = os.stat(pszDestDataSource).st_mode + if (mode & stat.S_IFDIR) == 0: + pszNewLayerName = os.path.splitext(os.path.basename(pszDestDataSource))[0] + except: + pass + + + psInfo = SetupTargetLayer( poDS, \ + poResultSet, + poODS, \ + papszLCO, \ + pszNewLayerName, \ + bTransform, \ + poOutputSRS, \ + bNullifyOutputSRS, \ + poSourceSRS, \ + papszSelFields, \ + bAppend, eGType, bPromoteToMulti, nCoordDim, bOverwrite, \ + papszFieldTypesToString, \ + bWrapDateline, \ + bExplodeCollections, \ + pszZField, \ + pszWHERE ) + + poResultSet.ResetReading() + + if psInfo is None or not TranslateLayer( psInfo, poDS, poResultSet, poODS, \ + poOutputSRS, bNullifyOutputSRS, \ + eGType, bPromoteToMulti, nCoordDim, \ + eGeomOp, dfGeomOpParam, \ + nCountLayerFeatures, \ + poClipSrc, poClipDst, \ + bExplodeCollections, \ + nSrcFileSize, None, \ + pfnProgress, pProgressArg ): + print( + "Terminating translation prematurely after failed\n" + \ + "translation from sql statement." ) + + return False + + poDS.ReleaseResultSet( poResultSet ) + + +# -------------------------------------------------------------------- +# Special case for layer interleaving mode. +# -------------------------------------------------------------------- + elif bSrcIsOSM and gdal.GetConfigOption("OGR_INTERLEAVED_READING", None) is None: + + gdal.SetConfigOption("OGR_INTERLEAVED_READING", "YES") + + #if (bSplitListFields) + #{ + # fprintf( stderr, "FAILURE: -splitlistfields not supported in this mode\n" ); + # exit( 1 ); + #} + + nSrcLayerCount = poDS.GetLayerCount() + pasAssocLayers = [ AssociatedLayers() for i in range(nSrcLayerCount) ] + +# -------------------------------------------------------------------- +# Special case to improve user experience when translating into +# single file shapefile and source has only one layer, and that +# the layer name isn't specified +# -------------------------------------------------------------------- + + if EQUAL(poDriver.GetName(), "ESRI Shapefile") and \ + (len(papszLayers) == 1 or nSrcLayerCount == 1) and pszNewLayerName is None: + try: + mode = os.stat(pszDestDataSource).st_mode + if (mode & stat.S_IFDIR) == 0: + pszNewLayerName = os.path.splitext(os.path.basename(pszDestDataSource))[0] + except: + pass + + if bDisplayProgress and bSrcIsOSM: + pfnProgress = progress_func + pProgressArg = progress_data + +# -------------------------------------------------------------------- +# If no target layer specified, use all source layers. +# -------------------------------------------------------------------- + if len(papszLayers) == 0: + papszLayers = [ None for i in range(nSrcLayerCount) ] + for iLayer in range(nSrcLayerCount): + poLayer = poDS.GetLayer(iLayer) + if poLayer is None: + print("FAILURE: Couldn't fetch advertised layer %d!" % iLayer) + return False + + papszLayers[iLayer] = poLayer.GetName() + else: + if bSrcIsOSM: + osInterestLayers = "SET interest_layers =" + for iLayer in range(len(papszLayers)): + if iLayer != 0: + osInterestLayers = osInterestLayers + "," + osInterestLayers = osInterestLayers + papszLayers[iLayer] + + poDS.ExecuteSQL(osInterestLayers, None, None) + +# -------------------------------------------------------------------- +# First pass to set filters and create target layers. +# -------------------------------------------------------------------- + for iLayer in range(nSrcLayerCount): + poLayer = poDS.GetLayer(iLayer) + if poLayer is None: + print("FAILURE: Couldn't fetch advertised layer %d!" % iLayer) + return False + + pasAssocLayers[iLayer].poSrcLayer = poLayer + + if CSLFindString(papszLayers, poLayer.GetName()) >= 0: + if pszWHERE is not None: + if poLayer.SetAttributeFilter( pszWHERE ) != 0: + print("FAILURE: SetAttributeFilter(%s) on layer '%s' failed.\n" % (pszWHERE, poLayer.GetName()) ) + if not bSkipFailures: + return False + + if poSpatialFilter is not None: + poLayer.SetSpatialFilter( poSpatialFilter ) + + psInfo = SetupTargetLayer( poDS, \ + poLayer, \ + poODS, \ + papszLCO, \ + pszNewLayerName, \ + bTransform, \ + poOutputSRS, \ + bNullifyOutputSRS, \ + poSourceSRS, \ + papszSelFields, \ + bAppend, eGType, bPromoteToMulti, nCoordDim, bOverwrite, \ + papszFieldTypesToString, \ + bWrapDateline, \ + bExplodeCollections, \ + pszZField, \ + pszWHERE ) + + if psInfo is None and not bSkipFailures: + return False + + pasAssocLayers[iLayer].psInfo = psInfo + else: + pasAssocLayers[iLayer].psInfo = None + +# -------------------------------------------------------------------- +# Second pass to process features in a interleaved layer mode. +# -------------------------------------------------------------------- + bHasLayersNonEmpty = True + while bHasLayersNonEmpty: + bHasLayersNonEmpty = False + + for iLayer in range(nSrcLayerCount): + poLayer = pasAssocLayers[iLayer].poSrcLayer + psInfo = pasAssocLayers[iLayer].psInfo + anReadFeatureCount = [0] + + if psInfo is not None: + if not TranslateLayer(psInfo, poDS, poLayer, poODS, \ + poOutputSRS, bNullifyOutputSRS, \ + eGType, bPromoteToMulti, nCoordDim, \ + eGeomOp, dfGeomOpParam, \ + 0, \ + poClipSrc, poClipDst, \ + bExplodeCollections, \ + nSrcFileSize, \ + anReadFeatureCount, \ + pfnProgress, pProgressArg ) \ + and not bSkipFailures: + print( + "Terminating translation prematurely after failed\n" + \ + "translation of layer " + poLayer.GetName() + " (use -skipfailures to skip errors)") + + return False + else: + # No matching target layer : just consumes the features + + poFeature = poLayer.GetNextFeature() + while poFeature is not None: + anReadFeatureCount[0] = anReadFeatureCount[0] + 1 + poFeature = poLayer.GetNextFeature() + + if anReadFeatureCount[0] != 0: + bHasLayersNonEmpty = True + + else: + + nLayerCount = 0 + papoLayers = [] + +# -------------------------------------------------------------------- +# Process each data source layer. +# -------------------------------------------------------------------- + if len(papszLayers) == 0: + nLayerCount = poDS.GetLayerCount() + papoLayers = [None for i in range(nLayerCount)] + iLayer = 0 + + for iLayer in range(nLayerCount): + poLayer = poDS.GetLayer(iLayer) + + if poLayer is None: + print("FAILURE: Couldn't fetch advertised layer %d!" % iLayer) + return False + + papoLayers[iLayer] = poLayer + iLayer = iLayer + 1 + +# -------------------------------------------------------------------- +# Process specified data source layers. +# -------------------------------------------------------------------- + else: + nLayerCount = len(papszLayers) + papoLayers = [None for i in range(nLayerCount)] + iLayer = 0 + + for layername in papszLayers: + poLayer = poDS.GetLayerByName(layername) + + if poLayer is None: + print("FAILURE: Couldn't fetch advertised layer %s!" % layername) + return False + + papoLayers[iLayer] = poLayer + iLayer = iLayer + 1 + + panLayerCountFeatures = [0 for i in range(nLayerCount)] + nCountLayersFeatures = 0 + nAccCountFeatures = 0 + + # First pass to apply filters and count all features if necessary + for iLayer in range(nLayerCount): + poLayer = papoLayers[iLayer] + + if pszWHERE is not None: + if poLayer.SetAttributeFilter( pszWHERE ) != 0: + print("FAILURE: SetAttributeFilter(%s) failed." % pszWHERE) + if not bSkipFailures: + return False + + if poSpatialFilter is not None: + poLayer.SetSpatialFilter( poSpatialFilter ) + + if bDisplayProgress and not bSrcIsOSM: + if not poLayer.TestCapability(ogr.OLCFastFeatureCount): + print("Progress turned off as fast feature count is not available.") + bDisplayProgress = False + else: + panLayerCountFeatures[iLayer] = poLayer.GetFeatureCount() + nCountLayersFeatures += panLayerCountFeatures[iLayer] + + # Second pass to do the real job + for iLayer in range(nLayerCount): + poLayer = papoLayers[iLayer] + + if bDisplayProgress: + if bSrcIsOSM: + pfnProgress = progress_func + pProgressArg = progress_data + else: + pfnProgress = ScaledProgressFunc + pProgressArg = ScaledProgressObject( \ + nAccCountFeatures * 1.0 / nCountLayersFeatures, \ + (nAccCountFeatures + panLayerCountFeatures[iLayer]) * 1.0 / nCountLayersFeatures, \ + progress_func, progress_data) + + nAccCountFeatures += panLayerCountFeatures[iLayer] + +# -------------------------------------------------------------------- +# Special case to improve user experience when translating into +# single file shapefile and source has only one layer, and that +# the layer name isn't specified +# -------------------------------------------------------------------- + if EQUAL(poDriver.GetName(), "ESRI Shapefile") and \ + nLayerCount == 1 and pszNewLayerName is None: + try: + mode = os.stat(pszDestDataSource).st_mode + if (mode & stat.S_IFDIR) == 0: + pszNewLayerName = os.path.splitext(os.path.basename(pszDestDataSource))[0] + except: + pass + + + psInfo = SetupTargetLayer( poDS, \ + poLayer, \ + poODS, \ + papszLCO, \ + pszNewLayerName, \ + bTransform, \ + poOutputSRS, \ + bNullifyOutputSRS, \ + poSourceSRS, \ + papszSelFields, \ + bAppend, eGType, bPromoteToMulti, nCoordDim, bOverwrite, \ + papszFieldTypesToString, \ + bWrapDateline, \ + bExplodeCollections, \ + pszZField, \ + pszWHERE ) + + poLayer.ResetReading() + + if (psInfo is None or \ + not TranslateLayer( psInfo, poDS, poLayer, poODS, \ + poOutputSRS, bNullifyOutputSRS, \ + eGType, bPromoteToMulti, nCoordDim, \ + eGeomOp, dfGeomOpParam, \ + panLayerCountFeatures[iLayer], \ + poClipSrc, poClipDst, \ + bExplodeCollections, \ + nSrcFileSize, None, \ + pfnProgress, pProgressArg )) \ + and not bSkipFailures: + print( + "Terminating translation prematurely after failed\n" + \ + "translation of layer " + poLayer.GetLayerDefn().GetName() + " (use -skipfailures to skip errors)") + + return False + +# -------------------------------------------------------------------- +# Close down. +# -------------------------------------------------------------------- + # We must explicitly destroy the output dataset in order the file + # to be properly closed ! + poODS.Destroy() + poDS.Destroy() + + return True + +#********************************************************************** +# Usage() +#********************************************************************** + +def Usage(): + + print( "Usage: ogr2ogr [--help-general] [-skipfailures] [-append] [-update] [-gt n]\n" + \ + " [-select field_list] [-where restricted_where] \n" + \ + " [-progress] [-sql ] \n" + \ + " [-spat xmin ymin xmax ymax] [-preserve_fid] [-fid FID]\n" + \ + " [-a_srs srs_def] [-t_srs srs_def] [-s_srs srs_def]\n" + \ + " [-f format_name] [-overwrite] [[-dsco NAME=VALUE] ...]\n" + \ + " [-simplify tolerance]\n" + \ + #// " [-segmentize max_dist] [-fieldTypeToString All|(type1[,type2]*)]\n" + \ + " [-fieldTypeToString All|(type1[,type2]*)] [-explodecollections] \n" + \ + " dst_datasource_name src_datasource_name\n" + \ + " [-lco NAME=VALUE] [-nln name] [-nlt type] [-dim 2|3] [layer [layer ...]]\n" + \ + "\n" + \ + " -f format_name: output file format name, possible values are:") + + for iDriver in range(ogr.GetDriverCount()): + poDriver = ogr.GetDriver(iDriver) + + if poDriver.TestCapability( ogr.ODrCCreateDataSource ): + print( " -f \"" + poDriver.GetName() + "\"" ) + + print( " -append: Append to existing layer instead of creating new if it exists\n" + \ + " -overwrite: delete the output layer and recreate it empty\n" + \ + " -update: Open existing output datasource in update mode\n" + \ + " -progress: Display progress on terminal. Only works if input layers have the \"fast feature count\" capability\n" + \ + " -select field_list: Comma-delimited list of fields from input layer to\n" + \ + " copy to the new layer (defaults to all)\n" + \ + " -where restricted_where: Attribute query (like SQL WHERE)\n" + \ + " -sql statement: Execute given SQL statement and save result.\n" + \ + " -skipfailures: skip features or layers that fail to convert\n" + \ + " -gt n: group n features per transaction (default 200)\n" + \ + " -spat xmin ymin xmax ymax: spatial query extents\n" + \ + " -simplify tolerance: distance tolerance for simplification.\n" + \ + #//" -segmentize max_dist: maximum distance between 2 nodes.\n" + \ + #//" Used to create intermediate points\n" + \ + " -dsco NAME=VALUE: Dataset creation option (format specific)\n" + \ + " -lco NAME=VALUE: Layer creation option (format specific)\n" + \ + " -nln name: Assign an alternate name to the new layer\n" + \ + " -nlt type: Force a geometry type for new layer. One of NONE, GEOMETRY,\n" + \ + " POINT, LINESTRING, POLYGON, GEOMETRYCOLLECTION, MULTIPOINT,\n" + \ + " MULTIPOLYGON, or MULTILINESTRING. Add \"25D\" for 3D layers.\n" + \ + " Default is type of source layer.\n" + \ + " -dim dimension: Force the coordinate dimension to the specified value.\n" + \ + " -fieldTypeToString type1,...: Converts fields of specified types to\n" + \ + " fields of type string in the new layer. Valid types are : \n" + \ + " Integer, Real, String, Date, Time, DateTime, Binary, IntegerList, RealList,\n" + \ + " StringList. Special value All can be used to convert all fields to strings.") + + print(" -a_srs srs_def: Assign an output SRS\n" + " -t_srs srs_def: Reproject/transform to this SRS on output\n" + " -s_srs srs_def: Override source SRS\n" + "\n" + " Srs_def can be a full WKT definition (hard to escape properly),\n" + " or a well known definition (i.e. EPSG:4326) or a file with a WKT\n" + " definition." ) + + return False + +def CSLFindString(v, mystr): + i = 0 + for strIter in v: + if EQUAL(strIter, mystr): + return i + i = i + 1 + return -1 + +def IsNumber( pszStr): + try: + (float)(pszStr) + return True + except: + return False + +def LoadGeometry( pszDS, pszSQL, pszLyr, pszWhere): + poGeom = None + + poDS = ogr.Open( pszDS, False ) + if poDS is None: + return None + + if pszSQL is not None: + poLyr = poDS.ExecuteSQL( pszSQL, None, None ) + elif pszLyr is not None: + poLyr = poDS.GetLayerByName(pszLyr) + else: + poLyr = poDS.GetLayer(0) + + if poLyr is None: + print("Failed to identify source layer from datasource.") + poDS.Destroy() + return None + + if pszWhere is not None: + poLyr.SetAttributeFilter(pszWhere) + + poFeat = poLyr.GetNextFeature() + while poFeat is not None: + poSrcGeom = poFeat.GetGeometryRef() + if poSrcGeom is not None: + eType = wkbFlatten(poSrcGeom.GetGeometryType()) + + if poGeom is None: + poGeom = ogr.Geometry( ogr.wkbMultiPolygon ) + + if eType == ogr.wkbPolygon: + poGeom.AddGeometry( poSrcGeom ) + elif eType == ogr.wkbMultiPolygon: + for iGeom in range(poSrcGeom.GetGeometryCount()): + poGeom.AddGeometry(poSrcGeom.GetGeometryRef(iGeom) ) + + else: + print("ERROR: Geometry not of polygon type." ) + if pszSQL is not None: + poDS.ReleaseResultSet( poLyr ) + poDS.Destroy() + return None + + poFeat = poLyr.GetNextFeature() + + if pszSQL is not None: + poDS.ReleaseResultSet( poLyr ) + poDS.Destroy() + + return poGeom + + +def wkbFlatten(x): + return x & (~ogr.wkb25DBit) + +#********************************************************************** +# SetZ() +#********************************************************************** + +def SetZ (poGeom, dfZ ): + + if poGeom is None: + return + + eGType = wkbFlatten(poGeom.GetGeometryType()) + if eGType == ogr.wkbPoint: + poGeom.SetPoint(0, poGeom.GetX(), poGeom.GetY(), dfZ) + + elif eGType == ogr.wkbLineString or \ + eGType == ogr.wkbLinearRing: + for i in range(poGeom.GetPointCount()): + poGeom.SetPoint(i, poGeom.GetX(i), poGeom.GetY(i), dfZ) + + elif eGType == ogr.wkbPolygon or \ + eGType == ogr.wkbMultiPoint or \ + eGType == ogr.wkbMultiLineString or \ + eGType == ogr.wkbMultiPolygon or \ + eGType == ogr.wkbGeometryCollection: + for i in range(poGeom.GetGeometryCount()): + SetZ(poGeom.GetGeometryRef(i), dfZ) + +#********************************************************************** +# SetupTargetLayer() +#********************************************************************** + +def SetupTargetLayer( poSrcDS, poSrcLayer, poDstDS, papszLCO, pszNewLayerName, \ + bTransform, poOutputSRS, bNullifyOutputSRS, poSourceSRS, papszSelFields, \ + bAppend, eGType, bPromoteToMulti, nCoordDim, bOverwrite, \ + papszFieldTypesToString, bWrapDateline, \ + bExplodeCollections, pszZField, pszWHERE) : + + if pszNewLayerName is None: + pszNewLayerName = poSrcLayer.GetLayerDefn().GetName() + +# -------------------------------------------------------------------- +# Setup coordinate transformation if we need it. +# -------------------------------------------------------------------- + poCT = None + + if bTransform: + if poSourceSRS is None: + poSourceSRS = poSrcLayer.GetSpatialRef() + + if poSourceSRS is None: + print("Can't transform coordinates, source layer has no\n" + \ + "coordinate system. Use -s_srs to set one." ) + return None + + poCT = osr.CoordinateTransformation( poSourceSRS, poOutputSRS ) + if gdal.GetLastErrorMsg().find( 'Unable to load PROJ.4 library' ) != -1: + poCT = None + + if poCT is None: + pszWKT = None + + print("Failed to create coordinate transformation between the\n" + \ + "following coordinate systems. This may be because they\n" + \ + "are not transformable, or because projection services\n" + \ + "(PROJ.4 DLL/.so) could not be loaded." ) + + pszWKT = poSourceSRS.ExportToPrettyWkt( 0 ) + print( "Source:\n" + pszWKT ) + + pszWKT = poOutputSRS.ExportToPrettyWkt( 0 ) + print( "Target:\n" + pszWKT ) + return None + +# -------------------------------------------------------------------- +# Get other info. +# -------------------------------------------------------------------- + poSrcFDefn = poSrcLayer.GetLayerDefn() + + if poOutputSRS is None and not bNullifyOutputSRS: + poOutputSRS = poSrcLayer.GetSpatialRef() + +# -------------------------------------------------------------------- +# Find the layer. +# -------------------------------------------------------------------- + + # GetLayerByName() can instantiate layers that would have been + # 'hidden' otherwise, for example, non-spatial tables in a + # PostGIS-enabled database, so this apparently useless command is + # not useless. (#4012) + gdal.PushErrorHandler('CPLQuietErrorHandler') + poDstLayer = poDstDS.GetLayerByName(pszNewLayerName) + gdal.PopErrorHandler() + gdal.ErrorReset() + + iLayer = -1 + if poDstLayer is not None: + nLayerCount = poDstDS.GetLayerCount() + for iLayer in range(nLayerCount): + poLayer = poDstDS.GetLayer(iLayer) + # The .cpp version compares on pointers directly, but we cannot + # do this with swig object, so just compare the names. + if poLayer is not None \ + and poLayer.GetName() == poDstLayer.GetName(): + break + + if (iLayer == nLayerCount): + # Shouldn't happen with an ideal driver + poDstLayer = None + +# -------------------------------------------------------------------- +# If the user requested overwrite, and we have the layer in +# question we need to delete it now so it will get recreated +# (overwritten). +# -------------------------------------------------------------------- + if poDstLayer is not None and bOverwrite: + if poDstDS.DeleteLayer( iLayer ) != 0: + print("DeleteLayer() failed when overwrite requested." ) + return None + + poDstLayer = None + +# -------------------------------------------------------------------- +# If the layer does not exist, then create it. +# -------------------------------------------------------------------- + if poDstLayer is None: + if eGType == -2: + eGType = poSrcFDefn.GetGeomType() + + n25DBit = eGType & ogr.wkb25DBit + if bPromoteToMulti: + if wkbFlatten(eGType) == ogr.wkbLineString: + eGType = ogr.wkbMultiLineString | n25DBit + elif wkbFlatten(eGType) == ogr.wkbPolygon: + eGType = ogr.wkbMultiPolygon | n25DBit + + if bExplodeCollections: + if wkbFlatten(eGType) == ogr.wkbMultiPoint: + eGType = ogr.wkbPoint | n25DBit + elif wkbFlatten(eGType) == ogr.wkbMultiLineString: + eGType = ogr.wkbLineString | n25DBit + elif wkbFlatten(eGType) == ogr.wkbMultiPolygon: + eGType = ogr.wkbPolygon | n25DBit + elif wkbFlatten(eGType) == ogr.wkbGeometryCollection: + eGType = ogr.wkbUnknown | n25DBit + + if pszZField is not None: + eGType = eGType | ogr.wkb25DBit + + if nCoordDim == 2: + eGType = eGType & ~ogr.wkb25DBit + elif nCoordDim == 3: + eGType = eGType | ogr.wkb25DBit + + if poDstDS.TestCapability( ogr.ODsCCreateLayer ) == False: + print("Layer " + pszNewLayerName + "not found, and CreateLayer not supported by driver.") + return None + + gdal.ErrorReset() + + poDstLayer = poDstDS.CreateLayer( pszNewLayerName, poOutputSRS, \ + eGType, papszLCO ) + + if poDstLayer is None: + return None + + bAppend = False + +# -------------------------------------------------------------------- +# Otherwise we will append to it, if append was requested. +# -------------------------------------------------------------------- + elif not bAppend: + print("FAILED: Layer " + pszNewLayerName + "already exists, and -append not specified.\n" + \ + " Consider using -append, or -overwrite.") + return None + else: + if len(papszLCO) > 0: + print("WARNING: Layer creation options ignored since an existing layer is\n" + \ + " being appended to." ) + +# -------------------------------------------------------------------- +# Add fields. Default to copy all field. +# If only a subset of all fields requested, then output only +# the selected fields, and in the order that they were +# selected. +# -------------------------------------------------------------------- + + # Initialize the index-to-index map to -1's + nSrcFieldCount = poSrcFDefn.GetFieldCount() + panMap = [ -1 for i in range(nSrcFieldCount) ] + + poDstFDefn = poDstLayer.GetLayerDefn() + + if papszSelFields is not None and not bAppend: + + nDstFieldCount = 0 + if poDstFDefn is not None: + nDstFieldCount = poDstFDefn.GetFieldCount() + + for iField in range(len(papszSelFields)): + + iSrcField = poSrcFDefn.GetFieldIndex(papszSelFields[iField]) + if iSrcField >= 0: + poSrcFieldDefn = poSrcFDefn.GetFieldDefn(iSrcField) + oFieldDefn = ogr.FieldDefn( poSrcFieldDefn.GetNameRef(), + poSrcFieldDefn.GetType() ) + oFieldDefn.SetWidth( poSrcFieldDefn.GetWidth() ) + oFieldDefn.SetPrecision( poSrcFieldDefn.GetPrecision() ) + + if papszFieldTypesToString is not None and \ + (CSLFindString(papszFieldTypesToString, "All") != -1 or \ + CSLFindString(papszFieldTypesToString, \ + ogr.GetFieldTypeName(poSrcFieldDefn.GetType())) != -1): + + oFieldDefn.SetType(ogr.OFTString) + + # The field may have been already created at layer creation + iDstField = -1 + if poDstFDefn is not None: + iDstField = poDstFDefn.GetFieldIndex(oFieldDefn.GetNameRef()) + if iDstField >= 0: + panMap[iSrcField] = iDstField + elif poDstLayer.CreateField( oFieldDefn ) == 0: + # now that we've created a field, GetLayerDefn() won't return NULL + if poDstFDefn is None: + poDstFDefn = poDstLayer.GetLayerDefn() + + # Sanity check : if it fails, the driver is buggy + if poDstFDefn is not None and \ + poDstFDefn.GetFieldCount() != nDstFieldCount + 1: + print("The output driver has claimed to have added the %s field, but it did not!" % oFieldDefn.GetNameRef() ) + else: + panMap[iSrcField] = nDstFieldCount + nDstFieldCount = nDstFieldCount + 1 + + else: + print("Field '" + papszSelFields[iField] + "' not found in source layer.") + if not bSkipFailures: + return None + + # -------------------------------------------------------------------- + # Use SetIgnoredFields() on source layer if available + # -------------------------------------------------------------------- + + # Here we differ from the ogr2ogr.cpp implementation since the OGRFeatureQuery + # isn't mapped to swig. So in that case just don't use SetIgnoredFields() + # to avoid issue raised in #4015 + if poSrcLayer.TestCapability(ogr.OLCIgnoreFields) and pszWHERE is None: + papszIgnoredFields = [] + for iSrcField in range(nSrcFieldCount): + pszFieldName = poSrcFDefn.GetFieldDefn(iSrcField).GetNameRef() + bFieldRequested = False + for iField in range(len(papszSelFields)): + if EQUAL(pszFieldName, papszSelFields[iField]): + bFieldRequested = True + break + + if pszZField is not None and EQUAL(pszFieldName, pszZField): + bFieldRequested = True + + # If source field not requested, add it to ignored files list + if not bFieldRequested: + papszIgnoredFields.append(pszFieldName) + + poSrcLayer.SetIgnoredFields(papszIgnoredFields) + + elif not bAppend: + + nDstFieldCount = 0 + if poDstFDefn is not None: + nDstFieldCount = poDstFDefn.GetFieldCount() + + for iField in range(nSrcFieldCount): + + poSrcFieldDefn = poSrcFDefn.GetFieldDefn(iField) + oFieldDefn = ogr.FieldDefn( poSrcFieldDefn.GetNameRef(), + poSrcFieldDefn.GetType() ) + oFieldDefn.SetWidth( poSrcFieldDefn.GetWidth() ) + oFieldDefn.SetPrecision( poSrcFieldDefn.GetPrecision() ) + + if papszFieldTypesToString is not None and \ + (CSLFindString(papszFieldTypesToString, "All") != -1 or \ + CSLFindString(papszFieldTypesToString, \ + ogr.GetFieldTypeName(poSrcFieldDefn.GetType())) != -1): + + oFieldDefn.SetType(ogr.OFTString) + + # The field may have been already created at layer creation + iDstField = -1 + if poDstFDefn is not None: + iDstField = poDstFDefn.GetFieldIndex(oFieldDefn.GetNameRef()) + if iDstField >= 0: + panMap[iField] = iDstField + elif poDstLayer.CreateField( oFieldDefn ) == 0: + # now that we've created a field, GetLayerDefn() won't return NULL + if poDstFDefn is None: + poDstFDefn = poDstLayer.GetLayerDefn() + + # Sanity check : if it fails, the driver is buggy + if poDstFDefn is not None and \ + poDstFDefn.GetFieldCount() != nDstFieldCount + 1: + print("The output driver has claimed to have added the %s field, but it did not!" % oFieldDefn.GetNameRef() ) + else: + panMap[iField] = nDstFieldCount + nDstFieldCount = nDstFieldCount + 1 + + else: + # For an existing layer, build the map by fetching the index in the destination + # layer for each source field + if poDstFDefn is None: + print( "poDstFDefn == NULL.\n" ) + return None + + for iField in range(nSrcFieldCount): + poSrcFieldDefn = poSrcFDefn.GetFieldDefn(iField) + iDstField = poDstFDefn.GetFieldIndex(poSrcFieldDefn.GetNameRef()) + if iDstField >= 0: + panMap[iField] = iDstField + + iSrcZField = -1 + if pszZField is not None: + iSrcZField = poSrcFDefn.GetFieldIndex(pszZField) + + psInfo = TargetLayerInfo() + psInfo.poDstLayer = poDstLayer + psInfo.poCT = poCT + #psInfo.papszTransformOptions = papszTransformOptions + psInfo.panMap = panMap + psInfo.iSrcZField = iSrcZField + + return psInfo + +#********************************************************************** +# TranslateLayer() +#********************************************************************** + +def TranslateLayer( psInfo, poSrcDS, poSrcLayer, poDstDS, \ + poOutputSRS, bNullifyOutputSRS, \ + eGType, bPromoteToMulti, nCoordDim, eGeomOp, dfGeomOpParam, \ + nCountLayerFeatures, \ + poClipSrc, poClipDst, bExplodeCollections, nSrcFileSize, \ + pnReadFeatureCount, pfnProgress, pProgressArg) : + + bForceToPolygon = False + bForceToMultiPolygon = False + bForceToMultiLineString = False + + poDstLayer = psInfo.poDstLayer + #papszTransformOptions = psInfo.papszTransformOptions + poCT = psInfo.poCT + panMap = psInfo.panMap + iSrcZField = psInfo.iSrcZField + + if poOutputSRS is None and not bNullifyOutputSRS: + poOutputSRS = poSrcLayer.GetSpatialRef() + + if wkbFlatten(eGType) == ogr.wkbPolygon: + bForceToPolygon = True + elif wkbFlatten(eGType) == ogr.wkbMultiPolygon: + bForceToMultiPolygon = True + elif wkbFlatten(eGType) == ogr.wkbMultiLineString: + bForceToMultiLineString = True + +# -------------------------------------------------------------------- +# Transfer features. +# -------------------------------------------------------------------- + nFeaturesInTransaction = 0 + nCount = 0 + + if nGroupTransactions > 0: + poDstLayer.StartTransaction() + + while True: + poDstFeature = None + + if nFIDToFetch != ogr.NullFID: + + #// Only fetch feature on first pass. + if nFeaturesInTransaction == 0: + poFeature = poSrcLayer.GetFeature(nFIDToFetch) + else: + poFeature = None + + else: + poFeature = poSrcLayer.GetNextFeature() + + if poFeature is None: + break + + nParts = 0 + nIters = 1 + if bExplodeCollections: + poSrcGeometry = poFeature.GetGeometryRef() + if poSrcGeometry is not None: + eSrcType = wkbFlatten(poSrcGeometry.GetGeometryType()) + if eSrcType == ogr.wkbMultiPoint or \ + eSrcType == ogr.wkbMultiLineString or \ + eSrcType == ogr.wkbMultiPolygon or \ + eSrcType == ogr.wkbGeometryCollection: + nParts = poSrcGeometry.GetGeometryCount() + nIters = nParts + if nIters == 0: + nIters = 1 + + for iPart in range(nIters): + nFeaturesInTransaction = nFeaturesInTransaction + 1 + if nFeaturesInTransaction == nGroupTransactions: + poDstLayer.CommitTransaction() + poDstLayer.StartTransaction() + nFeaturesInTransaction = 0 + + gdal.ErrorReset() + poDstFeature = ogr.Feature( poDstLayer.GetLayerDefn() ) + + if poDstFeature.SetFromWithMap( poFeature, 1, panMap ) != 0: + + if nGroupTransactions > 0: + poDstLayer.CommitTransaction() + + print("Unable to translate feature %d from layer %s" % (poFeature.GetFID() , poSrcLayer.GetName() )) + + return False + + if bPreserveFID: + poDstFeature.SetFID( poFeature.GetFID() ) + + poDstGeometry = poDstFeature.GetGeometryRef() + if poDstGeometry is not None: + + if nParts > 0: + # For -explodecollections, extract the iPart(th) of the geometry + poPart = poDstGeometry.GetGeometryRef(iPart).Clone() + poDstFeature.SetGeometryDirectly(poPart) + poDstGeometry = poPart + + if iSrcZField != -1: + SetZ(poDstGeometry, poFeature.GetFieldAsDouble(iSrcZField)) + # This will correct the coordinate dimension to 3 + poDupGeometry = poDstGeometry.Clone() + poDstFeature.SetGeometryDirectly(poDupGeometry) + poDstGeometry = poDupGeometry + + + if nCoordDim == 2 or nCoordDim == 3: + poDstGeometry.SetCoordinateDimension( nCoordDim ) + + if eGeomOp == GeomOperation.SEGMENTIZE: + pass + #if (poDstFeature.GetGeometryRef() is not None and dfGeomOpParam > 0) + # poDstFeature.GetGeometryRef().segmentize(dfGeomOpParam); + elif eGeomOp == GeomOperation.SIMPLIFY_PRESERVE_TOPOLOGY and dfGeomOpParam > 0: + poNewGeom = poDstGeometry.SimplifyPreserveTopology(dfGeomOpParam) + if poNewGeom is not None: + poDstFeature.SetGeometryDirectly(poNewGeom) + poDstGeometry = poNewGeom + + if poClipSrc is not None: + poClipped = poDstGeometry.Intersection(poClipSrc) + if poClipped is None or poClipped.IsEmpty(): + # Report progress + nCount = nCount +1 + if pfnProgress is not None: + pfnProgress(nCount * 1.0 / nCountLayerFeatures, "", pProgressArg) + continue + + poDstFeature.SetGeometryDirectly(poClipped) + poDstGeometry = poClipped + + if poCT is not None: + eErr = poDstGeometry.Transform( poCT ) + if eErr != 0: + if nGroupTransactions > 0: + poDstLayer.CommitTransaction() + + print("Failed to reproject feature %d (geometry probably out of source or destination SRS)." % poFeature.GetFID()) + if not bSkipFailures: + return False + + elif poOutputSRS is not None: + poDstGeometry.AssignSpatialReference(poOutputSRS) + + if poClipDst is not None: + poClipped = poDstGeometry.Intersection(poClipDst) + if poClipped is None or poClipped.IsEmpty(): + continue + + poDstFeature.SetGeometryDirectly(poClipped) + poDstGeometry = poClipped + + if bForceToPolygon: + poDstFeature.SetGeometryDirectly(ogr.ForceToPolygon(poDstGeometry)) + + elif bForceToMultiPolygon or \ + (bPromoteToMulti and wkbFlatten(poDstGeometry.GetGeometryType()) == ogr.wkbPolygon): + poDstFeature.SetGeometryDirectly(ogr.ForceToMultiPolygon(poDstGeometry)) + + elif bForceToMultiLineString or \ + (bPromoteToMulti and wkbFlatten(poDstGeometry.GetGeometryType()) == ogr.wkbLineString): + poDstFeature.SetGeometryDirectly(ogr.ForceToMultiLineString(poDstGeometry)) + + gdal.ErrorReset() + if poDstLayer.CreateFeature( poDstFeature ) != 0 and not bSkipFailures: + if nGroupTransactions > 0: + poDstLayer.RollbackTransaction() + + return False + + # Report progress + nCount = nCount + 1 + if pfnProgress is not None: + if nSrcFileSize != 0: + if (nCount % 1000) == 0: + poFCLayer = poSrcDS.ExecuteSQL("GetBytesRead()", None, None) + if poFCLayer is not None: + poFeat = poFCLayer.GetNextFeature() + if poFeat is not None: + pszReadSize = poFeat.GetFieldAsString(0) + nReadSize = int(pszReadSize) + pfnProgress(nReadSize * 1.0 / nSrcFileSize, "", pProgressArg) + poSrcDS.ReleaseResultSet(poFCLayer) + else: + pfnProgress(nCount * 1.0 / nCountLayerFeatures, "", pProgressArg) + + if pnReadFeatureCount is not None: + pnReadFeatureCount[0] = nCount + + if nGroupTransactions > 0: + poDstLayer.CommitTransaction() + + return True + +if __name__ == '__main__': + version_num = int(gdal.VersionInfo('VERSION_NUM')) + if version_num < 1800: # because of ogr.GetFieldTypeName + print('ERROR: Python bindings of GDAL 1.8.0 or later required') + sys.exit(1) + + if not main(sys.argv): + sys.exit(1) + else: + sys.exit(0) diff --git a/lydorn_utils/polygon_utils.py b/lydorn_utils/polygon_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f32d63685080c1c3e9eafff5bc203f89946c2865 --- /dev/null +++ b/lydorn_utils/polygon_utils.py @@ -0,0 +1,1811 @@ +import sys +import time +from functools import partial +import math +import random +import numpy as np +import scipy.spatial +from PIL import Image, ImageDraw, ImageFilter +import skimage.draw +import skimage +from descartes import PolygonPatch +from matplotlib.collections import PatchCollection +from multiprocess import Pool +import multiprocess +from tqdm import tqdm + +from lydorn_utils import python_utils + +if python_utils.module_exists("skimage.measure"): + from skimage.measure import approximate_polygon + +if python_utils.module_exists("shapely"): + import shapely.geometry + import shapely.affinity + import shapely.ops + import shapely.prepared + import shapely.validation + + +def is_polygon_clockwise(polygon): + rolled_polygon = np.roll(polygon, shift=1, axis=0) + double_signed_area = np.sum((rolled_polygon[:, 0] - polygon[:, 0]) * (rolled_polygon[:, 1] + polygon[:, 1])) + if 0 < double_signed_area: + return True + else: + return False + + +def orient_polygon(polygon, orientation="CW"): + poly_is_orientated_cw = is_polygon_clockwise(polygon) + if (poly_is_orientated_cw and orientation == "CCW") or (not poly_is_orientated_cw and orientation == "CW"): + return np.flip(polygon, axis=0) + else: + return polygon + + +def orient_polygons(polygons, orientation="CW"): + return [orient_polygon(polygon, orientation=orientation) for polygon in polygons] + + +def raster_to_polygon(image, vertex_count): + contours = skimage.measure.find_contours(image, 0.5) + contour = np.empty_like(contours[0]) + contour[:, 0] = contours[0][:, 1] + contour[:, 1] = contours[0][:, 0] + + # Simplify until vertex_count + tolerance = 0.1 + tolerance_step = 0.1 + simplified_contour = contour + while 1 + vertex_count < len(simplified_contour): + simplified_contour = approximate_polygon(contour, tolerance=tolerance) + tolerance += tolerance_step + + simplified_contour = simplified_contour[:-1] + + # plt.imshow(image, cmap="gray") + # plot_polygon(simplified_contour, draw_labels=False) + # plt.show() + + return simplified_contour + + +def l2diffs(polygon1, polygon2): + """ + Computes vertex-wise L2 difference between the two polygons. + As the two polygons may not have the same starting vertex, + all shifts are considred and the shift resulting in the minimum mean L2 difference is chosen + + :param polygon1: + :param polygon2: + :return: + """ + # Make polygons of equal length + if len(polygon1) != len(polygon2): + while len(polygon1) < len(polygon2): + polygon1 = np.append(polygon1, [polygon1[-1, :]], axis=0) + while len(polygon2) < len(polygon1): + polygon2 = np.append(polygon2, [polygon2[-1, :]], axis=0) + vertex_count = len(polygon1) + + def naive_l2diffs(polygon1, polygon2): + naive_l2diffs_result = np.sqrt(np.power(np.sum(polygon1 - polygon2, axis=1), 2)) + return naive_l2diffs_result + + min_l2_diffs = naive_l2diffs(polygon1, polygon2) + min_mean_l2_diffs = np.mean(min_l2_diffs, axis=0) + for i in range(1, vertex_count): + current_naive_l2diffs = naive_l2diffs(np.roll(polygon1, shift=i, axis=0), polygon2) + current_naive_mean_l2diffs = np.mean(current_naive_l2diffs, axis=0) + if current_naive_mean_l2diffs < min_mean_l2_diffs: + min_l2_diffs = current_naive_l2diffs + min_mean_l2_diffs = current_naive_mean_l2diffs + return min_l2_diffs + + +def intersect_polygons(simple_polygon, multi_polygon): + """ + + :param input_polygon: + :param target_polygon: + :return: List of a simple polygon: [poly1, poly2,...] with a multi polygon: [[(x1, y1), (x2, y2), ...], [...]] + """ + poly1 = shapely.geometry.Polygon(simple_polygon).buffer(0) + poly2 = shapely.geometry.MultiPolygon(shapely.geometry.Polygon(polygon) for polygon in multi_polygon).buffer(0) + intersection_poly = poly1.intersection(poly2) + if 0 < intersection_poly.area: + if intersection_poly.type == 'Polygon': + coords = intersection_poly.exterior.coords + return [coords] + elif intersection_poly.type == 'MultiPolygon': + ret_coords = [] + for poly in intersection_poly: + coords = poly.exterior.coords + ret_coords.append(coords) + return ret_coords + return None + + +def check_intersection_with_polygon(input_polygon, target_polygon): + poly1 = shapely.geometry.Polygon(input_polygon).buffer(0) + poly2 = shapely.geometry.Polygon(target_polygon).buffer(0) + intersection_poly = poly1.intersection(poly2) + intersection_area = intersection_poly.area + is_intersection = 0 < intersection_area + return is_intersection + + +def check_intersection_with_polygons(input_polygon, target_polygons): + """ + Returns True if there is an intersection with at least one polygon in target_polygons + :param input_polygon: + :param target_polygons: + :return: + """ + for target_polygon in target_polygons: + if check_intersection_with_polygon(input_polygon, target_polygon): + return True + return False + + +def polygon_area(polygon): + poly = shapely.geometry.Polygon(polygon).buffer(0) + return poly.area + + +def polygon_union(polygon1, polygon2): + poly1 = shapely.geometry.Polygon(polygon1).buffer(0) + poly2 = shapely.geometry.Polygon(polygon2).buffer(0) + union_poly = poly1.union(poly2) + return np.array(union_poly.exterior.coords) + + +def polygon_iou(polygon1, polygon2): + poly1 = shapely.geometry.Polygon(polygon1).buffer(0) + poly2 = shapely.geometry.Polygon(polygon2).buffer(0) + intersection_poly = poly1.intersection(poly2) + union_poly = poly1.union(poly2) + intersection_area = intersection_poly.area + union_area = union_poly.area + if union_area: + iou = intersection_area / union_area + else: + iou = 0 + return iou + + +def generate_polygon(cx, cy, ave_radius, irregularity, spikeyness, vertex_count): + """ + Start with the centre of the polygon at cx, cy, + then creates the polygon by sampling points on a circle around the centre. + Random noise is added by varying the angular spacing between sequential points, + and by varying the radial distance of each point from the centre. + + Params: + cx, cy - coordinates of the "centre" of the polygon + ave_radius - in px, the average radius of this polygon, this roughly controls how large the polygon is, + really only useful for order of magnitude. + irregularity - [0,1] indicating how much variance there is in the angular spacing of vertices. [0,1] will map to + [0, 2 * pi / vertex_count] + spikeyness - [0,1] indicating how much variance there is in each vertex from the circle of radius ave_radius. + [0,1] will map to [0, ave_radius] + vertex_count - self-explanatory + + Returns a list of vertices, in CCW order. + """ + + irregularity = clip(irregularity, 0, 1) * 2 * math.pi / vertex_count + spikeyness = clip(spikeyness, 0, 1) * ave_radius + + # generate n angle steps + angle_steps = [] + lower = (2 * math.pi / vertex_count) - irregularity + upper = (2 * math.pi / vertex_count) + irregularity + angle_sum = 0 + for i in range(vertex_count): + tmp = random.uniform(lower, upper) + angle_steps.append(tmp) + angle_sum = angle_sum + tmp + + # normalize the steps so that point 0 and point n+1 are the same + k = angle_sum / (2 * math.pi) + for i in range(vertex_count): + angle_steps[i] = angle_steps[i] / k + + # now generate the points + points = [] + angle = random.uniform(0, 2 * math.pi) + for i in range(vertex_count): + r_i = clip(random.gauss(ave_radius, spikeyness), 0, 2 * ave_radius) + x = cx + r_i * math.cos(angle) + y = cy + r_i * math.sin(angle) + points.append((x, y)) + + angle = angle + angle_steps[i] + + return points + + +def clip(x, mini, maxi): + if mini > maxi: + return x + elif x < mini: + return mini + elif x > maxi: + return maxi + else: + return x + + +def scale_bounding_box(bounding_box, scale): + half_width = math.ceil((bounding_box[2] - bounding_box[0]) * scale / 2) + half_height = math.ceil((bounding_box[3] - bounding_box[1]) * scale / 2) + center = [round((bounding_box[0] + bounding_box[2]) / 2), round((bounding_box[1] + bounding_box[3]) / 2)] + scaled_bounding_box = [int(center[0] - half_width), int(center[1] - half_height), int(center[0] + half_width), + int(center[1] + half_height)] + return scaled_bounding_box + + +def pad_bounding_box(bbox, pad): + return [bbox[0] + pad, bbox[1] + pad, bbox[2] - pad, bbox[3] - pad] + + +def compute_bounding_box(polygon, scale=1, boundingbox_margin=0, fit=None): + # Compute base bounding box + bounding_box = [np.min(polygon[:, 0]), np.min(polygon[:, 1]), np.max(polygon[:, 0]), np.max(polygon[:, 1])] + # Scale + half_width = math.ceil((bounding_box[2] - bounding_box[0]) * scale / 2) + half_height = math.ceil((bounding_box[3] - bounding_box[1]) * scale / 2) + # Add margin + half_width += boundingbox_margin + half_height += boundingbox_margin + # Compute square bounding box + if fit == "square": + half_width = half_height = max(half_width, half_height) + center = [round((bounding_box[0] + bounding_box[2]) / 2), round((bounding_box[1] + bounding_box[3]) / 2)] + bounding_box = [int(center[0] - half_width), int(center[1] - half_height), int(center[0] + half_width), + int(center[1] + half_height)] + return bounding_box + + +def compute_patch(polygon, patch_size): + centroid = np.mean(polygon, axis=0) + half_height = half_width = patch_size / 2 + bounding_box = [math.ceil(centroid[0] - half_width), math.ceil(centroid[1] - half_height), + math.ceil(centroid[0] + half_width), math.ceil(centroid[1] + half_height)] + return bounding_box + + +def bounding_box_within_bounds(bounding_box, bounds): + return bounds[0] <= bounding_box[0] and bounds[1] <= bounding_box[1] and bounding_box[2] <= bounds[2] and \ + bounding_box[3] <= bounds[3] + + +def vertex_within_bounds(vertex, bounds): + return bounds[0] <= vertex[0] <= bounds[2] and \ + bounds[1] <= vertex[1] <= bounds[3] + + +def edge_within_bounds(edge, bounds): + return vertex_within_bounds(edge[0], bounds) and vertex_within_bounds(edge[1], bounds) + + +def bounding_box_area(bounding_box): + return (bounding_box[2] - bounding_box[0]) * (bounding_box[3] - bounding_box[1]) + + +def convert_to_image_patch_space(polygon_image_space, bounding_box): + polygon_image_patch_space = np.empty_like(polygon_image_space) + polygon_image_patch_space[:, 0] = polygon_image_space[:, 0] - bounding_box[0] + polygon_image_patch_space[:, 1] = polygon_image_space[:, 1] - bounding_box[1] + return polygon_image_patch_space + + +def translate_polygons(polygons, translation): + for polygon in polygons: + polygon[:, 0] += translation[0] + polygon[:, 1] += translation[1] + return polygons + + +def strip_redundant_vertex(vertices, epsilon=1): + assert len(vertices.shape) == 2 # Is a polygon + new_vertices = vertices + if 1 < vertices.shape[0]: + if np.sum(np.absolute(vertices[0, :] - vertices[-1, :])) < epsilon: + new_vertices = vertices[:-1, :] + return new_vertices + + +def remove_doubles(vertices, epsilon=0.1): + dists = np.linalg.norm(np.roll(vertices, -1, axis=0) - vertices, axis=-1) + new_vertices = vertices[epsilon < dists] + return new_vertices + + +def simplify_polygon(polygon, tolerance=1): + approx_polygon = approximate_polygon(polygon, tolerance=tolerance) + return approx_polygon + + +def simplify_polygons(polygons, tolerance=1): + approx_polygons = [] + for polygon in polygons: + approx_polygon = approximate_polygon(polygon, tolerance=tolerance) + approx_polygons.append(approx_polygon) + return approx_polygons + + +def pad_polygon(vertices, target_length): + assert len(vertices.shape) == 2 # Is a polygon + assert vertices.shape[0] <= target_length + padding_length = target_length - vertices.shape[0] + padding = np.tile(vertices[-1], [padding_length, 1]) + padded_vertices = np.append(vertices, padding, axis=0) + return padded_vertices + + +def compute_diameter(polygon): + dist = scipy.spatial.distance.cdist(polygon, polygon) + return dist.max() + + +def plot_polygon(polygon, color=None, draw_labels=True, label_direction=1, indexing="xy", axis=None): + if python_utils.module_exists("matplotlib.pyplot"): + import matplotlib.pyplot as plt + + if axis is None: + axis = plt.gca() + + polygon_closed = np.append(polygon, [polygon[0, :]], axis=0) + if indexing == "xy=": + axis.plot(polygon_closed[:, 0], polygon_closed[:, 1], color=color, linewidth=3.0) + elif indexing == "ij": + axis.plot(polygon_closed[:, 1], polygon_closed[:, 0], color=color, linewidth=3.0) + else: + print("WARNING: Invalid indexing argument") + + if draw_labels: + labels = range(1, polygon.shape[0] + 1) + for label, x, y in zip(labels, polygon[:, 0], polygon[:, 1]): + axis.annotate( + label, + xy=(x, y), xytext=(-20 * label_direction, 20 * label_direction), + textcoords='offset points', ha='right', va='bottom', + bbox=dict(boxstyle='round,pad=0.25', fc=color, alpha=0.75), + arrowprops=dict(arrowstyle='->', color=color, connectionstyle='arc3,rad=0')) + + +def plot_polygons(polygons, color=None, draw_labels=True, label_direction=1, indexing="xy", axis=None): + for polygon in polygons: + plot_polygon(polygon, color=color, draw_labels=draw_labels, label_direction=label_direction, indexing=indexing, + axis=axis) + + +def compute_edge_normal(edge): + normal = np.array([- (edge[1][1] - edge[0][1]), + edge[1][0] - edge[0][0]]) + normal_norm = np.sqrt(np.sum(np.square(normal))) + normal /= normal_norm + return normal + + +def compute_vector_angle(x, y): + if x < 0.0: + slope = y / x + angle = np.pi + np.arctan(slope) + elif 0.0 < x: + slope = y / x + angle = np.arctan(slope) + else: + if 0 < y: + angle = np.pi / 2 + else: + angle = 3 * np.pi / 2 + if angle < 0.0: + angle += 2 * np.pi + return angle + + +def compute_edge_normal_angle_edge(edge): + normal = compute_edge_normal(edge) + normal_x = normal[1] + normal_y = normal[0] + angle = compute_vector_angle(normal_x, normal_y) + return angle + + +def polygon_in_bounding_box(polygon, bounding_box): + """ + Returns True if all vertices of polygons are inside bounding_box + :param polygon: [N, 2] + :param bounding_box: [row_min, col_min, row_max, col_max] + :return: + """ + result = np.all( + np.logical_and( + np.logical_and(bounding_box[0] <= polygon[:, 0], polygon[:, 0] <= bounding_box[2]), + np.logical_and(bounding_box[1] <= polygon[:, 1], polygon[:, 1] <= bounding_box[3]) + ) + ) + return result + + +def filter_polygons_in_bounding_box(polygons, bounding_box): + """ + Only keep polygons that are fully inside bounding_box + + :param polygons: [shape(N, 2), ...] + :param bounding_box: [row_min, col_min, row_max, col_max] + :return: + """ + filtered_polygons = [] + for polygon in polygons: + if polygon_in_bounding_box(polygon, bounding_box): + filtered_polygons.append(polygon) + return filtered_polygons + + +def transform_polygon_to_bounding_box_space(polygon, bounding_box): + """ + + :param polygon: shape(N, 2) + :param bounding_box: [row_min, col_min, row_max, col_max] + :return: + """ + assert len(polygon.shape) and polygon.shape[1] == 2, "polygon should have shape (N, 2), not shape {}".format( + polygon.shape) + assert len(bounding_box) == 4, "bounding_box should have 4 elements: [row_min, col_min, row_max, col_max]" + transformed_polygon = polygon.copy() + transformed_polygon[:, 0] -= bounding_box[0] + transformed_polygon[:, 1] -= bounding_box[1] + return transformed_polygon + + +def transform_polygons_to_bounding_box_space(polygons, bounding_box): + transformed_polygons = [] + for polygon in polygons: + transformed_polygons.append(transform_polygon_to_bounding_box_space(polygon, bounding_box)) + return transformed_polygons + + +def crop_polygon_to_patch(polygon, bounding_box): + return transform_polygon_to_bounding_box_space(polygon, bounding_box) + + +def crop_polygon_to_patch_if_touch(polygon, bounding_box): + assert type(polygon) == np.ndarray, "polygon should be a numpy array, not {}".format(type(polygon)) + assert len(polygon.shape) == 2 and polygon.shape[1] == 2, "polygon should be of shape (N, 2), not {}".format( + polygon.shape) + # Verify that at least one vertex is inside bounding_box + polygon_touches_patch = np.any( + np.logical_and( + np.logical_and(bounding_box[0] <= polygon[:, 0], polygon[:, 0] <= bounding_box[2]), + np.logical_and(bounding_box[1] <= polygon[:, 1], polygon[:, 1] <= bounding_box[3]) + ) + ) + if polygon_touches_patch: + return crop_polygon_to_patch(polygon, bounding_box) + else: + return None + + +def crop_polygons_to_patch_if_touch(polygons, bounding_box, return_indices=False): + assert type(polygons) == list, "polygons should be a list" + if return_indices: + indices = [] + cropped_polygons = [] + for i, polygon in enumerate(polygons): + cropped_polygon = crop_polygon_to_patch_if_touch(polygon, bounding_box) + if cropped_polygon is not None: + cropped_polygons.append(cropped_polygon) + if return_indices: + indices.append(i) + if return_indices: + return cropped_polygons, indices + else: + return cropped_polygons + + +def crop_polygons_to_patch(polygons, bounding_box): + cropped_polygons = [] + for polygon in polygons: + cropped_polygon = crop_polygon_to_patch(polygon, bounding_box) + if cropped_polygon is not None: + cropped_polygons.append(cropped_polygon) + return cropped_polygons + + +def patch_polygons(polygons, minx, miny, maxx, maxy): + """ + Filters out polygons that do not touch the bbox and translate those that do to the box's coordinate system. + + @param polygons: [shapely.geometry.Polygon, ...] + @param maxy: + @param maxx: + @param miny: + @param minx: + @return: [shapely.geometry.Polygon, ...] + """ + assert type(polygons) == list, "polygons should be a list" + if len(polygons) == 0: + return polygons + assert type(polygons[0]) == shapely.geometry.Polygon, \ + f"Items of the polygons list should be of type shapely.geometry.Polygon, not {type(polygons[0])}" + + box_polygon = shapely.geometry.box(minx, miny, maxx, maxy) + polygons = filter(box_polygon.intersects, polygons) + + polygons = map(partial(shapely.affinity.translate, xoff=-minx, yoff=-miny), polygons) + + return list(polygons) + + +def polygon_remove_holes(polygon): + polygon_no_holes = [] + for coords in polygon: + if not np.isnan(coords[0]) and not np.isnan(coords[1]): + polygon_no_holes.append(coords) + else: + break + return np.array(polygon_no_holes) + + +def polygons_remove_holes(polygons): + gt_polygons_no_holes = [] + for polygon in polygons: + gt_polygons_no_holes.append(polygon_remove_holes(polygon)) + return gt_polygons_no_holes + + +def apply_batch_disp_map_to_polygons(pred_disp_field_map_batch, disp_polygons_batch): + """ + + :param pred_disp_field_map_batch: shape(batch_size, height, width, 2) + :param disp_polygons_batch: shape(batch_size, polygon_count, vertex_count, 2) + :return: + """ + + # Apply all displacements at once + batch_count = pred_disp_field_map_batch.shape[0] + row_count = pred_disp_field_map_batch.shape[1] + col_count = pred_disp_field_map_batch.shape[2] + + disp_polygons_batch_int = np.round(disp_polygons_batch).astype(np.int) + # Clip coordinates to the field map: + disp_polygons_batch_int_nearest_valid_field = np.maximum(0, disp_polygons_batch_int) + disp_polygons_batch_int_nearest_valid_field[:, :, :, 0] = np.minimum( + disp_polygons_batch_int_nearest_valid_field[:, :, :, 0], row_count - 1) + disp_polygons_batch_int_nearest_valid_field[:, :, :, 1] = np.minimum( + disp_polygons_batch_int_nearest_valid_field[:, :, :, 1], col_count - 1) + + aligned_disp_polygons_batch = disp_polygons_batch.copy() + for batch_index in range(batch_count): + mask = ~np.isnan(disp_polygons_batch[batch_index, :, :, 0]) # Checking one coordinate is enough + aligned_disp_polygons_batch[batch_index, mask, 0] += pred_disp_field_map_batch[batch_index, + disp_polygons_batch_int_nearest_valid_field[ + batch_index, mask, 0], + disp_polygons_batch_int_nearest_valid_field[ + batch_index, mask, 1], 0].flatten() + aligned_disp_polygons_batch[batch_index, mask, 1] += pred_disp_field_map_batch[batch_index, + disp_polygons_batch_int_nearest_valid_field[ + batch_index, mask, 0], + disp_polygons_batch_int_nearest_valid_field[ + batch_index, mask, 1], 1].flatten() + return aligned_disp_polygons_batch + + +def apply_disp_map_to_polygons(disp_field_map, polygons): + """ + + :param disp_field_map: shape(height, width, 2) + :param polygon_list: [shape(N, 2), shape(M, 2), ...] + :return: + """ + disp_field_map_batch = np.expand_dims(disp_field_map, axis=0) + disp_polygons = [] + for polygon in polygons: + polygon_batch = np.expand_dims(np.expand_dims(polygon, axis=0), axis=0) # Add batch and polygon_count dims + disp_polygon_batch = apply_batch_disp_map_to_polygons(disp_field_map_batch, polygon_batch) + disp_polygon_batch = disp_polygon_batch[0, 0] # Remove batch and polygon_count dims + disp_polygons.append(disp_polygon_batch) + return disp_polygons + + +# This next function is somewhat redundant with apply_disp_map_to_polygons... (but displaces in the opposite direction) +def apply_displacement_field_to_polygons(polygons, disp_field_map): + disp_polygons = [] + for polygon in polygons: + mask_nans = np.isnan(polygon) # Will be necessary when polygons with holes are handled + polygon_int = np.round(polygon).astype(np.int) + polygon_int_clipped = np.maximum(0, polygon_int) + polygon_int_clipped[:, 0] = np.minimum(disp_field_map.shape[0] - 1, polygon_int_clipped[:, 0]) + polygon_int_clipped[:, 1] = np.minimum(disp_field_map.shape[1] - 1, polygon_int_clipped[:, 1]) + disp_polygon = polygon.copy() + disp_polygon[~mask_nans[:, 0], 0] -= disp_field_map[polygon_int_clipped[~mask_nans[:, 0], 0], + polygon_int_clipped[~mask_nans[:, 0], 1], 0] + disp_polygon[~mask_nans[:, 1], 1] -= disp_field_map[polygon_int_clipped[~mask_nans[:, 1], 0], + polygon_int_clipped[~mask_nans[:, 1], 1], 1] + disp_polygons.append(disp_polygon) + return disp_polygons + + +def apply_displacement_fields_to_polygons(polygons, disp_field_maps): + disp_field_map_count = disp_field_maps.shape[0] + disp_polygons_list = [] + for i in range(disp_field_map_count): + disp_polygons = apply_displacement_field_to_polygons(polygons, disp_field_maps[i, :, :, :]) + disp_polygons_list.append(disp_polygons) + return disp_polygons_list + + +def draw_line(shape, line, width, blur_radius=0): + im = Image.new("L", (shape[1], shape[0])) + # im_px_access = im.load() + draw = ImageDraw.Draw(im) + vertex_list = [] + for coords in line: + vertex = (coords[1], coords[0]) + vertex_list.append(vertex) + draw.line(vertex_list, fill=255, width=width) + if 0 < blur_radius: + im = im.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + array = np.array(im) / 255 + return array + + +def draw_triangle(shape, triangle, blur_radius=0): + im = Image.new("L", (shape[1], shape[0])) + # im_px_access = im.load() + draw = ImageDraw.Draw(im) + vertex_list = [] + for coords in triangle: + vertex = (coords[1], coords[0]) + vertex_list.append(vertex) + draw.polygon(vertex_list, fill=255) + if 0 < blur_radius: + im = im.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + array = np.array(im) / 255 + return array + + +def draw_polygon(polygon, shape, fill=True, edges=True, vertices=True, line_width=3): + # TODO: handle holes in polygons + im = Image.new("RGB", (shape[1], shape[0])) + im_px_access = im.load() + draw = ImageDraw.Draw(im) + + vertex_list = [] + for coords in polygon: + vertex = (coords[1], coords[0]) + if not np.isnan(vertex[0]) and not np.isnan(vertex[1]): + vertex_list.append(vertex) + else: + break + if edges: + draw.line(vertex_list, fill=(0, 255, 0), width=line_width) + if fill: + draw.polygon(vertex_list, fill=(255, 0, 0)) + if vertices: + draw.point(vertex_list, fill=(0, 0, 255)) + + # Convert image to numpy array with the right number of channels + array = np.array(im) + selection = [fill, edges, vertices] + selected_array = array[:, :, selection] + return selected_array + + +def _draw_circle(draw, center, radius, fill): + draw.ellipse([center[0] - radius, + center[1] - radius, + center[0] + radius, + center[1] + radius], fill=fill, outline=None) + + +def draw_polygons(polygons, shape, fill=True, edges=True, vertices=True, line_width=3, antialiasing=False): + # TODO: handle holes in polygons + polygons = polygons_remove_holes(polygons) + polygons = polygons_close(polygons) + + if antialiasing: + draw_shape = (2 * shape[0], 2 * shape[1]) + else: + draw_shape = shape + # Channels + fill_channel_index = 0 # Always first channel + edges_channel_index = fill # If fill == True, take second channel. If not then take first + vertices_channel_index = fill + edges # Same principle as above + channel_count = fill + edges + vertices + im_draw_list = [] + for channel_index in range(channel_count): + im = Image.new("L", (draw_shape[1], draw_shape[0])) + im_px_access = im.load() + draw = ImageDraw.Draw(im) + im_draw_list.append((im, draw)) + + for polygon in polygons: + if antialiasing: + polygon *= 2 + vertex_list = [] + for coords in polygon: + vertex_list.append((coords[1], coords[0])) + if fill: + draw = im_draw_list[fill_channel_index][1] + draw.polygon(vertex_list, fill=255) + if edges: + draw = im_draw_list[edges_channel_index][1] + draw.line(vertex_list, fill=255, width=line_width) + if vertices: + draw = im_draw_list[vertices_channel_index][1] + for vertex in vertex_list: + _draw_circle(draw, vertex, line_width / 2, fill=255) + + im_list = [] + if antialiasing: + # resize images: + for im_draw in im_draw_list: + resize_shape = (shape[1], shape[0]) + im_list.append(im_draw[0].resize(resize_shape, Image.BILINEAR)) + else: + for im_draw in im_draw_list: + im_list.append(im_draw[0]) + + # Convert image to numpy array with the right number of channels + array_list = [np.array(im) for im in im_list] + array = np.stack(array_list, axis=-1) + return array + + +def draw_polygon_map(polygons, shape, fill=True, edges=True, vertices=True, line_width=3): + """ + Alias for draw_polygon function + + :param polygons: + :param shape: + :param fill: + :param edges: + :param vertices: + :param line_width: + :return: + """ + return draw_polygons(polygons, shape, fill=fill, edges=edges, vertices=vertices, line_width=line_width) + + +def draw_polygon_maps(polygons_list, shape, fill=True, edges=True, vertices=True, line_width=3): + polygon_maps_list = [] + for polygons in polygons_list: + polygon_map = draw_polygon_map(polygons, shape, fill=fill, edges=edges, vertices=vertices, + line_width=line_width) + polygon_maps_list.append(polygon_map) + disp_field_maps = np.stack(polygon_maps_list, axis=0) + return disp_field_maps + + +def swap_coords(polygon): + polygon_new = polygon.copy() + polygon_new[..., 0] = polygon[..., 1] + polygon_new[..., 1] = polygon[..., 0] + return polygon_new + + +def prepare_polygons_for_tfrecord(gt_polygons, disp_polygons_list, boundingbox=None): + assert len(gt_polygons) + + # print("Starting to crop polygons") + # start = time.time() + + dtype = gt_polygons[0].dtype + cropped_gt_polygons = [] + cropped_disp_polygons_list = [[] for i in range(len(disp_polygons_list))] + polygon_length = 0 + for polygon_index, gt_polygon in enumerate(gt_polygons): + if boundingbox is not None: + cropped_gt_polygon = crop_polygon_to_patch_if_touch(gt_polygon, boundingbox) + else: + cropped_gt_polygon = gt_polygon + if cropped_gt_polygon is not None: + cropped_gt_polygons.append(cropped_gt_polygon) + if polygon_length < cropped_gt_polygon.shape[0]: + polygon_length = cropped_gt_polygon.shape[0] + # Crop disp polygons + for disp_index, disp_polygons in enumerate(disp_polygons_list): + disp_polygon = disp_polygons[polygon_index] + if boundingbox is not None: + cropped_disp_polygon = crop_polygon_to_patch(disp_polygon, boundingbox) + else: + cropped_disp_polygon = disp_polygon + cropped_disp_polygons_list[disp_index].append(cropped_disp_polygon) + + # end = time.time() + # print("Finished cropping polygons in in {}s".format(end - start)) + # + # print("Starting to pad polygons") + # start = time.time() + + polygon_count = len(cropped_gt_polygons) + if polygon_count: + # Add +1 to both dimensions for end-of-item NaNs + padded_gt_polygons = np.empty((polygon_count + 1, polygon_length + 1, 2), dtype=dtype) + padded_gt_polygons[:, :, :] = np.nan + padded_disp_polygons_array = np.empty((len(disp_polygons_list), polygon_count + 1, polygon_length + 1, 2), + dtype=dtype) + padded_disp_polygons_array[:, :, :] = np.nan + for i, polygon in enumerate(cropped_gt_polygons): + padded_gt_polygons[i, 0:polygon.shape[0], :] = polygon + for j, polygons in enumerate(cropped_disp_polygons_list): + for i, polygon in enumerate(polygons): + padded_disp_polygons_array[j, i, 0:polygon.shape[0], :] = polygon + else: + padded_gt_polygons = padded_disp_polygons_array = None + + # end = time.time() + # print("Finished padding polygons in in {}s".format(end - start)) + + return padded_gt_polygons, padded_disp_polygons_array + + +def prepare_stages_polygons_for_tfrecord(gt_polygons, disp_polygons_list_list, boundingbox): + assert len(gt_polygons) + + print(gt_polygons) + print(disp_polygons_list_list) + + exit() + + # print("Starting to crop polygons") + # start = time.time() + + dtype = gt_polygons[0].dtype + cropped_gt_polygons = [] + cropped_disp_polygons_list_list = [[[] for i in range(len(disp_polygons_list))] for disp_polygons_list in + disp_polygons_list_list] + polygon_length = 0 + for polygon_index, gt_polygon in enumerate(gt_polygons): + cropped_gt_polygon = crop_polygon_to_patch_if_touch(gt_polygon, boundingbox) + if cropped_gt_polygon is not None: + cropped_gt_polygons.append(cropped_gt_polygon) + if polygon_length < cropped_gt_polygon.shape[0]: + polygon_length = cropped_gt_polygon.shape[0] + # Crop disp polygons + for stage_index, disp_polygons_list in enumerate(disp_polygons_list_list): + for disp_index, disp_polygons in enumerate(disp_polygons_list): + disp_polygon = disp_polygons[polygon_index] + cropped_disp_polygon = crop_polygon_to_patch(disp_polygon, boundingbox) + cropped_disp_polygons_list_list[stage_index][disp_index].append(cropped_disp_polygon) + + # end = time.time() + # print("Finished cropping polygons in in {}s".format(end - start)) + # + # print("Starting to pad polygons") + # start = time.time() + + polygon_count = len(cropped_gt_polygons) + if polygon_count: + # Add +1 to both dimensions for end-of-item NaNs + padded_gt_polygons = np.empty((polygon_count + 1, polygon_length + 1, 2), dtype=dtype) + padded_gt_polygons[:, :, :] = np.nan + padded_disp_polygons_array = np.empty( + (len(disp_polygons_list_list), len(disp_polygons_list_list[0]), polygon_count + 1, polygon_length + 1, 2), + dtype=dtype) + padded_disp_polygons_array[:, :, :] = np.nan + for i, polygon in enumerate(cropped_gt_polygons): + padded_gt_polygons[i, 0:polygon.shape[0], :] = polygon + for k, cropped_disp_polygons_list in enumerate(cropped_disp_polygons_list_list): + for j, polygons in enumerate(cropped_disp_polygons_list): + for i, polygon in enumerate(polygons): + padded_disp_polygons_array[k, j, i, 0:polygon.shape[0], :] = polygon + else: + padded_gt_polygons = padded_disp_polygons_array = None + + # end = time.time() + # print("Finished padding polygons in in {}s".format(end - start)) + + return padded_gt_polygons, padded_disp_polygons_array + + +def rescale_polygon(polygons, scaling_factor): + """ + + :param polygons: + :return: scaling_factor + """ + if len(polygons): + rescaled_polygons = [polygon * scaling_factor for polygon in polygons] + return rescaled_polygons + else: + return polygons + + +def get_edge_center(edge): + return np.mean(edge, axis=0) + + +def get_edge_length(edge): + return np.sqrt(np.sum(np.square(edge[0] - edge[1]))) + + +def get_edges_angle(edge1, edge2): + x1 = edge1[1, 0] - edge1[0, 0] + y1 = edge1[1, 1] - edge1[0, 1] + x2 = edge2[1, 0] - edge2[0, 0] + y2 = edge2[1, 1] - edge2[0, 1] + angle1 = compute_vector_angle(x1, y1) + angle2 = compute_vector_angle(x2, y2) + edges_angle = math.fabs(angle1 - angle2) % (2 * math.pi) + if math.pi < edges_angle: + edges_angle = 2 * math.pi - edges_angle + return edges_angle + + +def compute_angle_two_points(point_source, point_target): + vector = point_target - point_source + angle = compute_vector_angle(vector[0], vector[1]) + return angle + + +def compute_angle_three_points(point_source, point_target1, point_target2): + squared_dist_source_target1 = math.pow((point_source[0] - point_target1[0]), 2) + math.pow( + (point_source[1] - point_target1[1]), 2) + squared_dist_source_target2 = math.pow((point_source[0] - point_target2[0]), 2) + math.pow( + (point_source[1] - point_target2[1]), 2) + squared_dist_target1_target2 = math.pow((point_target1[0] - point_target2[0]), 2) + math.pow( + (point_target1[1] - point_target2[1]), 2) + dist_source_target1 = math.sqrt(squared_dist_source_target1) + dist_source_target2 = math.sqrt(squared_dist_source_target2) + try: + cos = (squared_dist_source_target1 + squared_dist_source_target2 - squared_dist_target1_target2) / ( + 2 * dist_source_target1 * dist_source_target2) + except ZeroDivisionError: + return float('inf') + cos = max(min(cos, 1), + -1) # Avoid some math domain error due to cos being slightly bigger than 1 (from floating point operations) + angle = math.acos(cos) + return angle + + +def are_edges_overlapping(edge1, edge2, threshold): + """ + Checks if at least 2 different vertices of either edge lies on the other edge: it characterizes an overlap + :param edge1: + :param edge2: + :param threshold: + :return: + """ + count_list = [ + is_vertex_on_edge(edge1[0], edge2, threshold), + is_vertex_on_edge(edge1[1], edge2, threshold), + is_vertex_on_edge(edge2[0], edge1, threshold), + is_vertex_on_edge(edge2[1], edge1, threshold), + ] + # Count number of identical vertices + identical_vertex_list = [ + np.array_equal(edge1[0], edge2[0]), + np.array_equal(edge1[0], edge2[1]), + np.array_equal(edge1[1], edge2[0]), + np.array_equal(edge1[1], edge2[1]), + ] + adjusted_count = np.sum(count_list) - np.sum(identical_vertex_list) + return 2 <= adjusted_count + + +# def are_edges_collinear(edge1, edge2, angle_threshold): +# edges_angle = get_edges_angle(edge1, edge2) +# return edges_angle < angle_threshold + + +def get_line_intersect(a1, a2, b1, b2): + """ + Returns the point of intersection of the lines passing through a2,a1 and b2,b1. + a1: [x, y] a point on the first line + a2: [x, y] another point on the first line + b1: [x, y] a point on the second line + b2: [x, y] another point on the second line + """ + s = np.vstack([a1, a2, b1, b2]) # s for stacked + h = np.hstack((s, np.ones((4, 1)))) # h for homogeneous + l1 = np.cross(h[0], h[1]) # get first line + l2 = np.cross(h[2], h[3]) # get second line + x, y, z = np.cross(l1, l2) # point of intersection + if z == 0: # lines are parallel + return float('inf'), float('inf') + return x / z, y / z + + +def are_edges_intersecting(edge1, edge2, epsilon=1e-6): + """ + edge1 and edge2 should not have a common vertex between them + :param edge1: + :param edge2: + :return: + """ + intersect = get_line_intersect(edge1[0], edge1[1], edge2[0], edge2[1]) + # print("---") + # print(edge1) + # print(edge2) + # print(intersect) + if intersect[0] == float('inf') or intersect[1] == float('inf'): + # Lines don't intersect + return False + else: + # Lines intersect + # Check if intersect point belongs to both edges + angle1 = compute_angle_three_points(intersect, edge1[0], edge1[1]) + angle2 = compute_angle_three_points(intersect, edge2[0], edge2[1]) + intersect_belongs_to_edges = (math.pi - epsilon) < angle1 and (math.pi - epsilon) < angle2 + return intersect_belongs_to_edges + + +def shorten_edge(edge, length_to_cut1, length_to_cut2, min_length): + center = get_edge_center(edge) + total_length = get_edge_length(edge) + new_length = total_length - length_to_cut1 - length_to_cut2 + if min_length <= new_length: + scale = new_length / total_length + new_edge = (edge.copy() - center) * scale + center + return new_edge + else: + return None + + +def is_edge_in_triangle(edge, triangle): + return edge[0] in triangle and edge[1] in triangle + + +def get_connectivity_of_edge(edge, triangles): + connectivity = 0 + for triangle in triangles: + connectivity += is_edge_in_triangle(edge, triangle) + return connectivity + + +def get_connectivity_of_edges(edges, triangles): + connectivity_of_edges = [] + for edge in edges: + connectivity_of_edge = get_connectivity_of_edge(edge, triangles) + connectivity_of_edges.append(connectivity_of_edge) + return connectivity_of_edges + + +def polygon_to_closest_int(polygons): + int_polygons = [] + for polygon in polygons: + int_polygon = np.round(polygon) + int_polygons.append(int_polygon) + return int_polygons + + +def is_vertex_on_edge(vertex, edge, threshold): + """ + :param vertex: + :param edge: + :param threshold: + :return: + """ + # Compare distances sum to edge length + edge_length = get_edge_length(edge) + dist1 = get_edge_length([vertex, edge[0]]) + dist2 = get_edge_length([vertex, edge[1]]) + vertex_on_edge = (dist1 + dist2) < (edge_length + threshold) + return vertex_on_edge + + +def get_face_edges(face_vertices): + edges = [] + prev_vertex = face_vertices[0] + for vertex in face_vertices[1:]: + edge = (prev_vertex, vertex) + edges.append(edge) + + # For next iteration: + prev_vertex = vertex + return edges + + +def find_edge_in_face(edge, face_vertices): + # Copy inputs list so that we don't modify it + face_vertices = face_vertices[:] + face_vertices.append(face_vertices[0]) # Close face (does not matter if it is already closed) + edges = get_face_edges(face_vertices) + index = edges.index(edge) + return index + + +def clean_degenerate_face_edges(face_vertices): + def recursive_clean_degenerate_face_edges(open_face_vertices): + face_vertex_count = len(open_face_vertices) + cleaned_open_face_vertices = [] + skip = False + for index in range(face_vertex_count): + if skip: + skip = False + else: + prev_vertex = open_face_vertices[(index - 1) % face_vertex_count] + vertex = open_face_vertices[index] + next_vertex = open_face_vertices[(index + 1) % face_vertex_count] + if prev_vertex != next_vertex: + cleaned_open_face_vertices.append(vertex) + else: + skip = True + if len(cleaned_open_face_vertices) < face_vertex_count: + return recursive_clean_degenerate_face_edges(cleaned_open_face_vertices) + else: + return cleaned_open_face_vertices + + open_face_vertices = face_vertices[:-1] + cleaned_face_vertices = recursive_clean_degenerate_face_edges(open_face_vertices) + # Close cleaned_face_vertices + cleaned_face_vertices.append(cleaned_face_vertices[0]) + return cleaned_face_vertices + + +def merge_vertices(main_face_vertices, extra_face_vertices, common_edge): + sorted_common_edge = tuple(sorted(common_edge)) + open_face_vertices_pair = (main_face_vertices[:-1], extra_face_vertices[:-1]) + face_index = 0 # 0: current_face == main_face, 1: current_face == extra_face + vertex_index = 0 + start_vertex = vertex = open_face_vertices_pair[face_index][vertex_index] + merged_face_vertices = [start_vertex] + faces_merged = False + while not faces_merged: + # Get next vertex + next_vertex_index = (vertex_index + 1) % len(open_face_vertices_pair[face_index]) + next_vertex = open_face_vertices_pair[face_index][next_vertex_index] + edge = (vertex, next_vertex) + sorted_edge = tuple(sorted(edge)) + if sorted_edge == sorted_common_edge: + # Switch current face + face_index = 1 - face_index + # Find vertex_index in new current face + reverse_edge = (edge[1], edge[0]) # Because we are now on the other face + edge_index = find_edge_in_face(reverse_edge, open_face_vertices_pair[face_index]) + vertex_index = edge_index + 1 # Index of the second vertex of edge + # vertex_index = open_face_vertices_pair[face_index].index(vertex) + vertex_index = (vertex_index + 1) % len(open_face_vertices_pair[face_index]) + vertex = open_face_vertices_pair[face_index][vertex_index] + merged_face_vertices.append(vertex) + faces_merged = vertex == start_vertex # This also makes the merged_face closed + # Remove degenerate face edges (edges where the face if on both sides of it) + cleaned_merged_face_vertices = clean_degenerate_face_edges(merged_face_vertices) + return cleaned_merged_face_vertices + + +def polygon_close(polygon): + return np.concatenate((polygon, polygon[0:1, :]), axis=0) + + +def polygons_close(polygons): + return [polygon_close(polygon) for polygon in polygons] + + +# def init_cross_field(polygons, shape): +# """ +# Cross field: {v_1, v_2, -v_1, -v_2} encoded as {v_1, v_2}. +# This is not invariant to symmetries. +# +# :param polygons: +# :param shape: +# :return: cross_field_array (shape[0], shape[1], 2), dtype=np.int8 +# """ +# def draw_edge(edge, v1): +# rr, cc = skimage.draw.line(edge[0][0], edge[0][1], edge[1][0], edge[1][1]) +# mask = (0 <= rr) & (rr < shape[0]) & (0 <= cc) & (cc < shape[1]) +# cross_field_array[rr[mask], cc[mask], 0] = v1.real +# cross_field_array[rr[mask], cc[mask], 1] = v1.imag +# +# polygons = polygons_remove_holes(polygons) +# polygons = polygons_close(polygons) +# +# cross_field_array = np.zeros(shape + (4,), dtype=np.float) +# +# for polygon in polygons: +# # --- edges: +# edge_vect_array = np.diff(polygon, axis=0) +# norm = np.linalg.norm(edge_vect_array, axis=1, keepdims=True) +# # if not np.all(0 < norm): +# # print("WARNING: one of the norms is zero, which cannot be used to divide") +# # print("polygon that raised this warning:") +# # print(polygon) +# # exit() +# edge_dir_array = edge_vect_array / norm +# edge_v1_array = edge_dir_array.view(np.complex)[..., 0] +# # edge_v2_array is zero +# +# # --- vertices: +# vertex_v1_array = edge_v1_array +# vertex_v2_array = - np.roll(edge_v1_array, 1, axis=0) +# +# # --- Draw values +# polygon = polygon.astype(np.int) +# +# for i in range(polygon.shape[0] - 1): +# edge = (polygon[i], polygon[i+1]) +# v1 = edge_v1_array[i] +# draw_edge(edge, v1) +# +# vertex_array = polygon[:-1] +# mask = (0 <= vertex_array[:, 0]) & (vertex_array[:, 0] < shape[0])\ +# & (0 <= vertex_array[:, 1]) & (vertex_array[:, 1] < shape[1]) +# cross_field_array[vertex_array[mask, 0], vertex_array[mask, 1], 0] = vertex_v1_array[mask].real +# cross_field_array[vertex_array[mask, 0], vertex_array[mask, 1], 1] = vertex_v1_array[mask].imag +# cross_field_array[vertex_array[mask, 0], vertex_array[mask, 1], 2] = vertex_v2_array[mask].real +# cross_field_array[vertex_array[mask, 0], vertex_array[mask, 1], 3] = vertex_v2_array[mask].imag +# +# # --- Encode cross-field with integer complex to save memory because abs(cross_field_array) <= 1 anyway. +# cross_field_array = (127*cross_field_array).astype(np.int8) +# +# return cross_field_array + + +# def init_angle_field(polygons, shape): +# """ +# Angle field {\theta_1} the tangent vector's angle for every pixel, specified on the polygon edges. +# Angle between 0 and pi. +# Also indices of those angle values. +# This is not invariant to symmetries. +# +# :param polygons: +# :param shape: +# :return: (angles: np.array((num_edge_pixels, ), dtype=np.uint8), +# indices: np.array((num_edge_pixels, 2), dtype=np.int)) +# """ +# def draw_edge(edge, angle): +# rr, cc = skimage.draw.line(edge[0][0], edge[0][1], edge[1][0], edge[1][1]) +# edge_mask = (0 <= rr) & (rr < shape[0]) & (0 <= cc) & (cc < shape[1]) +# angle_field_array[rr[edge_mask], cc[edge_mask]] = angle +# mask[rr[edge_mask], cc[edge_mask]] = True +# +# polygons = polygons_remove_holes(polygons) +# polygons = polygons_close(polygons) +# +# angle_field_array = np.zeros(shape, dtype=np.float) +# mask = np.zeros(shape, dtype=np.bool) +# +# for polygon in polygons: +# # --- edges: +# edge_vect_array = np.diff(polygon, axis=0) +# edge_angle_array = np.angle(edge_vect_array[:, 0] + 1j * edge_vect_array[:, 1]) +# neg_indices = np.where(edge_angle_array < 0) +# edge_angle_array[neg_indices] += np.pi +# +# # --- Draw values +# polygon = polygon.astype(np.int) +# +# for i in range(polygon.shape[0] - 1): +# edge = (polygon[i], polygon[i+1]) +# angle = edge_angle_array[i] +# draw_edge(edge, angle) +# +# # --- Encode angle-field with positive integers to save memory because angle is between 0 and pi. +# indices = np.stack(np.where(mask), axis=-1) +# angles = angle_field_array[indices[:, 0], indices[:, 1]] +# angles = (255*angles/np.pi).round().astype(np.uint8) +# +# return angles, indices + + +def init_angle_field(polygons, shape, line_width=1): + """ + Angle field {\theta_1} the tangent vector's angle for every pixel, specified on the polygon edges. + Angle between 0 and pi. + This is not invariant to symmetries. + + :param polygons: + :param shape: + :return: (angles: np.array((num_edge_pixels, ), dtype=np.uint8), + mask: np.array((num_edge_pixels, 2), dtype=np.int)) + """ + assert type(polygons) == list, "polygons should be a list" + + polygons = polygons_remove_holes(polygons) + polygons = polygons_close(polygons) + + im = Image.new("L", (shape[1], shape[0])) + im_px_access = im.load() + draw = ImageDraw.Draw(im) + + for polygon in polygons: + # --- edges: + edge_vect_array = np.diff(polygon, axis=0) + edge_angle_array = np.angle(edge_vect_array[:, 0] + 1j * edge_vect_array[:, 1]) + neg_indices = np.where(edge_angle_array < 0) + edge_angle_array[neg_indices] += np.pi + + for i in range(polygon.shape[0] - 1): + edge = (polygon[i], polygon[i + 1]) + angle = edge_angle_array[i] + uint8_angle = int((255 * angle / np.pi).round()) + line = [(edge[0][1], edge[0][0]), (edge[1][1], edge[1][0])] + draw.line(line, fill=uint8_angle, width=line_width) + _draw_circle(draw, line[0], radius=line_width / 2, fill=uint8_angle) + _draw_circle(draw, line[1], radius=line_width / 2, fill=uint8_angle) + + # Convert image to numpy array + array = np.array(im) + return array + + +def plot_geometries(axis, geometries, linewidths=1, markersize=3): + if len(geometries): + patches = [] + for i, geometry in enumerate(geometries): + if geometry.geom_type == "Polygon": + polygon = shapely.geometry.Polygon(geometry) + if not polygon.is_empty: + patch = PolygonPatch(polygon) + patches.append(patch) + axis.plot(*polygon.exterior.xy, marker="o", markersize=markersize) + for interior in polygon.interiors: + axis.plot(*interior.xy, marker="o", markersize=markersize) + elif geometry.geom_type == "LineString" or geometry.geom_type == "LinearRing": + axis.plot(*geometry.xy, marker="o", markersize=markersize) + else: + raise NotImplementedError(f"Geom type {geometry.geom_type} not recognized.") + random.seed(1) + colors = random.choices([ + [0, 0, 1, 1], + [0, 1, 0, 1], + [1, 0, 0, 1], + [1, 1, 0, 1], + [1, 0, 1, 1], + [0, 1, 1, 1], + [0.5, 1, 0, 1], + [1, 0.5, 0, 1], + [0.5, 0, 1, 1], + [1, 0, 0.5, 1], + [0, 0.5, 1, 1], + [0, 1, 0.5, 1], + ], k=len(patches)) + edgecolors = np.array(colors) + facecolors = edgecolors.copy() + p = PatchCollection(patches, facecolors=facecolors, edgecolors=edgecolors, linewidths=linewidths) + axis.add_collection(p) + + +def sample_geometry(geom, density): + """ + Sample edges of geom with a homogeneous density. + + @param geom: + @param density: + @return: + """ + if isinstance(geom, shapely.geometry.GeometryCollection): + # tic = time.time() + + sampled_geom = shapely.geometry.GeometryCollection([sample_geometry(g, density) for g in geom]) + + # toc = time.time() + # print(f"sample_geometry: {toc - tic}s") + elif isinstance(geom, shapely.geometry.Polygon): + sampled_exterior = sample_geometry(geom.exterior, density) + sampled_interiors = [sample_geometry(interior, density) for interior in geom.interiors] + sampled_geom = shapely.geometry.Polygon(sampled_exterior, sampled_interiors) + elif isinstance(geom, shapely.geometry.LineString): + sampled_x = [] + sampled_y = [] + coords = np.array(geom.coords[:]) + lengths = np.linalg.norm(coords[:-1] - coords[1:], axis=1) + for i in range(len(lengths)): + start = geom.coords[i] + end = geom.coords[i + 1] + length = lengths[i] + num = max(1, int(round(length / density))) + 1 + x_seq = np.linspace(start[0], end[0], num) + y_seq = np.linspace(start[1], end[1], num) + if 0 < i: + x_seq = x_seq[1:] + y_seq = y_seq[1:] + sampled_x.append(x_seq) + sampled_y.append(y_seq) + sampled_x = np.concatenate(sampled_x) + sampled_y = np.concatenate(sampled_y) + sampled_coords = zip(sampled_x, sampled_y) + sampled_geom = shapely.geometry.LineString(sampled_coords) + else: + raise TypeError(f"geom of type {type(geom)} not supported!") + return sampled_geom + +# +# def sample_half_tangent_endpoints(geom, length=0.1): +# """ +# Add 2 vertices per edge, very close to the edge's endpoints. They represent both half-tangent endpoints +# @param geom: +# @param length: +# @return: +# """ +# if isinstance(geom, shapely.geometry.GeometryCollection): +# sampled_geom = shapely.geometry.GeometryCollection([sample_half_tangent_endpoints(g, length) for g in geom]) +# elif isinstance(geom, shapely.geometry.Polygon): +# sampled_exterior = sample_half_tangent_endpoints(geom.exterior, length) +# sampled_interiors = [sample_half_tangent_endpoints(interior, length) for interior in geom.interiors] +# sampled_geom = shapely.geometry.Polygon(sampled_exterior, sampled_interiors) +# elif isinstance(geom, shapely.geometry.LineString): +# coords = np.array(geom.coords[:]) +# edge_vecs = coords[1:] - coords[:-1] +# norms = np.linalg.norm(edge_vecs, axis=1) +# edge_dirs = edge_vecs / norms[:, None] +# sampled_coords = [coords[0]] # Init with first vertex +# for edge_i in range(edge_dirs.shape[0]): +# first_half_tangent_endpoint = coords[edge_i] + length * edge_dirs[edge_i] +# sampled_coords.append(first_half_tangent_endpoint) +# second_half_tangent_endpoint = coords[edge_i + 1] - length * edge_dirs[edge_i] +# sampled_coords.append(second_half_tangent_endpoint) +# sampled_coords.append(coords[edge_i + 1]) # Next vertex +# sampled_geom = shapely.geometry.LineString(sampled_coords) +# else: +# raise TypeError(f"geom of type {type(geom)} not supported!") +# return sampled_geom + + +def point_project_onto_geometry(coord, target): + point = shapely.geometry.Point(coord) + _, projected_point = shapely.ops.nearest_points(point, target) + # dist = point.distance(projected_point) + return projected_point.coords[0] + + +def project_onto_geometry(geom, target, pool: Pool=None): + """ + Projects all points from line_string onto target. + @param geom: + @param target: + @param pool: + @return: + """ + if isinstance(geom, shapely.geometry.GeometryCollection): + # tic = time.time() + + if pool is None: + projected_geom = [project_onto_geometry(g, target, pool=pool) for g in geom] + else: + partial_project_onto_geometry = partial(project_onto_geometry, target=target) + projected_geom = pool.map(partial_project_onto_geometry, geom) + projected_geom = shapely.geometry.GeometryCollection(projected_geom) + + # toc = time.time() + # print(f"project_onto_geometry: {toc - tic}s") + elif isinstance(geom, shapely.geometry.Polygon): + projected_exterior = project_onto_geometry(geom.exterior, target) + projected_interiors = [project_onto_geometry(interior, target) for interior in geom.interiors] + try: + projected_geom = shapely.geometry.Polygon(projected_exterior, projected_interiors) + except shapely.errors.TopologicalError as e: + import matplotlib.pyplot as plt + fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(8, 4), sharex=True, sharey=True) + ax = axes.ravel() + plot_geometries(ax[0], [geom]) + plot_geometries(ax[1], target) + plot_geometries(ax[2], [projected_exterior, *projected_interiors]) + fig.tight_layout() + plt.show() + raise e + elif isinstance(geom, shapely.geometry.LineString): + projected_coords = [point_project_onto_geometry(coord, target) for coord in geom.coords] + projected_geom = shapely.geometry.LineString(projected_coords) + else: + raise TypeError(f"geom of type {type(geom)} not supported!") + return projected_geom + +# +# def compute_edge_measures(geom1, geom2, max_stretch, metric_name="cosine"): +# """ +# +# @param geom1: +# @param geom2: +# @param max_stretch: Edges of geom2 than are longer than those of geom1 with a factor greater than max_stretch are ignored +# @param metric_name: +# @return: +# """ +# assert type(geom1) == type(geom2), f"geom1 and geom2 must be of the same type, not {type(geom1)} and {type(geom2)}" +# if isinstance(geom1, shapely.geometry.GeometryCollection): +# # tic = time.time() +# +# edge_measures_edge_dists_list = [compute_edge_measures(_geom1, _geom2, max_stretch, metric_name=metric_name) for _geom1, _geom2 in zip(geom1, geom2)] +# if len(edge_measures_edge_dists_list): +# edge_measures_list, edge_dists_list = zip(*edge_measures_edge_dists_list) +# edge_measures = np.concatenate(edge_measures_list) +# edge_dists = np.concatenate(edge_dists_list) +# else: +# edge_measures = np.array([]) +# edge_dists = np.array([]) +# +# # toc = time.time() +# # print(f"compute_edge_distance: {toc - tic}s") +# # elif isinstance(geom1, shapely.geometry.Polygon): +# # distances_exterior = compute_edge_distance(geom1.exterior, geom2.exterior, tolerance, max_stretch, dist=dist) +# # distances_interiors = [compute_edge_distance(interior1, interior2, tolerance, max_stretch, dist=dist) for interior1, interior2 in zip(geom1.interiors, geom2.interiors)] +# # distances = [distances_exterior, *distances_interiors] +# # distances = np.concatenate(distances) +# elif isinstance(geom1, shapely.geometry.LineString): +# assert len(geom1.coords) == len(geom2.coords), "geom1 and geom2 must have the same length" +# points1 = np.array(geom1.coords) +# points2 = np.array(geom2.coords) +# # Mark points that are farther away than tolerance between points1 and points2 to remove then from further computation +# point_dists = np.linalg.norm(points1 - points2, axis=1) +# if metric_name == "cosine": +# edges1 = points1[1:] - points1[:-1] +# edges2 = points2[1:] - points2[:-1] +# edge_dists = (point_dists[1:] + point_dists[:-1]) / 2 +# # Remove edges with a norm of zero +# norm1 = np.linalg.norm(edges1, axis=1) +# norm2 = np.linalg.norm(edges2, axis=1) +# norm_valid_mask = 0 < norm1 * norm2 +# edges1 = edges1[norm_valid_mask] +# edges2 = edges2[norm_valid_mask] +# norm1 = norm1[norm_valid_mask] +# norm2 = norm2[norm_valid_mask] +# edge_dists = edge_dists[norm_valid_mask] +# # Remove edges that have been stretched more than max_stretch +# stretch = norm2 / norm1 +# stretch_valid_mask = np.logical_and(1 / max_stretch < stretch, stretch < max_stretch) +# edges1 = edges1[stretch_valid_mask] +# edges2 = edges2[stretch_valid_mask] +# norm1 = norm1[stretch_valid_mask] +# norm2 = norm2[stretch_valid_mask] +# edge_dists = edge_dists[stretch_valid_mask] +# # Compute +# edge_measures = np.sum(np.multiply(edges1, edges2), axis=1) / (norm1 * norm2) +# else: +# raise NotImplemented(f"Metric '{metric_name}' is not implemented") +# else: +# raise TypeError(f"geom of type {type(geom1)} not supported!") +# return edge_measures, edge_dists + + +def compute_contour_measure(pred_polygon, gt_contours, sampling_spacing, max_stretch, metric_name="cosine"): + pred_contours = shapely.geometry.GeometryCollection([pred_polygon.exterior, *pred_polygon.interiors]) + sampled_pred_contours = sample_geometry(pred_contours, sampling_spacing) + # Project sampled contour points to ground truth contours + projected_pred_contours = project_onto_geometry(sampled_pred_contours, gt_contours) + contour_measures = [] + for contour, proj_contour in zip(sampled_pred_contours, projected_pred_contours): + coords = np.array(contour.coords[:]) + proj_coords = np.array(proj_contour.coords[:]) + edges = coords[1:] - coords[:-1] + proj_edges = proj_coords[1:] - proj_coords[:-1] + # Remove edges with a norm of zero + edge_norms = np.linalg.norm(edges, axis=1) + proj_edge_norms = np.linalg.norm(proj_edges, axis=1) + norm_valid_mask = 0 < edge_norms * proj_edge_norms + edges = edges[norm_valid_mask] + proj_edges = proj_edges[norm_valid_mask] + edge_norms = edge_norms[norm_valid_mask] + proj_edge_norms = proj_edge_norms[norm_valid_mask] + # Remove edge that have stretched more than max_stretch (invalid projection) + stretch = edge_norms / proj_edge_norms + stretch_valid_mask = np.logical_and(1 / max_stretch < stretch, stretch < max_stretch) + edges = edges[stretch_valid_mask] + if edges.shape[0] == 0: + # Invalid projection for the whole contour, skip it + continue + proj_edges = proj_edges[stretch_valid_mask] + edge_norms = edge_norms[stretch_valid_mask] + proj_edge_norms = proj_edge_norms[stretch_valid_mask] + scalar_products = np.abs(np.sum(np.multiply(edges, proj_edges), axis=1) / (edge_norms * proj_edge_norms)) + try: + contour_measures.append(scalar_products.min()) + except ValueError: + import matplotlib.pyplot as plt + fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(8, 4), sharex=True, sharey=True) + ax = axes.ravel() + plot_geometries(ax[0], [contour]) + plot_geometries(ax[1], [proj_contour]) + plot_geometries(ax[2], gt_contours) + fig.tight_layout() + plt.show() + if len(contour_measures): + min_scalar_product = min(contour_measures) + measure = np.arccos(min_scalar_product) + return measure + else: + return None + + +def compute_polygon_contour_measures(pred_polygons: list, gt_polygons: list, sampling_spacing: float, min_precision: float, max_stretch: float, metric_name: str="cosine", progressbar=False): + """ + pred_polygons are sampled with sampling_spacing before projecting those sampled points to gt_polygons. + Then the + + @param pred_polygons: + @param gt_polygons: + @param sampling_spacing: + @param min_precision: Polygons in pred_polygons must have a precision with gt_polygons above min_precision to be included in further computations + @param max_stretch: Exclude edges that have been stretched by the projection more than max_stretch from further computation + @param metric_name: Metric type, can be "cosine" or ... + @return: + """ + assert isinstance(pred_polygons, list), "pred_polygons should be a list" + assert isinstance(gt_polygons, list), "gt_polygons should be a list" + if len(pred_polygons) == 0 or len(gt_polygons) == 0: + return np.array([]), [], [] + assert isinstance(pred_polygons[0], shapely.geometry.Polygon), \ + f"Items of pred_polygons should be of type shapely.geometry.Polygon, not {type(pred_polygons[0])}" + assert isinstance(gt_polygons[0], shapely.geometry.Polygon), \ + f"Items of gt_polygons should be of type shapely.geometry.Polygon, not {type(gt_polygons[0])}" + gt_polygons = shapely.geometry.collection.GeometryCollection(gt_polygons) + pred_polygons = shapely.geometry.collection.GeometryCollection(pred_polygons) + # Filter pred_polygons to have at least a precision with gt_polygons of min_precision + filtered_pred_polygons = [pred_polygon for pred_polygon in pred_polygons if min_precision < pred_polygon.intersection(gt_polygons).area / pred_polygon.area] + # Extract contours of gt polygons + gt_contours = shapely.geometry.collection.GeometryCollection([contour for polygon in gt_polygons for contour in [polygon.exterior, *polygon.interiors]]) + # Measure metric for each pred polygon + if progressbar: + process_id = int(multiprocess.current_process().name[-1]) + iterator = tqdm(filtered_pred_polygons, desc="Contour measure", leave=False, position=process_id) + else: + iterator = filtered_pred_polygons + half_tangent_max_angles = [compute_contour_measure(pred_polygon, gt_contours, sampling_spacing=sampling_spacing, max_stretch=max_stretch, metric_name=metric_name) + for pred_polygon in iterator] + return half_tangent_max_angles + + +def fix_polygons(polygons, buffer=0.0): + polygons_geom = shapely.ops.unary_union(polygons) # Fix overlapping polygons + polygons_geom = polygons_geom.buffer(buffer) # Fix self-intersecting polygons and other things + fixed_polygons = [] + if polygons_geom.geom_type == "MultiPolygon": + for poly in polygons_geom: + fixed_polygons.append(poly) + elif polygons_geom.geom_type == "Polygon": + fixed_polygons.append(polygons_geom) + else: + raise TypeError(f"Geom type {polygons_geom.geom_type} not recognized.") + return fixed_polygons + + +POINTS = [] + +# +# def compute_half_tangent_measure(pred_polygon, gt_contours, step=0.1, metric_name="angle"): +# """ +# For each vertex in pred_polygon, find the closest gt contour and the closest point on that contour. From that point, compute both half-tangents. +# measure angle difference between half-tangents of pred and corresponding gt points. +# @param pred_polygon: +# @param gt_contours: +# @param metric_name: +# @return: +# """ +# assert isinstance(pred_polygon, shapely.geometry.Polygon), "pred_polygon should be a shapely Polygon" +# pred_contours = [pred_polygon.exterior, *pred_polygon.interiors] +# tangent_measures_list = [] +# for pred_contour in pred_contours: +# pos_array = np.array(pred_contour.coords[:]) +# pred_tangents = pos_array[1:] - pos_array[:-1] +# gt_tangent_1_list = [] +# gt_tangent_2_list = [] +# for i, pos in enumerate(pos_array[:-1]): +# pred_point = shapely.geometry.Point(pos) +# dist_to_gt = np.inf +# closest_gt_contour = None +# for gt_contour in gt_contours: +# d = pred_point.distance(gt_contour) +# if d < dist_to_gt: +# dist_to_gt = d +# closest_gt_contour = gt_contour +# gt_point_t = closest_gt_contour.project(pred_point) # References the projection of pred_point onto closest_gt_contour with a 1d referencing coordinate t +# # --- Compute tangents of projected point on gt: +# gt_point_tangent_1 = closest_gt_contour.interpolate(gt_point_t - step) +# POINTS.append(gt_point_tangent_1) +# gt_point = closest_gt_contour.interpolate(gt_point_t) +# POINTS.append(gt_point) +# gt_point_tangent_2 = closest_gt_contour.interpolate(gt_point_t + step) +# POINTS.append(gt_point_tangent_2) +# gt_pos_tangent_1 = np.array(gt_point_tangent_1.coords[0]) +# gt_pos_tangent_2 = np.array(gt_point_tangent_2.coords[0]) +# gt_pos = np.array(gt_point.coords[0]) +# gt_tangent_1 = gt_pos_tangent_1 - gt_pos +# gt_tangent_2 = gt_pos_tangent_2 - gt_pos +# gt_tangent_1_list.append(gt_tangent_1) +# gt_tangent_2_list.append(gt_tangent_2) +# gt_tangents_1 = np.stack(gt_tangent_1_list, axis=0) +# gt_tangents_2 = np.stack(gt_tangent_2_list, axis=0) +# # Measure dist between pred_tangents and gt_tangents +# pred_norms = np.linalg.norm(pred_tangents, axis=1) +# tangent_1_measures = np.abs(np.sum(np.multiply(np.roll(pred_tangents, 1, axis=0), gt_tangents_1), axis=1) / (np.roll(pred_norms, 1, axis=0) * step)) +# tangent_2_measures = np.abs(np.sum(np.multiply(pred_tangents, gt_tangents_2), axis=1) / (pred_norms * step)) +# print(tangent_1_measures) +# print(tangent_2_measures) +# tangent_measures_list.append(tangent_1_measures) +# tangent_measures_list.append(tangent_2_measures) +# tangent_measures = np.concatenate(tangent_measures_list) +# min_scalar_product = np.min(tangent_measures) +# max_angle = np.arccos(min_scalar_product) +# return max_angle + +# +# def compute_vertex_measures(pred_polygons: list, gt_polygons: list, min_precision: float, metric_name: str="angle", pool: Pool=None): +# """ +# Computes measure for each pred_polygon +# @param pred_polygons: +# @param gt_polygons: +# @param min_precision: +# @param metric_name: +# @param pool: +# @return: +# """ +# assert isinstance(pred_polygons, list), "pred_polygons should be a list" +# assert isinstance(gt_polygons, list), "gt_polygons should be a list" +# if len(pred_polygons) == 0 or len(gt_polygons) == 0: +# return np.array([]), [], [] +# assert isinstance(pred_polygons[0], shapely.geometry.Polygon), \ +# f"Items of pred_polygons should be of type shapely.geometry.Polygon, not {type(pred_polygons[0])}" +# assert isinstance(gt_polygons[0], shapely.geometry.Polygon), \ +# f"Items of gt_polygons should be of type shapely.geometry.Polygon, not {type(gt_polygons[0])}" +# gt_polygons = shapely.geometry.collection.GeometryCollection(gt_polygons) +# pred_polygons = shapely.geometry.collection.GeometryCollection(pred_polygons) +# # Filter pred_polygons to have at least a precision with gt_polygons of min_precision +# filtered_pred_polygons = [pred_polygon for pred_polygon in pred_polygons if min_precision < pred_polygon.intersection(gt_polygons).area / pred_polygon.area] +# # Extract contours of gt polygons +# gt_contours = shapely.geometry.collection.GeometryCollection([contour for polygon in gt_polygons for contour in [polygon.exterior, *polygon.interiors]]) +# # Measure metric for each pre polygon +# half_tangent_max_angles = [compute_half_tangent_measure(pred_polygon, gt_contours, metric_name=metric_name) +# for pred_polygon in filtered_pred_polygons] +# return half_tangent_max_angles + + +def main(): + import matplotlib.pyplot as plt + + gt_polygon_1 = shapely.geometry.Polygon( + [ + [0, 0], + [10, 0], + [10, 10], + [0, 10] + ], + # [[ + # [0.1, 0.1], + # [0.9, 0.1], + # [0.9, 0.9], + # [0.1, 0.9] + # ]] + ) + # gt_polygon_2 = shapely.geometry.Polygon([ + # [2, 2], + # [5, 0], + # [5, 6], + # [0, 4] + # ]) + pred_polygon_1 = shapely.geometry.Polygon( + [ + [0.1, 0.1], + [10.1, 0], + [9.9, 9], + [9, 10.1], + [0.1, 10] + ], + # [ + # [0, 0], + # [10, 0], + # [10, 9], + # [10, 10], + # [9, 10], + # [0, 10] + # ], + ) + pred_polygons = [pred_polygon_1] + gt_polygons = [gt_polygon_1] + + max_angle_diffs = compute_polygon_contour_measures(pred_polygons, gt_polygons, sampling_spacing=0.1, min_precision=0.5, max_stretch=2) + # half_tangent_max_angles = compute_vertex_measures(pred_polygons, gt_polygons, min_precision=0.5) + + # print(cosine_similarities.mean()) + print(max_angle_diffs[0] * 180 / np.pi) + + fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(8, 4), sharex=True, sharey=True) + ax = axes.ravel() + + plot_geometries(ax[0], gt_polygons) + plot_geometries(ax[1], pred_polygons) + # plot_geometries(ax[2], projected_pred_contours) + for point in POINTS: + ax[2].plot(*point.xy, marker="o", markersize=1) + + fig.tight_layout() + plt.show() + + +if __name__ == "__main__": + main() diff --git a/lydorn_utils/print_utils.py b/lydorn_utils/print_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e27493700a5abc758cbcc0370d7093545da7d9ab --- /dev/null +++ b/lydorn_utils/print_utils.py @@ -0,0 +1,63 @@ +class bcolors: + HEADER = '\033[95m' + OKBLUE = '\033[94m' + OKGREEN = '\033[92m' + WARNING = '\033[93m' + FAIL = '\033[91m' + DEBUG = '\033[31;40m' + ENDC = '\033[0m' + BOLD = '\033[1m' + UNDERLINE = '\033[4m' + + +def print_info(*args): + print(bcolors.OKBLUE + " ".join(map(str, args)) + bcolors.ENDC) + + +def print_success(*args): + print(bcolors.OKGREEN + " ".join(map(str, args)) + bcolors.ENDC) + + +def print_failure(*args): + print(bcolors.FAIL + " ".join(map(str, args)) + bcolors.ENDC) + + +def print_error(*args): + print_failure(*args) + + +def print_warning(*args): + print(bcolors.WARNING + " ".join(map(str, args)) + bcolors.ENDC) + + +def print_debug(*args): + print(bcolors.DEBUG + " ".join(map(str, args)) + bcolors.ENDC) + + +def print_format_table(): + """ + prints table of formatted text format options + """ + for style in range(8): + for fg in range(30, 38): + s1 = '' + for bg in range(40, 48): + format = ';'.join([str(style), str(fg), str(bg)]) + s1 += '\x1b[%sm %s \x1b[0m' % (format, format) + print(s1) + print('\n') + + +def main(): + print_format_table() + + print_info("Info") + print_success("Success") + print_failure("Failure") + print_error("ERROR") + print_warning("WARNING") + print_debug("Debug") + + +if __name__ == '__main__': + main() diff --git a/lydorn_utils/python_utils.py b/lydorn_utils/python_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..06373abf6bdc4ffa44f577031ff16bb70374aedc --- /dev/null +++ b/lydorn_utils/python_utils.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +import os +import errno +import json +from jsmin import jsmin +from lydorn_utils import print_utils + + +def module_exists(module_name): + try: + __import__(module_name) + except ImportError: + return False + else: + return True + + +def choose_first_existing_path(path_list, return_tried_paths=False): + paths_tried = [] + for path in path_list: + path = os.path.expanduser(path) + paths_tried.append(path) + if os.path.exists(path): + if return_tried_paths: + return path, paths_tried + else: + return path + + if return_tried_paths: + return None, paths_tried + else: + return None + + +def get_display_availability(): + return "DISPLAY" in os.environ + + +def get_filepaths(dir_path, endswith_str="", startswith_str="", not_endswith_str=None, not_startswith_str=None): + if os.path.isdir(dir_path): + image_filepaths = [] + for path, dnames, fnames in os.walk(dir_path): + fnames = sorted(fnames) + image_filepaths.extend([os.path.join(path, x) for x in fnames + if x.endswith(endswith_str) + and x.startswith(startswith_str) + and (not_endswith_str is None or not x.endswith(not_endswith_str)) + and (not_startswith_str is None or not x.startswith(not_startswith_str))]) + return image_filepaths + else: + raise NotADirectoryError(errno.ENOENT, os.strerror(errno.ENOENT), dir_path) + + +def get_dir_list_filepaths(dir_path_list, endswith_str="", startswith_str=""): + image_filepaths = [] + for dir_path in dir_path_list: + image_filepaths.extend(get_filepaths(dir_path, endswith_str=endswith_str, startswith_str=startswith_str)) + return image_filepaths + + +def save_json(filepath, data): + dirpath = os.path.dirname(filepath) + os.makedirs(dirpath, exist_ok=True) + with open(filepath, 'w') as outfile: + json.dump(data, outfile) + return True + + +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 wipe_dir(dirpath): + filepaths = get_filepaths(dirpath) + for filepath in filepaths: + os.remove(filepath) + + +def split_list_into_chunks(l, n, pad=False): + """Yield successive n-sized chunks from l.""" + for i in range(0, len(l), n): + if pad: + chunk = l[i:i + n] + if len(chunk) < n: + chunk.extend([chunk[-1]]*(n - len(chunk))) + yield chunk + else: + yield l[i:i + n] + + +def params_to_str(params): + def to_str(value): + if type(value) == float and value == int(value): + return str(int(value)) + return str(value) + + return "_".join(["{}_{}".format(key, to_str(params[key])) for key in sorted(params.keys())]) + + +def main(): + l = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + batches = split_list_into_chunks(l, 4, pad=True) + for batch in batches: + print(batch) + + +if __name__ == '__main__': + main() diff --git a/lydorn_utils/run_utils.py b/lydorn_utils/run_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5e932ac49c25d983f3e6915500ce0899fd9c6b6b --- /dev/null +++ b/lydorn_utils/run_utils.py @@ -0,0 +1,766 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +import os +import time +import datetime +from jsmin import jsmin +import json + +from . import print_utils +from . import python_utils + +# Stolen from Docker: +NAME_SET = set([ + # Muhammad ibn Jābir al-Ḥarrānī al-Battānī was a founding father of astronomy. https://en.wikipedia.org/wiki/Mu%E1%B8%A5ammad_ibn_J%C4%81bir_al-%E1%B8%A4arr%C4%81n%C4%AB_al-Batt%C4%81n%C4%AB + "albattani", + + # Frances E. Allen, became the first female IBM Fellow in 1989. In 2006, she became the first female recipient of the ACM's Turing Award. https://en.wikipedia.org/wiki/Frances_E._Allen + "allen", + + # June Almeida - Scottish virologist who took the first pictures of the rubella virus - https://en.wikipedia.org/wiki/June_Almeida + "almeida", + + # Maria Gaetana Agnesi - Italian mathematician, philosopher, theologian and humanitarian. She was the first woman to write a mathematics handbook and the first woman appointed as a Mathematics Professor at a University. https://en.wikipedia.org/wiki/Maria_Gaetana_Agnesi + "agnesi", + + # Archimedes was a physicist, engineer and mathematician who invented too many things to list them here. https://en.wikipedia.org/wiki/Archimedes + "archimedes", + + # Maria Ardinghelli - Italian translator, mathematician and physicist - https://en.wikipedia.org/wiki/Maria_Ardinghelli + "ardinghelli", + + # Aryabhata - Ancient Indian mathematician-astronomer during 476-550 CE https://en.wikipedia.org/wiki/Aryabhata + "aryabhata", + + # Wanda Austin - Wanda Austin is the President and CEO of The Aerospace Corporation, a leading architect for the US security space programs. https://en.wikipedia.org/wiki/Wanda_Austin + "austin", + + # Charles Babbage invented the concept of a programmable computer. https://en.wikipedia.org/wiki/Charles_Babbage. + "babbage", + + # Stefan Banach - Polish mathematician, was one of the founders of modern functional analysis. https://en.wikipedia.org/wiki/Stefan_Banach + "banach", + + # John Bardeen co-invented the transistor - https://en.wikipedia.org/wiki/John_Bardeen + "bardeen", + + # Jean Bartik, born Betty Jean Jennings, was one of the original programmers for the ENIAC computer. https://en.wikipedia.org/wiki/Jean_Bartik + "bartik", + + # Laura Bassi, the world's first female professor https://en.wikipedia.org/wiki/Laura_Bassi + "bassi", + + # Hugh Beaver, British engineer, founder of the Guinness Book of World Records https://en.wikipedia.org/wiki/Hugh_Beaver + "beaver", + + # Alexander Graham Bell - an eminent Scottish-born scientist, inventor, engineer and innovator who is credited with inventing the first practical telephone - https://en.wikipedia.org/wiki/Alexander_Graham_Bell + "bell", + + # Karl Friedrich Benz - a German automobile engineer. Inventor of the first practical motorcar. https://en.wikipedia.org/wiki/Karl_Benz + "benz", + + # Homi J Bhabha - was an Indian nuclear physicist, founding director, and professor of physics at the Tata Institute of Fundamental Research. Colloquially known as "father of Indian nuclear programme"- https://en.wikipedia.org/wiki/Homi_J._Bhabha + "bhabha", + + # Bhaskara II - Ancient Indian mathematician-astronomer whose work on calculus predates Newton and Leibniz by over half a millennium - https://en.wikipedia.org/wiki/Bh%C4%81skara_II#Calculus + "bhaskara", + + # Elizabeth Blackwell - American doctor and first American woman to receive a medical degree - https://en.wikipedia.org/wiki/Elizabeth_Blackwell + "blackwell", + + # Niels Bohr is the father of quantum theory. https://en.wikipedia.org/wiki/Niels_Bohr. + "bohr", + + # Kathleen Booth, she's credited with writing the first assembly language. https://en.wikipedia.org/wiki/Kathleen_Booth + "booth", + + # Anita Borg - Anita Borg was the founding director of the Institute for Women and Technology (IWT). https://en.wikipedia.org/wiki/Anita_Borg + "borg", + + # Satyendra Nath Bose - He provided the foundation for Bose–Einstein statistics and the theory of the Bose–Einstein condensate. - https://en.wikipedia.org/wiki/Satyendra_Nath_Bose + "bose", + + # Evelyn Boyd Granville - She was one of the first African-American woman to receive a Ph.D. in mathematics; she earned it in 1949 from Yale University. https://en.wikipedia.org/wiki/Evelyn_Boyd_Granville + "boyd", + + # Brahmagupta - Ancient Indian mathematician during 598-670 CE who gave rules to compute with zero - https://en.wikipedia.org/wiki/Brahmagupta#Zero + "brahmagupta", + + # Walter Houser Brattain co-invented the transistor - https://en.wikipedia.org/wiki/Walter_Houser_Brattain + "brattain", + + # Emmett Brown invented time travel. https://en.wikipedia.org/wiki/Emmett_Brown (thanks Brian Goff) + "brown", + + # Rachel Carson - American marine biologist and conservationist, her book Silent Spring and other writings are credited with advancing the global environmental movement. https://en.wikipedia.org/wiki/Rachel_Carson + "carson", + + # Subrahmanyan Chandrasekhar - Astrophysicist known for his mathematical theory on different stages and evolution in structures of the stars. He has won nobel prize for physics - https://en.wikipedia.org/wiki/Subrahmanyan_Chandrasekhar + "chandrasekhar", + + # Sergey Alexeyevich Chaplygin (Russian: Серге́й Алексе́евич Чаплы́гин; April 5, 1869 – October 8, 1942) was a Russian and Soviet physicist, mathematician, and mechanical engineer. He is known for mathematical formulas such as Chaplygin's equation and for a hypothetical substance in cosmology called Chaplygin gas, named after him. https://en.wikipedia.org/wiki/Sergey_Chaplygin + "chaplygin", + + # Asima Chatterjee was an indian organic chemist noted for her research on vinca alkaloids, development of drugs for treatment of epilepsy and malaria - https://en.wikipedia.org/wiki/Asima_Chatterjee + "chatterjee", + + # Pafnuty Chebyshev - Russian mathematitian. He is known fo his works on probability, statistics, mechanics, analytical geometry and number theory https://en.wikipedia.org/wiki/Pafnuty_Chebyshev + "chebyshev", + + # Claude Shannon - The father of information theory and founder of digital circuit design theory. (https://en.wikipedia.org/wiki/Claude_Shannon) + "shannon", + + # Joan Clarke - Bletchley Park code breaker during the Second World War who pioneered techniques that remained top secret for decades. Also an accomplished numismatist https://en.wikipedia.org/wiki/Joan_Clarke + "clarke", + + # Jane Colden - American botanist widely considered the first female American botanist - https://en.wikipedia.org/wiki/Jane_Colden + "colden", + + # Gerty Theresa Cori - American biochemist who became the third woman—and first American woman—to win a Nobel Prize in science, and the first woman to be awarded the Nobel Prize in Physiology or Medicine. Cori was born in Prague. https://en.wikipedia.org/wiki/Gerty_Cori + "cori", + + # Seymour Roger Cray was an American electrical engineer and supercomputer architect who designed a series of computers that were the fastest in the world for decades. https://en.wikipedia.org/wiki/Seymour_Cray + "cray", + + # This entry reflects a husband and wife team who worked together: + # Joan Curran was a Welsh scientist who developed radar and invented chaff, a radar countermeasure. https://en.wikipedia.org/wiki/Joan_Curran + # Samuel Curran was an Irish physicist who worked alongside his wife during WWII and invented the proximity fuse. https://en.wikipedia.org/wiki/Samuel_Curran + "curran", + + # Marie Curie discovered radioactivity. https://en.wikipedia.org/wiki/Marie_Curie. + "curie", + + # Charles Darwin established the principles of natural evolution. https://en.wikipedia.org/wiki/Charles_Darwin. + "darwin", + + # Leonardo Da Vinci invented too many things to list here. https://en.wikipedia.org/wiki/Leonardo_da_Vinci. + "davinci", + + # Edsger Wybe Dijkstra was a Dutch computer scientist and mathematical scientist. https://en.wikipedia.org/wiki/Edsger_W._Dijkstra. + "dijkstra", + + # Donna Dubinsky - played an integral role in the development of personal digital assistants (PDAs) serving as CEO of Palm, Inc. and co-founding Handspring. https://en.wikipedia.org/wiki/Donna_Dubinsky + "dubinsky", + + # Annie Easley - She was a leading member of the team which developed software for the Centaur rocket stage and one of the first African-Americans in her field. https://en.wikipedia.org/wiki/Annie_Easley + "easley", + + # Thomas Alva Edison, prolific inventor https://en.wikipedia.org/wiki/Thomas_Edison + "edison", + + # Albert Einstein invented the general theory of relativity. https://en.wikipedia.org/wiki/Albert_Einstein + "einstein", + + # Gertrude Elion - American biochemist, pharmacologist and the 1988 recipient of the Nobel Prize in Medicine - https://en.wikipedia.org/wiki/Gertrude_Elion + "elion", + + # Alexandra Asanovna Elbakyan (Russian: Алекса́ндра Аса́новна Элбакя́н) is a Kazakhstani graduate student, computer programmer, internet pirate in hiding, and the creator of the site Sci-Hub. Nature has listed her in 2016 in the top ten people that mattered in science, and Ars Technica has compared her to Aaron Swartz. - https://en.wikipedia.org/wiki/Alexandra_Elbakyan + "elbakyan", + + # Douglas Engelbart gave the mother of all demos: https://en.wikipedia.org/wiki/Douglas_Engelbart + "engelbart", + + # Euclid invented geometry. https://en.wikipedia.org/wiki/Euclid + "euclid", + + # Leonhard Euler invented large parts of modern mathematics. https://de.wikipedia.org/wiki/Leonhard_Euler + "euler", + + # Pierre de Fermat pioneered several aspects of modern mathematics. https://en.wikipedia.org/wiki/Pierre_de_Fermat + "fermat", + + # Enrico Fermi invented the first nuclear reactor. https://en.wikipedia.org/wiki/Enrico_Fermi. + "fermi", + + # Richard Feynman was a key contributor to quantum mechanics and particle physics. https://en.wikipedia.org/wiki/Richard_Feynman + "feynman", + + # Benjamin Franklin is famous for his experiments in electricity and the invention of the lightning rod. + "franklin", + + # Galileo was a founding father of modern astronomy, and faced politics and obscurantism to establish scientific truth. https://en.wikipedia.org/wiki/Galileo_Galilei + "galileo", + + # William Henry "Bill" Gates III is an American business magnate, philanthropist, investor, computer programmer, and inventor. https://en.wikipedia.org/wiki/Bill_Gates + "gates", + + # Adele Goldberg, was one of the designers and developers of the Smalltalk language. https://en.wikipedia.org/wiki/Adele_Goldberg_(computer_scientist) + "goldberg", + + # Adele Goldstine, born Adele Katz, wrote the complete technical description for the first electronic digital computer, ENIAC. https://en.wikipedia.org/wiki/Adele_Goldstine + "goldstine", + + # Shafi Goldwasser is a computer scientist known for creating theoretical foundations of modern cryptography. Winner of 2012 ACM Turing Award. https://en.wikipedia.org/wiki/Shafi_Goldwasser + "goldwasser", + + # James Golick, all around gangster. + "golick", + + # Jane Goodall - British primatologist, ethologist, and anthropologist who is considered to be the world's foremost expert on chimpanzees - https://en.wikipedia.org/wiki/Jane_Goodall + "goodall", + + # Lois Haibt - American computer scientist, part of the team at IBM that developed FORTRAN - https://en.wikipedia.org/wiki/Lois_Haibt + "haibt", + + # Margaret Hamilton - Director of the Software Engineering Division of the MIT Instrumentation Laboratory, which developed on-board flight software for the Apollo space program. https://en.wikipedia.org/wiki/Margaret_Hamilton_(scientist) + "hamilton", + + # Stephen Hawking pioneered the field of cosmology by combining general relativity and quantum mechanics. https://en.wikipedia.org/wiki/Stephen_Hawking + "hawking", + + # Werner Heisenberg was a founding father of quantum mechanics. https://en.wikipedia.org/wiki/Werner_Heisenberg + "heisenberg", + + # Grete Hermann was a German philosopher noted for her philosophical work on the foundations of quantum mechanics. https://en.wikipedia.org/wiki/Grete_Hermann + "hermann", + + # Jaroslav Heyrovský was the inventor of the polarographic method, father of the electroanalytical method, and recipient of the Nobel Prize in 1959. His main field of work was polarography. https://en.wikipedia.org/wiki/Jaroslav_Heyrovsk%C3%BD + "heyrovsky", + + # Dorothy Hodgkin was a British biochemist, credited with the development of protein crystallography. She was awarded the Nobel Prize in Chemistry in 1964. https://en.wikipedia.org/wiki/Dorothy_Hodgkin + "hodgkin", + + # Erna Schneider Hoover revolutionized modern communication by inventing a computerized telephone switching method. https://en.wikipedia.org/wiki/Erna_Schneider_Hoover + "hoover", + + # Grace Hopper developed the first compiler for a computer programming language and is credited with popularizing the term "debugging" for fixing computer glitches. https://en.wikipedia.org/wiki/Grace_Hopper + "hopper", + + # Frances Hugle, she was an American scientist, engineer, and inventor who contributed to the understanding of semiconductors, integrated circuitry, and the unique electrical principles of microscopic materials. https://en.wikipedia.org/wiki/Frances_Hugle + "hugle", + + # Hypatia - Greek Alexandrine Neoplatonist philosopher in Egypt who was one of the earliest mothers of mathematics - https://en.wikipedia.org/wiki/Hypatia + "hypatia", + + # Mary Jackson, American mathematician and aerospace engineer who earned the highest title within NASA's engineering department - https://en.wikipedia.org/wiki/Mary_Jackson_(engineer) + "jackson", + + # Yeong-Sil Jang was a Korean scientist and astronomer during the Joseon Dynasty; he invented the first metal printing press and water gauge. https://en.wikipedia.org/wiki/Jang_Yeong-sil + "jang", + + # Betty Jennings - one of the original programmers of the ENIAC. https://en.wikipedia.org/wiki/ENIAC - https://en.wikipedia.org/wiki/Jean_Bartik + "jennings", + + # Mary Lou Jepsen, was the founder and chief technology officer of One Laptop Per Child (OLPC), and the founder of Pixel Qi. https://en.wikipedia.org/wiki/Mary_Lou_Jepsen + "jepsen", + + # Katherine Coleman Goble Johnson - American physicist and mathematician contributed to the NASA. https://en.wikipedia.org/wiki/Katherine_Johnson + "johnson", + + # Irène Joliot-Curie - French scientist who was awarded the Nobel Prize for Chemistry in 1935. Daughter of Marie and Pierre Curie. https://en.wikipedia.org/wiki/Ir%C3%A8ne_Joliot-Curie + "joliot", + + # Karen Spärck Jones came up with the concept of inverse document frequency, which is used in most search engines today. https://en.wikipedia.org/wiki/Karen_Sp%C3%A4rck_Jones + "jones", + + # A. P. J. Abdul Kalam - is an Indian scientist aka Missile Man of India for his work on the development of ballistic missile and launch vehicle technology - https://en.wikipedia.org/wiki/A._P._J._Abdul_Kalam + "kalam", + + # Sergey Petrovich Kapitsa (Russian: Серге́й Петро́вич Капи́ца; 14 February 1928 – 14 August 2012) was a Russian physicist and demographer. He was best known as host of the popular and long-running Russian scientific TV show, Evident, but Incredible. His father was the Nobel laureate Soviet-era physicist Pyotr Kapitsa, and his brother was the geographer and Antarctic explorer Andrey Kapitsa. - https://en.wikipedia.org/wiki/Sergey_Kapitsa + "kapitsa", + + # Susan Kare, created the icons and many of the interface elements for the original Apple Macintosh in the 1980s, and was an original employee of NeXT, working as the Creative Director. https://en.wikipedia.org/wiki/Susan_Kare + "kare", + + # Mstislav Keldysh - a Soviet scientist in the field of mathematics and mechanics, academician of the USSR Academy of Sciences (1946), President of the USSR Academy of Sciences (1961–1975), three times Hero of Socialist Labor (1956, 1961, 1971), fellow of the Royal Society of Edinburgh (1968). https://en.wikipedia.org/wiki/Mstislav_Keldysh + "keldysh", + + # Mary Kenneth Keller, Sister Mary Kenneth Keller became the first American woman to earn a PhD in Computer Science in 1965. https://en.wikipedia.org/wiki/Mary_Kenneth_Keller + "keller", + + # Johannes Kepler, German astronomer known for his three laws of planetary motion - https://en.wikipedia.org/wiki/Johannes_Kepler + "kepler", + + # Har Gobind Khorana - Indian-American biochemist who shared the 1968 Nobel Prize for Physiology - https://en.wikipedia.org/wiki/Har_Gobind_Khorana + "khorana", + + # Jack Kilby invented silicone integrated circuits and gave Silicon Valley its name. - https://en.wikipedia.org/wiki/Jack_Kilby + "kilby", + + # Maria Kirch - German astronomer and first woman to discover a comet - https://en.wikipedia.org/wiki/Maria_Margarethe_Kirch + "kirch", + + # Donald Knuth - American computer scientist, author of "The Art of Computer Programming" and creator of the TeX typesetting system. https://en.wikipedia.org/wiki/Donald_Knuth + "knuth", + + # Sophie Kowalevski - Russian mathematician responsible for important original contributions to analysis, differential equations and mechanics - https://en.wikipedia.org/wiki/Sofia_Kovalevskaya + "kowalevski", + + # Marie-Jeanne de Lalande - French astronomer, mathematician and cataloguer of stars - https://en.wikipedia.org/wiki/Marie-Jeanne_de_Lalande + "lalande", + + # Hedy Lamarr - Actress and inventor. The principles of her work are now incorporated into modern Wi-Fi, CDMA and Bluetooth technology. https://en.wikipedia.org/wiki/Hedy_Lamarr + "lamarr", + + # Leslie B. Lamport - American computer scientist. Lamport is best known for his seminal work in distributed systems and was the winner of the 2013 Turing Award. https://en.wikipedia.org/wiki/Leslie_Lamport + "lamport", + + # Mary Leakey - British paleoanthropologist who discovered the first fossilized Proconsul skull - https://en.wikipedia.org/wiki/Mary_Leakey + "leakey", + + # Henrietta Swan Leavitt - she was an American astronomer who discovered the relation between the luminosity and the period of Cepheid variable stars. https://en.wikipedia.org/wiki/Henrietta_Swan_Leavitt + "leavitt", + + # Daniel Lewin - Mathematician, Akamai co-founder, soldier, 9/11 victim-- Developed optimization techniques for routing traffic on the internet. Died attempting to stop the 9-11 hijackers. https://en.wikipedia.org/wiki/Daniel_Lewin + "lewin", + + # Ruth Lichterman - one of the original programmers of the ENIAC. https://en.wikipedia.org/wiki/ENIAC - https://en.wikipedia.org/wiki/Ruth_Teitelbaum + "lichterman", + + # Barbara Liskov - co-developed the Liskov substitution principle. Liskov was also the winner of the Turing Prize in 2008. - https://en.wikipedia.org/wiki/Barbara_Liskov + "liskov", + + # Ada Lovelace invented the first algorithm. https://en.wikipedia.org/wiki/Ada_Lovelace (thanks James Turnbull) + "lovelace", + + # Auguste and Louis Lumière - the first filmmakers in history - https://en.wikipedia.org/wiki/Auguste_and_Louis_Lumi%C3%A8re + "lumiere", + + # Mahavira - Ancient Indian mathematician during 9th century AD who discovered basic algebraic identities - https://en.wikipedia.org/wiki/Mah%C4%81v%C4%ABra_(mathematician) + "mahavira", + + # Maria Mayer - American theoretical physicist and Nobel laureate in Physics for proposing the nuclear shell model of the atomic nucleus - https://en.wikipedia.org/wiki/Maria_Mayer + "mayer", + + # John McCarthy invented LISP: https://en.wikipedia.org/wiki/John_McCarthy_(computer_scientist) + "mccarthy", + + # Barbara McClintock - a distinguished American cytogeneticist, 1983 Nobel Laureate in Physiology or Medicine for discovering transposons. https://en.wikipedia.org/wiki/Barbara_McClintock + "mcclintock", + + # Malcolm McLean invented the modern shipping container: https://en.wikipedia.org/wiki/Malcom_McLean + "mclean", + + # Kay McNulty - one of the original programmers of the ENIAC. https://en.wikipedia.org/wiki/ENIAC - https://en.wikipedia.org/wiki/Kathleen_Antonelli + "mcnulty", + + # Dmitri Mendeleev - a chemist and inventor. He formulated the Periodic Law, created a farsighted version of the periodic table of elements, and used it to correct the properties of some already discovered elements and also to predict the properties of eight elements yet to be discovered. https://en.wikipedia.org/wiki/Dmitri_Mendeleev + "mendeleev", + + # Lise Meitner - Austrian/Swedish physicist who was involved in the discovery of nuclear fission. The element meitnerium is named after her - https://en.wikipedia.org/wiki/Lise_Meitner + "meitner", + + # Carla Meninsky, was the game designer and programmer for Atari 2600 games Dodge 'Em and Warlords. https://en.wikipedia.org/wiki/Carla_Meninsky + "meninsky", + + # Johanna Mestorf - German prehistoric archaeologist and first female museum director in Germany - https://en.wikipedia.org/wiki/Johanna_Mestorf + "mestorf", + + # Marvin Minsky - Pioneer in Artificial Intelligence, co-founder of the MIT's AI Lab, won the Turing Award in 1969. https://en.wikipedia.org/wiki/Marvin_Minsky + "minsky", + + # Maryam Mirzakhani - an Iranian mathematician and the first woman to win the Fields Medal. https://en.wikipedia.org/wiki/Maryam_Mirzakhani + "mirzakhani", + + # Samuel Morse - contributed to the invention of a single-wire telegraph system based on European telegraphs and was a co-developer of the Morse code - https://en.wikipedia.org/wiki/Samuel_Morse + "morse", + + # Ian Murdock - founder of the Debian project - https://en.wikipedia.org/wiki/Ian_Murdock + "murdock", + + # John von Neumann - todays computer architectures are based on the von Neumann architecture. https://en.wikipedia.org/wiki/Von_Neumann_architecture + "neumann", + + # Isaac Newton invented classic mechanics and modern optics. https://en.wikipedia.org/wiki/Isaac_Newton + "newton", + + # Florence Nightingale, more prominently known as a nurse, was also the first female member of the Royal Statistical Society and a pioneer in statistical graphics https://en.wikipedia.org/wiki/Florence_Nightingale#Statistics_and_sanitary_reform + "nightingale", + + # Alfred Nobel - a Swedish chemist, engineer, innovator, and armaments manufacturer (inventor of dynamite) - https://en.wikipedia.org/wiki/Alfred_Nobel + "nobel", + + # Emmy Noether, German mathematician. Noether's Theorem is named after her. https://en.wikipedia.org/wiki/Emmy_Noether + "noether", + + # Poppy Northcutt. Poppy Northcutt was the first woman to work as part of NASA’s Mission Control. http://www.businessinsider.com/poppy-northcutt-helped-apollo-astronauts-2014-12?op=1 + "northcutt", + + # Robert Noyce invented silicone integrated circuits and gave Silicon Valley its name. - https://en.wikipedia.org/wiki/Robert_Noyce + "noyce", + + # Panini - Ancient Indian linguist and grammarian from 4th century CE who worked on the world's first formal system - https://en.wikipedia.org/wiki/P%C4%81%E1%B9%87ini#Comparison_with_modern_formal_systems + "panini", + + # Ambroise Pare invented modern surgery. https://en.wikipedia.org/wiki/Ambroise_Par%C3%A9 + "pare", + + # Louis Pasteur discovered vaccination, fermentation and pasteurization. https://en.wikipedia.org/wiki/Louis_Pasteur. + "pasteur", + + # Cecilia Payne-Gaposchkin was an astronomer and astrophysicist who, in 1925, proposed in her Ph.D. thesis an explanation for the composition of stars in terms of the relative abundances of hydrogen and helium. https://en.wikipedia.org/wiki/Cecilia_Payne-Gaposchkin + "payne", + + # Radia Perlman is a software designer and network engineer and most famous for her invention of the spanning-tree protocol (STP). https://en.wikipedia.org/wiki/Radia_Perlman + "perlman", + + # Rob Pike was a key contributor to Unix, Plan 9, the X graphic system, utf-8, and the Go programming language. https://en.wikipedia.org/wiki/Rob_Pike + "pike", + + # Henri Poincaré made fundamental contributions in several fields of mathematics. https://en.wikipedia.org/wiki/Henri_Poincar%C3%A9 + "poincare", + + # Laura Poitras is a director and producer whose work, made possible by open source crypto tools, advances the causes of truth and freedom of information by reporting disclosures by whistleblowers such as Edward Snowden. https://en.wikipedia.org/wiki/Laura_Poitras + "poitras", + + # Tat’yana Avenirovna Proskuriakova (Russian: Татья́на Авени́ровна Проскуряко́ва) (January 23 [O.S. January 10] 1909 – August 30, 1985) was a Russian-American Mayanist scholar and archaeologist who contributed significantly to the deciphering of Maya hieroglyphs, the writing system of the pre-Columbian Maya civilization of Mesoamerica. https://en.wikipedia.org/wiki/Tatiana_Proskouriakoff + "proskuriakova", + + # Claudius Ptolemy - a Greco-Egyptian writer of Alexandria, known as a mathematician, astronomer, geographer, astrologer, and poet of a single epigram in the Greek Anthology - https://en.wikipedia.org/wiki/Ptolemy + "ptolemy", + + # C. V. Raman - Indian physicist who won the Nobel Prize in 1930 for proposing the Raman effect. - https://en.wikipedia.org/wiki/C._V._Raman + "raman", + + # Srinivasa Ramanujan - Indian mathematician and autodidact who made extraordinary contributions to mathematical analysis, number theory, infinite series, and continued fractions. - https://en.wikipedia.org/wiki/Srinivasa_Ramanujan + "ramanujan", + + # Sally Kristen Ride was an American physicist and astronaut. She was the first American woman in space, and the youngest American astronaut. https://en.wikipedia.org/wiki/Sally_Ride + "ride", + + # Rita Levi-Montalcini - Won Nobel Prize in Physiology or Medicine jointly with colleague Stanley Cohen for the discovery of nerve growth factor (https://en.wikipedia.org/wiki/Rita_Levi-Montalcini) + "montalcini", + + # Dennis Ritchie - co-creator of UNIX and the C programming language. - https://en.wikipedia.org/wiki/Dennis_Ritchie + "ritchie", + + # Wilhelm Conrad Röntgen - German physicist who was awarded the first Nobel Prize in Physics in 1901 for the discovery of X-rays (Röntgen rays). https://en.wikipedia.org/wiki/Wilhelm_R%C3%B6ntgen + "roentgen", + + # Rosalind Franklin - British biophysicist and X-ray crystallographer whose research was critical to the understanding of DNA - https://en.wikipedia.org/wiki/Rosalind_Franklin + "rosalind", + + # Meghnad Saha - Indian astrophysicist best known for his development of the Saha equation, used to describe chemical and physical conditions in stars - https://en.wikipedia.org/wiki/Meghnad_Saha + "saha", + + # Jean E. Sammet developed FORMAC, the first widely used computer language for symbolic manipulation of mathematical formulas. https://en.wikipedia.org/wiki/Jean_E._Sammet + "sammet", + + # Carol Shaw - Originally an Atari employee, Carol Shaw is said to be the first female video game designer. https://en.wikipedia.org/wiki/Carol_Shaw_(video_game_designer) + "shaw", + + # Dame Stephanie "Steve" Shirley - Founded a software company in 1962 employing women working from home. https://en.wikipedia.org/wiki/Steve_Shirley + "shirley", + + # William Shockley co-invented the transistor - https://en.wikipedia.org/wiki/William_Shockley + "shockley", + + # Françoise Barré-Sinoussi - French virologist and Nobel Prize Laureate in Physiology or Medicine; her work was fundamental in identifying HIV as the cause of AIDS. https://en.wikipedia.org/wiki/Fran%C3%A7oise_Barr%C3%A9-Sinoussi + "sinoussi", + + # Betty Snyder - one of the original programmers of the ENIAC. https://en.wikipedia.org/wiki/ENIAC - https://en.wikipedia.org/wiki/Betty_Holberton + "snyder", + + # Frances Spence - one of the original programmers of the ENIAC. https://en.wikipedia.org/wiki/ENIAC - https://en.wikipedia.org/wiki/Frances_Spence + "spence", + + # Richard Matthew Stallman - the founder of the Free Software movement, the GNU project, the Free Software Foundation, and the League for Programming Freedom. He also invented the concept of copyleft to protect the ideals of this movement, and enshrined this concept in the widely-used GPL (General Public License) for software. https://en.wikiquote.org/wiki/Richard_Stallman + "stallman", + + # Lina Solomonovna Stern (or Shtern; Russian: Лина Соломоновна Штерн; 26 August 1878 – 7 March 1968) was a Soviet biochemist, physiologist and humanist whose medical discoveries saved thousands of lives at the fronts of World War II. She is best known for her pioneering work on blood–brain barrier, which she described as hemato-encephalic barrier in 1921. https://en.wikipedia.org/wiki/Lina_Stern + "shtern", + + # Michael Stonebraker is a database research pioneer and architect of Ingres, Postgres, VoltDB and SciDB. Winner of 2014 ACM Turing Award. https://en.wikipedia.org/wiki/Michael_Stonebraker + "stonebraker", + + # Janese Swanson (with others) developed the first of the Carmen Sandiego games. She went on to found Girl Tech. https://en.wikipedia.org/wiki/Janese_Swanson + "swanson", + + # Aaron Swartz was influential in creating RSS, Markdown, Creative Commons, Reddit, and much of the internet as we know it today. He was devoted to freedom of information on the web. https://en.wikiquote.org/wiki/Aaron_Swartz + "swartz", + + # Bertha Swirles was a theoretical physicist who made a number of contributions to early quantum theory. https://en.wikipedia.org/wiki/Bertha_Swirles + "swirles", + + # Valentina Tereshkova is a russian engineer, cosmonaut and politician. She was the first woman flying to space in 1963. In 2013, at the age of 76, she offered to go on a one-way mission to mars. https://en.wikipedia.org/wiki/Valentina_Tereshkova + "tereshkova", + + # Nikola Tesla invented the AC electric system and every gadget ever used by a James Bond villain. https://en.wikipedia.org/wiki/Nikola_Tesla + "tesla", + + # Ken Thompson - co-creator of UNIX and the C programming language - https://en.wikipedia.org/wiki/Ken_Thompson + "thompson", + + # Linus Torvalds invented Linux and Git. https://en.wikipedia.org/wiki/Linus_Torvalds + "torvalds", + + # Alan Turing was a founding father of computer science. https://en.wikipedia.org/wiki/Alan_Turing. + "turing", + + # Varahamihira - Ancient Indian mathematician who discovered trigonometric formulae during 505-587 CE - https://en.wikipedia.org/wiki/Var%C4%81hamihira#Contributions + "varahamihira", + + # Dorothy Vaughan was a NASA mathematician and computer programmer on the SCOUT launch vehicle program that put America's first satellites into space - https://en.wikipedia.org/wiki/Dorothy_Vaughan + "vaughan", + + # Sir Mokshagundam Visvesvaraya - is a notable Indian engineer. He is a recipient of the Indian Republic's highest honour, the Bharat Ratna, in 1955. On his birthday, 15 September is celebrated as Engineer's Day in India in his memory - https://en.wikipedia.org/wiki/Visvesvaraya + "visvesvaraya", + + # Christiane Nüsslein-Volhard - German biologist, won Nobel Prize in Physiology or Medicine in 1995 for research on the genetic control of embryonic development. https://en.wikipedia.org/wiki/Christiane_N%C3%BCsslein-Volhard + "volhard", + + # Cédric Villani - French mathematician, won Fields Medal, Fermat Prize and Poincaré Price for his work in differential geometry and statistical mechanics. https://en.wikipedia.org/wiki/C%C3%A9dric_Villani + "villani", + + # Marlyn Wescoff - one of the original programmers of the ENIAC. https://en.wikipedia.org/wiki/ENIAC - https://en.wikipedia.org/wiki/Marlyn_Meltzer + "wescoff", + + # Andrew Wiles - Notable British mathematician who proved the enigmatic Fermat's Last Theorem - https://en.wikipedia.org/wiki/Andrew_Wiles + "wiles", + + # Roberta Williams, did pioneering work in graphical adventure games for personal computers, particularly the King's Quest series. https://en.wikipedia.org/wiki/Roberta_Williams + "williams", + + # Sophie Wilson designed the first Acorn Micro-Computer and the instruction set for ARM processors. https://en.wikipedia.org/wiki/Sophie_Wilson + "wilson", + + # Jeannette Wing - co-developed the Liskov substitution principle. - https://en.wikipedia.org/wiki/Jeannette_Wing + "wing", + + # Steve Wozniak invented the Apple I and Apple II. https://en.wikipedia.org/wiki/Steve_Wozniak + "wozniak", + + # The Wright brothers, Orville and Wilbur - credited with inventing and building the world's first successful airplane and making the first controlled, powered and sustained heavier-than-air human flight - https://en.wikipedia.org/wiki/Wright_brothers + "wright", + + # Rosalyn Sussman Yalow - Rosalyn Sussman Yalow was an American medical physicist, and a co-winner of the 1977 Nobel Prize in Physiology or Medicine for development of the radioimmunoassay technique. https://en.wikipedia.org/wiki/Rosalyn_Sussman_Yalow + "yalow", + + # Ada Yonath - an Israeli crystallographer, the first woman from the Middle East to win a Nobel prize in the sciences. https://en.wikipedia.org/wiki/Ada_Yonath + "yonath", + + # Nikolay Yegorovich Zhukovsky (Russian: Никола́й Его́рович Жуко́вский, January 17 1847 – March 17, 1921) was a Russian scientist, mathematician and engineer, and a founding father of modern aero- and hydrodynamics. Whereas contemporary scientists scoffed at the idea of human flight, Zhukovsky was the first to undertake the study of airflow. He is often called the Father of Russian Aviation. https://en.wikipedia.org/wiki/Nikolay_Yegorovich_Zhukovsky + "zhukovsky", +]) + + +def setup_run_dir(runs_dirpath, run_name=None, new_run=False, check_exists=False): + """ + If new_run is True, creates a new directory: + If run_name is None, generate a random name + else build the created directory name with run_name + + If new_run is False, return an existing directory: + if run_name is None, return the last created directory (from timestamp) + else return the last created directory (from timestamp) whose name starts with run_name, + if that does not exist and check_exists is False create a new run with run_name, + if check_exists is True, then raise an error. + + Special case: if there is no existing runs, the new_run is not taken into account and the function behaves like new_run is True. + + :param runs_dirpath: Parent directory path of all the runs + :param run_name: + :param new_run: + :param check_exists: + :return: Run directory path. The directory name is in the form "run_name _ timestamp" + """ + # Create runs directory of it does not exist + if not os.path.exists(runs_dirpath): + os.makedirs(runs_dirpath, exist_ok=True) + + existing_run_dirnames = os.listdir(runs_dirpath) + if new_run or (not new_run and not 0 < len(existing_run_dirnames)): + if run_name is not None: + # Create another directory name for the run, with its name starting with run_name + name_timestamped = create_name_timestamped(run_name) + else: + # Create another directory name for the run, excluding the existing names + existing_run_names = [existing_run_dirname.split(" _ ")[0] for existing_run_dirname in + existing_run_dirnames] + name_timestamped = create_free_name_timestamped(exclude_list=existing_run_names) + current_run_dirpath = os.path.join(runs_dirpath, name_timestamped) + os.mkdir(current_run_dirpath) + else: + if run_name is not None: + # Pick run dir based on run_name + filtered_existing_run_dirnames = [existing_run_dirname for existing_run_dirname in existing_run_dirnames if + existing_run_dirname.split(" _ ")[0] == run_name] + if filtered_existing_run_dirnames: + filtered_existing_run_timestamps = [filtered_existing_run_dirname.split(" _ ")[1] for + filtered_existing_run_dirname in + filtered_existing_run_dirnames] + filtered_last_index = filtered_existing_run_timestamps.index(max(filtered_existing_run_timestamps)) + current_run_dirname = filtered_existing_run_dirnames[filtered_last_index] + else: + if check_exists: + raise FileNotFoundError("Run '{}' does not exist.".format(run_name)) + else: + return setup_run_dir(runs_dirpath, run_name=run_name, new_run=True) + else: + # Pick last run dir based on timestamp + existing_run_timestamps = [existing_run_dirname.split(" _ ")[1] for existing_run_dirname in + existing_run_dirnames] + last_index = existing_run_timestamps.index(max(existing_run_timestamps)) + current_run_dirname = existing_run_dirnames[last_index] + current_run_dirpath = os.path.join(runs_dirpath, current_run_dirname) + return current_run_dirpath + + +def create_name_timestamped(name): + timestamp = time.time() + formatted_timestamp = datetime.datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S') + name_timestamped = name + " _ " + formatted_timestamp + return name_timestamped + + +def create_free_name_timestamped(exclude_list=None): + if exclude_list is not None: + names = list(NAME_SET - set(exclude_list)) + else: + names = list(NAME_SET) + assert 0 < len(names), \ + "In create_free_name_timestamped(), all possible names have been used. " \ + "Cannot create a new name without a collision! Delete some runs to continue..." + sorted_names = sorted(names) + name = sorted_names[0] + name_timestamped = create_name_timestamped(name) + return name_timestamped + + +def setup_run_subdir(run_dir, subdirname): + subdirpath = os.path.join(run_dir, subdirname) + os.makedirs(subdirpath, exist_ok=True) + return subdirpath + + +def setup_run_subdirs(run_dir, logs_dirname="logs", checkpoints_dirname="checkpoints"): + logs_dir = os.path.join(run_dir, logs_dirname) + checkpoints_dir = os.path.join(run_dir, checkpoints_dirname) + os.makedirs(logs_dir, exist_ok=True) + os.makedirs(checkpoints_dir, exist_ok=True) + return logs_dir, checkpoints_dir + + +def wipe_run_subdirs(run_dir, logs_dirname="logs", checkpoints_dirname="checkpoints"): + logs_dir = os.path.join(run_dir, logs_dirname) + checkpoints_dir = os.path.join(run_dir, checkpoints_dirname) + python_utils.wipe_dir(logs_dir) + python_utils.wipe_dir(checkpoints_dir) + + +def save_config(config, config_dirpath): + filepath = os.path.join(config_dirpath, 'config.json') + with open(filepath, 'w') as outfile: + json.dump(config, outfile) + # shutil.copyfile(os.path.join(project_dir, "config.py"), os.path.join(current_logs_dir, "config.py")) + + +def load_config(config_name="config", config_dirpath="", try_default=False): + if os.path.splitext(config_name)[1] == ".json": + config_filepath = os.path.join(config_dirpath, config_name) + else: + config_filepath = os.path.join(config_dirpath, config_name + ".json") + + try: + with open(config_filepath, 'r') as f: + minified = jsmin(f.read()) + try: + config = json.loads(minified) + except json.decoder.JSONDecodeError as e: + print_utils.print_error("ERROR: Parsing config failed:") + print(e) + print_utils.print_info("Minified JSON causing the problem:") + print(str(minified)) + exit() + return config + except FileNotFoundError: + if config_name == "config" and config_dirpath == "": + print_utils.print_warning( + "WARNING: the default config file was not found....") + return None + elif try_default: + print_utils.print_warning( + "WARNING: config file {} was not found, opening default config file config.defaults.json instead.".format( + config_filepath)) + return load_config() + else: + print_utils.print_warning( + "WARNING: config file {} was not found.".format(config_filepath)) + return None + + +def _merge_dictionaries(dict1, dict2): + """ + Recursive merge dictionaries. + + :param dict1: Base dictionary to merge. + :param dict2: Dictionary to merge on top of base dictionary. + :return: Merged dictionary + """ + for key, val in dict1.items(): + if isinstance(val, dict): + dict2_node = dict2.setdefault(key, {}) + _merge_dictionaries(val, dict2_node) + else: + if key not in dict2: + dict2[key] = val + + return dict2 + + +def load_defaults_in_config(config: dict, filepath_key: str="defaults_filepath", depth: int=0) -> dict: + """ + Searches in the config dict for keys equal to "defaults_filepath". + When one is found, read the json at the defaults_filepath and add the defaults for the current params. + + Example: + config = { + "some_params": { + filepath_key: "path/to/default_params.json", + "param_2": { + "sub_param_1": 0, + "sub_param_2": 0, + } + } + } + and there is a file at path/to/default_params.json which reads: + { + "param_1": 0, + "param_2": { + "sub_param_2": 1, + } + } + + The returned config dict will be: + { + "some_params": { + "param_1": 0, + "param_2": { + "sub_param_1": 0, + "sub_param_2": 1, + } + } + } + Note the value of param_2.sub_param_2 whose default was overwritten while param_2.sub_param_1 was left to default + + @param config: + @param filepath_key: + @return: config with all defaults loaded + """ + assert isinstance(config, dict), f"config should be of type dict, not {type(config)}" + + if filepath_key in config and isinstance(config[filepath_key], str): + defaults = python_utils.load_json(config[filepath_key]) + if defaults: + tabs = "\t"*depth + print_utils.print_info(f"{tabs}INFO: Loading defaults from {config[filepath_key]}") + # Recursively process defaults + defaults = load_defaults_in_config(defaults, filepath_key=filepath_key, depth=depth+1) + # Copy defaults in config if the key is not already there + config = _merge_dictionaries(defaults, config) + # for default_name, default_value in defaults.items(): + # if default_name not in config: + # config[default_name] = default_value + # Delete filepath_key key + del config[filepath_key] + else: + print_utils.print_error(f"ERROR: Could not load defaults from {config[filepath_key]}!") + raise ValueError + + # Check items of all other keys + for key, item in config.items(): + if type(item) == dict: + config[key] = load_defaults_in_config(item, filepath_key=filepath_key, depth=depth+1) + + return config diff --git a/paper/README.md b/paper/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ca95116dc753a8ddf8bfa6371f2269ef13981c82 --- /dev/null +++ b/paper/README.md @@ -0,0 +1,64 @@ +# Introduction + +This "paper" folder contains scripts used to generate certain figures/results for our paper. +These scripts use data saved by launching the ```main.py``` with ```--mode=eval``` or with ```--mode=eval_coco```. + +We introduce each script in the following. + +## convert_stats_to_latex.py + +When launching ```main.py``` with ```--mode=eval_coco```, it will compute COCO measures (AP/AR and their variants) and save them in .json files with filenames "*.stats.*.annotation.*.json". +Use this script to print those metrics in LaTex code ready to be copy-pasted into tables. + +Run as(any filename matching "*.stats.*.annotation.*.json" under the given ```dirpath``` will be printed): +``` +convert_stats_to_latex.py --dirpath +``` + +## plot_complexity_fidelity.py + +Use this script to plot the complexity vs fidelity figures. + +Run as (no arguments, will have to change paths directly in the "main" function of the script): +``` +plot_complexity_fidelity.py +``` + +## plot_contour_angle_hist.py + +This script is used to compute and plot the "relative angle distribution" histograms from contours detected in segmentation probability maps. +It is an additional measure for building regularity that we propose. It does not need any +ground truth annotation so that results are evaluated on their own. +We use the simple polygonization method to obtain contours from the segmentation map. The +minimum rotated rectangle is computed for each building. Then the relative angle between each +contour edge and the principal axis of the associated minimum rotated rectangle is computed. For +a collection of contours, we aggregate the data in the form of a distribution of relative angles. +If the distribution is more homogeneous, it means buildings are less regular, i.e. smoother. +Conversely, if the distribution has peaks around certain relative angle values (which are expected +to be 0°, 90°, and 180° for buildings), it means buildings are more regular, with sharper corners +having similar angles. + +Run as (no arguments, will have to change paths directly in the "main" function of the script): +``` +plot_contour_angle_hist.py +``` + +## plot_contour_metrics.py + +This script is used to plot our "max tangent angle error" metric. +All computations are done when launching ```main.py``` with ```--mode=eval_coco``` which generates .json files with filenames "test.metrics.test.annotation.poly.*.json". +This script takes as input the path to the "eval directory" where the results for all runs are saved and where these .json files will be found. + +Example use (filenames to specific runs and .json result metrics should be changed directly in the "main" function of the script): +``` +plot_contour_metrics.py --dirpath +``` + +## show_result_image.py + +This script uses the pycocotools API to plot result contours in [polygon] format on top of test images. + +Run as (no arguments, will have to change paths directly in the "main" function of the script): +``` +plot_contour_angle_hist.py +``` \ No newline at end of file diff --git a/paper/convert_stats_to_latex.py b/paper/convert_stats_to_latex.py new file mode 100644 index 0000000000000000000000000000000000000000..4a0953a3e0b9b2d97ddd0e697cad2aabe102ee28 --- /dev/null +++ b/paper/convert_stats_to_latex.py @@ -0,0 +1,61 @@ +import os.path +import fnmatch +import argparse + +from lydorn_utils import python_utils + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--dirpath', + default="/home/lydorn/data/mapping_challenge_dataset/eval_runs", + type=str, + help='Path to eval directory') + + args = argparser.parse_args() + return args + + +def convert(in_filepath, stat_names): + stats = python_utils.load_json(in_filepath) + if stats: + string = "" + for stat_name in stat_names: + string += str(round(100 * stats[stat_name], 1)) + " & " + return string[:-2] + "\\\\" + else: + print("File not found!") + return "" + + +def main(): + args = get_args() + + stat_names = ["AP", "AP_50", "AP_75", "AP_S", "AP_M", "AP_L", "AR", "AR_50", "AR_75", "AR_S", "AR_M", "AR_L"] + + dirname_list = next(os.walk(args.dirpath))[1] + dirname_list = sorted(dirname_list) + + run_file_latex_list = [("Run name", "filename", "Latex")] + run_max_len = 0 + file_max_len = 0 + for dirname in dirname_list: + dirpath = os.path.join(args.dirpath, dirname) + in_filename_list = fnmatch.filter(os.listdir(dirpath), "*.stats.*.annotation.*.json") + in_filename_list = sorted(in_filename_list) + for in_filename in in_filename_list: + in_filepath = os.path.join(dirpath, in_filename) + latex = convert(in_filepath, stat_names) + run = dirname[:-len(" | 0000-00-00 00:00:00")] + run_file_latex_list.append((run, in_filename, latex)) + run_max_len = max(run_max_len, len(run)) + file_max_len = max(file_max_len, len(in_filename)) + + # print + for run, file, latex in run_file_latex_list: + print(run.ljust(run_max_len, ' '), file.ljust(file_max_len, ' '), latex) + + +if __name__ == '__main__': + main() diff --git a/paper/plot_complexity_fidelity.py b/paper/plot_complexity_fidelity.py new file mode 100644 index 0000000000000000000000000000000000000000..c19dcb776eb523485894961192d2138f416bb67e --- /dev/null +++ b/paper/plot_complexity_fidelity.py @@ -0,0 +1,76 @@ +import os +import matplotlib.pyplot as plt + +from lydorn_utils import python_utils +from lydorn_utils import print_utils + + +def get_stat_from_all(stat_filepath_format, method_info, tolerances, stat_name): + stat_list = [0 for _ in tolerances] + for i, tolerance in enumerate(tolerances): + filepath = stat_filepath_format.format(method_info["name"], tolerance) + stats = python_utils.load_json(filepath) + if stats: + stat_list[i] = stats[stat_name] + else: + print_utils.print_warning("WARNING: could not open {}".format(filepath)) + return stat_list + + +def plot_stat(stat_filepath_format, method_info_list, tolerances, stat_name, exp_name): + legend = [] + for method_info in method_info_list: + ap_list = get_stat_from_all(stat_filepath_format, method_info, tolerances, stat_name) + legend.append(method_info["title"]) + + plt.plot(tolerances, ap_list) + + plt.legend(legend, loc='lower left') + plt.xlabel("Tolerance") + plt.ylabel(stat_name) + plt.title(exp_name + ": " + stat_name + " vs tolerance") + plt.savefig(exp_name.replace(" ", "_") + "_" + stat_name + "_vs_tolerance.pdf") + plt.show() + + +def main(): + method_info_list = [ + { + "title": "Baseline polygonization", + "name": "simple" + }, + { + "title": "Our polygonization", + "name": "acm" + }, + ] + tolerances = [0.125, 0.25, 0.5, 1, 2, 4, 8, 16] + eval_runs_dirpath = "/data/data/mapping_challenge_dataset/eval_runs_cluster" + + info_list = [ + { + "exp_name": "U-Net16 full method", + "run_dirname": "mapping_dataset.unet16.train_val | 2020-02-21 03:09:03", + }, + { + "exp_name": "U-Net16 field off", + "run_dirname": "mapping_dataset.unet16.field_off.train_val | 2020-02-28 23:51:16", + }, + { + "exp_name": "DeepLab101 full method", + "run_dirname": "mapping_dataset.deeplab101.train_val | 2020-02-24 23:57:19", + }, + { + "exp_name": "DeepLab101 field off", + "run_dirname": "mapping_dataset.deeplab101.field_off.train_val | 2020-03-02 00:03:45", + }, + ] + for info in info_list: + stat_filepath_format = os.path.join(eval_runs_dirpath, info["run_dirname"], "test.stats.test.annotation.poly.{}.tol_{}.json") + + plot_stat(stat_filepath_format, method_info_list, tolerances, "AP", info["exp_name"]) + plot_stat(stat_filepath_format, method_info_list, tolerances, "AR", info["exp_name"]) + + +if __name__ == '__main__': + main() diff --git a/paper/plot_contour_angle_hist.py b/paper/plot_contour_angle_hist.py new file mode 100644 index 0000000000000000000000000000000000000000..ca69f106a727526ab3e49a97234655ec71a1a194 --- /dev/null +++ b/paper/plot_contour_angle_hist.py @@ -0,0 +1,90 @@ +import matplotlib.pyplot as plt + +import numpy as np +import shapely.geometry +import shapely.affinity +import skimage.io +import skimage.measure +from matplotlib.ticker import StrMethodFormatter + + +def compute_polygon_angles(polygon): + # --- Rotate polygon so that main axis is aligned with the x-axis + min_rot_rect = polygon.minimum_rotated_rectangle + min_rot_rect_contour = np.array(min_rot_rect.exterior.coords[:]) + min_rot_rect_edges = min_rot_rect_contour[1:] - min_rot_rect_contour[:-1] + min_rot_rect_norms = np.linalg.norm(min_rot_rect_edges, axis=1) + max_norms_index = np.argmax(min_rot_rect_norms) + longest_edge = min_rot_rect_edges[max_norms_index] + main_angle = np.angle(longest_edge[0] + 1j*longest_edge[1]) + polygon = shapely.affinity.rotate(polygon, -main_angle, use_radians=True) + + contour = np.array(polygon.exterior) + edges = contour[1:] - contour[:-1] + edges = edges[:, 1] + 1j * edges[:, 0] + angles = np.angle(edges) + angles[angles < 0] += np.pi # Don't care about direction of edge + + return angles + + +def get_angles(mask_filepath, level=0.5, tol=0.1): + # Read images + mask = skimage.io.imread(mask_filepath) / 255 + + # Compute contours + contours = skimage.measure.find_contours(mask, level, fully_connected='low', positive_orientation='high') + polygons = [shapely.geometry.Polygon(contour[:, ::-1]) for contour in contours] + # Filter out really small polylines + polygons = [polygon for polygon in polygons if 2 < polygon.area] + # Simplify + polygons = [polygon.simplify(tol, preserve_topology=True) for polygon in polygons] + + # Compute angles + contours_angles = [compute_polygon_angles(polygon) for polygon in polygons] + + angles = np.concatenate(contours_angles) + relative_degrees = angles * 180 / np.pi + return relative_degrees + + +def plot_contour_angle_hist(list_info, level=0.5, tol=0.1): + start = 0 + stop = 180 + bin_count = 100 + bin_edges = np.linspace(start, stop, bin_count + 1) + bin_width = (stop - start) / bin_count + for i, info in enumerate(list_info): + degrees = get_angles(info["mask_filepath"], level, tol) + hist, bin_edges = np.histogram(degrees, bins=bin_edges) + freq = hist / np.sum(hist) + plt.bar(bin_edges[1:] - bin_width/2, freq, width=bin_width, alpha=0.5, label=info["name"]) + + plt.title("Histogram of relative contour angles") + plt.xlabel("Relative angle") + plt.ylabel("Freq") + plt.gca().xaxis.set_major_formatter(StrMethodFormatter(u"{x:.0f}°")) + plt.legend(loc="upper left") + plt.xlim(0, 180) + plt.savefig("histogram_of_relative_contour_angles.pdf", transparent=True) + plt.show() + + +def main(): + + list_info = [ + { + "name": "ICTNet", + "mask_filepath": "inria_dataset_test_sample_result.ictlab.jpg" + }, + { + "name": "Ours", + "mask_filepath": "inria_dataset_test_sample_result.ours.tif" + }, + ] + + plot_contour_angle_hist(list_info, tol=1) + + +if __name__ == '__main__': + main() diff --git a/paper/plot_contour_metrics.py b/paper/plot_contour_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..cf58cd6141db92dfaa81db03c4c80113876c2b1f --- /dev/null +++ b/paper/plot_contour_metrics.py @@ -0,0 +1,161 @@ +import argparse +import os + +import matplotlib.pyplot as plt + +import numpy as np + +from lydorn_utils import python_utils +from lydorn_utils import print_utils + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--dirpath', + default="/home/lydorn/data/mapping_challenge_dataset/eval_runs", + type=str, + help='Path to eval directory') + + args = argparser.parse_args() + return args + + +def plot_metric(dirpath, info_list): + legend = [] + for info in info_list: + metrics_filepath = os.path.join(dirpath, info["metrics_filepath"]) + metrics = python_utils.load_json(metrics_filepath) + if metrics: + max_angle_diffs = np.array(metrics["max_angle_diffs"]) + total = len(max_angle_diffs) + angle_thresholds = range(0, 91) + fraction_under_threshold_list = [] + for angle_threshold in angle_thresholds: + fraction_under_threshold = np.sum(max_angle_diffs < angle_threshold) / total + fraction_under_threshold_list.append(fraction_under_threshold) + # Plot + plt.plot(angle_thresholds, fraction_under_threshold_list) + + # Compute mean + mean_error = np.mean(max_angle_diffs) + + legend.append(f"{info['name']}: {mean_error:.1f}°") + + else: + print_utils.print_warning("WARNING: could not open {}".format(info["metrics_filepath"])) + + plt.legend(legend, loc='lower right') + plt.xlabel("Threshold (degrees)") + plt.ylabel("Fraction of detections") + axes = plt.gca() + axes.set_xlim([0, 90]) + axes.set_ylim([0, 1]) + title = f"Cumulative max tangent angle error per detection" + plt.title(title) + plt.savefig(title.lower().replace(" ", "_") + ".pdf") + plt.show() + + +def main(): + args = get_args() + + + # Mapping challenge: + info_list = [ + { + "name": "UResNet101 (no field), simple poly.", + "metrics_filepath": "mapping_dataset.unet_resnet101_pretrained.field_off.train_val | 2020-05-21 08:33:20/test.metrics.test.annotation.poly.simple.tol_0.125.json" + }, + { + "name": "UResNet101 (with field), simple poly.", + "metrics_filepath": "mapping_dataset.unet_resnet101_pretrained.train_val | 2020-05-21 08:32:48/test.metrics.test.annotation.poly.simple.tol_0.125.json" + }, + { + "name": "UResNet101 (with field), our poly.", + "metrics_filepath": "mapping_dataset.unet_resnet101_pretrained.train_val | 2020-05-21 08:32:48/test.metrics.test.annotation.poly.acm.tol_0.125.json" + }, + { + "name": "UResNet101 (no $L_{align90}$), our poly.", + "metrics_filepath": "mapping_dataset.unet_resnet101_pretrained.align90_off.train_val | 2020-11-02 07:34:43/test.metrics.test.annotation.poly.acm.tol_0.125.json" + }, + { + "name": "UResNet101 (no $L_{int edge}$), our poly.", + "metrics_filepath": "mapping_dataset.unet_resnet101_pretrained.edge_int_off.train_val | 2020-11-02 07:34:54/test.metrics.test.annotation.poly.acm.tol_0.125.json" + }, + { + "name": "UResNet101 (no $L_{int align}$ and $L_{edge align}$), our poly.", + "metrics_filepath": "mapping_dataset.unet_resnet101_pretrained.seg_framefield_off.train_val | 2020-10-29 11:27:52/test.metrics.test.annotation.poly.acm.tol_0.125.json" + }, + { + "name": "UResNet101 (no $L_{smooth}$), our poly.", + "metrics_filepath": "mapping_dataset.unet_resnet101_pretrained.smooth_off.train_val | 2020-10-29 11:18:33/test.metrics.test.annotation.poly.acm.tol_0.125.json" + }, + + { + "name": "PolyMapper", + "metrics_filepath": "mapping_dataset.polymapper | 0000-00-00 00:00:00/test.metrics.test.annotation.poly.json" + }, + { + "name": "U-Net variant, ASIP poly.", + "metrics_filepath": "mapping_dataset.asip | 0000-00-00 00:00:00/test.metrics.test.annotation.poly.json" + }, + { + "name": "Zorzi et al.", + "metrics_filepath": "mapping_dataset.zorzi | 0000-00-00 00:00:00/test.metrics.test.annotation.poly.json" + }, + { + "name": "U-Net variant, UResNet101 poly", + "metrics_filepath": "mapping_dataset.open_solution_full | 0000-00-00 00:00:00/test.metrics.test.annotation.seg_cleaned.poly.json" + } + ] + + # Inria Polygonized Dataset + # info_list = [ + # { + # "name": "UResNet101 (no field), simple poly.", + # "metrics_filepath": "/home/lydorn/data/AerialImageDataset/raw/test/pred_ours_leaderboard_new_losses.field_off/poly_shapefile.simple.tol_1/aggr_metrics.json" + # }, + # { + # "name": "UResNet101 (with field), our poly.", + # "metrics_filepath": "/home/lydorn/data/AerialImageDataset/raw/test/pred_ours_leaderboard/poly_shapefile.acm.tol_0.125/aggr_metrics.json" + # }, + # { + # "name": "Zorzi et al.", + # "metrics_filepath": "/home/lydorn/data/AerialImageDataset/raw/test/pred_zorzi/shapes/aggr_metrics.json" + # }, + # { + # "name": "ICTNet, simple poly.", + # "metrics_filepath": "/home/lydorn/data/AerialImageDataset/raw/test/pred_ictnet/shp/aggr_metrics.json" + # }, + # { + # "name": "Khvedchenya, simple poly.", + # "metrics_filepath": "/home/lydorn/data/AerialImageDataset/raw/test/pred_khvedchenya/shp/aggr_metrics.json" + # }, + # ] + + # LuxCarta's Bangkok image + # info_list = [ + # { + # "name": "ACM", + # "metrics_filepath": "/home/lydorn/repos/lydorn/frame_field_learning/frame_field_learning/test_images/Bangkok/Bangkok3bands.poly_acm.metrics.json" + # }, + # { + # "name": "ASM", + # "metrics_filepath": "/home/lydorn/repos/lydorn/frame_field_learning/frame_field_learning/test_images/Bangkok/Bangkok3bands.poly_asm.metrics.json" + # }, + # { + # "name": "ASM regularized", + # "metrics_filepath": "/home/lydorn/repos/lydorn/frame_field_learning/frame_field_learning/test_images/Bangkok/Bangkok3bands.reg.metrics.json" + # }, + # { + # "name": "Company", + # "metrics_filepath": "/home/lydorn/repos/lydorn/frame_field_learning/frame_field_learning/test_images/Bangkok/Luxcarta/Building_Thailand_Bangkok_pansharpened25.metrics.json" + # }, + # ] + + plot_metric(args.dirpath, info_list) + + +if __name__ == '__main__': + main() diff --git a/paper/rounded_corners.py b/paper/rounded_corners.py new file mode 100644 index 0000000000000000000000000000000000000000..8b2c0c62e6f0bb6d1b59768b84df6f5508eb472d --- /dev/null +++ b/paper/rounded_corners.py @@ -0,0 +1,96 @@ +import numpy as np +import matplotlib.pyplot as plt + +import shapely.geometry + +import torch_lydorn.torchvision +from numpy.core._multiarray_umath import ndarray +from scipy import ndimage + +import vectorization_ambiguities + + +def create_polygons(): + polygons = [ + np.array([ + [25, 25], + [75, 25], + [75, 50], + [50, 50], + [50, 75], + [25, 75], + [25, 25], + ]) + ] + return polygons + + +def displace_polygons(polygons, max_global, max_polygon, max_vertex, max_rot_deg): + new_polygons = [] + global_disp = np.random.uniform(-1, 1, 2) * max_global + for polygon in polygons: + polygon_disp = np.random.uniform(-1, 1, 2) * max_polygon + vertex_disp = np.random.uniform(-1, 1, polygon.shape) * max_vertex + new_polygon = polygon + global_disp + polygon_disp + vertex_disp + + # Rotation + geom = shapely.geometry.Polygon(new_polygon) + angle = np.random.uniform(-1, 1, 1) * max_rot_deg + geom = shapely.affinity.rotate(geom, angle, origin='center', use_radians=False) + new_polygon = geom.exterior.coords[:] + + new_polygons.append(new_polygon) + return new_polygons + + +def rasterize(image, polygons): + polygons = [shapely.geometry.Polygon(polygon) for polygon in polygons] + raster = torch_lydorn.torchvision.transforms.Rasterize(fill=True, edges=False, vertices=False, line_width=4, antialiasing=True)(image, polygons) + raster = raster[:, :, 0] + raster = ndimage.gaussian_filter(raster, sigma=1) # Simulates blurriness of overhead image, which leads to blurriness of segmentation + return raster + + +def plot(image, out_filepath, dpi=300): + height = image.shape[0] + width = image.shape[1] + f, axis = plt.subplots(1, 1, figsize=(width, height), dpi=dpi) + + # Plot image + axis.imshow(image, cmap="gray") + + axis.autoscale(False) + axis.axis('equal') + axis.axis('off') + plt.subplots_adjust(left=0, right=1, top=1, bottom=0) # Plot without margins + plt.savefig(out_filepath, transparent=True, dpi=dpi) + plt.close() + + +def main(): + shape = (100, 100) + samples = 10 + all_rasters: ndarray = np.empty((*shape, samples)) + methods = [ + "marching_squares", + "border_following", + "rasterio" + ] + + for s in range(samples): + polygons = create_polygons() + polygons = displace_polygons(polygons, max_global=0, max_polygon=3, max_vertex=0.5, max_rot_deg=1) # Simulates imperfect ground truth annotations + raster = rasterize(np.zeros(shape), polygons) / 255 + all_rasters[:, :, s] = raster + plot(raster, f"rounded_corners_sample_{s:02d}.png", dpi=1) + + mean_raster_s = np.mean(all_rasters[:, :, :(s+1)], axis=-1) # Simulates training to reduce average loss over all (noisy) ground truth + plot(mean_raster_s, f"rounded_corners_avg_{s:02d}.png", dpi=1) + + for m in methods: + contours = vectorization_ambiguities.detect_contours(mean_raster_s, method=m) + vectorization_ambiguities.plot(mean_raster_s, contours, f"rounded_corners_avg_contour_{m}.pdf", linewidth=60, dpi=1, grid=False) + + +if __name__ == "__main__": + main() diff --git a/paper/show_result_image.py b/paper/show_result_image.py new file mode 100644 index 0000000000000000000000000000000000000000..8e36976c5624ef5b4f7bf8c570084dbde5d8fb5f --- /dev/null +++ b/paper/show_result_image.py @@ -0,0 +1,96 @@ +import matplotlib.pyplot as plt + +from pycocotools.coco import COCO +import skimage.io as io +import os + + +def plot_result(output_filename_format, im, image_id, coco): + print("Plotting image" + output_filename_format.format(image_id)) + annotation_ids = coco.getAnnIds(imgIds=image_id) + annotations = coco.loadAnns(annotation_ids) + dpi = 100 + f, axis = plt.subplots(1, 1, figsize=(im.shape[1] / dpi, im.shape[0] / dpi), dpi=dpi) + axis.imshow(im) + coco.showAnns(annotations) + axis.autoscale(False) + axis.axis('equal') + axis.axis('off') + axis.set_xlim([0, im.shape[1]]) + axis.set_ylim([im.shape[0], 0]) + plt.subplots_adjust(left=0, right=1, top=1, bottom=0) # Plot without margins + plt.savefig(output_filename_format.format(image_id)) + plt.show() + + +def show_result_image(val_annotations_filepath, val_images_dirpath, info_list, image_id_list): + coco_gt = COCO(val_annotations_filepath) + image_ids = coco_gt.getImgIds(catIds=coco_gt.getCatIds()) + + for info in info_list: + if info["annotations_filepath"] == "gt_annotations": + coco_dt = coco_gt + else: + coco_dt = coco_gt.loadRes(info["annotations_filepath"]) + + for image_id in image_id_list: + assert image_id in image_ids, "image_id is invalid" + img_info = coco_gt.loadImgs(image_id)[0] + image_path = os.path.join(val_images_dirpath, img_info["file_name"]) + im = io.imread(image_path) + + plot_result(info["output_filename_format"], im, image_id, coco_dt) + + +def main(): + val_annotations_filepath = "/data/data/mapping_challenge_dataset/raw/val/annotation.json" + val_images_dirpath = "/data/data/mapping_challenge_dataset/raw/val/images" + + # Mapping challenge: + info_list = [ + { + "output_filename_format": "crowdai_gt_{:012d}.pdf", + "annotations_filepath": "gt_annotations" + }, + # { + # "output_filename_format": "crowdai_unet_resnet101_{:012d}.poly_viz.acm.tol_1.pdf", + # "annotations_filepath": "/data/data/mapping_challenge_dataset/eval_runs/mapping_dataset.unet_resnet101_pretrained.train_val | 2020-05-21 08:32:48/test.annotation.poly.acm.tol_1.json" + # }, + # { + # "output_filename_format": "crowdai_polymapper_{:012d}.poly_viz.pdf", + # "annotations_filepath": "/data/data/mapping_challenge_dataset/eval_runs/mapping_dataset.polymapper | 0000-00-00 00:00:00/test.annotation.poly.json" + # }, + # { + # "output_filename_format": "crowdai_li_{:012d}.poly_viz.pdf", + # "annotations_filepath": "/data/data/mapping_challenge_dataset/eval_runs/mapping_dataset.mu | 0000-00-00 00:00:00/test.annotation.poly.json" + # }, + + { + "output_filename_format": "crowdai_unet_resnet101_full_{:012d}.poly_viz.acm.tol_0.125.pdf", + "annotations_filepath": "/data/data/mapping_challenge_dataset/eval_runs/mapping_dataset.unet_resnet101_pretrained.train_val | 2020-05-21 08:32:48/test.annotation.poly.acm.tol_0.125.json" + }, + { + "output_filename_format": "crowdai_unet_resnet101_full_{:012d}.poly_viz.simple.tol_0.125.pdf", + "annotations_filepath": "/data/data/mapping_challenge_dataset/eval_runs/mapping_dataset.unet_resnet101_pretrained.train_val | 2020-05-21 08:32:48/test.annotation.poly.simple.tol_0.125.json" + }, + { + "output_filename_format": "crowdai_unet_resnet101_no_field_{:012d}.poly_viz.simple.tol_0.125.pdf", + "annotations_filepath": "/data/data/mapping_challenge_dataset/eval_runs/mapping_dataset.unet_resnet101_pretrained.field_off.train_val | 2020-05-21 08:33:20/test.annotation.poly.simple.tol_0.125.json" + }, + { + "output_filename_format": "crowdai_unet16_no_coupling_losses_{:012d}.poly_viz.simple.tol_0.125.pdf", + "annotations_filepath": "/data/data/mapping_challenge_dataset/eval_runs/mapping_dataset.unet16.coupling_losses_off.train_val | 2020-03-01 13:27:45/test.annotation.poly.simple.tol_0.125.json" + }, + { + "output_filename_format": "crowdai_unet16_no_coupling_losses_{:012d}.poly_viz.acm.tol_0.125.pdf", + "annotations_filepath": "/data/data/mapping_challenge_dataset/eval_runs/mapping_dataset.unet16.coupling_losses_off.train_val | 2020-03-01 13:27:45/test.annotation.poly.acm.tol_0.125.json" + }, + ] + image_id_list = [21219, 443, 371, 265, 18205, 1, 2, 3, 4] + # image_id_list = [443] + + show_result_image(val_annotations_filepath, val_images_dirpath, info_list, image_id_list) + + +if __name__ == '__main__': + main() diff --git a/paper/vectorization_ambiguities.py b/paper/vectorization_ambiguities.py new file mode 100644 index 0000000000000000000000000000000000000000..7468716bf902531e402422c8f14ee5b02cd614ea --- /dev/null +++ b/paper/vectorization_ambiguities.py @@ -0,0 +1,80 @@ +import numpy as np +import matplotlib.pyplot as plt +import skimage.measure +import cv2 + + +def create_seg(): + seg = np.zeros((6, 8)) + # Triangle: + seg[1, 4] = 1 + seg[2, 3:5] = 1 + seg[3, 2:5] = 1 + seg[4, 1:5] = 1 + # L extension: + seg[3:5, 5:7] = 1 + return seg + + +def detect_contours(seg, method): + if method == "marching_squares": + contours = skimage.measure.find_contours(seg, 0.5, fully_connected='low', positive_orientation='high') + elif method == "border_following": + u, contours, _ = cv2.findContours((0.5 < seg).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) + contours = [contour[:, 0, ::-1] for contour in contours] + contours = [np.concatenate((contour, contour[0:1, :]), axis=0) for contour in contours] + elif method == "rasterio": + import rasterio.features + shapes = rasterio.features.shapes((0.5 < seg).astype(np.uint8)) + contours = [] + for shape in shapes: + for coords in shape[0]["coordinates"][1:]: + contours.append(np.array(coords)[:, ::-1] - 0.5) + else: + raise ValueError(f"Method {method} not recognized!") + return contours + + +def plot(image, contours, out_filepath, linewidth=6, dpi=300, grid=True): + height = image.shape[0] + width = image.shape[1] + f, axis = plt.subplots(1, 1, figsize=(width, height), dpi=dpi) + + # Plot image + axis.imshow(image, cmap="gray") + + # Grid lines + if grid: + for p in range(image.shape[1]): + plt.axvline(p + 0.5, color=[0.5]*3, linewidth=0.5) + for p in range(image.shape[0]): + plt.axhline(p + 0.5, color=[0.5]*3, linewidth=0.5) + + # Plot contours + for contour in contours: + plt.plot(contour[:, 1], contour[:, 0], linewidth=linewidth) + + axis.autoscale(False) + axis.axis('equal') + axis.axis('off') + plt.subplots_adjust(left=0, right=1, top=1, bottom=0) # Plot without margins + plt.savefig(out_filepath, transparent=True, dpi=dpi) + plt.close() + + +def main(): + seg = create_seg() + + methods = [ + "marching_squares", + "border_following", + "rasterio" + ] + + for m in methods: + contours = detect_contours(seg, method=m) + plot(seg, contours, f"vectorization_ambiguities_{m}.pdf", dpi=300) + + +if __name__ == "__main__": + main() diff --git a/precidt.py b/precidt.py new file mode 100644 index 0000000000000000000000000000000000000000..197054bb916ff4213220ea373723f4ea687fddb1 --- /dev/null +++ b/precidt.py @@ -0,0 +1,14 @@ +''' +Author: Egrt +Date: 2022-03-18 10:12:19 +LastEditors: Egrt +LastEditTime: 2022-03-18 11:07:27 +FilePath: \Polygonization-by-Frame-Field-Learning\precidt.py +''' +import os +images_dir = 'images/val/images' +images_path = os.listdir(images_dir) +for index, image_name in enumerate(images_path): + if index<100: + image_path = os.path.join(images_dir, image_name) + os.system('') \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..a61b1a0928489b87f5611c16e4b8d7e115e7b7c5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,23 @@ +scipy==1.4.1 +numpy==1.22.3 +matplotlib==3.3.2 +opencv_python==4.5.4.60 +torch==1.4.0 +torchvision==0.5.0 +tqdm==4.63.0 +Pillow==8.2.0 +h5py==2.10.0 +gradio==2.5.3 +jsmin==3.0.1 +kornia==0.5.0 +shapely==1.8.1 +skan==0.10.0 +descartes==1.1.0 +multiprocess==0.70.12.2 +dill==0.3.4 +gdal==3.4.1 +rasterio==1.2.10 +overpy==0.6 +pyproj==2.5.0 +fiona==1.8.21 +zipp==3.4.0 \ No newline at end of file diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 0000000000000000000000000000000000000000..9b4cd5825f1adfb000d1f099b6052385442c621c --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,32 @@ +# Introduction + +This "scripts" folder contains stand-alone scripts for some useful tasks detailed below. + +## mask_to_json.py + +Use this script to convert .png segmentation masks from the Open Solution from the CrowdAI challenge +(https://github.com/neptune-ai/open-solution-mapping-challenge) +to the COCO .json format with RLE mask encododing. +Run as: +``` +mask_to_json.py --mask_dirpath --output_filepath +``` + +## plot_framefield.py + +Use this script to plot a framefield saved as a .npy file. Can be useful for visualization. +Explanation about its arguments can be accessed with: +``` +mask_to_json.py --help +``` + +## ply_to_json.py + +Use this script to convert .ply segmentation polygons from the paper +"Li, M., Lafarge, F., Marlet, R.: Approximating shapes in images with low-complexity polygons. In: CVPR (2020)" +to the COCO .json format with [polygon] mask encoding. In order to fill the score field of each annotation in the COCO format, we also need access to segmentation masks. + +Run as +``` +ply_to_json.py --ply_dirpath --mask_dirpath --output_filepath +``` \ No newline at end of file diff --git a/scripts/clean_coco.py b/scripts/clean_coco.py new file mode 100644 index 0000000000000000000000000000000000000000..80964107368a3c75db31f4c34c81659a0285a0e9 --- /dev/null +++ b/scripts/clean_coco.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 + +################################################################### +# Use this script to clean COCO mask detection. It merges detections that touch each other taking into account their scores. +# Example use: +# python clean_coco.py --gt_filepath /home/lydorn/data/mapping_challenge_dataset/raw/val/annotation.json --in_filepath "/home/lydorn/data/mapping_challenge_dataset/eval_runs/mapping_dataset.open_solution_full | 0000-00-00 00:00:00/test.annotation.seg.json" --out_filepath "/home/lydorn/data/mapping_challenge_dataset/eval_runs/mapping_dataset.open_solution_full | 0000-00-00 00:00:00/test.annotation.seg_cleaned.json" +################################################################### + +import argparse +from pycocotools.coco import COCO +from pycocotools import mask as cocomask +from multiprocess import Pool +import numpy as np + +import json +from tqdm import tqdm + +import torch + +from frame_field_learning import inference, polygonize_acm, save_utils + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--gt_filepath', + required=True, + type=str, + help='Filepath of the ground truth annotations in COCO format (.json file).') + argparser.add_argument( + '--in_filepath', + required=True, + type=str, + help='Filepath of the input mask annotations in COCO format (.json file).') + argparser.add_argument( + '--out_filepath', + required=True, + type=str, + help='Filepath of the output polygon annotations in COCO format (.json file).') + args = argparser.parse_args() + return args + + +def clean_one(im_data): + img, dts = im_data + seg_image = np.zeros((img["height"], img["width"]), dtype=np.float) + mask_image = np.zeros((img["height"], img["width"]), dtype=np.uint8) + + # Rank detections by score first + dts = sorted(dts, key=lambda k: k['score'], reverse=True) + for dt in dts: + if 0 < dt["score"]: + dt_mask = cocomask.decode(dt["segmentation"]) + intersection = mask_image * dt_mask + if intersection.sum() == 0: + mask_image = np.maximum(mask_image, dt_mask) + seg_image = np.maximum(seg_image, dt_mask * dt["score"]) + + sample = { + "seg": torch.tensor(seg_image[None, :, :]), + "seg_mask": torch.tensor(mask_image).int(), + "image_id": torch.tensor(img["id"]) + } + coco_ann = save_utils.seg_coco(sample) + return coco_ann + + +def clean_masks(gt_filepath, in_filepath, out_filepath): + coco_gt = COCO(gt_filepath) + coco_dt = coco_gt.loadRes(in_filepath) + + # --- Clean input COCO mask detections --- # + img_ids = sorted(coco_dt.getImgIds()) + im_data_list = [] + for img_id in img_ids: + img = coco_gt.loadImgs(img_id)[0] + dts = coco_dt.loadAnns(coco_dt.getAnnIds(imgIds=img_id)) + im_data_list.append((img, dts)) + + pool = Pool() + output_annotations_list = list(tqdm(pool.imap(clean_one, im_data_list), desc="Clean detections", total=len(im_data_list))) + output_annotations = [output_annotation for output_annotations in output_annotations_list for output_annotation in output_annotations] + + print("Saving output...") + with open(out_filepath, 'w') as outfile: + json.dump(output_annotations, outfile) + + +if __name__ == "__main__": + args = get_args() + clean_masks(args.gt_filepath, args.in_filepath, args.out_filepath) diff --git a/scripts/convert_seg_single_channel.py b/scripts/convert_seg_single_channel.py new file mode 100644 index 0000000000000000000000000000000000000000..b7b32180721c795b4ec86f3d6ae3340de3fdad56 --- /dev/null +++ b/scripts/convert_seg_single_channel.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +import os + +import argparse +import skimage.io +import skimage.external.tifffile +from multiprocess import Pool +from functools import partial +from tqdm import tqdm +import cv2 + +try: + __import__("frame_field_learning.local_utils") +except ImportError: + print("ERROR: The frame_field_learning package is not installed! " + "Execute script setup.sh to install local dependencies such as frame_field_learning in develop mode.") + exit() + +from lydorn_utils import print_utils + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--filepath', + required=True, + type=str, + nargs='*', + help='Path(s) to tiff seg RGB images to convert to single channel segmentation map (only keep the first channel).') + argparser.add_argument( + '--out_dirpath', + type=str, + help='Path to the output directory for the converted images.') + + args = argparser.parse_args() + return args + + +def convert_one(filepath, out_dirpath): + image = skimage.io.imread(filepath) + gray_image = image[:, :, 0] + + basename = os.path.basename(filepath) + name = basename.split(".")[0] + out_filepath = os.path.join(out_dirpath, name + ".png") + os.makedirs(os.path.dirname(out_filepath), exist_ok=True) + + cv2.imwrite(out_filepath, gray_image, [cv2.IMWRITE_PNG_COMPRESSION, 9]) + + +def main(): + args = get_args() + print_utils.print_info(f"INFO: converting {len(args.filepath)} seg images.") + + pool = Pool() + list(tqdm(pool.imap(partial(convert_one, out_dirpath=args.out_dirpath), args.filepath), desc="RGB to Gray", total=len(args.filepath))) + + +if __name__ == '__main__': + main() diff --git a/scripts/eval_shapefiles.py b/scripts/eval_shapefiles.py new file mode 100644 index 0000000000000000000000000000000000000000..927e2ff091aa18b80984e452ac0c5688ee9c461a --- /dev/null +++ b/scripts/eval_shapefiles.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python3 +import os + +import fiona +import numpy as np +import shapely.geometry +import argparse +from multiprocess import Pool +from functools import partial +from tqdm import tqdm + +try: + __import__("frame_field_learning.local_utils") +except ImportError: + print("ERROR: The frame_field_learning package is not installed! " + "Execute script setup.sh to install local dependencies such as frame_field_learning in develop mode.") + exit() + +from lydorn_utils import polygon_utils, python_utils, print_utils, geo_utils + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--im_filepath', + required=True, + type=str, + nargs='*', + help='Path(s) to tiff images (to get projection info).') + argparser.add_argument( + '--gt_filepath', + required=True, + type=str, + nargs='*', + help='Path(s) to ground truth.') + argparser.add_argument( + '--pred_filepath', + required=True, + type=str, + nargs='*', + help='Path(s) to predictions.') + argparser.add_argument( + '--overwrite', + action='store_true', + help='The default behavior is to not re-compute a metric if it already has been computed. ' + 'Add the --overwrite flag if you want to overwrite existing metric results.') + + args = argparser.parse_args() + return args + + +def load_geom(geom_filepath, im_filepath): + ext = os.path.splitext(geom_filepath)[-1] + if ext == ".geojson": + file = fiona.open(geom_filepath) + assert len(file) == 1, "There should be only one feature per file" + feat = file[0] + geoms = shapely.geometry.shape(feat["geometry"]) + elif ext == ".shp": + polygons, _ = geo_utils.get_polygons_from_shapefile(im_filepath, geom_filepath, progressbar=False) + geoms = shapely.geometry.collection.GeometryCollection([shapely.geometry.Polygon(polygon[:, ::-1]) for polygon in polygons]) + else: + raise ValueError(f"Geometry can not be loaded from a {ext} file.") + return geoms + + +def extract_name(filepath): + basename = os.path.basename(filepath) + name = basename.split(".")[0] + return name + + +def match_im_gt_pred(im_filepaths, gt_filepaths, pred_filepaths): + im_names = list(map(extract_name, im_filepaths)) + gt_names = list(map(extract_name, gt_filepaths)) + + im_gt_pred_filepaths = [] + for pred_filepath in pred_filepaths: + name = extract_name(pred_filepath) + im_index = im_names.index(name) + gt_index = gt_names.index(name) + im_gt_pred_filepaths.append((im_filepaths[im_index], gt_filepaths[gt_index], pred_filepath)) + + return im_gt_pred_filepaths + + +def eval_one(im_gt_pred_filepath, overwrite=False): + im_filepath, gt_filepath, pred_filepath = im_gt_pred_filepath + metrics_filepath = os.path.splitext(pred_filepath)[0] + ".metrics.json" + iou_filepath = os.path.splitext(pred_filepath)[0] + ".iou.json" + + metrics = False + iou = False + if not overwrite: + # Try reading metrics and iou json + metrics = python_utils.load_json(metrics_filepath) + iou = python_utils.load_json(iou_filepath) + + if not metrics or not iou: + # Have to compute at least one so load geometries + gt_polygons = load_geom(gt_filepath, im_filepath) + fixed_gt_polygons = polygon_utils.fix_polygons(gt_polygons, + buffer=0.0001) # Buffer adds vertices but is needed to repair some geometries + pred_polygons = load_geom(pred_filepath, im_filepath) + fixed_pred_polygons = polygon_utils.fix_polygons(pred_polygons) + + if not metrics: + # Compute and save metrics + max_angle_diffs = polygon_utils.compute_polygon_contour_measures(fixed_pred_polygons, fixed_gt_polygons, + sampling_spacing=1.0, min_precision=0.5, + max_stretch=2, progressbar=True) + max_angle_diffs = [value for value in max_angle_diffs if value is not None] + max_angle_diffs = np.array(max_angle_diffs) + max_angle_diffs = max_angle_diffs * 180 / np.pi # Convert to degrees + metrics = { + "max_angle_diffs": list(max_angle_diffs) + } + python_utils.save_json(metrics_filepath, metrics) + + if not iou: + fixed_gt_polygon_collection = shapely.geometry.collection.GeometryCollection(fixed_gt_polygons) + fixed_pred_polygon_collection = shapely.geometry.collection.GeometryCollection(fixed_pred_polygons) + intersection = fixed_gt_polygon_collection.intersection(fixed_pred_polygon_collection).area + union = fixed_gt_polygon_collection.union(fixed_pred_polygon_collection).area + iou = { + "intersection": intersection, + "union": union + } + python_utils.save_json(iou_filepath, iou) + + return { + "metrics": metrics, + "iou": iou, + } + + +def main(): + args = get_args() + print_utils.print_info(f"INFO: evaluating {len(args.pred_filepath)} predictions.") + + # Match files together + im_gt_pred_filepaths = match_im_gt_pred(args.im_filepath, args.gt_filepath, args.pred_filepath) + + pool = Pool() + metrics_iou_list = list(tqdm(pool.imap(partial(eval_one, overwrite=args.overwrite), im_gt_pred_filepaths), desc="Compute eval metrics", total=len(im_gt_pred_filepaths))) + + # Aggregate metrics and IoU + aggr_metrics = { + "max_angle_diffs": [] + } + aggr_iou = { + "intersection": 0, + "union": 0 + } + for metrics_iou in metrics_iou_list: + if metrics_iou["metrics"]: + aggr_metrics["max_angle_diffs"] += metrics_iou["metrics"]["max_angle_diffs"] + if metrics_iou["iou"]: + aggr_iou["intersection"] += metrics_iou["iou"]["intersection"] + aggr_iou["union"] += metrics_iou["iou"]["union"] + aggr_iou["iou"] = aggr_iou["intersection"] / aggr_iou["union"] + + aggr_metrics_filepath = os.path.join(os.path.dirname(args.pred_filepath[0]), "aggr_metrics.json") + aggr_iou_filepath = os.path.join(os.path.dirname(args.pred_filepath[0]), "aggr_iou.json") + python_utils.save_json(aggr_metrics_filepath, aggr_metrics) + python_utils.save_json(aggr_iou_filepath, aggr_iou) + + +if __name__ == '__main__': + main() diff --git a/scripts/masks_to_json.py b/scripts/masks_to_json.py new file mode 100644 index 0000000000000000000000000000000000000000..85debdf7e5aa0507f27b97b8a3662195d105eb01 --- /dev/null +++ b/scripts/masks_to_json.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 + +################################################################### +# Use this script to convert .png masks from the Open Solution from the CrowdAI challenge: +# https://github.com/neptune-ai/open-solution-mapping-challenge +# to the COCO .json format +################################################################### + +import fnmatch +import os +import argparse + +import skimage.io +import skimage.morphology +import pycocotools.mask +import numpy as np +from tqdm import tqdm +import json +import skimage.measure + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--mask_dirpath', + required=True, + type=str, + help='Path to the directory where the mask .png files are.') + argparser.add_argument( + '--output_filepath', + required=True, + type=str, + help='Filepath of the final .json.') + args = argparser.parse_args() + return args + + +def masks_to_json(mask_dirpath, output_filepath): + filenames = fnmatch.filter(os.listdir(mask_dirpath), "*.png") + + annotations = [] + for filename in tqdm(filenames, desc="Process masks:"): + image_id = int(os.path.splitext(filename)[0]) + seg = skimage.io.imread(os.path.join(mask_dirpath, filename)) + labels = skimage.morphology.label(seg) + properties = skimage.measure.regionprops(labels, cache=True) + for i, contour_props in enumerate(properties): + skimage_bbox = contour_props["bbox"] + coco_bbox = [skimage_bbox[1], skimage_bbox[0], + skimage_bbox[3] - skimage_bbox[1], skimage_bbox[2] - skimage_bbox[0]] + + image_mask = labels == (i + 1) # The mask has to span the whole image + rle = pycocotools.mask.encode(np.asfortranarray(image_mask)) + rle["counts"] = rle["counts"].decode("utf-8") + annotation = { + "category_id": 100, # Building + "bbox": coco_bbox, + "segmentation": rle, + "score": 1.0, + "image_id": image_id} + annotations.append(annotation) + + with open(output_filepath, 'w') as outfile: + json.dump(annotations, outfile) + + +if __name__ == "__main__": + args = get_args() + mask_dirpath = args.mask_dirpath + output_filepath = args.output_filepath + masks_to_json(mask_dirpath, output_filepath) diff --git a/scripts/plot_framefield.py b/scripts/plot_framefield.py new file mode 100644 index 0000000000000000000000000000000000000000..4f2a1f78741b4f68fa4090a250d79300ac512aec --- /dev/null +++ b/scripts/plot_framefield.py @@ -0,0 +1,65 @@ +import argparse + +import numpy as np +import matplotlib.pyplot as plt + +from frame_field_learning import plot_utils + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '-f', '--filepath', + required=True, + type=str, + help='Path to the .npy to plot the framefield from.') + argparser.add_argument( + '-o', '--out_filepath', + required=True, + type=str, + help='Path to save the image.') + + args = argparser.parse_args() + return args + + +def save_plot_framefield(framefield, out_filepath): + # Setup plot + height = framefield.shape[0] + width = framefield.shape[1] + f, axis = plt.subplots(1, 1, figsize=(width // 10, height // 10)) + + # axis.imshow(im) + transparent_im = np.zeros((height, width, 4)) + axis.imshow(transparent_im) + + framefield_stride = 8 + plot_utils.plot_framefield(axis, framefield, framefield_stride, alpha=1, width=2) + + axis.autoscale(False) + axis.axis('equal') + axis.axis('off') + + # f.tight_layout() + plt.subplots_adjust(left=0, right=1, top=1, bottom=0) # Plot without margins + plt.savefig(out_filepath, transparent=True) + plt.close() + # plt.show() + + +def main(): + # --- Process args --- # + args = get_args() + + # Read + framefield = np.load(args.filepath) + + # Cut bbox + # bbox = [2700, 2300, 3200, 2900] + # framefield = framefield[bbox[0]:bbox[2], bbox[1]:bbox[3]] + + save_plot_framefield(framefield, args.out_filepath) + + +if __name__ == '__main__': + main() diff --git a/scripts/ply_to_json.py b/scripts/ply_to_json.py new file mode 100644 index 0000000000000000000000000000000000000000..379724cb7bb120f07f71cda0b951fed974e5c3f8 --- /dev/null +++ b/scripts/ply_to_json.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 + +################################################################### +# Use this script to convert .ply results from +# Li, M., Lafarge, F., Marlet, R.: Approximating shapes in images with low-complexitypolygons. In: CVPR (2020) +# to COCO .json format +################################################################### + + +import fnmatch +import os +import argparse + +import skimage.io +import skimage.morphology +import numpy as np +from tqdm import tqdm +import json +from plyfile import PlyData, PlyElement +import shapely.geometry +import shapely.ops + +from frame_field_learning import polygonize_utils, save_utils + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--ply_dirpath', + required=True, + type=str, + help='Path to the directory where the .ply are.') + argparser.add_argument( + '--mask_dirpath', + required=True, + type=str, + help='Path to the directory where the masks are (used to compute probability of each polygonal partition.') + argparser.add_argument( + '--output_filepath', + required=True, + type=str, + help='Filepath of the final .json.') + args = argparser.parse_args() + return args + + +def ply_to_json(ply_dirpath, mask_dirpath, output_filepath, mask_filename_format="{:012d}.png"): + filenames = fnmatch.filter(os.listdir(ply_dirpath), "*.ply") + + all_annotations = [] + for filename in tqdm(filenames, desc="Process ply files:"): + image_id = int(os.path.splitext(filename)[0]) + + # Load .ply + plydata = PlyData.read(os.path.join(ply_dirpath, filename)) + x = plydata['vertex']['x'] + y = 299 - plydata['vertex']['y'] + pos = np.stack([x, y], axis=1) + vertex1 = plydata['edge'].data["vertex1"] + vertex2 = plydata['edge'].data["vertex2"] + edge_index = np.stack([vertex1, vertex2], axis=1) + edge = pos[edge_index] + linestrings = [] + for e in edge: + linestrings.append(shapely.geometry.LineString(e)) + + # Load mask + mask_filename = mask_filename_format.format(image_id) + mask = 0 < skimage.io.imread(os.path.join(mask_dirpath, mask_filename)) + + # Convert to polygons + polygons = shapely.ops.polygonize(linestrings) + + # Remove low prob polygons + filtered_polygons = [] + filtered_polygon_probs = [] + for polygon in polygons: + prob = polygonize_utils.compute_geom_prob(polygon, mask) + # print("simple:", np_indicator.min(), np_indicator.mean(), np_indicator.max(), prob) + if 0.5 < prob: + filtered_polygons.append(polygon) + filtered_polygon_probs.append(prob) + + annotations = save_utils.poly_coco(filtered_polygons, filtered_polygon_probs, image_id) + all_annotations.extend(annotations) + + with open(output_filepath, 'w') as outfile: + json.dump(all_annotations, outfile) + + +if __name__ == "__main__": + args = get_args() + ply_dirpath = args.ply_dirpath + mask_dirpath = args.mask_dirpath + output_filepath = args.output_filepath + ply_to_json(ply_dirpath, mask_dirpath, output_filepath) diff --git a/scripts/polygonize_coco.py b/scripts/polygonize_coco.py new file mode 100644 index 0000000000000000000000000000000000000000..24a0e89b7c0448d97fbf7f16e7d81168f0a98083 --- /dev/null +++ b/scripts/polygonize_coco.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 + +################################################################### +# Use this script to polygonize binary mask detection in COCO format (for example frome the Open Solution from the CrowdAI challenge: +# https://github.com/neptune-ai/open-solution-mapping-challenge) +# using the frame field polygonization method and save the output in COCO format. +# Example use: +# python polygonize_coco.py --run_dirpath "/home/lydorn/repos/lydorn/Polygonization-by-Frame-Field-Learning/frame_field_learning/runs/mapping_dataset.unet_resnet101_pretrained.train_val | 2020-09-07 11:28:51" --images_dirpath "/home/lydorn/data/mapping_challenge_dataset/raw/val/images" --gt_filepath /home/lydorn/data/mapping_challenge_dataset/raw/val/annotation.json --in_filepath "/home/lydorn/data/mapping_challenge_dataset/eval_runs/mapping_dataset.open_solution | 0000-00-00 00:00:00/test.annotation.seg.json" --out_filepath "/home/lydorn/data/mapping_challenge_dataset/eval_runs/mapping_dataset.open_solution | 0000-00-00 00:00:00/test.annotation.poly.json" +################################################################### + +import argparse +from pycocotools.coco import COCO +from pycocotools import mask as cocomask +from pycocotools.cocoeval import COCOeval +import numpy as np +import skimage.io +import matplotlib.pyplot as plt +import pylab +import random +import os +import json +import sys +from tqdm import tqdm + +import torch + +from frame_field_learning.model import FrameFieldModel +from frame_field_learning import inference, polygonize_acm, data_transforms, save_utils, polygonize_utils +from lydorn_utils import run_utils, print_utils, python_utils +from backbone import get_backbone +import torch_lydorn + +pylab.rcParams['figure.figsize'] = (8.0, 10.0) + +# Image stats for the open challenge dataset: +image_mean = [0.30483739, 0.35143595, 0.3973895] +image_std = [0.16362707, 0.15187606, 0.14273278] + +polygonize_config = { + "steps": 500, + "data_level": 0.5, + "data_coef": 0.1, + "length_coef": 0.4, + "crossfield_coef": 0.5, + "poly_lr": 0.001, + "warmup_iters": 499, + "warmup_factor": 0.1, + "device": "cuda", + "tolerance": 0.125, + "seg_threshold": 0.5, + "min_area": 10 +} + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--run_dirpath', + required=True, + type=str, + help='Full path to the run directory to use for frame field prediction (needed for frame field polygonization).') + argparser.add_argument( + '--images_dirpath', + required=True, + type=str, + help='Path to the images directory to use for frame field prediction (needed for frame field polygonization).') + argparser.add_argument( + '--gt_filepath', + required=True, + type=str, + help='Filepath of the ground truth annotations in COCO format (.json file).') + argparser.add_argument( + '--in_filepath', + required=True, + type=str, + help='Filepath of the input mask annotations in COCO format (.json file).') + argparser.add_argument( + '--out_filepath', + required=True, + type=str, + help='Filepath of the output polygon annotations in COCO format (.json file).') + argparser.add_argument( + '--batch_size', + default=16, + type=int, + help='Batch size for running inference on the model.') + argparser.add_argument( + '--batch_size_mult', + default=64, + type=int, + help='Multiply batch_size by this factor for polygonization.') + args = argparser.parse_args() + return args + + +def list_to_batch(sample_data_list): + tile_data = {} + for key in sample_data_list[0].keys(): + if isinstance(sample_data_list[0][key], list): + tile_data[key] = [item for _tile_data in sample_data_list for item in _tile_data[key]] + elif isinstance(sample_data_list[0][key], torch.Tensor): + tile_data[key] = torch.cat([_tile_data[key] for _tile_data in sample_data_list], dim=0) + else: + raise TypeError(f"Type {type(sample_data_list[0][key])} is not handled!") + return tile_data + + +def run_model(config, model, sample_data_list): + tile_data = list_to_batch(sample_data_list) + tile_data = inference.inference(config, model, tile_data, compute_polygonization=False) + return tile_data + + +def run_polygonization(sample_data_list): + tile_data = list_to_batch(sample_data_list) + # Polygonize input mask with predicted frame field + seg_batch = tile_data["mask_image"] + crossfield_batch = tile_data["crossfield"] + + polygons_batch, _ = polygonize_acm.polygonize(seg_batch, crossfield_batch, polygonize_config) + # Discard the probs computed by polygonize(). They will be computed next using the score_image + + # Convert to COCO format + coco_ann_list = [] + for polygons, img_id, score_image in zip(polygons_batch, tile_data["img_id"], tile_data["score_image"]): + scores = polygonize_utils.compute_geom_prob(polygons, score_image[0, :, :].numpy()) + coco_ann = save_utils.poly_coco(polygons, scores, image_id=img_id) + coco_ann_list.extend(coco_ann) + + return coco_ann_list + + +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) + + +if __name__ == "__main__": + args = get_args() + polygonize_masks(args.run_dirpath, args.images_dirpath, args.gt_filepath, args.in_filepath, args.out_filepath, args.batch_size, args.batch_size_mult) diff --git a/scripts/polygonize_mask.py b/scripts/polygonize_mask.py new file mode 100644 index 0000000000000000000000000000000000000000..86e9e77c92e180d74e4272f7566226efa9369cb0 --- /dev/null +++ b/scripts/polygonize_mask.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 + +################################################################### +# Use this script to extract polygons from binary masks using a model trained for this task. +# I used it to polygonize the original ground truth masks from the Inria Aerial Image Labeling Dataset. +# The first step is to train a network whose input is a binary mask and output is a segmentation + frame field. +# I did this on rasterized OSM annotation corresponding to the Inria dataset +# (so that there is a ground truth for the frame field). +################################################################### + + +import argparse +import sys +import os +import numpy as np +import torch_lydorn +from tqdm import tqdm +import skimage.io +import torch + +try: + __import__("frame_field_learning.local_utils") +except ImportError: + print("ERROR: The frame_field_learning package is not installed! " + "Execute script setup.sh to install local dependencies such as frame_field_learning in develop mode.") + exit() + +from frame_field_learning import data_transforms, polygonize_asm, save_utils, polygonize_acm, measures +from frame_field_learning.model import FrameFieldModel +from frame_field_learning import inference +from frame_field_learning import local_utils + +from torch_lydorn import torchvision +from lydorn_utils import run_utils, geo_utils, polygon_utils +from lydorn_utils import print_utils + +from backbone import get_backbone + +# polygonize_config = { +# "data_level": 0.5, +# "step_thresholds": [0, 500], # From 0 to 500: gradually go from coefs[0] to coefs[1] +# "data_coefs": [0.9, 0.09], +# "length_coefs": [0.1, 0.01], +# "crossfield_coefs": [0.0, 0.05], +# "poly_lr": 0.1, +# "device": "cuda", +# "tolerance": 0.001, +# "seg_threshold": 0.5, +# "min_area": 10, +# } +polygonize_config = { + "steps": 500, + "data_level": 0.5, + "data_coef": 0.1, + "length_coef": 0.4, + "crossfield_coef": 0.5, + "poly_lr": 0.01, + "warmup_iters": 100, + "warmup_factor": 0.1, + "device": "cuda", + "tolerance": 0.5, + "seg_threshold": 0.5, + "min_area": 10 +} + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '-f', '--filepath', + required=True, + type=str, + nargs='*', + help='Filepaths to the binary images to polygonize.') + + argparser.add_argument( + '-r', '--runs_dirpath', + default="runs", + type=str, + help='Directory where runs are recorded (model saves and logs).') + + argparser.add_argument( + '--run_name', + required=True, + type=str, + help='Name of the run to use for predicting the frame field needed by the polygonization algorithm.' + 'That name does not include the timestamp of the folder name: | .') + argparser.add_argument( + '--eval_patch_size', + type=int, + help='When evaluating, patch size the tile split into.') + argparser.add_argument( + '--eval_patch_overlap', + type=int, + help='When evaluating, patch the tile with the specified overlap to reduce edge artifacts when reconstructing ' + 'the whole tile') + argparser.add_argument( + '--out_ext', + type=str, + default="geojson", + choices=['geojson', 'shp'], + help="File extension of the output geometry. 'geojson': GeoJSON, 'shp': shapefile") + + args = argparser.parse_args() + return args + + +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) + + +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) + + +if __name__ == '__main__': + main() diff --git a/scripts/shp_to_json.py b/scripts/shp_to_json.py new file mode 100644 index 0000000000000000000000000000000000000000..1391f575e5018a679f2322f5d1ae6247aaecdaa0 --- /dev/null +++ b/scripts/shp_to_json.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 + +################################################################### +# Use this script to convert shapefiles to the COCO .json format. +################################################################### + +import fnmatch +import os +import argparse +import fiona +import shapely.geometry + + +import skimage.io +import skimage.morphology +import pycocotools.mask +import numpy as np +from tqdm import tqdm +import json +import skimage.measure + + +def get_args(): + argparser = argparse.ArgumentParser(description=__doc__) + argparser.add_argument( + '--shp_dirpath', + required=True, + type=str, + help='Path to the directory where the shapefiles are.') + argparser.add_argument( + '--output_filepath', + required=True, + type=str, + help='Filepath of the final .json.') + args = argparser.parse_args() + return args + + +def shp_to_json(shp_dirpath, output_filepath): + filenames = fnmatch.filter(os.listdir(shp_dirpath), "*.shp") + filenames = sorted(filenames) + + annotations = [] + for filename in tqdm(filenames, desc="Process shapefiles:"): + shapefile = fiona.open(os.path.join(shp_dirpath, filename)) + + polygons = [] + for feature in shapefile: + geometry = shapely.geometry.shape(feature["geometry"]) + if geometry.type == "MultiPolygon": + for polygon in geometry.geoms: + polygons.append(polygon) + elif geometry.type == "Polygon": + polygons.append(geometry) + else: + raise TypeError(f"geometry.type should be either Polygon or MultiPolygon, not {geometry.type}.") + + image_id = int(filename.split(".")[0]) + for polygon in polygons: + bbox = np.round([polygon.bounds[0], polygon.bounds[1], + polygon.bounds[2] - polygon.bounds[0], polygon.bounds[3] - polygon.bounds[1]], 2) + contour = np.array(polygon.exterior.coords) + contour[:, 1] *= -1 # Shapefiles have inverted y axis... + exterior = list(np.round(contour.reshape(-1), 2)) + segmentation = [exterior] + annotation = { + "category_id": 100, # Building + "bbox": list(bbox), + "segmentation": segmentation, + "score": 1.0, + "image_id": image_id} + annotations.append(annotation) + + # seg = skimage.io.imread() + # labels = skimage.morphology.label(seg) + # properties = skimage.measure.regionprops(labels, cache=True) + # for i, contour_props in enumerate(properties): + # skimage_bbox = contour_props["bbox"] + # coco_bbox = [skimage_bbox[1], skimage_bbox[0], + # skimage_bbox[3] - skimage_bbox[1], skimage_bbox[2] - skimage_bbox[0]] + # + # image_mask = labels == (i + 1) # The mask has to span the whole image + # rle = pycocotools.mask.encode(np.asfortranarray(image_mask)) + # rle["counts"] = rle["counts"].decode("utf-8") + # annotation = { + # "category_id": 100, # Building + # "bbox": coco_bbox, + # "segmentation": rle, + # "score": 1.0, + # "image_id": image_id} + # annotations.append(annotation) + + with open(output_filepath, 'w') as outfile: + json.dump(annotations, outfile) + + +if __name__ == "__main__": + args = get_args() + shp_dirpath = args.shp_dirpath + output_filepath = args.output_filepath + shp_to_json(shp_dirpath, output_filepath) diff --git a/singularity/README.md b/singularity/README.md new file mode 100644 index 0000000000000000000000000000000000000000..0fd1272ca7713718fa9c77853fca156e4d016145 --- /dev/null +++ b/singularity/README.md @@ -0,0 +1,65 @@ +# Introduction + +I explain here how to use Singularity with Docker images (on Fedora) if ti is needed. + +# Install Singularity on Fedora + +I had to install these dependencies: +``` +dnf install libarchive-devel +dnf install squashfs-tools +``` + +Then follow [Quick installaton steps](https://www.sylabs.io/guides/2.6/user-guide/quick_start.html#quick-installation-steps) from Singularity 2.6 docs: +``` +git clone https://github.com/sylabs/singularity.git + +cd singularity + +git fetch --all + +git checkout 2.6.0 + +./autogen.sh + +./configure --prefix=/usr/local + +make + +sudo make install +``` + +# Build a Singularity image from a local Docker image: + +The Docker image has to put on a registry for Singularity to use it. Usually this is the Docker hub registry but you can use a local one too: + +Create local registry: +``` +docker run -d -p 5000:5000 --restart=always --name registry registry:2 +``` + +Push local docker container to it: +``` +docker tag localhost:5000/ +docker push localhost:5000/ +``` + +Create a Singularity def file: +``` +Bootstrap: docker +Registry: http://localhost:5000 +Namespace: +From: +``` + +Build singularity image: +``` +sudo SINGULARITY_NOHTTPS=1 singularity build .simg Singularity +``` + +# Send image to cluster + +``` +scp .simg nef-devel: +scp frame-field-learning_1.2.simg nef-devel:frame_field_learning/singularity/ +``` \ No newline at end of file diff --git a/singularity/Singularity b/singularity/Singularity new file mode 100644 index 0000000000000000000000000000000000000000..2103138b645e6bde4d3c291eb21c1f5e2a4b6741 --- /dev/null +++ b/singularity/Singularity @@ -0,0 +1,4 @@ +Bootstrap: docker +Registry: http://localhost:5000 +Namespace: +From: frame-field-learning:1.2 \ No newline at end of file diff --git a/torch_lydorn/LICENSE.txt b/torch_lydorn/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..1c9dd054c47da23697cec2006b97b63290449a21 --- /dev/null +++ b/torch_lydorn/LICENSE.txt @@ -0,0 +1,94 @@ +*** +CONTRAT DE LICENCE DE LOGICIEL + +Logiciel MAPALIGNMENT ©Inria 2018, tout droit réservé, ci-après dénommé "le Logiciel". + +Le Logiciel a été conçu et réalisé par des chercheurs de l’équipe-projet TITANE d’Inria (Institut National de Recherche en Informatique et Automatique). + +Inria, Domaine de Voluceau, Rocquencourt - BP 105 +78153 Le Chesnay Cedex, France + +Inria détient tous les droits de propriété sur le Logiciel. + +Le Logiciel a été déposé auprès de l'Agence pour la Protection des Programmes (APP) sous le numéro . + +Le Logiciel est en cours de développement et Inria souhaite qu'il soit utilisé par la communauté scientifique de façon à le tester et l'évaluer, et afin qu’Inria puisse le cas échéant le faire évoluer. + +A cette fin, Inria a décidé de distribuer le Logiciel. + +Inria concède à l'utilisateur académique, gratuitement, sans droit de sous-licence, pour une période de un (1) an à compter du téléchargement du code source, le droit non-exclusif d'utiliser le Logiciel à fins de recherche. Toute autre utilisation sans l’accord préalable d’Inria est exclue. + +L’utilisateur académique reconnaît expressément avoir reçu d’Inria toutes les informations lui permettant d’apprécier l’adéquation du Logiciel à ses besoins et de prendre toutes les précautions utiles pour sa mise en œuvre et son utilisation. + +Le Logiciel est distribué sous forme d'un fichier source. + +Si le Logiciel est utilisé pour la publication de résultats, l’utilisateur devra citer le Logiciel de la façon suivante : + +@misc{girard2020polygonal, + title={Polygonal Building Segmentation by Frame Field Learning}, + author={Nicolas Girard and Dmitriy Smirnov and Justin Solomon and Yuliya Tarabalka}, + year={2020}, + eprint={2004.14875}, + archivePrefix={arXiv}, + primaryClass={cs.CV} +} + + +Tout utilisateur du Logiciel pourra communiquer ses remarques d'utilisation du Logiciel aux développeurs de MAPALIGNMENT : nicolas.girard@inria.fr + + +L'UTILISATEUR NE PEUT FAIRE NI UTILISATION NI EXPLOITATION NI DISTRIBUTION COMMERCIALE DU LOGICIEL SANS L'ACCORD EXPRÈS PRÉALABLE d’INRIA (stip-sam@inria.fr). +TOUT ACTE CONTRAIRE CONSTITUERAIT UNE CONTREFAÇON. + +LE LOGICIEL EST FOURNI "TEL QU'EN L'ÉTAT" SANS AUCUNE GARANTIE DE QUELQUE NATURE, IMPLICITE OU EXPLICITE, QUANT À SON UTILISATION COMMERCIALE, PROFESSIONNELLE, LÉGALE OU NON, OU AUTRE, SA COMMERCIALISATION OU SON ADAPTATION. + +SAUF LORSQU'EXPLICITEMENT PRÉVU PAR LA LOI, INRIA NE POURRA ÊTRE TENU POUR RESPONSABLE DE TOUT DOMMAGE OU PRÉJUDICE DIRECT,INDIRECT, (PERTES FINANCIÈRES DUES AU MANQUE À GAGNER, À L'INTERRUPTION D'ACTIVITÉS OU À LA PERTE DE DONNÉES, ETC...) DÉCOULANT DE L'UTILISATION DE TOUT OU PARTIE DU LOGICIEL OU DE L'IMPOSSIBILITÉ D'UTILISER CELUI-CI. +  +*** + +SOFTWARE LICENSE AGREEMENT + + +Software MAPALIGNMENT ©Inria – 2018, all rights reserved, hereinafter "the Software". + +This software has been developed by researchers of TITANE project team of Inria (Institut National de Recherche en Informatique et Automatique). + +Inria, Domaine de Voluceau, Rocquencourt - BP 105 +78153 Le Chesnay Cedex, FRANCE + + +Inria holds all the ownership rights on the Software. + +The Software has been registered with the Agence pour la Protection des Programmes (APP) under . + +The Software is still being currently developed. It is the Inria’s aim for the Software to be used by the scientific community so as to test it and, evaluate it so that Inria may improve it. + +For these reasons Inria has decided to distribute the Software. + +Inria grants to the academic user, a free of charge, without right to sublicense non-exclusive right to use the Software for research purposes for a period of one (1) year from the date of the download of the source code. Any other use without of prior consent of Inria is prohibited. + +The academic user explicitly acknowledges having received from Inria all information allowing him to appreciate the adequacy between of the Software and his needs and to undertake all necessary precautions for his execution and use. + +The Software is provided only as a source. + +In case of using the Software for a publication or other results obtained through the use of the Software, user should cite the Software as follows : + +@misc{girard2020polygonal, + title={Polygonal Building Segmentation by Frame Field Learning}, + author={Nicolas Girard and Dmitriy Smirnov and Justin Solomon and Yuliya Tarabalka}, + year={2020}, + eprint={2004.14875}, + archivePrefix={arXiv}, + primaryClass={cs.CV} +} + + +Every user of the Software could communicate to the developers of MAPALIGNMENT [nicolas.girard@inria.fr] his or her remarks as to the use of the Software. + +THE USER CANNOT USE, EXPLOIT OR COMMERCIALY DISTRIBUTE THE SOFTWARE WITHOUT PRIOR AND EXPLICIT CONSENT OF INRIA (stip-sam@inria.fr). ANY SUCH ACTION WILL CONSTITUTE A FORGERY. + +THIS SOFTWARE IS PROVIDED "AS IS" WITHOUT ANY WARRANTIES OF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES,WITH REGARDS TO COMMERCIAL USE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR ADAPTATION. + +UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL INRIA OR THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + + diff --git a/torch_lydorn/README.md b/torch_lydorn/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ca403926fa2cdf1aad27b69456215b931a42e000 --- /dev/null +++ b/torch_lydorn/README.md @@ -0,0 +1,8 @@ +My own extensions to several torch libraries: + +- PyTorch +- Torchvision +- PyTorch Ggeometric +- kornia + +And additional utils related to torch.setup.py \ No newline at end of file diff --git a/torch_lydorn/__init__.py b/torch_lydorn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db91421141251ef0daf1228c530da14288cd715f --- /dev/null +++ b/torch_lydorn/__init__.py @@ -0,0 +1,12 @@ +''' +Author: Egrt +Date: 2022-03-17 17:29:56 +LastEditors: Egrt +LastEditTime: 2022-03-17 17:47:09 +FilePath: \Polygonization-by-Frame-Field-Learning\torch_lydorn\__init__.py +''' +from .torch.nn.functionnal import * +from .torch.utils.complex import * +from .torch.utils.data import * +from .torchvision import * +from .kornia import * \ No newline at end of file diff --git a/torch_lydorn/kornia/__init__.py b/torch_lydorn/kornia/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1a21d1efea612a7bbcc676514a2633818c6e418b --- /dev/null +++ b/torch_lydorn/kornia/__init__.py @@ -0,0 +1,3 @@ +from .filters import * +from .geometry import * +from .augmentation import * diff --git a/torch_lydorn/kornia/augmentation/__init__.py b/torch_lydorn/kornia/augmentation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a98102745d87e8067dcab2ea95587062fde9cb7c --- /dev/null +++ b/torch_lydorn/kornia/augmentation/__init__.py @@ -0,0 +1,3 @@ +from .augmentations import * + +__all__ = ["random_vflip"] diff --git a/torch_lydorn/kornia/augmentation/augmentations.py b/torch_lydorn/kornia/augmentation/augmentations.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ff9af2a65c3f0d64845ebb69520fcf922f127c --- /dev/null +++ b/torch_lydorn/kornia/augmentation/augmentations.py @@ -0,0 +1,78 @@ +from typing import Tuple, List, Union, cast +import torch + +from kornia.geometry.transform import vflip, rotate + +UnionType = Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] + + +def random_rotate(input: torch.Tensor) -> UnionType: + r"""Rotate a tensor image or a batch of tensor images randomly. + Input should be a tensor of shape (C, H, W) or a batch of tensors :math:`(*, C, H, W)`. + Args: + input tensor. + Returns: + torch.Tensor: The rotated input + """ + + if not torch.is_tensor(input): + raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}") + + device: torch.device = input.device + input = input.unsqueeze(0) + input = input.view((-1, (*input.shape[-3:]))) + angle: torch.Tensor = torch.empty(input.shape[0], device=device).uniform_(-180, -180) + + rotated = rotate(input, angle) + return rotated + + +def random_vflip(input: torch.Tensor, p: float = 0.5, return_transform: bool = False) -> UnionType: + r"""Vertically flip a tensor image or a batch of tensor images randomly with a given probability. + Input should be a tensor of shape (C, H, W) or a batch of tensors :math:`(*, C, H, W)`. + Args: + p (float): probability of the image being flipped. Default value is 0.5 + return_transform (bool): if ``True`` return the matrix describing the transformation applied to each + input tensor. + Returns: + torch.Tensor: The vertically flipped input + torch.Tensor: The applied transformation matrix :math: `(*, 3, 3)` if return_transform flag + is set to ``True`` + """ + + if not torch.is_tensor(input): + raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}") + + if not isinstance(p, float): + raise TypeError(f"The probability should be a float number. Got {type(p)}") + + if not isinstance(return_transform, bool): + raise TypeError(f"The return_transform flag must be a bool. Got {type(return_transform)}") + + device: torch.device = input.device + dtype: torch.dtype = input.dtype + + input = input.unsqueeze(0) + input = input.view((-1, (*input.shape[-3:]))) + + probs: torch.Tensor = torch.empty(input.shape[0], device=device).uniform_(0, 1) + + to_flip: torch.Tensor = probs < p + flipped: torch.Tensor = input.clone() + + flipped[to_flip] = vflip(input[to_flip]) + + if return_transform: + + trans_mat: torch.Tensor = torch.eye(3, device=device, dtype=dtype).expand(input.shape[0], -1, -1) + + w: int = input.shape[-2] + flip_mat: torch.Tensor = torch.tensor([[-1, 0, w], + [0, 1, 0], + [0, 0, 1]]) + + trans_mat[to_flip] = flip_mat.to(device).to(dtype) + + return flipped, trans_mat + + return flipped diff --git a/torch_lydorn/kornia/filters/__init__.py b/torch_lydorn/kornia/filters/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cb9274dd3a442f330dbed64be0c6d64e8c8f33c3 --- /dev/null +++ b/torch_lydorn/kornia/filters/__init__.py @@ -0,0 +1,7 @@ +from .sobel import SpatialGradient +from .kernels import get_spatial_gradient_kernel2d + +__all__ = [ + "SpatialGradient", + "get_spatial_gradient_kernel2d", +] \ No newline at end of file diff --git a/torch_lydorn/kornia/filters/kernels.py b/torch_lydorn/kornia/filters/kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..2286807c3b62f78fc8ae19e77cca23917d0abc4e --- /dev/null +++ b/torch_lydorn/kornia/filters/kernels.py @@ -0,0 +1,56 @@ +import torch + +from kornia.filters.kernels import get_sobel_kernel2d, get_sobel_kernel2d_2nd_order, get_diff_kernel2d, get_diff_kernel2d_2nd_order + + +def get_scharr_kernel_3x3() -> torch.Tensor: + """Utility function that returns a sobel kernel of 3x3""" + return torch.tensor([ + [-47., 0., 47.], + [-162., 0., 162.], + [-47., 0., 47.], + ]) + + +def get_scharr_kernel2d(coord: str = "xy") -> torch.Tensor: + assert coord == "xy" or coord == "ij" + kernel_x: torch.Tensor = get_scharr_kernel_3x3() + kernel_y: torch.Tensor = kernel_x.transpose(0, 1) + if coord == "xy": + return torch.stack([kernel_x, kernel_y]) + elif coord == "ij": + return torch.stack([kernel_y, kernel_x]) + + +def get_spatial_gradient_kernel2d(mode: str, order: int, coord: str = "xy") -> torch.Tensor: + r"""Function that returns kernel for 1st or 2nd order image gradients, + using one of the following operators: sobel, diff""" + if mode not in ['sobel', 'diff', 'scharr']: + raise TypeError("mode should be either sobel, diff or scharr. Got {}".format(mode)) + if order not in [1, 2]: + raise TypeError("order should be either 1 or 2\ + Got {}".format(order)) + if mode == 'sobel' and order == 1: + kernel: torch.Tensor = get_sobel_kernel2d() + elif mode == 'sobel' and order == 2: + kernel = get_sobel_kernel2d_2nd_order() + elif mode == 'diff' and order == 1: + kernel = get_diff_kernel2d() + elif mode == 'diff' and order == 2: + kernel = get_diff_kernel2d_2nd_order() + elif mode == 'scharr' and order == 1: + kernel = get_scharr_kernel2d(coord) + else: + raise NotImplementedError("") + return kernel + + +def main(): + import cv2 + k_x, k_y = cv2.getDerivKernels(kx=3, ky=3, dx=1, dy=1, ksize=-1) + print(k_x) + print(k_y) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/torch_lydorn/kornia/filters/sobel.py b/torch_lydorn/kornia/filters/sobel.py new file mode 100644 index 0000000000000000000000000000000000000000..101ac46b2d13038be343ac2609115a1ad2227067 --- /dev/null +++ b/torch_lydorn/kornia/filters/sobel.py @@ -0,0 +1,71 @@ +import kornia + +import torch + +from kornia.filters.kernels import normalize_kernel2d + +from .kernels import get_spatial_gradient_kernel2d + + +class SpatialGradient(torch.nn.Module): + r"""Computes the first order image derivative in both x and y using a Sobel or Scharr + operator. + + Return: + torch.Tensor: the sobel edges of the input feature map. + + Shape: + - Input: :math:`(B, C, H, W)` + - Output: :math:`(B, C, 2, H, W)` + + Examples: + input = torch.rand(1, 3, 4, 4) + output = kornia.filters.SpatialGradient()(input) # 1x3x2x4x4 + """ + + def __init__(self, + mode: str = 'sobel', + order: int = 1, + normalized: bool = True, + coord: str = "xy", + device: str = "cpu", + dtype: torch.dtype = torch.float) -> None: + super(SpatialGradient, self).__init__() + self.normalized: bool = normalized + self.order: int = order + self.mode: str = mode + self.kernel: torch.Tensor = get_spatial_gradient_kernel2d(mode, order, coord) + if self.normalized: + self.kernel = normalize_kernel2d(self.kernel) + # Pad with "replicate for spatial dims, but with zeros for channel + self.spatial_pad = [self.kernel.size(1) // 2, + self.kernel.size(1) // 2, + self.kernel.size(2) // 2, + self.kernel.size(2) // 2] + # Prepare kernel + self.kernel: torch.Tensor = self.kernel.to(device).to(dtype).detach() + self.kernel: torch.Tensor = self.kernel.unsqueeze(1).unsqueeze(1) + self.kernel: torch.Tensor = self.kernel.flip(-3) + return + + def __repr__(self) -> str: + return self.__class__.__name__ + '('\ + 'order=' + str(self.order) + ', ' + \ + 'normalized=' + str(self.normalized) + ', ' + \ + 'mode=' + self.mode + ')' + + def forward(self, inp: torch.Tensor) -> torch.Tensor: # type: ignore + if not torch.is_tensor(inp): + raise TypeError("Input type is not a torch.Tensor. Got {}" + .format(type(inp))) + if not len(inp.shape) == 4: + raise ValueError("Invalid input shape, we expect BxCxHxW. Got: {}" + .format(inp.shape)) + # prepare kernel + b, c, h, w = inp.shape + + # convolve inp tensor with sobel kernel + out_channels: int = 3 if self.order == 2 else 2 + padded_inp: torch.Tensor = torch.nn.functional.pad(inp.reshape(b * c, 1, h, w), + self.spatial_pad, 'replicate')[:, :, None] + return torch.nn.functional.conv3d(padded_inp, self.kernel, padding=0).view(b, c, out_channels, h, w) diff --git a/torch_lydorn/kornia/geometry/__init__.py b/torch_lydorn/kornia/geometry/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..29621e921633da78ab9b159b5b2abdb701e49899 --- /dev/null +++ b/torch_lydorn/kornia/geometry/__init__.py @@ -0,0 +1 @@ +from .transform import * diff --git a/torch_lydorn/kornia/geometry/transform/__init__.py b/torch_lydorn/kornia/geometry/transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f0a52c46fc7d670c9058014fed0ac24606690693 --- /dev/null +++ b/torch_lydorn/kornia/geometry/transform/__init__.py @@ -0,0 +1 @@ +from .imgwarp import * diff --git a/torch_lydorn/kornia/geometry/transform/imgwarp.py b/torch_lydorn/kornia/geometry/transform/imgwarp.py new file mode 100644 index 0000000000000000000000000000000000000000..534d50421cb8d6b1ab171d9a0e32b29ee5a20355 --- /dev/null +++ b/torch_lydorn/kornia/geometry/transform/imgwarp.py @@ -0,0 +1,40 @@ +import numpy as np +import torch + +__all__ = [ + "get_affine_grid", +] + + +def get_affine_grid(tensor: torch.Tensor, angle: torch.Tensor, offset: torch.Tensor, scale: torch.Tensor=None) -> torch.Tensor: + r"""Get rotation sample grid to rotate the image anti-clockwise. + """ + assert len(tensor.shape) == 4, "tensor should have shape (N, C, H, W)" + assert len(angle.shape) == 1, "tensor should have shape (N,)" + assert len(offset.shape) == 2 and offset.shape[1] == 2, "tensor should have shape (N, 2)" + assert tensor.shape[0] == angle.shape[0] == offset.shape[0], "tensor, angle and offset should have the same batch_size" + rad = np.pi * angle / 180 + affine_mat = torch.zeros((rad.shape[0], 2, 3), device=rad.device) + sin_rad = torch.sin(rad) + cos_rad = torch.cos(rad) + if scale is not None: + assert len(scale.shape) == 1, "tensor should have shape (N,)" + affine_mat[:, 0, 0] = cos_rad * scale + affine_mat[:, 1, 1] = cos_rad * scale + else: + affine_mat[:, 0, 0] = cos_rad + affine_mat[:, 1, 1] = cos_rad + affine_mat[:, 0, 1] = - sin_rad + affine_mat[:, 1, 0] = sin_rad + affine_mat[:, :, 2] = offset + # affine_grid = torch.nn.functional.affine_grid(affine_mat, tensor.shape, align_corners=False) # uncomment when using more recent version of PyTorch + affine_grid = torch.nn.functional.affine_grid(affine_mat, tensor.shape) + return affine_grid + + +def main(): + pass + + +if __name__ == '__main__': + main() diff --git a/torch_lydorn/torch/__init__.py b/torch_lydorn/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ffe234098c208f35fee87f0767b1d67a4f9d6154 --- /dev/null +++ b/torch_lydorn/torch/__init__.py @@ -0,0 +1,10 @@ +''' +Author: Egrt +Date: 2022-03-17 17:45:47 +LastEditors: Egrt +LastEditTime: 2022-03-17 17:45:47 +FilePath: \Polygonization-by-Frame-Field-Learning\torch_lydorn\torch\__init__.py +''' +from .nn.functionnal import * +from .utils.complex import * +from .utils.data import * \ No newline at end of file diff --git a/torch_lydorn/torch/nn/functionnal.py b/torch_lydorn/torch/nn/functionnal.py new file mode 100644 index 0000000000000000000000000000000000000000..7831143a6ac207651605150a3e41ffc5d8111264 --- /dev/null +++ b/torch_lydorn/torch/nn/functionnal.py @@ -0,0 +1,75 @@ +import torch + + +def bilinear_interpolate(im, pos, batch=None): + # From https://gist.github.com/peteflorence/a1da2c759ca1ac2b74af9a83f69ce20e + x = pos[:, 1] + y = pos[:, 0] + x0 = torch.floor(x).long() + x1 = x0 + 1 + + y0 = torch.floor(y).long() + y1 = y0 + 1 + + x0_int = torch.clamp(x0, 0, im.shape[-1] - 1) + x1_int = torch.clamp(x1, 0, im.shape[-1] - 1) + y0_int = torch.clamp(y0, 0, im.shape[-2] - 1) + y1_int = torch.clamp(y1, 0, im.shape[-2] - 1) + + if batch is not None: + Ia = im[batch, :, y0_int, x0_int] + Ib = im[batch, :, y1_int, x0_int] + Ic = im[batch, :, y0_int, x1_int] + Id = im[batch, :, y1_int, x1_int] + else: + Ia = im[..., y0_int, x0_int].t() + Ib = im[..., y1_int, x0_int].t() + Ic = im[..., y0_int, x1_int].t() + Id = im[..., y1_int, x1_int].t() + + wa = (x1.float() - x) * (y1.float() - y) + wb = (x1.float() - x) * (y - y0.float()) + wc = (x - x0.float()) * (y1.float() - y) + wd = (x - x0.float()) * (y - y0.float()) + + wa = wa.unsqueeze(1) + wb = wb.unsqueeze(1) + wc = wc.unsqueeze(1) + wd = wd.unsqueeze(1) + + out = wa * Ia + wb * Ib + wc * Ic + wd * Id + + return out + + +def main(): + im = torch.tensor([ + [ + [0, 0.5, 0, 0], + [0.25, 1, 0, 0], + ], + [ + [1, 1, 1, 1], + [1, 1, 1, 1], + ], + [ + [2, 2, 2, 2], + [2, 2, 2, 2], + ] + ], dtype=torch.float) + im = im[:, None, ...] + print(im.shape) + print(im) + pos = torch.tensor([ + [1.0, 0], + [0.5, 0.5], + [0.5, 0.5], + ]) + batch = torch.tensor([0, 1, 2], dtype=torch.long) + + val = bilinear_interpolate(im, pos, batch=batch) + print(val) + + +if __name__ == '__main__': + main() diff --git a/torch_lydorn/torch/utils/complex.py b/torch_lydorn/torch/utils/complex.py new file mode 100644 index 0000000000000000000000000000000000000000..e659e31eebbfd4c92294b0473a8e3aed9e89a111 --- /dev/null +++ b/torch_lydorn/torch/utils/complex.py @@ -0,0 +1,77 @@ +import torch + + +def get_real(t, complex_dim=-1): + return t.select(complex_dim, 0) + + +def get_imag(t, complex_dim=-1): + return t.select(complex_dim, 1) + + +def complex_mul(t1, t2, complex_dim=-1): + t1_real = get_real(t1, complex_dim) + t1_imag = get_imag(t1, complex_dim) + t2_real = get_real(t2, complex_dim) + t2_imag = get_imag(t2, complex_dim) + + ac = t1_real * t2_real + bd = t1_imag * t2_imag + ad = t1_real * t2_imag + bc = t1_imag * t2_real + tr_real = ac - bd + tr_imag = ad + bc + + tr = torch.stack([tr_real, tr_imag], dim=complex_dim) + + return tr + + +def complex_sqrt(t, complex_dim=-1): + sqrt_t_abs = torch.sqrt(complex_abs(t, complex_dim)) + sqrt_t_arg = complex_arg(t, complex_dim) / 2 + # Overwrite t with cos(\theta / 2) + i sin(\theta / 2): + sqrt_t = sqrt_t_abs.unsqueeze(complex_dim) * torch.stack([torch.cos(sqrt_t_arg), torch.sin(sqrt_t_arg)], dim=complex_dim) + return sqrt_t + + +def complex_abs_squared(t, complex_dim=-1): + return get_real(t, complex_dim)**2 + get_imag(t, complex_dim)**2 + + +def complex_abs(t, complex_dim=-1): + return torch.sqrt(complex_abs_squared(t, complex_dim=complex_dim)) + + +def complex_arg(t, complex_dim=-1): + return torch.atan2(get_imag(t, complex_dim), get_real(t, complex_dim)) + + +def main(): + device = None + + t1 = torch.Tensor([ + [2, 0], + [0, 2], + [-1, 0], + [0, -1], + ]).to(device) + t2 = torch.Tensor([ + [2, 0], + [0, 2], + [-1, 0], + [0, -1], + ]).to(device) + complex_dim = -1 + + print(t1.int()) + print(t2.int()) + t1_mul_t2 = complex_mul(t1, t2, complex_dim) + print(t1_mul_t2.int()) + + sqrt_t1_mul_t2 = complex_sqrt(t1_mul_t2) + print(sqrt_t1_mul_t2.int()) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/torch_lydorn/torch/utils/data/__init__.py b/torch_lydorn/torch/utils/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e199d71210370fb52304af5f4b00a9ab7c79513d --- /dev/null +++ b/torch_lydorn/torch/utils/data/__init__.py @@ -0,0 +1,2 @@ +from .dataset import Dataset, __repr__, files_exist +from .makedirs import makedirs diff --git a/torch_lydorn/torch/utils/data/dataset.py b/torch_lydorn/torch/utils/data/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..4a1ad2bb5c7e9fedd7f50ad8b0cd2817cdc742b5 --- /dev/null +++ b/torch_lydorn/torch/utils/data/dataset.py @@ -0,0 +1,212 @@ +import copy +import collections +import os.path as osp +import warnings +import re + +import torch.utils.data + +from .makedirs import makedirs + + +def to_list(x): + if not isinstance(x, collections.Iterable) or isinstance(x, str): + x = [x] + return x + + +def files_exist(files): + return all([osp.exists(f) for f in files]) + + +def __repr__(obj): + if obj is None: + return 'None' + return re.sub('(<.*?)\\s.*(>)', r'\1\2', obj.__repr__()) + + +class Dataset(torch.utils.data.Dataset): + r"""Dataset base class for creating datasets with pre-processing. + Based on Dataset from pytorch-geometric: see `here `__ for the accompanying tutorial. + Args: + root (string, optional): Root directory where the dataset should be + saved. (optional: :obj:`None`) + transform (callable, optional): A function/transform that takes in an + :obj:`torch_geometric.data.Data` object and returns a transformed + version. The data object will be transformed before every access. + (default: :obj:`None`) + pre_transform (callable, optional): A function/transform that takes in + an :obj:`torch_geometric.data.Data` object and returns a + transformed version. The data object will be transformed before + being saved to disk. (default: :obj:`None`) + pre_filter (callable, optional): A function that takes in an + :obj:`torch_geometric.data.Data` object and returns a boolean + value, indicating whether the data object should be included in the + final dataset. (default: :obj:`None`) + """ + @property + def raw_file_names(self): + r"""The name of the files to find in the :obj:`self.raw_dir` folder in + order to skip the download.""" + raise NotImplementedError + + @property + def processed_file_names(self): + r"""The name of the files to find in the :obj:`self.processed_dir` + folder in order to skip the processing.""" + raise NotImplementedError + + def download(self): + r"""Downloads the dataset to the :obj:`self.raw_dir` folder.""" + raise NotImplementedError + + def process(self): + r"""Processes the dataset to the :obj:`self.processed_dir` folder.""" + raise NotImplementedError + + def len(self): + raise NotImplementedError + + def get(self, idx): + r"""Gets the data object at index :obj:`idx`.""" + raise NotImplementedError + + def __init__(self, root=None, transform=None, pre_transform=None, + pre_filter=None): + super(Dataset, self).__init__() + + if isinstance(root, str): + root = osp.expanduser(osp.normpath(root)) + + self.root = root + self.transform = transform + self.pre_transform = pre_transform + self.pre_filter = pre_filter + self.__indices__ = None + + if 'download' in self.__class__.__dict__.keys(): + self._download() + + if 'process' in self.__class__.__dict__.keys(): + self._process() + + def indices(self): + if self.__indices__ is not None: + return self.__indices__ + else: + return range(len(self)) + + @property + def raw_dir(self): + return osp.join(self.root, 'raw') + + @property + def processed_dir(self): + return osp.join(self.root, 'processed') + + @property + def raw_paths(self): + r"""The filepaths to find in order to skip the download.""" + files = to_list(self.raw_file_names) + return [osp.join(self.raw_dir, f) for f in files] + + @property + def processed_paths(self): + r"""The filepaths to find in the :obj:`self.processed_dir` + folder in order to skip the processing.""" + files = to_list(self.processed_file_names) + return [osp.join(self.processed_dir, f) for f in files] + + def _download(self): + if files_exist(self.raw_paths): # pragma: no cover + return + + makedirs(self.raw_dir) + self.download() + + def _process(self): + f = osp.join(self.processed_dir, 'pre_transform.pt') + if osp.exists(f) and torch.load(f) != __repr__(self.pre_transform): + warnings.warn( + 'The `pre_transform` argument differs from the one used in ' + 'the pre-processed version of this dataset. If you really ' + 'want to make use of another pre-processing technique, make ' + 'sure to delete `{}` first.'.format(self.processed_dir)) + f = osp.join(self.processed_dir, 'pre_filter.pt') + if osp.exists(f) and torch.load(f) != __repr__(self.pre_filter): + warnings.warn( + 'The `pre_filter` argument differs from the one used in the ' + 'pre-processed version of this dataset. If you really want to ' + 'make use of another pre-fitering technique, make sure to ' + 'delete `{}` first.'.format(self.processed_dir)) + + if files_exist(self.processed_paths): # pragma: no cover + return + + print('Processing...') + + makedirs(self.processed_dir) + self.process() + + path = osp.join(self.processed_dir, 'pre_transform.pt') + torch.save(__repr__(self.pre_transform), path) + path = osp.join(self.processed_dir, 'pre_filter.pt') + torch.save(__repr__(self.pre_filter), path) + + print('Done!') + + def __len__(self): + r"""The number of examples in the dataset.""" + if self.__indices__ is not None: + return len(self.__indices__) + return self.len() + + def __getitem__(self, idx): + r"""Gets the data object at index :obj:`idx` and transforms it (in case + a :obj:`self.transform` is given). + In case :obj:`idx` is a slicing object, *e.g.*, :obj:`[2:5]`, a list, a + tuple, a LongTensor or a BoolTensor, will return a subset of the + dataset at the specified indices.""" + if isinstance(idx, int): + data = self.get(self.indices()[idx]) + data = data if self.transform is None else self.transform(data) + return data + else: + return self.index_select(idx) + + def index_select(self, idx): + indices = self.indices() + + if isinstance(idx, slice): + indices = indices[idx] + elif torch.is_tensor(idx): + if idx.dtype == torch.long: + return self.index_select(idx.tolist()) + elif idx.dtype == torch.bool or idx.dtype == torch.uint8: + return self.index_select(idx.nonzero().flatten().tolist()) + elif isinstance(idx, list) or isinstance(idx, tuple): + indices = [indices[i] for i in idx] + else: + raise IndexError( + 'Only integers, slices (`:`), list, tuples, and long or bool ' + 'tensors are valid indices (got {}).'.format( + type(idx).__name__)) + + dataset = copy.copy(self) + dataset.__indices__ = indices + return dataset + + def shuffle(self, return_perm=False): + r"""Randomly shuffles the examples in the dataset. + Args: + return_perm (bool, optional): If set to :obj:`True`, will + additionally return the random permutation used to shuffle the + dataset. (default: :obj:`False`) + """ + perm = torch.randperm(len(self)) + dataset = self.index_select(perm) + return (dataset, perm) if return_perm is True else dataset + + def __repr__(self): # pragma: no cover + return f'{self.__class__.__name__}({len(self)})' diff --git a/torch_lydorn/torch/utils/data/makedirs.py b/torch_lydorn/torch/utils/data/makedirs.py new file mode 100644 index 0000000000000000000000000000000000000000..159f4e9d46247816cbb80b15d65c13cb0c10afbe --- /dev/null +++ b/torch_lydorn/torch/utils/data/makedirs.py @@ -0,0 +1,11 @@ +import os +import os.path as osp +import errno + + +def makedirs(path): + try: + os.makedirs(osp.expanduser(osp.normpath(path))) + except OSError as e: + if e.errno != errno.EEXIST and osp.isdir(path): + raise e diff --git a/torch_lydorn/torchvision/__init__.py b/torch_lydorn/torchvision/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aebf39701637fd30d47317a7743cfdf7f749c74c --- /dev/null +++ b/torch_lydorn/torchvision/__init__.py @@ -0,0 +1 @@ +from . import transforms diff --git a/torch_lydorn/torchvision/datasets/__init__.py b/torch_lydorn/torchvision/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c96f7315ac7c73d57f48e5759b0932f48de2e3d4 --- /dev/null +++ b/torch_lydorn/torchvision/datasets/__init__.py @@ -0,0 +1,6 @@ +from .inria_aerial import InriaAerial +from .luxcarta_buildings import LuxcartaBuildings +from .mapping_challenge import MappingChallenge +from .open_cities_competition import RasterizedOpenCities, OpenCitiesTestDataset +from . import utils +from .xview2_dataset import xView2Dataset diff --git a/torch_lydorn/torchvision/datasets/config.luxcarta_dataset.json b/torch_lydorn/torchvision/datasets/config.luxcarta_dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..26dab5e54c7f9039dc26b1602c53cd2b973d5654 --- /dev/null +++ b/torch_lydorn/torchvision/datasets/config.luxcarta_dataset.json @@ -0,0 +1,28 @@ +{ + "data_dir_candidates": [ + "/local/shared/data", // try cluster local node first + "/data/titane/user/nigirard/data", // Try cluster /data directory + "~/data", // In home directory (docker) + "/data" // In landsat's /data volume (docker) + ], + "data_root_partial_dirpath": "luxcarta_precise_buildings", + "num_workers": 10, + "dataset_params": { + "data_patch_size": 725, // Size of patch saved on disk if data aug is True (allows for rotating patches for the train split) + "input_patch_size": 512 // Size of patch fed to the model + }, + "data_split_params": { + "seed": 0, // Change this to change the random splitting of data in train/val/test + "train_fraction": 0.9, + "val_fraction": 0.1 // test_fraction is the rest + }, + "data_aug_params": { + "enable": true, + "vflip": true, + "rotate": true, + "color_jitter": true, + "device": "cuda" + }, + + "device": "cuda" // Only has effects when mode is val or test. When mode is train, always use CUDA +}, diff --git a/torch_lydorn/torchvision/datasets/inria_aerial.py b/torch_lydorn/torchvision/datasets/inria_aerial.py new file mode 100644 index 0000000000000000000000000000000000000000..7f18456c2cbc898c0dd09fb6670306e3234f4510 --- /dev/null +++ b/torch_lydorn/torchvision/datasets/inria_aerial.py @@ -0,0 +1,592 @@ +import fnmatch +import os.path +import pathlib +import sys +import time + +import shapely.geometry +import multiprocess +import itertools +import skimage.io +import numpy as np + +from tqdm import tqdm + +import torch +import torch.utils.data +import torchvision + +from lydorn_utils import run_utils, image_utils, polygon_utils, geo_utils +from lydorn_utils import print_utils +from lydorn_utils import python_utils + +from torch_lydorn.torchvision.datasets import utils + +CITY_METADATA_DICT = { + "bloomington": { + "fold": "test", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.44583929, 0.46205078, 0.35783887], + "std": [0.18212699, 0.17152641, 0.16157062], + }, + "bellingham": { + "fold": "test", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.3766195, 0.391402, 0.32659722], + "std": [0.18134978, 0.16412577, 0.16369793], + }, + "innsbruck": { + "fold": "test", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.41375683, 0.41818116, 0.38940192], + "std": [0.16616156, 0.14364722, 0.13317743], + }, + "sfo": { + "fold": "test", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.59388761, 0.61522012, 0.54348289], + "std": [0.25730708, 0.23301019, 0.23707742], + }, + "tyrol-e": { + "fold": "test", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.44171042, 0.48147037, 0.44642358], + "std": [0.1808623, 0.15437789, 0.15102051], + }, + "austin": { + "fold": "train", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.39584444, 0.40599795, 0.38298687], + "std": [0.17341954, 0.16856597, 0.16360443], + }, + "chicago": { + "fold": "train", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.4055142, 0.42844002, 0.38229637], + "std": [0.2133328, 0.20827106, 0.20132315], + }, + "kitsap": { + "fold": "train", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.34717916, 0.37854108, 0.32571001], + "std": [0.17048794, 0.14537676, 0.13466496], + }, + "tyrol-w": { + "fold": "train", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.39704218, 0.4545488, 0.4321427], + "std": [0.19484766, 0.1742585, 0.15186383], + }, + "vienna": { + "fold": "train", + "pixelsize": 0.3, + "numbers": list(range(1, 37)), + "mean": [0.47861977, 0.46878486, 0.44043111], + "std": [0.22614806, 0.19949128, 0.19524506], + }, +} + +IMAGE_DIRNAME = "images" +IMAGE_NAME_FORMAT = "{city}{number}" +IMAGE_FILENAME_FORMAT = IMAGE_NAME_FORMAT + ".tif" # City name, number + + +class InriaAerial(torch.utils.data.Dataset): + """ + Inria Aerial Image Dataset + """ + + def __init__(self, root: str, fold: str="train", pre_process: bool=True, tile_filter=None, patch_size: int=None, patch_stride: int=None, + pre_transform=None, transform=None, small: bool=False, pool_size: int=1, raw_dirname: str="raw", processed_dirname: str="processed", + gt_source: str="disk", gt_type: str="npy", gt_dirname: str="gt_polygons", mask_only: bool=False): + """ + + @param root: + @param fold: + @param pre_process: If True, the dataset will be pre-processed first, saving training patches on disk. If False, data will be serve on-the-fly without any patching. + @param tile_filter: Function to call on tile_info, if returns True, include that tile. If returns False, exclude that tile. Does not affect pre-processing. + @param patch_size: + @param patch_stride: + @param pre_transform: + @param transform: + @param small: If True, use a small subset of the dataset (for testing) + @param pool_size: + @param processed_dirname: + @param gt_source: Can be "disk" for annotation that are on disk or "osm" to download from OSM (not implemented) + @param gt_type: Type of annotation files on disk: can be "npy", "geojson" or "tif" + @param gt_dirname: Name of directory with annotation files + @param mask_only: If True, discard the RGB image, sample's "image" field is a single-channel binary mask of the polygons and there is no ground truth segmentation. + This is to allow learning only the frame field from binary masks in order to polygonize binary masks + """ + assert gt_source in {"disk", "osm"}, "gt_source should be disk or osm" + assert gt_type in {"npy", "geojson", "tif"}, f"gt_type should be npy, geojson or tif, not {gt_type}" + self.root = root + self.fold = fold + self.pre_process = pre_process + self.tile_filter = tile_filter + self.patch_size = patch_size + self.patch_stride = patch_stride + self.pre_transform = pre_transform + self.transform = transform + self.small = small + if self.small: + print_utils.print_info("INFO: Using small version of the Inria dataset.") + self.pool_size = pool_size + self.raw_dirname = raw_dirname + self.gt_source = gt_source + self.gt_type = gt_type + self.gt_dirname = gt_dirname + self.mask_only = mask_only + + # Fill default values + if self.gt_source == "disk": + print_utils.print_info("INFO: annotations will be loaded from disk") + elif self.gt_source == "osm": + print_utils.print_info("INFO: annotations will be downloaded from OSM. " + "Make sure you have an internet connection to the OSM servers!") + + if self.pre_process: + # Setup of pre-process + processed_dirname_extention = f"{processed_dirname}.source_{self.gt_source}.type_{self.gt_type}" + if self.gt_dirname is not None: + processed_dirname_extention += f".dirname_{self.gt_dirname}" + if self.mask_only: + processed_dirname_extention += f".mask_only_{int(self.mask_only)}" + processed_dirname_extention += f".patch_size_{int(self.patch_size)}" + self.processed_dirpath = os.path.join(self.root, processed_dirname_extention, self.fold) + self.stats_filepath = os.path.join(self.processed_dirpath, "stats-small.pt" if self.small else "stats.pt") + self.processed_flag_filepath = os.path.join(self.processed_dirpath, + "processed_flag-small" if self.small else "processed_flag") + + # Check if dataset has finished pre-processing by checking flag: + if os.path.exists(self.processed_flag_filepath): + # Process done, load stats + self.stats = torch.load(self.stats_filepath) + else: + # Pre-process not finished, launch it: + tile_info_list = self.get_tile_info_list(tile_filter=None) + self.stats = self.process(tile_info_list) + # Save stats + torch.save(self.stats, self.stats_filepath) + # Mark dataset as processed with flag + pathlib.Path(self.processed_flag_filepath).touch() + + # Get processed_relative_paths with filter + tile_info_list = self.get_tile_info_list(tile_filter=self.tile_filter) + self.processed_relative_paths = self.get_processed_relative_paths(tile_info_list) + else: + # Setup data sample list + self.tile_info_list = self.get_tile_info_list(tile_filter=self.tile_filter) + + def get_tile_info_list(self, tile_filter=None): + tile_info_list = [] + for city, info in CITY_METADATA_DICT.items(): + if not info["fold"] == self.fold: + continue + if self.small: + numbers = [*info["numbers"][:5], info["numbers"][-1]] + else: + numbers = info["numbers"] + for number in numbers: + image_info = { + "city": city, + "number": number, + "pixelsize": info["pixelsize"], + "mean": np.array(info["mean"]), + "std": np.array(info["std"]), + } + tile_info_list.append(image_info) + if tile_filter is not None: + tile_info_list = list(filter(self.tile_filter, tile_info_list)) + return tile_info_list + + def get_processed_relative_paths(self, tile_info_list): + processed_relative_paths = [] + for tile_info in tile_info_list: + processed_tile_relative_dirpath = os.path.join(tile_info['city'], f"{tile_info['number']:02d}") + processed_tile_dirpath = os.path.join(self.processed_dirpath, processed_tile_relative_dirpath) + sample_filenames = fnmatch.filter(os.listdir(processed_tile_dirpath), "data.*.pt") + processed_tile_relative_paths = [os.path.join(processed_tile_relative_dirpath, sample_filename) for sample_filename + in sample_filenames] + processed_relative_paths.extend(processed_tile_relative_paths) + return sorted(processed_relative_paths) + + def process(self, tile_info_list): + # os.makedirs(os.path.join(self.root, self.processed_dirname), exist_ok=True) + with multiprocess.Pool(self.pool_size) as p: + stats_all = list( + tqdm(p.imap(self._process_one, tile_info_list), total=len(tile_info_list), desc="Process")) + + stats = {} + if not self.mask_only: + stats_all = list(filter(None.__ne__, stats_all)) + stat_lists = {} + for stats_one in stats_all: + for key, stat in stats_one.items(): + if key in stat_lists: + stat_lists[key].append(stat) + else: + stat_lists[key] = [stat] + + # Aggregate stats + if "class_freq" in stat_lists and "num" in stat_lists: + class_freq_array = np.stack(stat_lists["class_freq"], axis=0) + num_array = np.stack(stat_lists["num"], axis=0) + if num_array.min() == 0: + raise ZeroDivisionError("num_array has some zeros values, cannot divide!") + stats["class_freq"] = np.sum(class_freq_array*num_array[:, None], axis=0) / np.sum(num_array) + + return stats + + def load_raw_data(self, tile_info): + raw_data = {} + + # Image: + raw_data["image_filepath"] = os.path.join(self.root, self.raw_dirname, self.fold, IMAGE_DIRNAME, + IMAGE_FILENAME_FORMAT.format(city=tile_info["city"], number=tile_info["number"])) + raw_data["image"] = skimage.io.imread(raw_data["image_filepath"]) + assert len(raw_data["image"].shape) == 3 and raw_data["image"].shape[2] == 3, f"image should have shape (H, W, 3), not {raw_data['image'].shape}..." + + # Annotations: + if self.gt_source == "disk": + gt_base_filepath = os.path.join(self.root, self.raw_dirname, self.fold, self.gt_dirname, + IMAGE_NAME_FORMAT.format(city=tile_info["city"], + number=tile_info["number"])) + gt_filepath = gt_base_filepath + "." + self.gt_type + if not os.path.exists(gt_filepath): + raw_data["gt_polygons"] = [] + return raw_data + if self.gt_type == "npy": + np_gt_polygons = np.load(gt_filepath, allow_pickle=True) + gt_polygons = [] + for np_gt_polygon in np_gt_polygons: + try: + gt_polygons.append(shapely.geometry.Polygon(np_gt_polygon[:, ::-1])) + except ValueError: + # Invalid polygon, continue without it + continue + raw_data["gt_polygons"] = gt_polygons + elif self.gt_type == "geojson": + geojson = python_utils.load_json(gt_filepath) + raw_data["gt_polygons"] = list(shapely.geometry.shape(geojson)) + elif self.gt_type == "tif": + raw_data["gt_polygons_image"] = skimage.io.imread(gt_filepath)[:, :, None] + assert len(raw_data["gt_polygons_image"].shape) == 3 and raw_data["gt_polygons_image"].shape[2] == 1, \ + f"Mask should have shape (H, W, 1), not {raw_data['gt_polygons_image'].shape}..." + elif self.gt_source == "osm": + raise NotImplementedError( + "Downloading from OSM is not implemented (takes too long to download, better download to disk first...).") + # np_gt_polygons = geo_utils.get_polygons_from_osm(image_filepath, tag="building", ij_coords=False) + + return raw_data + + def _process_one(self, tile_info): + process_id = int(multiprocess.current_process().name[-1]) + # print(f"\n--- {process_id} ---\n") + + # --- Init + tile_name = IMAGE_NAME_FORMAT.format(city=tile_info["city"], number=tile_info["number"]) + processed_tile_relative_dirpath = os.path.join(tile_info['city'], f"{tile_info['number']:02d}") + processed_tile_dirpath = os.path.join(self.processed_dirpath, processed_tile_relative_dirpath) + processed_flag_filepath = os.path.join(processed_tile_dirpath, "processed_flag") + stats_filepath = os.path.join(processed_tile_dirpath, "stats.pt") + os.makedirs(processed_tile_dirpath, exist_ok=True) + stats = {} + + # --- Check if tile has been processed already + if os.path.exists(processed_flag_filepath): + if not self.mask_only: + stats = torch.load(stats_filepath) + return stats + + # --- Read data: + raw_data = self.load_raw_data(tile_info) + + # --- Patch tiles + if self.patch_size is not None: + patch_stride = self.patch_stride if self.patch_stride is not None else self.patch_size + patch_boundingboxes = image_utils.compute_patch_boundingboxes(raw_data["image"].shape[0:2], + stride=patch_stride, + patch_res=self.patch_size) + class_freq_list = [] + for i, bbox in enumerate(tqdm(patch_boundingboxes, desc=f"Patching {tile_name}", leave=False, position=process_id)): + sample = { + "image_filepath": raw_data["image_filepath"], + "name": f"{tile_name}.rowmin_{bbox[0]}_colmin_{bbox[1]}_rowmax_{bbox[2]}_colmax_{bbox[3]}", + "bbox": bbox, + "city": tile_info["city"], + "number": tile_info["number"], + } + + if self.gt_type == "npy" or self.gt_type == "geojson": + patch_gt_polygons = polygon_utils.patch_polygons(raw_data["gt_polygons"], minx=bbox[1], miny=bbox[0], + maxx=bbox[3], maxy=bbox[2]) + sample["gt_polygons"] = patch_gt_polygons + elif self.gt_type == "tif": + patch_gt_mask = raw_data["gt_polygons_image"][bbox[0]:bbox[2], bbox[1]:bbox[3], :] + sample["gt_polygons_image"] = patch_gt_mask + + sample["image"] = raw_data["image"][bbox[0]:bbox[2], bbox[1]:bbox[3], :] + + sample = self.pre_transform(sample) # Needs "image" to infer shape even if mask_only is True + if self.mask_only: + del sample["image"] # Don't need RGB image anymore + + relative_filepath = os.path.join(processed_tile_relative_dirpath, "data.{:06d}.pt".format(i)) + filepath = os.path.join(self.processed_dirpath, relative_filepath) + torch.save(sample, filepath) + + # Compute stats + if not self.mask_only: + if self.gt_type == "npy" or self.gt_type == "geojson": + class_freq_list.append(np.mean(sample["gt_polygons_image"], axis=(0, 1)) / 255) + elif self.gt_type == "mask": + raise NotImplementedError("mask class freq") + else: + raise NotImplementedError(f"gt_type={self.gt_type} not implemented for computing stats") + + # Aggregate stats + if not self.mask_only: + if len(class_freq_list): + class_freq_array = np.stack(class_freq_list, axis=0) + stats["class_freq"] = np.mean(class_freq_array, axis=0) + stats["num"] = len(class_freq_list) + else: + print("Empty tile:", tile_info["city"], tile_info["number"], "polygons:", len(raw_data["gt_polygons"])) + else: + raise NotImplemented("patch_size is None") + + # Save stats + if not self.mask_only: + torch.save(stats, stats_filepath) + + # Mark tile as processed with flag + pathlib.Path(processed_flag_filepath).touch() + + return stats + + def __len__(self): + if self.pre_process: + return len(self.processed_relative_paths) + else: + return len(self.tile_info_list) + + def __getitem__(self, idx): + if self.pre_process: + filepath = os.path.join(self.processed_dirpath, self.processed_relative_paths[idx]) + data = torch.load(filepath) + if self.mask_only: + data["image"] = np.repeat(data["gt_polygons_image"][:, :, 0:1], 3, axis=-1) # Fill image slot + data["image_mean"] = np.array([0.5, 0.5, 0.5]) + data["image_std"] = np.array([1, 1, 1]) + else: + data["image_mean"] = np.array(CITY_METADATA_DICT[data["city"]]["mean"]) + data["image_std"] = np.array(CITY_METADATA_DICT[data["city"]]["std"]) + data["class_freq"] = self.stats["class_freq"] + else: + tile_info = self.tile_info_list[idx] + # Load raw data + data = self.load_raw_data(tile_info) + data["name"] = IMAGE_NAME_FORMAT.format(city=tile_info["city"], number=tile_info["number"]) + data["image_mean"] = np.array(tile_info["mean"]) + data["image_std"] = np.array(tile_info["std"]) + data = self.transform(data) + return 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": "AerialImageDataset", + "pre_process": False, + "gt_source": "disk", + "gt_type": "tif", + "gt_dirname": "gt", + "mask_only": False, + "small": True, + "data_patch_size": 425, + "input_patch_size": 300, + + "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"]) + 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": config["dataset_params"]["mask_only"], + } + train_val_split_point = config["dataset_params"]["train_fraction"] * 36 + def train_tile_filter(tile): return tile["number"] <= train_val_split_point + def val_tile_filter(tile): return train_val_split_point < tile["number"] + # --- --- # + fold = "train" + if fold == "train": + dataset = InriaAerial(root_dir, fold="train", tile_filter=train_tile_filter, **kwargs) + elif fold == "val": + dataset = InriaAerial(root_dir, fold="train", tile_filter=val_tile_filter, **kwargs) + elif fold == "test": + dataset = InriaAerial(root_dir, fold="test", **kwargs) + + 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"].min().item(), batch["distances"].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"].min().item(), batch["sizes"].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) + + 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...") + + +if __name__ == '__main__': + main() diff --git a/torch_lydorn/torchvision/datasets/luxcarta_buildings.py b/torch_lydorn/torchvision/datasets/luxcarta_buildings.py new file mode 100644 index 0000000000000000000000000000000000000000..1e3b5aa48ece577221dcb23b94a67aa001dd19c2 --- /dev/null +++ b/torch_lydorn/torchvision/datasets/luxcarta_buildings.py @@ -0,0 +1,398 @@ +import skimage.io +from functools import partial +from multiprocess import Pool +import itertools + +import numpy as np +import rasterio +import rasterio.rio.insp +import fiona +from pyproj import Proj, transform, CRS + +import torch_lydorn.torch.utils.data + +import os + +from tqdm import tqdm + +import torch + +from frame_field_learning import data_transforms + +from lydorn_utils import run_utils +from lydorn_utils import print_utils +from lydorn_utils import python_utils +from lydorn_utils import polygon_utils +from lydorn_utils import ogr2ogr +from lydorn_utils import image_utils + + +class LuxcartaBuildings(torch_lydorn.torch.utils.data.Dataset): + def __init__(self, root, transform=None, pre_transform=None, fold="train", patch_size=None, patch_stride=None, + pool_size=1): + assert fold in {"train", "test"}, "fold should be either train of test" + self.fold = fold + self.patch_size = patch_size + self.patch_stride = patch_stride + self.pool_size = pool_size + self._processed_filepaths = None + # TODO: implement pool_size option + super(LuxcartaBuildings, self).__init__(root, transform, pre_transform) + + @property + def raw_dir(self): + return os.path.join(self.root, 'raw', self.fold) + + @property + def processed_dir(self): + return os.path.join(self.root, 'processed', self.fold) + + @property + def raw_file_names(self): + return [] + + @property + def raw_sample_metadata_list(self): + returned_list = [] + for dirname in os.listdir(self.raw_dir): + dirpath = os.path.join(self.raw_dir, dirname) + if os.path.isdir(dirpath): + image_filepath_list = python_utils.get_filepaths(dirpath, endswith_str="_crop.tif", + startswith_str="Ortho", + not_endswith_str=".tif_crop.tif") + gt_polygons_filepath_list = python_utils.get_filepaths(dirpath, endswith_str=".shp", + startswith_str="Building") + mask_filepath_list = python_utils.get_filepaths(dirpath, endswith_str=".kml", startswith_str="zone") + if len(image_filepath_list) and len(gt_polygons_filepath_list): + metadata = { + "dirname": dirname, + "image_filepath": image_filepath_list[0], + "gt_polygons_filepath": gt_polygons_filepath_list[0], + } + if len(mask_filepath_list): + metadata["mask_filepath"] = mask_filepath_list[0] + returned_list.append(metadata) + return returned_list + + @property + def processed_filepaths(self): + if self._processed_filepaths is None: + self._processed_filepaths = [] + for dirname in os.listdir(self.processed_dir): + dirpath = os.path.join(self.processed_dir, dirname) + if os.path.isdir(dirpath): + filepath_list = python_utils.get_filepaths(dirpath, endswith_str=".pt", startswith_str="data.") + self._processed_filepaths.extend(filepath_list) + return self._processed_filepaths + + def __len__(self): + return len(self.processed_filepaths) + + def _download(self): + pass + + def download(self): + pass + + def _get_mask_multi_polygon(self, mask_filepath, shp_srs): + # Read mask and convert to shapefile's crs: + mask_shp_filepath = os.path.join(os.path.dirname(mask_filepath), "mask.shp") + + # TODO: see what's up with the warnings: + ogr2ogr.main(["", "-f", "ESRI Shapefile", mask_shp_filepath, + mask_filepath]) # Convert .kml into .shp + mask = fiona.open(mask_shp_filepath, "r") + mask_srs = Proj(mask.crs["init"]) + + mask_multi_polygon = [] + for feat in mask: + mask_polygon = [] + for point in feat['geometry']['coordinates'][0]: + long, lat = point[:2] # one 2D point of the LinearRing + x, y = transform(mask_srs, shp_srs, long, lat, always_xy=True) # transform the point + mask_polygon.append((x, y)) + mask_multi_polygon.append(mask_polygon) + # mask_polygon is now in UTM proj, ready to compare to shapefile proj + + mask.close() + + return mask_multi_polygon + + def _read_annotations(self, annotations_filepath, raster, mask_filepath=None): + shapefile = fiona.open(annotations_filepath, "r") + shp_srs = Proj(shapefile.crs["init"]) + raster_srs = Proj(raster.crs) + + # --- Read and crop shapefile with mask if specified --- # + if mask_filepath is not None: + mask_multi_polygon = self._get_mask_multi_polygon(mask_filepath, shp_srs) + else: + mask_multi_polygon = None + + process_feat_partial = partial(process_feat, shp_srs=shp_srs, raster_srs=raster_srs, + raster_transform=raster.transform, mask_multi_polygon=mask_multi_polygon) + with Pool() as pool: + out_polygons = list( + tqdm(pool.imap(process_feat_partial, shapefile), desc="Process shp feature", total=len(shapefile), + leave=True)) + out_polygons = list(itertools.chain.from_iterable(out_polygons)) + + shapefile.close() + + return out_polygons + + def process(self, metadata_list): + progress_bar = tqdm(metadata_list, desc="Pre-process") + for metadata in progress_bar: + progress_bar.set_postfix(image=metadata["dirname"], status="Loading image") + # Load image + # image = skimage.io.imread(metadata["image_filepath"]) + raster = rasterio.open(metadata["image_filepath"]) + # print(raster) + # print(dir(raster)) + # print(raster.meta) + # exit() + + progress_bar.set_postfix(image=metadata["dirname"], status="Process shapefile") + mask_filepath = metadata["mask_filepath"] if "mask_filepath" in metadata else None + gt_polygons = self._read_annotations(metadata["gt_polygons_filepath"], raster, mask_filepath=mask_filepath) + + # Compute image mean and std + b1, b2, b3 = raster.read() + image = np.stack([b1, b2, b3], axis=-1) + progress_bar.set_postfix(image=metadata["dirname"], status="Compute mean and std") + 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) + + if self.patch_size is not None: + # Patch the tile + progress_bar.set_postfix(image=metadata["dirname"], status="Patching") + patch_stride = self.patch_stride if self.patch_stride is not None else self.patch_size + patch_boundingboxes = image_utils.compute_patch_boundingboxes(image.shape[0:2], + stride=patch_stride, + patch_res=self.patch_size) + for i, patch_boundingbox in enumerate(tqdm(patch_boundingboxes, desc="Process patches", leave=False)): + patch_gt_polygons = polygon_utils.crop_polygons_to_patch_if_touch(gt_polygons, patch_boundingbox) + if len(patch_gt_polygons) == 0: + # Do not save patches empty of polygons # TODO: keep empty patches? + break + patch_image = image[patch_boundingbox[0]:patch_boundingbox[2], + patch_boundingbox[1]:patch_boundingbox[3], :] + sample = { + "dirname": metadata["dirname"], + "image": patch_image, + "image_mean": mean, + "image_std": std, + "gt_polygons": patch_gt_polygons, + "image_filepath": metadata["image_filepath"], + } + + if self.pre_transform: + sample = self.pre_transform(sample) + + filepath = os.path.join(self.processed_dir, sample["dirname"], "data.{:06d}.pt".format(i)) + os.makedirs(os.path.dirname(filepath), exist_ok=True) + torch.save(sample, filepath) + else: + # Tile is saved as is + sample = { + "dirname": metadata["dirname"], + "image": image, + "image_mean": mean, + "image_std": std, + "gt_polygons": gt_polygons, + "image_filepath": metadata["image_filepath"], + } + + if self.pre_transform: + sample = self.pre_transform(sample) + + filepath = os.path.join(self.processed_dir, sample["dirname"], "data.{:06d}.pt".format(0)) + os.makedirs(os.path.dirname(filepath), exist_ok=True) + torch.save(sample, filepath) + + flag_filepath = os.path.join(self.processed_dir, metadata["dirname"], "flag.json") + flag = python_utils.load_json(flag_filepath) + flag["done"] = True + python_utils.save_json(flag_filepath, flag) + + raster.close() + + def _process(self): + to_process_metadata_list = [] + for metadata in self.raw_sample_metadata_list: + flag_filepath = os.path.join(self.processed_dir, metadata["dirname"], "flag.json") + if os.path.exists(flag_filepath): + flag = python_utils.load_json(flag_filepath) + if not flag["done"]: + to_process_metadata_list.append(metadata) + else: + flag = { + "done": False + } + python_utils.save_json(flag_filepath, flag) + to_process_metadata_list.append(metadata) + + if len(to_process_metadata_list) == 0: + return + + print('Processing...') + + torch_lydorn.torch.utils.data.makedirs(self.processed_dir) + self.process(to_process_metadata_list) + + path = os.path.join(self.processed_dir, 'pre_transform.pt') + torch.save(torch_lydorn.torch.utils.data.__repr__(self.pre_transform), path) + path = os.path.join(self.processed_dir, 'pre_filter.pt') + torch.save(torch_lydorn.torch.utils.data.__repr__(self.pre_filter), path) + + print('Done!') + + def get(self, idx): + filepath = self.processed_filepaths[idx] + data = torch.load(filepath) + data["patch_bbox"] = torch.tensor([0, 0, 0, 0]) # TODO: implement in pre-processing + tile_name = os.path.basename(os.path.dirname(filepath)) + patch_name = os.path.basename(filepath) + patch_name = patch_name[len("data."):-len(".pt")] + data["name"] = tile_name + "." + patch_name + return data + + +def process_feat(feat, shp_srs, raster_srs, raster_transform, mask_multi_polygon=None): + out_polygons = [] + + if feat['geometry']["type"] == "Polygon": + poly = feat['geometry']['coordinates'] + polygons = process_polygon_feat(poly, shp_srs, raster_srs, raster_transform, mask_multi_polygon) + out_polygons.extend(polygons) + + elif feat['geometry']["type"] == "MultiPolygon": + for poly in feat['geometry']["coordinates"]: + polygons = process_polygon_feat(poly, shp_srs, raster_srs, raster_transform, mask_multi_polygon) + out_polygons.extend(polygons) + + return out_polygons + + +def process_polygon_feat(in_polygon, shp_srs, raster_srs, raster_transform, mask_multi_polygon=None): + out_polygons = [] + points = in_polygon[0] # TODO: handle holes + # Intersect with mask_polygon if specified + if mask_multi_polygon is not None: + multi_polygon_simple = polygon_utils.intersect_polygons(points, mask_multi_polygon) + if multi_polygon_simple is None: + return out_polygons + else: + multi_polygon_simple = [points] + + for polygon_simple in multi_polygon_simple: + new_poly = [] + for point in polygon_simple: + x, y = point[:2] # 740524.429227941 7175355.263524155 + x, y = transform(shp_srs, raster_srs, x, y) # transform the point # 740520.728530676 7175320.732711278 + j, i = ~raster_transform * (x, y) + new_poly.append((i, j)) # 2962.534577447921 2457.457061359659 + out_polygons.append(np.array(new_poly)) + + return out_polygons + + +def get_seg_display(seg): + seg_display = np.zeros([seg.shape[0], seg.shape[1], 4], dtype=np.float) + if len(seg.shape) == 2: + seg_display[..., 0] = seg + seg_display[..., 3] = seg + else: + for i in range(seg.shape[-1]): + seg_display[..., i] = seg[..., i] + seg_display[..., 3] = np.clip(np.sum(seg, axis=-1), 0, 1) + return seg_display + + +def main(): + # --- Params --- # + config_name = "config.luxcarta_dataset" + # --- --- # + + # Load config + config = run_utils.load_config(config_name) + if config is None: + print_utils.print_error( + "ERROR: cannot continue without a config file. Exiting now...") + exit() + + # 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): + train_pre_transform = data_transforms.get_offline_transform(config, + augmentations=config["data_aug_params"]["enable"]) + eval_pre_transform = data_transforms.get_offline_transform(config, augmentations=False) + # --- Online transform done on the host (CPU): + train_online_cpu_transform = data_transforms.get_online_cpu_transform(config, + augmentations=config["data_aug_params"][ + "enable"]) + eval_online_cpu_transform = data_transforms.get_online_cpu_transform(config, augmentations=False) + # --- Online transform performed on the device (GPU): + train_online_cuda_transform = data_transforms.get_online_cuda_transform(config, + augmentations=config["data_aug_params"][ + "enable"]) + eval_online_cuda_transform = data_transforms.get_online_cuda_transform(config, augmentations=False) + # --- --- # + + data_patch_size = config["dataset_params"]["data_patch_size"] if config["data_aug_params"]["enable"] else config["dataset_params"]["input_patch_size"] + fold = "test" + if fold == "train": + dataset = LuxcartaBuildings(root_dir, + transform=train_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"]) + elif fold == "test": + dataset = LuxcartaBuildings(root_dir, + transform=train_online_cpu_transform, + pre_transform=data_transforms.get_offline_transform_patch(), + fold="test", + pool_size=config["num_workers"]) + + print("# --- Sample 0 --- #") + sample = dataset[0] + print(sample["image"].shape) + print(sample["gt_polygons_image"].shape) + print("# --- Samples --- #") + # for data in tqdm(dataset): + # pass + + data_loader = torch.utils.data.DataLoader(dataset, batch_size=10, shuffle=True, num_workers=0) + print("# --- Batches --- #") + for batch in tqdm(data_loader): + print(batch["image"].shape) + print(batch["gt_polygons_image"].shape) + + # Save output to visualize + seg = np.array(batch["gt_polygons_image"][0]) / 255 # First batch + seg = np.moveaxis(seg, 0, -1) + seg_display = 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) + + input("Enter to continue...") + + +if __name__ == '__main__': + main() diff --git a/torch_lydorn/torchvision/datasets/mapping_challenge.py b/torch_lydorn/torchvision/datasets/mapping_challenge.py new file mode 100644 index 0000000000000000000000000000000000000000..f4daee2bee6ebb4af17c092f199fe16939ee6a7f --- /dev/null +++ b/torch_lydorn/torchvision/datasets/mapping_challenge.py @@ -0,0 +1,349 @@ +import os +import pathlib +import warnings + +import skimage.io +from multiprocess import Pool +from functools import partial + +import numpy as np +from pycocotools.coco import COCO +import shapely.geometry + +from tqdm import tqdm + +import torch + +from lydorn_utils import print_utils +from lydorn_utils import python_utils + +from torch_lydorn.torch.utils.data import Dataset as LydornDataset, makedirs, files_exist, __repr__ + +from torch_lydorn.torchvision.datasets import utils + + +class MappingChallenge(LydornDataset): + 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 load_image_ids(self): + image_id_list_filepath = os.path.join(self.processed_dir, "image_id_list-small.json" if self.small else "image_id_list.json") + if os.path.exists(image_id_list_filepath): + image_id_list = python_utils.load_json(image_id_list_filepath) + else: + coco = self.get_coco() + image_id_list = coco.getImgIds(catIds=coco.getCatIds()) + # Save for later so that the whole coco object doesn't have to be instantiated when just reading processed samples with multiple workers: + python_utils.save_json(image_id_list_filepath, image_id_list) + return image_id_list + + def get_coco(self): + if self.coco is None: + annotation_filename = "annotation-small.json" if self.small else "annotation.json" + annotations_filepath = os.path.join(self.raw_dir, self.fold, annotation_filename) + self.coco = COCO(annotations_filepath) + return self.coco + + @property + def processed_dir(self): + return os.path.join(self.root, 'processed', self.fold) + + @property + def processed_file_names(self): + l = [] + for image_id in self.image_id_list: + l.append(os.path.join("data_{:012d}.pt".format(image_id))) + return l + + def __len__(self): + return len(self.image_id_list) + + def _download(self): + pass + + def download(self): + pass + + def _process(self): + f = os.path.join(self.processed_dir, 'pre_transform.pt') + if os.path.exists(f) and torch.load(f) != __repr__(self.pre_transform): + warnings.warn( + 'The `pre_transform` argument differs from the one used in ' + 'the pre-processed version of this dataset. If you really ' + 'want to make use of another pre-processing technique, make ' + 'sure to delete `{}` first.'.format(self.processed_dir)) + f = os.path.join(self.processed_dir, 'pre_filter.pt') + if os.path.exists(f) and torch.load(f) != __repr__(self.pre_filter): + warnings.warn( + 'The `pre_filter` argument differs from the one used in the ' + 'pre-processed version of this dataset. If you really want to ' + 'make use of another pre-fitering technique, make sure to ' + 'delete `{}` first.'.format(self.processed_dir)) + + if os.path.exists(self.processed_flag_filepath): + return + + print('Processing...') + + makedirs(self.processed_dir) + self.process() + + path = os.path.join(self.processed_dir, 'pre_transform.pt') + torch.save(__repr__(self.pre_transform), path) + path = os.path.join(self.processed_dir, 'pre_filter.pt') + torch.save(__repr__(self.pre_filter), path) + + print('Done!') + + def process(self): + images_relative_dirpath = os.path.join("raw", self.fold, "images") + + image_info_list = [] + coco = self.get_coco() + for image_id in self.image_id_list: + filename = coco.loadImgs(image_id)[0]["file_name"] + annotation_ids = coco.getAnnIds(imgIds=image_id) + annotation_list = coco.loadAnns(annotation_ids) + image_info = { + "image_id": image_id, + "image_filepath": os.path.join(self.root, images_relative_dirpath, filename), + "image_relative_filepath": os.path.join(images_relative_dirpath, filename), + "annotation_list": annotation_list + } + image_info_list.append(image_info) + + partial_preprocess_one = partial(preprocess_one, pre_filter=self.pre_filter, pre_transform=self.pre_transform, + processed_dir=self.processed_dir) + with Pool(self.pool_size) as p: + sample_stats_list = list(tqdm(p.imap(partial_preprocess_one, image_info_list), total=len(image_info_list))) + + # Aggregate sample_stats_list + image_s0_list, image_s1_list, image_s2_list, class_freq_list = zip(*sample_stats_list) + image_s0_array = np.stack(image_s0_list, axis=0) + image_s1_array = np.stack(image_s1_list, axis=0) + image_s2_array = np.stack(image_s2_list, axis=0) + class_freq_array = np.stack(class_freq_list, axis=0) + + image_s0_total = np.sum(image_s0_array, axis=0) + image_s1_total = np.sum(image_s1_array, axis=0) + image_s2_total = np.sum(image_s2_array, axis=0) + + image_mean = image_s1_total / image_s0_total + image_std = np.sqrt(image_s2_total/image_s0_total - np.power(image_mean, 2)) + class_freq = np.sum(class_freq_array*image_s0_array[:, None], axis=0) / image_s0_total + + # Save aggregated stats + self.stats = { + "image_mean": image_mean, + "image_std": image_std, + "class_freq": class_freq, + } + torch.save(self.stats, self.stats_filepath) + + # Indicates that processing has been performed: + pathlib.Path(self.processed_flag_filepath).touch() + + def get(self, idx): + image_id = self.image_id_list[idx] + data = torch.load(os.path.join(self.processed_dir, "data_{:012d}.pt".format(image_id))) + data["image_mean"] = self.stats["image_mean"] + data["image_std"] = self.stats["image_std"] + data["class_freq"] = self.stats["class_freq"] + return data + + +def preprocess_one(image_info, pre_filter, pre_transform, processed_dir): + out_filepath = os.path.join(processed_dir, "data_{:012d}.pt".format(image_info["image_id"])) + data = None + if os.path.exists(out_filepath): + # Load already-processed sample + try: + data = torch.load(out_filepath) + except EOFError: + pass + if data is None: + # Process sample: + image = skimage.io.imread(image_info["image_filepath"]) + gt_polygons = [] + for annotation in image_info["annotation_list"]: + flattened_segmentation_list = annotation["segmentation"] + if len(flattened_segmentation_list) != 1: + print("WHAT!?!, len(flattened_segmentation_list = {}".format(len(flattened_segmentation_list))) + print("To implement: if more than one segmentation in flattened_segmentation_list (MS COCO format), does it mean it is a MultiPolygon or a Polygon with holes?") + raise NotImplementedError + flattened_arrays = np.array(flattened_segmentation_list) + coords = np.reshape(flattened_arrays, (-1, 2)) + polygon = shapely.geometry.Polygon(coords) + + # Filter out degenerate polygons (area is lower than 2.0) + if 2.0 < polygon.area: + gt_polygons.append(polygon) + + data = { + "image": image, + "gt_polygons": gt_polygons, + "image_relative_filepath": image_info["image_relative_filepath"], + "name": os.path.splitext(os.path.basename(image_info["image_relative_filepath"]))[0], + "image_id": image_info["image_id"] + } + + if pre_filter is not None and not pre_filter(data): + return + + if pre_transform is not None: + data = pre_transform(data) + + # masked_angles = data["gt_crossfield_angle"].astype(np.float) * data["gt_polygons_image"][:, :, 1].astype(np.float) + # skimage.io.imsave("gt_crossfield_angle.png", data["gt_crossfield_angle"]) + # skimage.io.imsave("masked_angles.png", masked_angles) + # exit() + + torch.save(data, out_filepath) + + # Compute stats for later aggregation for the whole dataset + normed_image = data["image"] / 255 + image_s0 = data["image"].shape[0] * data["image"].shape[1] # Number of pixels + image_s1 = np.sum(normed_image, axis=(0, 1)) # Sum of pixel normalized values + image_s2 = np.sum(np.power(normed_image, 2), axis=(0, 1)) + class_freq = np.mean(data["gt_polygons_image"], axis=(0, 1)) / 255 + + return image_s0, image_s1, image_s2, class_freq + + +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": { + "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...") + + +if __name__ == '__main__': + main() diff --git a/torch_lydorn/torchvision/datasets/open_cities_competition.py b/torch_lydorn/torchvision/datasets/open_cities_competition.py new file mode 100644 index 0000000000000000000000000000000000000000..84c71ac45e6c3fb6a09df4506e627b20968b75b4 --- /dev/null +++ b/torch_lydorn/torchvision/datasets/open_cities_competition.py @@ -0,0 +1,466 @@ +#!/usr/bin/env python3 + +import sys, csv, random, glob + +import numpy as np +import torch +from torch.utils.data import Dataset +from torchvision import transforms + +from PIL import Image +from tqdm import tqdm + +import json +import geojson +import rasterio +import rasterio.mask +from rasterio.windows import Window +from pyproj import CRS, Transformer + +from skimage.transform import resize + +from matplotlib import pyplot as plt + +from lydorn_utils import polygon_utils + + +class BuildingDataset(Dataset): + + def __init__(self, tier=1, show_mode=False, augment=False, small_subset=False, crop_size=1024, resize_size=224, + window_random_shift=2048, data_dir="./", baseline_mode=False, transform=None, val=False, val_split=0.1, split_seed=42, sampling_mode="polygons"): + super().__init__() + + self.crop_size = crop_size + self.resize_size = resize_size + + self.data_dir = data_dir + + random.seed(42) + + # Load TIF and geojson file names + # from train_metadata.csv + self.img_ids_to_tif = dict() + self.img_ids_to_geojson = dict() + with open(data_dir + "/train_metadata.csv", 'r') as f: + csvreader = csv.reader(f) + header = next(csvreader) + imgs=[] + for row in csvreader: + if int(row[3]) <= tier: + imgs.append(row) + random.shuffle(imgs) + n_train = int(len(imgs)*(1.0-val_split)) + if val: + imgs = imgs[n_train:] + else: + imgs = imgs[:n_train] + for row in imgs: + imgid = row[0].split("/")[2] + self.img_ids_to_tif[imgid] = row[0] + self.img_ids_to_geojson[imgid] = row[1] + + # Load all geojson files + print("Loading geojson files") + + self.geojson_data = dict() + self.all_polygons = [] + self.feat_to_img_id = [] + for ids, geojson_f in tqdm(self.img_ids_to_geojson.items()): + self.geojson_data[ids] = self._load_geojsonfile(data_dir + "/" + geojson_f) + self.all_polygons += self.geojson_data[ids]["features"] + for f in self.geojson_data[ids]["features"]: + self.feat_to_img_id.append(ids) + + print("Number of training polygons:", len(self.all_polygons)) + + # Open all tif files + # and get corresponding transformers + print("Opening rasters") + #self.rasters = dict() + img_to_transformer_dict = dict() + for ids, img_f in tqdm(self.img_ids_to_tif.items()): + with rasterio.open(data_dir + "/" + img_f) as raster: + img_to_transformer_dict[ids] = Transformer.from_crs(CRS.from_proj4("+proj=latlon"), raster.crs) + + print("Reading means and std") + self.stats = dict() + for ids, tifs in self.img_ids_to_tif.items(): + with open(data_dir + "/" + ".".join([tifs.split(".")[0], "tif.stats.json"]), "r") as statfile: + self.stats[ids] = json.load(statfile) + + + print("Reprojecting polygons") + + self.img_id_to_polys = dict() + for k, v in self.img_ids_to_tif.items(): + self.img_id_to_polys[k] = [] + + for i in tqdm(range(len(self.all_polygons))): + self.convert_poly(i, img_to_transformer_dict) + self.img_id_to_polys[self.feat_to_img_id[i]].append(self.all_polygons[i]) + + if small_subset: + self.all_polygons = self.all_polygons[:200] + + self.show_mode = show_mode + self.baseline_mode = baseline_mode + + self.aug_transforms = transforms.Compose([transforms.RandomVerticalFlip(), + transforms.RandomHorizontalFlip(), + transforms.RandomAffine((180), (0.2, 0.2), (0.5, 2)), + transforms.ColorJitter(0.2, 0.2, 0.5, 0.2), + transforms.Resize((resize_size, resize_size)), + transforms.ToTensor(), + # transforms.RandomErasing(), + transforms.ToPILImage()]) + + self.noaug_transforms = transforms.Compose([transforms.Resize((resize_size, resize_size))]) + + self.mask_transforms = transforms.Compose([transforms.Resize((resize_size, resize_size)), + transforms.ToTensor()]) + + self.img_transforms = transforms.Compose([transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225])]) + + self.wow_transforms = transform + + self.augment = augment + + self.window_random_shift = window_random_shift + + self.sampling_mode = sampling_mode + + def __len__(self): + if self.sampling_mode == "polygons": + return len(self.all_polygons) + else: + return 1000 + + def __getitem__(self, i): + img_id = self.feat_to_img_id[i] + + raster = rasterio.open(self.data_dir + "/" + self.img_ids_to_tif[img_id]) + + if self.sampling_mode == "random": + window = self._get_random_window() + else: + window = self._get_window(raster, self.all_polygons[i]) + + b1, b2, b3, b4 = raster.read(window=window) + out_image = np.stack([b1, b2, b3]) + out_image = np.swapaxes(out_image, 0, 2)[:, :, 0:3] + + + try: + img = Image.fromarray(out_image, 'RGB') + except: + print(window) + print(raster.width, raster.height) + + + # create rasterized edges + polygons = [] + for feat in self.img_id_to_polys[img_id]: + try: + # polywindow = rasterio.features.geometry_window(raster, [feat]) + polywindow = self._get_window(raster, feat, padded=False) + + if rasterio.windows.intersect(polywindow, window): + if feat["geometry"]["type"] == "MultiPolygon": + local_poly = [] + for p in feat["geomtry"]["coordinates"]: + local_poly += self._polygon_window_pixel_coords(raster, window, p[0]) + elif feat["geometry"]["type"] == "Polygon": + local_poly = self._polygon_window_pixel_coords(raster, window, + feat["geometry"]["coordinates"][0]) + ratio = self.resize_size / self.crop_size + scaled_poly = [] + for point in local_poly: + x, y = point + scaled_poly.append((x * ratio, y * ratio)) + polygons.append(scaled_poly) + except: + pass + + masks = polygon_utils.draw_polygons(polygons, (self.resize_size, self.resize_size), line_width=2, antialiasing=True) + masks = np.moveaxis(np.array(masks), 2, 0) + + angles = polygon_utils.init_angle_field(polygons, (self.resize_size, self.resize_size), line_width=4) + angles = np.expand_dims(angles, 0) + + if self.augment: + trans = self.aug_transforms + else: + trans = self.noaug_transforms + + mean = [self.stats[img_id]["stats"]["mean_r"], self.stats[img_id]["stats"]["mean_g"], self.stats[img_id]["stats"]["mean_b"]] + std = [self.stats[img_id]["stats"]["std_r"], self.stats[img_id]["stats"]["std_g"], self.stats[img_id]["stats"]["std_b"]] + + out_img = np.array(trans(img)) + out_img = torch.Tensor(np.moveaxis(out_img, -1, 0)) + + + if self.show_mode: + return {"image": trans(img), "gt_polygons_image": masks, "gt_crossfield_angle": angles, "image_mean": mean, + "image_std": std} + elif self.baseline_mode: + masks = np.moveaxis(np.array(masks), 0,2) + masks = masks / 255 + return {"image": self.img_transforms(trans(img)), "gt_polygons_image": masks} + else: + return {"image": out_img, + "gt_polygons_image": torch.Tensor(masks), + "gt_crossfield_angle": torch.Tensor(angles), + "image_mean": torch.Tensor(mean), + "image_std": torch.Tensor(std)} + + def _get_window(self, raster, feat, padded=True): + + poly = feat["geometry"]["coordinates"] + + box = [] + + if feat["geometry"]["type"] == "MultiPolygon": + local_poly = [] + for p in feat["geometry"]["coordinates"]: + local_poly += self._polygon_pixel_coords(raster, p[0]) + elif feat["geometry"]["type"] == "Polygon": + local_poly = self._polygon_pixel_coords(raster, feat["geometry"]["coordinates"][0]) + + for i in (0, 1): + res = sorted(local_poly, key=lambda x: x[i]) + box += list((res[0][i], res[-1][i])) + + if padded: + toplx, toply = (box[0] + box[1]) / 2 - self.crop_size / 2, (box[2] + box[3]) / 2 - self.crop_size / 2 + + while True: + shift_y = random.randint(-self.window_random_shift, self.window_random_shift) + shift_x = random.randint(-self.window_random_shift, self.window_random_shift) + win = Window(toplx + shift_x, toply + shift_y, self.crop_size, self.crop_size) + + # check that the window is in the image + if win.col_off + win.width < raster.width and win.col_off > 0: + if win.row_off + win.height < raster.height and win.row_off > 0: + break + return win + else: + return Window(box[0], box[1], box[2] - box[0], box[3] - box[1]) + + def _get_random_window(self): + + rastern = random.randint(0, len(self.img_ids_to_tif)-1) + + img_id = list(self.img_ids_to_tif)[rastern] + + tiff = self.img_ids_to_tif[img_id] + + raster = rasterio.open(self.data_dir + "/" + tiff) + + while True: + x = random.randint(0, raster.width-1) + y = random.randint(0, raster.height-1) + + sample = [val for val in raster.sample([raster.transform * (x,y)])][0] + if sample[3] == 255: + win = Window(x - self.crop_size/2, y - self.crop_size/2, self.crop_size, self.crop_size) + if win.col_off + win.width < raster.width and win.col_off > 0: + if win.row_off + win.height < raster.height and win.row_off > 0: + return win, img_id + + + def _polygon_pixel_coords(self, raster, poly): + # converts to raster pixel coords from utm + local_poly = [] + for point in poly: + local_poly.append(~raster.transform * (point[0], point[1])) + + return local_poly + + def _polygon_window_pixel_coords(self, raster, window, poly): + # converts to window coords from utm + local_poly = [] + for point in poly: + local_poly.append(~raster.window_transform(window) * (point[0], point[1])) + + return local_poly + + def _load_geojsonfile(self, filename): + with open(filename, 'r') as f: + content = f.read() + json = geojson.loads(content) + return json + + def convert_poly(self, i, transformers): + img_id = self.feat_to_img_id[i] + transformer = transformers[img_id] + + geom = self.all_polygons[i]["geometry"] + + if geom["type"] == "MultiPolygon": + for poly in geom["coordinates"]: + for part in poly: + for point in part: + point[0], point[1] = list(transformer.transform(point[0], point[1])) + elif geom["type"] == "Polygon": + for poly in geom["coordinates"]: + for point in poly: + point[0], point[1] = list(transformer.transform(point[0], point[1])) + + def augmentation(self, enable=True): + self.augment = enable + +class RasterizedOpenCities(BuildingDataset): + + def __init__(self, + tier=1, + show_mode=False, + augment=False, + small_subset=False, + crop_size=1024, + resize_size=224, + window_random_shift=2048, + data_dir="./", + baseline_mode=False, + transform=None, + val=False, + val_split=0.1, + split_seed=42, + sampling_mode="polygons"): + + super().__init__(tier, show_mode, augment, small_subset, crop_size, + resize_size, window_random_shift, data_dir, baseline_mode, transform, val, val_split, split_seed, sampling_mode) + + self.img_ids_to_label_raster = dict() + with open(data_dir + "/train_metadata.csv", 'r') as f: + csvreader = csv.reader(f) + header = next(csvreader) + for row in csvreader: + if int(row[3]) <= tier: + imgid = row[0].split("/")[2] + self.img_ids_to_label_raster[imgid] = row[1][:-7] + "tif" + + def __getitem__(self, i): + + + if self.sampling_mode == "random": + window, raster_id = self._get_random_window() + img_id=raster_id + raster = rasterio.open(self.data_dir + "/" + self.img_ids_to_tif[raster_id]) + label_raster = rasterio.open(self.data_dir + "/" + + self.img_ids_to_label_raster[raster_id]) + else: + img_id = self.feat_to_img_id[i] + raster = rasterio.open(self.data_dir + "/" + self.img_ids_to_tif[img_id]) + label_raster = rasterio.open(self.data_dir + "/" + + self.img_ids_to_label_raster[img_id]) + window = self._get_window(raster, self.all_polygons[i]) + + b1, b2, b3, b4 = raster.read(window=window) + out_image = np.stack([b1, b2, b3]) + out_image = np.moveaxis(out_image, 0, 2)[:, :, 0:3] + img = Image.fromarray(out_image, 'RGB') + + if self.augment: + trans = self.aug_transforms + else: + trans = self.noaug_transforms + + fill, edges, vertices, angles = label_raster.read(window=window) + masks = np.stack([fill, edges, vertices]) + + mean = [self.stats[img_id]["stats"]["mean_r"], self.stats[img_id]["stats"]["mean_g"], self.stats[img_id]["stats"]["mean_b"]] + std = [self.stats[img_id]["stats"]["std_r"], self.stats[img_id]["stats"]["std_g"], self.stats[img_id]["stats"]["std_b"]] + + masks = np.moveaxis(masks, 0, 2) + masks = Image.fromarray(masks, "RGB") + masks = masks.resize((self.resize_size, self.resize_size)) + angles = Image.fromarray(angles) + angles = angles.resize((self.resize_size, self.resize_size)) + + if self.show_mode: + masks = np.moveaxis(np.array(masks), 0, 2) + return {"image": trans(img), "gt_polygons_image": masks, "gt_crossfield_angle": angles, "image_mean": mean, + "image_std": std} + elif self.baseline_mode: + masks = np.array(masks, dtype=np.float) / 255 + return {"image": self.img_transforms(trans(img)), "gt_polygons_image": masks} + else: + out_img = np.array(trans(img)) + out_img = torch.Tensor(np.moveaxis(out_img, -1, 0)) + masks = np.moveaxis(np.array(masks), 2, 0) + masks = torch.Tensor(np.array(masks)) + angles = np.expand_dims(angles, 0) + angles = torch.Tensor(np.array(angles)) + return {"image": out_img, + "gt_polygons_image": masks, + "gt_crossfield_angle": angles, + "image_mean": torch.Tensor(mean), + "image_std": torch.Tensor(std), + "name":str(i), + "original_image": img_id} + + +class OpenCitiesTestDataset(Dataset): + + def __init__(self, img_dir, output_size): + super().__init__() + self.tif_files = glob.glob(img_dir + "/*/*.tif") + print("Found", str(len(self.tif_files)), "test images") + + self.transforms = transforms.Compose([transforms.Resize((output_size, output_size))]) + + def __len__(self): + return len(self.tif_files) + + def __getitem__(self, i): + img = self.get_img(i) + imgscaled = img / 255.0 + return {"image": img, + "image_filepath": self.tif_files[i], + "name":self.get_id(i), + "image_mean": torch.mean(imgscaled, (1,2)), + "image_std" : torch.std(imgscaled, (1,2))} + + def get_id(self, i): + + path = self.tif_files[i] + return path.split("/")[-2] + + def get_img(self, i): + + img = Image.open(self.tif_files[i]) + img = img.convert("RGB") + img = np.moveaxis(np.array(self.transforms(img)),2,0) + + return torch.Tensor(img) + + + + +if __name__ == "__main__": + + ds = RasterizedOpenCities(show_mode=True, tier=1, small_subset=False, sampling_mode="random") + + n_samples = int(sys.argv[1]) if len(sys.argv) > 1 else 1 + + for i in range(n_samples): + break + n = random.randint(0, len(ds)) + + sample = ds[n] + + fig = plt.figure() + + fig.add_subplot(1, 3, 1) + plt.imshow(sample["image"]) + fig.add_subplot(1, 3, 2) + plt.imshow(sample["gt_polygons_image"]) + fig.add_subplot(1, 3, 3) + plt.imshow(sample["gt_crossfield_angle"]) + plt.show() + + print("Testing whole dataset") + for i in tqdm(range(len(ds))): + sample = ds[i] diff --git a/torch_lydorn/torchvision/datasets/utils.py b/torch_lydorn/torchvision/datasets/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..07c4280f9bd1804693563c293a19b5453bf24f44 --- /dev/null +++ b/torch_lydorn/torchvision/datasets/utils.py @@ -0,0 +1,29 @@ +import numpy as np + + +def get_seg_display(seg): + seg_display = np.zeros([seg.shape[0], seg.shape[1], 4], dtype=np.float) + if len(seg.shape) == 2: + seg_display[..., 0] = seg + seg_display[..., 3] = seg + else: + for i in range(seg.shape[-1]): + seg_display[..., i] = seg[..., i] + seg_display[..., 3] = np.clip(np.sum(seg, axis=-1), 0, 1) + return seg_display + + +def batch_to_cuda(batch): + # Send data to computing device: + for key, item in batch.items(): + if hasattr(item, "cuda"): + batch[key] = item.cuda(non_blocking=True) + return batch + + +def batch_to_cpu(batch): + # Send data to computing device: + for key, item in batch.items(): + if hasattr(item, "cpu"): + batch[key] = item.cpu() + return batch diff --git a/torch_lydorn/torchvision/datasets/xview2_dataset.py b/torch_lydorn/torchvision/datasets/xview2_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..9c9682c99813d77e5405c155a330bcd9c72240dc --- /dev/null +++ b/torch_lydorn/torchvision/datasets/xview2_dataset.py @@ -0,0 +1,405 @@ +import fnmatch +import os.path +import pathlib +import random +import sys +import time +from collections import defaultdict + +import shapely.geometry +import shapely.wkt +import multiprocess +import itertools +import skimage.io +import numpy as np + +from tqdm import tqdm + +import torch +import torch.utils.data +import torchvision + +from lydorn_utils import run_utils, image_utils, polygon_utils, geo_utils +from lydorn_utils import print_utils +from lydorn_utils import python_utils + +from torch_lydorn.torchvision.datasets import utils + + +class xView2Dataset(torch.utils.data.Dataset): + """ + xView2 xBD dataset: https://xview2.org/ + """ + + def __init__(self, root: str, fold: str = "train", pre_process: bool = True, + patch_size: int = None, + pre_transform=None, transform=None, small: bool = False, pool_size: int = 1, raw_dirname: str = "raw", + processed_dirname: str = "processed"): + """ + + @param root: + @param fold: + @param pre_process: If True, the dataset will be pre-processed first, saving training patches on disk. If False, data will be serve on-the-fly without any patching. + @param patch_size: + @param pre_transform: + @param transform: + @param small: If True, use a small subset of the dataset (for testing) + @param pool_size: + @param processed_dirname: + """ + self.root = root + self.fold = fold + self.pre_process = pre_process + self.patch_size = patch_size + self.pre_transform = pre_transform + self.transform = transform + self.small = small + if self.small: + print_utils.print_info("INFO: Using small version of the xView2 xBD dataset.") + self.pool_size = pool_size + self.raw_dirname = raw_dirname + + if self.pre_process: + # Setup of pre-process + self.processed_dirpath = os.path.join(self.root, processed_dirname, self.fold) + stats_filepath = os.path.join(self.processed_dirpath, "stats-small.pt" if self.small else "stats.pt") + processed_relative_paths_filepath = os.path.join(self.processed_dirpath, + "processed_paths-small.json" if self.small else "processed_paths.json") + + # Check if dataset has finished pre-processing by checking processed_relative_paths_filepath: + if os.path.exists(processed_relative_paths_filepath): + # Process done, load stats and processed_relative_paths + self.stats = torch.load(stats_filepath) + self.processed_relative_paths = python_utils.load_json(processed_relative_paths_filepath) + else: + # Pre-process not finished, launch it: + tile_info_list = self.get_tile_info_list() + self.stats = self.process(tile_info_list) + # Save stats + torch.save(self.stats, stats_filepath) + # Save processed_relative_paths + self.processed_relative_paths = [tile_info["processed_relative_filepath"] for tile_info in tile_info_list] + python_utils.save_json(processed_relative_paths_filepath, self.processed_relative_paths) + else: + # Setup data sample list + self.tile_info_list = self.get_tile_info_list() + + def get_tile_info_list(self): + tile_info_list = [] + fold_dirpath = os.path.join(self.root, self.raw_dirname, self.fold) + images_dirpath = os.path.join(fold_dirpath, "images") + image_filenames = fnmatch.filter(os.listdir(images_dirpath), "*_pre_disaster.png") + image_filenames = sorted(image_filenames) + disaster_samples_dict = defaultdict(int) + for image_filename in image_filenames: + name_split = image_filename.split("_") + disaster = name_split[0] + if self.small: + disaster_samples_dict[disaster] += 1 + if 10 < disaster_samples_dict[disaster]: + continue # Skip this sample as there is already enough for the small dataset + number = int(name_split[1]) + tile_info = { + "name": f"{disaster}_{number:08d}", + "disaster": disaster, + "number": number, + "image_filepath": os.path.join(fold_dirpath, "images", f"{disaster}_{number:08d}_pre_disaster.png"), + "label_filepath": os.path.join(fold_dirpath, "labels", f"{disaster}_{number:08d}_pre_disaster.json"), + "processed_relative_filepath": os.path.join(disaster, f"{number:08d}.pt") + } + tile_info_list.append(tile_info) + return tile_info_list + + def process(self, tile_info_list): + # os.makedirs(os.path.join(self.root, self.processed_dirname), exist_ok=True) + with multiprocess.Pool(self.pool_size) as p: + list_of_stats = list( + tqdm(p.imap(self._process_one, tile_info_list), total=len(tile_info_list), desc="Process")) + + # Aggregate stats + mean_per_disaster = defaultdict(list) + var_per_disaster = defaultdict(list) + class_freq = [] + for stats in list_of_stats: + mean_per_disaster[stats["disaster"]].append(stats["mean"]) + var_per_disaster[stats["disaster"]].append(stats["var"]) + class_freq.append(stats["class_freq"]) + stats = { + "mean": {}, + "std": {}, + "class_freq": None + } + for disaster in mean_per_disaster.keys(): + stats["mean"][disaster] = np.mean(np.stack(mean_per_disaster[disaster], axis=0), axis=0) + stats["std"][disaster] = np.sqrt(np.mean(np.stack(var_per_disaster[disaster], axis=0), axis=0)) + stats["class_freq"] = np.mean(np.stack(class_freq, axis=0), axis=0) + + return stats + + def load_raw_data(self, tile_info): + # Image: + tile_info["image"] = skimage.io.imread(tile_info["image_filepath"]) + assert len(tile_info["image"].shape) == 3 and tile_info["image"].shape[ + 2] == 3, f"image should have shape (H, W, 3), not {tile_info['image'].shape}..." + + # Annotations: + label_json = python_utils.load_json(tile_info["label_filepath"]) + features_xy = label_json["features"]["xy"] + tile_info["gt_polygons"] = [shapely.wkt.loads(obj["wkt"]) for obj in features_xy] + + return tile_info + + def _process_one(self, tile_info): + # --- Init + processed_tile_filepath = os.path.join(self.processed_dirpath, tile_info["processed_relative_filepath"]) + processed_tile_dirpath = os.path.dirname(processed_tile_filepath) + stats_filepath = os.path.join(processed_tile_dirpath, f"{tile_info['number']:08d}.stats.pt") + os.makedirs(processed_tile_dirpath, exist_ok=True) + + # --- Check if tile has been processed already + if os.path.exists(stats_filepath): + stats = torch.load(stats_filepath) + return stats + + tile_info = self.load_raw_data(tile_info) + + tile_info = self.pre_transform(tile_info) + + # Compute stats + stats = { + "mean": np.mean(tile_info["image"].reshape(-1, tile_info["image"].shape[-1]), axis=0) / 255, + "var": np.var(tile_info["image"].reshape(-1, tile_info["image"].shape[-1]), axis=0) / 255, + "class_freq": np.mean(tile_info["gt_polygons_image"], axis=(0, 1)) / 255, + "disaster": tile_info["disaster"] # Add disaster name to stats for aggregating per disaster + } + + # Save data + torch.save(tile_info, processed_tile_filepath) + torch.save(stats, stats_filepath) + + return stats + + def __len__(self): + if self.pre_process: + return len(self.processed_relative_paths) + else: + return len(self.tile_info_list) + + def __getitem__(self, idx): + if self.pre_process: + filepath = os.path.join(self.processed_dirpath, self.processed_relative_paths[idx]) + data = torch.load(filepath) + data["image_mean"] = self.stats["mean"][data["disaster"]] + data["image_std"] = self.stats["std"][data["disaster"]] + data["class_freq"] = self.stats["class_freq"] + else: + tile_info = self.tile_info_list[idx] + # Load raw data + data = self.load_raw_data(tile_info) + raise NotImplementedError("Need to implement mean and std computation") + + # --- Crop to path_size + height, width, _ = data["image"].shape + pre_crop_image_norm = data["image"].shape[0] + data["image"].shape[1] + crop_i = random.randint(0, height - self.patch_size) + crop_j = random.randint(0, width - self.patch_size) + data["image"] = data["image"][crop_i:crop_i + self.patch_size, crop_j:crop_j + self.patch_size] + data["gt_polygons_image"] = data["gt_polygons_image"][crop_i:crop_i + self.patch_size, crop_j:crop_j + self.patch_size] + data["gt_crossfield_angle"] = data["gt_crossfield_angle"][crop_i:crop_i + self.patch_size, crop_j:crop_j + self.patch_size] + data["distances"] = data["distances"][crop_i:crop_i + self.patch_size, crop_j:crop_j + self.patch_size] + data["sizes"] = data["sizes"][crop_i:crop_i + self.patch_size, crop_j:crop_j + self.patch_size] + post_crop_image_norm = data["image"].shape[0] + data["image"].shape[1] + # Sizes and distances are affected by cropping because they are relative to the image's norm (height + width). + # All non-one pixels have to be renormalized: + size_ratio = pre_crop_image_norm / post_crop_image_norm + data["distances"][data["distances"] != 1] *= size_ratio + data["sizes"][data["sizes"] != 1] *= size_ratio + # --- + + data = self.transform(data) + return 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...") + + +if __name__ == '__main__': + main() diff --git a/torch_lydorn/torchvision/transforms/__init__.py b/torch_lydorn/torchvision/transforms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f89076b73e7e58f618c4bb44a75a17638d50ef1f --- /dev/null +++ b/torch_lydorn/torchvision/transforms/__init__.py @@ -0,0 +1,40 @@ +from .transforms import * + +from .angle_field_init import AngleFieldInit +from .approximate_polygon import ApproximatePolygon +from .filter_empty_polygons import FilterEmptyPolygons +from .filter_poly_vertex_count import FilterPolyVertexCount +from .keep_keys import KeepKeys +from .map import Map +from .tensorpoly import polygons_to_tensorpoly, tensorpoly_pad +from .tensorskeleton import Paths, Skeleton, TensorSkeleton, skeletons_to_tensorskeleton, tensorskeleton_to_skeletons +from .rasterize import Rasterize +from .remove_doubles import RemoveDoubles +from .remove_keys import RemoveKeys +from .sample_uniform import SampleUniform +from .to_patches import ToPatches +from .transform_by_key import TransformByKey + + +__all__ = [ + 'functional', + 'AngleFieldInit', + 'ApproximatePolygon', + 'FilterEmptyPolygons', + 'FilterPolyVertexCount', + 'KeepKeys', + 'Map', + 'polygons_to_tensorpoly', + 'tensorpoly_pad', + 'Paths', + 'Skeleton', + 'TensorSkeleton', + 'skeletons_to_tensorskeleton', + 'tensorskeleton_to_skeletons', + 'Rasterize', + 'RemoveDoubles', + 'RemoveKeys', + 'SampleUniform', + 'ToPatches', + 'TransformByKey', +] diff --git a/torch_lydorn/torchvision/transforms/angle_field_init.py b/torch_lydorn/torchvision/transforms/angle_field_init.py new file mode 100644 index 0000000000000000000000000000000000000000..36203b9ce43dd5144f33b3db8aa8072c1136a1d3 --- /dev/null +++ b/torch_lydorn/torchvision/transforms/angle_field_init.py @@ -0,0 +1,70 @@ +from PIL import Image, ImageDraw, ImageFilter +import numpy as np +import shapely.geometry +import shapely.affinity +from scipy.ndimage.morphology import distance_transform_edt + +from functools import partial + +from . import functional + + +class AngleFieldInit(object): + def __init__(self, line_width=1): + self.line_width = line_width + + def __call__(self, image, polygons): + size = (image.shape[0], image.shape[1]) + return init_angle_field(polygons, size, self.line_width) + + +def init_angle_field(polygons, shape, line_width=1): + """ + Angle field {\theta_1} the tangent vector's angle for every pixel, specified on the polygon edges. + Angle between 0 and pi. + This is not invariant to symmetries. + + :param polygons: + :param shape: + :return: (angles: np.array((num_edge_pixels, ), dtype=np.uint8), + mask: np.array((num_edge_pixels, 2), dtype=np.int)) + """ + assert type(polygons) == list, "polygons should be a list" + if len(polygons): + assert type(polygons[0]) == shapely.geometry.Polygon, "polygon should be a shapely.geometry.Polygon" + + im = Image.new("L", (shape[1], shape[0])) + im_px_access = im.load() + draw = ImageDraw.Draw(im) + + for polygon in polygons: + draw_linear_ring(draw, polygon.exterior, line_width) + for interior in polygon.interiors: + draw_linear_ring(draw, interior, line_width) + + # Convert image to numpy array + array = np.array(im) + return array + + +def draw_linear_ring(draw, linear_ring, line_width): + # --- edges: + coords = np.array(linear_ring) + edge_vect_array = np.diff(coords, axis=0) + edge_angle_array = np.angle(edge_vect_array[:, 1] + 1j * edge_vect_array[:, 0]) # ij coord sys + neg_indices = np.where(edge_angle_array < 0) + edge_angle_array[neg_indices] += np.pi + + first_uint8_angle = None + for i in range(coords.shape[0] - 1): + edge = (coords[i], coords[i + 1]) + angle = edge_angle_array[i] + uint8_angle = int((255 * angle / np.pi).round()) + if first_uint8_angle is None: + first_uint8_angle = uint8_angle + line = [(edge[0][0], edge[0][1]), (edge[1][0], edge[1][1])] + draw.line(line, fill=uint8_angle, width=line_width) + functional.draw_circle(draw, line[0], radius=line_width / 2, fill=uint8_angle) + + # Add first vertex back on top (equals to last vertex too): + functional.draw_circle(draw, line[1], radius=line_width / 2, fill=first_uint8_angle) diff --git a/torch_lydorn/torchvision/transforms/approximate_polygon.py b/torch_lydorn/torchvision/transforms/approximate_polygon.py new file mode 100644 index 0000000000000000000000000000000000000000..e35561dfc2a5e6f3dcd120607fa36acf026940f0 --- /dev/null +++ b/torch_lydorn/torchvision/transforms/approximate_polygon.py @@ -0,0 +1,11 @@ +from skimage.measure import approximate_polygon + + +class ApproximatePolygon(object): + """Simplifies polygons""" + + def __init__(self, tolerance=0.1): + self.tolerance = tolerance + + def __call__(self, polygons): + return [approximate_polygon(polygon, tolerance=self.tolerance) for polygon in polygons] diff --git a/torch_lydorn/torchvision/transforms/filter_empty_polygons.py b/torch_lydorn/torchvision/transforms/filter_empty_polygons.py new file mode 100644 index 0000000000000000000000000000000000000000..d9da3be17983c494000cca16fd1a643ca13e74d6 --- /dev/null +++ b/torch_lydorn/torchvision/transforms/filter_empty_polygons.py @@ -0,0 +1,14 @@ +class FilterEmptyPolygons(object): + """Removes None elements of input data_list""" + + def __init__(self, key): + self.key = key + + def _filter(self, item): + if item[self.key] is not None: + return len(item[self.key]) + else: + return False + + def __call__(self, data_list): + return [item for item in data_list if self._filter(item)] diff --git a/torch_lydorn/torchvision/transforms/filter_poly_vertex_count.py b/torch_lydorn/torchvision/transforms/filter_poly_vertex_count.py new file mode 100644 index 0000000000000000000000000000000000000000..3a1f82ca692049345c2ad7bc4da70aa33032f658 --- /dev/null +++ b/torch_lydorn/torchvision/transforms/filter_poly_vertex_count.py @@ -0,0 +1,23 @@ +class FilterPolyVertexCount(object): + """ + if min is set, only keep polygons with at least min vertices + if max is set, only keep polygons with at most max vertices + """ + + def __init__(self, min=None, max=None): + self.min = min + self.max = max + if self.min is not None and self.max is not None: + if self.max < self.min: + print("WARNING: min and max of FilterPolyVertexCount() are {} and {}," + " which creates an impossible-to_satisfy condition.".format(self.min, self.max)) + + def __call__(self, polygons): + new_polygons = [] + for polygon in polygons: + if self.min is not None and polygon.shape[0] < self.min: + continue + if self.max is not None and self.max < polygon.shape[0]: + continue + new_polygons.append(polygon) + return new_polygons \ No newline at end of file diff --git a/torch_lydorn/torchvision/transforms/functional.py b/torch_lydorn/torchvision/transforms/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..552b482e3592b491f10f7b8d0f2907e4859d622d --- /dev/null +++ b/torch_lydorn/torchvision/transforms/functional.py @@ -0,0 +1,227 @@ +import numbers +import numpy as np +import torch +from torchvision.transforms.functional import _is_pil_image, _is_numpy_image + +from torch_lydorn.torch.utils.complex import complex_mul, complex_abs_squared + +try: + import accimage +except ImportError: + accimage = None + + +__all__ = ["to_tensor", "batch_normalize"] + + +def _is_numpy(img): + return isinstance(img, np.ndarray) + + +def to_tensor(pic): + """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor without typecasting and rescaling + + See ``ToTensor`` for more details. + + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + if not(_is_pil_image(pic) or _is_numpy(pic)): + raise TypeError('pic should be PIL Image or ndarray. Got {}'.format(type(pic))) + + if _is_numpy(pic) and not _is_numpy_image(pic): + raise ValueError('pic should be 2/3 dimensional. Got {} dimensions.'.format(pic.ndim)) + + if isinstance(pic, np.ndarray): + # handle numpy array + if pic.ndim == 2: + pic = pic[:, :, None] + + img = torch.from_numpy(pic.transpose((2, 0, 1))) + return img + + if accimage is not None and isinstance(pic, accimage.Image): + nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.uint8) + pic.copyto(nppic) + return torch.from_numpy(nppic) + + # handle PIL Image + if pic.mode == 'I': + img = torch.from_numpy(np.array(pic, np.int32, copy=False)) + elif pic.mode == 'I;16': + img = torch.from_numpy(np.array(pic, np.int16, copy=False)) + elif pic.mode == 'F': + img = torch.from_numpy(np.array(pic, np.float32, copy=False)) + elif pic.mode == '1': + img = 255 * torch.from_numpy(np.array(pic, np.uint8, copy=False)) + else: + img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes())) + # PIL image mode: L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK + if pic.mode == 'YCbCr': + nchannel = 3 + elif pic.mode == 'I;16': + nchannel = 1 + else: + nchannel = len(pic.mode) + img = img.view(pic.size[1], pic.size[0], nchannel) + # put it from HWC to CHW format + # yikes, this transpose takes 80% of the loading time/CPU + img = img.transpose(0, 1).transpose(0, 2).contiguous() + return img + + +def batch_normalize(tensor, mean, std, inplace=False): + """Normalize a batched tensor image with batched mean and batched standard deviation. + .. note:: + This transform acts out of place by default, i.e., it does not mutates the input tensor. + Args: + tensor (Tensor): Tensor image of size (B, C, H, W) to be normalized. + mean (sequence): Tensor means of size (B, C). + std (sequence): Tensor standard deviations of size (B, C). + inplace(bool,optional): Bool to make this operation inplace. + Returns: + Tensor: Normalized Tensor image. + """ + assert len(tensor.shape) == 4, \ + "tensor should have 4 dims (B, H, W, C) , not {}".format(len(tensor.shape)) + assert len(mean.shape) == len(std.shape) == 2, \ + "mean and std should have 2 dims (B, C) , not {} and {}".format(len(mean.shape), len(std.shape)) + assert tensor.shape[1] == mean.shape[1] == std.shape[1], \ + "tensor, mean and std should have the same number of channels, not {}, {} and {}".format(tensor.shape[1], mean.shape[1], std.shape[1]) + + if not inplace: + tensor = tensor.clone() + + mean = mean.to(tensor.dtype) + std = std.to(tensor.dtype) + + tensor.sub_(mean[..., None, None]).div_(std[..., None, None]) + return tensor + + +def batch_denormalize(tensor, mean, std, inplace=False): + """Denormalize a batched tensor image with batched mean and batched standard deviation. + .. note:: + This transform acts out of place by default, i.e., it does not mutates the input tensor. + Args: + tensor (Tensor): Tensor image of size (B, C, H, W) to be normalized. + mean (sequence): Tensor means of size (B, C). + std (sequence): Tensor standard deviations of size (B, C). + inplace(bool,optional): Bool to make this operation inplace. + Returns: + Tensor: Normalized Tensor image. + """ + assert len(tensor.shape) == 4, \ + "tensor should have 4 dims (B, H, W, C) , not {}".format(len(tensor.shape)) + assert len(mean.shape) == len(std.shape) == 2, \ + "mean and std should have 2 dims (B, C) , not {} and {}".format(len(mean.shape), len(std.shape)) + assert tensor.shape[1] == mean.shape[1] == std.shape[1], \ + "tensor, mean and std should have the same number of channels, not {}, {} and {}".format( tensor.shape[-1], mean.shape[-1], std.shape[-1]) + + if not inplace: + tensor = tensor.clone() + + mean = mean.to(tensor.dtype) + std = std.to(tensor.dtype) + + tensor.mul_(std[..., None, None]).add_(mean[..., None, None]) + return tensor + + +def crop(tensor, top, left, height, width): + """Crop the given Tensor batch of images. + Args: + tensor (B, C, H, W): Tensor to be cropped. (0,0) denotes the top left corner of the image. + top (int): Vertical component of the top left corner of the crop box. + left (int): Horizontal component of the top left corner of the crop box. + height (int): Height of the crop box. + width (int): Width of the crop box. + Returns: + Tensor: Cropped image. + """ + return tensor[..., top:top+height, left:left+width] + + +def center_crop(tensor, output_size): + """Crop the given tensor batch of images and resize it to desired size. + + Args: + tensor (B, C, H, W): Tensor to be cropped. + output_size (sequence or int): (height, width) of the crop box. If int, + it is used for both directions + Returns: + Tensor: Cropped tensor. + """ + if isinstance(output_size, numbers.Number): + output_size = (int(output_size), int(output_size)) + tensor_height, tensor_width = tensor.shape[-2:] + crop_height, crop_width = output_size + crop_top = int(round((tensor_height - crop_height) / 2.)) + crop_left = int(round((tensor_width - crop_width) / 2.)) + return crop(tensor, crop_top, crop_left, crop_height, crop_width) + + +def rotate_anglefield(angle_field, angle): + """ + :param angle_field: (B, 1, H, W), in radians + :param angle_deg: (B) in degrees + :return: + """ + assert len(angle_field.shape) == 4, "angle_field should have shape (B, 1, H, W)" + assert len(angle.shape) == 1, "angle should have shape (B)" + assert angle_field.shape[0] == angle.shape[0], "angle_field and angle should have the same batch size" + angle_field += np.pi * angle[:, None, None, None] / 180 + return angle_field + + +def vflip_anglefield(angle_field): + """ + + :param angle_field: (B, 1, H, W), in radians + :return: + """ + assert len(angle_field.shape) == 4, "angle_field should have shape (B, 1, H, W)" + angle_field = np.pi - angle_field # Angle is expressed in ij coordinate (it's a horizontal flip in xy) + return angle_field + + +def rotate_framefield(framefield, angle): + """ + ONly rotates values of the framefield, does not rotate the spatial domain (use already-made torch functions to do that). + + @param framefield: shape (B, 4, H, W). The 4 channels represent the c_0 and c_2 complex coefficients. + @param angle: in degrees + @return: + """ + assert framefield.shape[1] == 4, f"framefield should have shape (B, 4, H, W), not {framefield.shape}" + rad = np.pi * angle / 180 + z_4angle = torch.tensor([np.cos(4*rad), np.sin(4*rad)], dtype=framefield.dtype, device=framefield.device) + z_2angle = torch.tensor([np.cos(2*rad), np.sin(2*rad)], dtype=framefield.dtype, device=framefield.device) + framefield[:, :2, :, :] = complex_mul(framefield[:, :2, :, :], z_4angle[None, :, None, None], complex_dim=1) + framefield[:, 2:, :, :] = complex_mul(framefield[:, 2:, :, :], z_2angle[None, :, None, None], complex_dim=1) + return framefield + + +def vflip_framefield(framefield): + """ + Flips the framefield vertically. This means switching the signs of the real part of u and v + (this is because the framefield is in ij coordinates: it's a horizontal flip in xy), which translates to + switching the signs of the imaginary parts of c_0 and c_2. + + @param framefield: shape (B, 4, H, W). The 4 channels represent the c_0 and c_2 complex coefficients. + @return: + """ + assert framefield.shape[1] == 4, f"framefield should have shape (B, 4, H, W), not {framefield.shape}" + framefield[:, 1, :, :] = - framefield[:, 1, :, :] + framefield[:, 3, :, :] = - framefield[:, 3, :, :] + return framefield + + +def draw_circle(draw, center, radius, fill): + draw.ellipse([center[0] - radius, + center[1] - radius, + center[0] + radius, + center[1] + radius], fill=fill, outline=None) diff --git a/torch_lydorn/torchvision/transforms/keep_keys.py b/torch_lydorn/torchvision/transforms/keep_keys.py new file mode 100644 index 0000000000000000000000000000000000000000..842a893e66fb68586fc54c9a012addb8efe42ec3 --- /dev/null +++ b/torch_lydorn/torchvision/transforms/keep_keys.py @@ -0,0 +1,7 @@ +class KeepKeys(object): + + def __init__(self, keys): + self.keys = keys + + def __call__(self, data): + return {key: data[key] for key in self.keys if key in data} diff --git a/torch_lydorn/torchvision/transforms/map.py b/torch_lydorn/torchvision/transforms/map.py new file mode 100644 index 0000000000000000000000000000000000000000..ccb29920d811ed5bac057b749a2f4f919f9773e1 --- /dev/null +++ b/torch_lydorn/torchvision/transforms/map.py @@ -0,0 +1,28 @@ +# class Map(object): +# """Applies parameter transform to all items of input list""" +# +# def __init__(self, transform, multi_outs=False): +# self.transform = transform +# if multi_outs: +# self.__call__ = self._map_multi_outs +# else: +# self.__call__ = self._map_single_out +# +# def _map_single_out(self, data_list): +# assert type(data_list) == list, "data_list should be a list" +# return [self.transform(item) for item in data_list] +# +# def _map_multi_outs(self, data_list): +# assert type(data_list) == list, "data_list should be a list" +# return tuple(zip(*[self.transform(item) for item in data_list])) + + +class Map(object): + """Applies parameter transform to all items of input list""" + + def __init__(self, transform): + self.transform = transform + + def __call__(self, data_list): + assert type(data_list) == list, "data_list should be a list" + return [self.transform(item) for item in data_list] diff --git a/torch_lydorn/torchvision/transforms/rasterize.py b/torch_lydorn/torchvision/transforms/rasterize.py new file mode 100644 index 0000000000000000000000000000000000000000..2eb6e2925681fdbd64f7713dd7a4082a847fcdd2 --- /dev/null +++ b/torch_lydorn/torchvision/transforms/rasterize.py @@ -0,0 +1,230 @@ +import math +import sys +import time + +import skimage.morphology +import skimage.io +from PIL import Image, ImageDraw, ImageFilter +import numpy as np +import shapely.geometry +import shapely.affinity +from lydorn_utils import print_utils +from scipy.ndimage.morphology import distance_transform_edt +import cv2 as cv + +from functools import partial + +import torch_lydorn.torchvision + + +class Rasterize(object): + """Rasterize polygons""" + + def __init__(self, fill=True, edges=True, vertices=True, line_width=3, antialiasing=False, return_distances=False, + return_sizes=False): + self.fill = fill + self.edges = edges + self.vertices = vertices + self.line_width = line_width + self.antialiasing = antialiasing + + if not return_distances and not return_sizes: + self.raster_func = partial(draw_polygons, fill=self.fill, edges=self.edges, vertices=self.vertices, + line_width=self.line_width, antialiasing=self.antialiasing) + elif return_distances and return_sizes: + self.raster_func = partial(compute_raster_distances_sizes, fill=self.fill, edges=self.edges, vertices=self.vertices, + line_width=self.line_width, antialiasing=self.antialiasing) + else: + raise NotImplementedError + + def __call__(self, image, polygons): + """ + If distances is True, also returns distances image + (sum of distance to closest and second-closest annotation for each pixel). + Same for sizes (size of annotation the pixel belongs to). + + """ + size = (image.shape[0], image.shape[1]) + out = self.raster_func(polygons, size) + return out + + +def compute_raster_distances_sizes(polygons, shape, fill=True, edges=True, vertices=True, line_width=3, antialiasing=False): + """ + Returns: + - distances: sum of distance to closest and second-closest annotation for each pixel. + - size_weights: relative size (normalized by image area) of annotation the pixel belongs to. + """ + assert type(polygons) == list, "polygons should be a list" + + # Filter out zero-area polygons + polygons = [polygon for polygon in polygons if 0 < polygon.area] + + # tic = time.time() + + channel_count = fill + edges + vertices + polygons_raster = np.zeros((*shape, channel_count), dtype=np.uint8) + distance_maps = np.ones((*shape, len(polygons))) # Init with max value (distances are normed) + sizes = np.ones(shape) # Init with max value (sizes are normed) + image_area = shape[0] * shape[1] + for i, polygon in enumerate(polygons): + minx, miny, maxx, maxy = polygon.bounds + mini = max(0, math.floor(miny) - 2*line_width) + minj = max(0, math.floor(minx) - 2*line_width) + maxi = min(polygons_raster.shape[0], math.ceil(maxy) + 2*line_width) + maxj = min(polygons_raster.shape[1], math.ceil(maxx) + 2*line_width) + bbox_shape = (maxi - mini, maxj - minj) + bbox_polygon = shapely.affinity.translate(polygon, xoff=-minj, yoff=-mini) + bbox_raster = draw_polygons([bbox_polygon], bbox_shape, fill, edges, vertices, line_width, antialiasing) + polygons_raster[mini:maxi, minj:maxj] = np.maximum(polygons_raster[mini:maxi, minj:maxj], bbox_raster) + bbox_mask = 0 < np.sum(bbox_raster, axis=2) # Polygon interior + edge + vertex + if bbox_mask.max(): # Make sure mask is not empty + polygon_mask = np.zeros(shape, dtype=np.bool) + polygon_mask[mini:maxi, minj:maxj] = bbox_mask + polygon_dist = cv.distanceTransform(1 - polygon_mask.astype(np.uint8), distanceType=cv.DIST_L2, maskSize=cv.DIST_MASK_5, + dstType=cv.CV_64F) + polygon_dist /= (polygon_mask.shape[0] + polygon_mask.shape[1]) # Normalize dist + distance_maps[:, :, i] = polygon_dist + + selem = skimage.morphology.disk(line_width) + bbox_dilated_mask = skimage.morphology.binary_dilation(bbox_mask, selem=selem) + sizes[mini:maxi, minj:maxj][bbox_dilated_mask] = polygon.area / image_area + + polygons_raster = np.clip(polygons_raster, 0, 255) + # skimage.io.imsave("polygons_raster.png", polygons_raster) + + if edges: + edge_channels = -1 + fill + edges + # Remove border edges because they correspond to cut buildings: + polygons_raster[:line_width, :, edge_channels] = 0 + polygons_raster[-line_width:, :, edge_channels] = 0 + polygons_raster[:, :line_width, edge_channels] = 0 + polygons_raster[:, -line_width:, edge_channels] = 0 + + distances = compute_distances(distance_maps) + # skimage.io.imsave("distances.png", distances) + + distances = distances.astype(np.float16) + sizes = sizes.astype(np.float16) + + # toc = time.time() + # print(f"Rasterize {len(polygons)} polygons: {toc - tic}s") + + return polygons_raster, distances, sizes + + +def compute_distances(distance_maps): + distance_maps.sort(axis=2) + distance_maps = distance_maps[:, :, :2] + distances = np.sum(distance_maps, axis=2) + return distances + + +def draw_polygons(polygons, shape, fill=True, edges=True, vertices=True, line_width=3, antialiasing=False): + assert type(polygons) == list, "polygons should be a list" + assert type(polygons[0]) == shapely.geometry.Polygon, "polygon should be a shapely.geometry.Polygon" + + if antialiasing: + draw_shape = (2 * shape[0], 2 * shape[1]) + polygons = [shapely.affinity.scale(polygon, xfact=2.0, yfact=2.0, origin=(0, 0)) for polygon in polygons] + line_width *= 2 + else: + draw_shape = shape + # Channels + fill_channel_index = 0 # Always first channel + edges_channel_index = fill # If fill == True, take second channel. If not then take first + vertices_channel_index = fill + edges # Same principle as above + channel_count = fill + edges + vertices + im_draw_list = [] + for channel_index in range(channel_count): + im = Image.new("L", (draw_shape[1], draw_shape[0])) + im_px_access = im.load() + draw = ImageDraw.Draw(im) + im_draw_list.append((im, draw)) + + for polygon in polygons: + if fill: + draw = im_draw_list[fill_channel_index][1] + draw.polygon(polygon.exterior.coords, fill=255) + for interior in polygon.interiors: + draw.polygon(interior.coords, fill=0) + if edges: + draw = im_draw_list[edges_channel_index][1] + draw.line(polygon.exterior.coords, fill=255, width=line_width) + for interior in polygon.interiors: + draw.line(interior.coords, fill=255, width=line_width) + if vertices: + draw = im_draw_list[vertices_channel_index][1] + for vertex in polygon.exterior.coords: + torch_lydorn.torchvision.transforms.functional.draw_circle(draw, vertex, line_width / 2, fill=255) + for interior in polygon.interiors: + for vertex in interior.coords: + torch_lydorn.torchvision.transforms.functional.draw_circle(draw, vertex, line_width / 2, fill=255) + + im_list = [] + if antialiasing: + # resize images: + for im_draw in im_draw_list: + resize_shape = (shape[1], shape[0]) + im_list.append(im_draw[0].resize(resize_shape, Image.BILINEAR)) + else: + for im_draw in im_draw_list: + im_list.append(im_draw[0]) + + # Convert image to numpy array with the right number of channels + array_list = [np.array(im) for im in im_list] + array = np.stack(array_list, axis=-1) + return array + + +def _rasterize_coco(image, polygons): + import pycocotools.mask as cocomask + + image_size = image.shape[:2] + mask = np.zeros(image_size) + for polygon in polygons: + rle = cocomask.frPyObjects([np.array(polygon.exterior.coords).reshape(-1)], image_size[0], image_size[1]) + m = cocomask.decode(rle) + + for i in range(m.shape[-1]): + mi = m[:, :, i] + mi = mi.reshape(image_size) + mask += mi + return mask + + +def _test(): + import skimage.io + + rasterize = Rasterize(fill=True, edges=False, vertices=False, line_width=2, antialiasing=True, return_distances=True, return_sizes=True) + + image = np.zeros((300, 300)) + polygons = [ + shapely.geometry.Polygon([ + [10.5, 10.5], + [100, 10], + [100, 150], + [10, 100], + [10, 10], + ]), + shapely.geometry.Polygon([ + [10+150, 10], + [100+150, 10], + [100+150, 100], + [10+150, 100], + [10+150, 10], + ]), + ] + polygons_raster, distances, size_weights = rasterize(image, polygons) + + skimage.io.imsave('rasterize.polygons_raster.png', polygons_raster) + skimage.io.imsave('rasterize.distances.png', distances) + skimage.io.imsave('rasterize.size_weights.png', size_weights) + + # Rasterize with pycocotools + coco_mask = _rasterize_coco(image, polygons) + skimage.io.imsave('rasterize.coco_mask.png', coco_mask) + + +if __name__ == "__main__": + _test() \ No newline at end of file diff --git a/torch_lydorn/torchvision/transforms/remove_doubles.py b/torch_lydorn/torchvision/transforms/remove_doubles.py new file mode 100644 index 0000000000000000000000000000000000000000..38d98ac582ed434a3e42ba92f89808c6a970eb06 --- /dev/null +++ b/torch_lydorn/torchvision/transforms/remove_doubles.py @@ -0,0 +1,11 @@ +from lydorn_utils import polygon_utils + + +class RemoveDoubles(object): + """Removes redundant vertices of all input polygons""" + + def __init__(self, epsilon=0.1): + self.epsilon = epsilon + + def __call__(self, polygons): + return [polygon_utils.remove_doubles(polygon, epsilon=self.epsilon) for polygon in polygons] diff --git a/torch_lydorn/torchvision/transforms/remove_keys.py b/torch_lydorn/torchvision/transforms/remove_keys.py new file mode 100644 index 0000000000000000000000000000000000000000..b431b14550441d680f4825eb4a99679bb74019ff --- /dev/null +++ b/torch_lydorn/torchvision/transforms/remove_keys.py @@ -0,0 +1,7 @@ +class RemoveKeys(object): + + def __init__(self, keys): + self.keys = keys + + def __call__(self, data): + return {key: data[key] for key in data.keys() if key not in self.keys} diff --git a/torch_lydorn/torchvision/transforms/sample_uniform.py b/torch_lydorn/torchvision/transforms/sample_uniform.py new file mode 100644 index 0000000000000000000000000000000000000000..c1071ec5c42236e89511e6edf472935844892b5e --- /dev/null +++ b/torch_lydorn/torchvision/transforms/sample_uniform.py @@ -0,0 +1,10 @@ +import torch + + +class SampleUniform(object): + + def __init__(self, low, high): + self.m = torch.distributions.uniform.Uniform(low, high) + + def __call__(self): + return self.m.sample() diff --git a/torch_lydorn/torchvision/transforms/tensorpoly.py b/torch_lydorn/torchvision/transforms/tensorpoly.py new file mode 100644 index 0000000000000000000000000000000000000000..1eef1bce9296f7d2e8561ba5a34e93d2c7d5c6fe --- /dev/null +++ b/torch_lydorn/torchvision/transforms/tensorpoly.py @@ -0,0 +1,181 @@ +import numpy as np + +import torch + + +class TensorPoly(object): + def __init__(self, pos, poly_slice, batch, batch_size, is_endpoint=None): + """ + + :param pos: + :param poly_slice: + :param batch: Batch index for each node + :param is_endpoint: One value per node. If true, that node is an endpoint and is thus part of an open polyline + """ + assert pos.shape[0] == batch.shape[0] + self.pos = pos + self.poly_slice = poly_slice + self.batch = batch + self.batch_size = batch_size + self.is_endpoint = is_endpoint + + self.to_padded_index = None # No pad initially + self.to_unpadded_poly_slice = None + + @property + def num_nodes(self): + return self.pos.shape[0] + + def to(self, device): + self.pos = self.pos.to(device) + self.poly_slice = self.poly_slice.to(device) + self.batch = self.batch.to(device) + if self.is_endpoint is not None: + self.is_endpoint = self.is_endpoint.to(device) + if self.to_padded_index is not None: + self.to_padded_index = self.to_padded_index.to(device) + if self.to_unpadded_poly_slice is not None: + self.to_unpadded_poly_slice = self.to_unpadded_poly_slice.to(device) + + +def polygons_to_tensorpoly(polygons_batch): + """ + Parametrizes N polygons into a 1d grid to be used in 1d conv later on: + - pos (n1+n2+..., 2) concatenation of polygons vertex positions + - poly_slice (poly_count, 2) polygon vertex slices [start, end) + + :param polygons_batch: Batch of polygons: [[(n1, 2), (n2, 2), ...], ...] + :return: TensorPoly(pos, poly_slice, batch, batch_size, is_endpoint) + """ + # TODO: If there are no polygons + batch_size = len(polygons_batch) + is_endpoint_list = [] + batch_list = [] + polygon_list = [] + for i, polygons in enumerate(polygons_batch): + for polygon in polygons: + if not np.max(np.abs(polygon[0] - polygon[-1])) < 1e-6: + # Polygon is open + is_endpoint = np.zeros(polygon.shape[0], dtype=np.bool) + is_endpoint[0] = True + is_endpoint[-1] = True + else: + # Polygon is closed, remove last redundant point + polygon = polygon[:-1, :] + is_endpoint = np.zeros(polygon.shape[0], dtype=np.bool) + batch = i * np.ones(polygon.shape[0], dtype=np.long) + is_endpoint_list.append(is_endpoint) + batch_list.append(batch) + polygon_list.append(polygon) + pos = np.concatenate(polygon_list, axis=0) + is_endpoint = np.concatenate(is_endpoint_list, axis=0) + batch = np.concatenate(batch_list, axis=0) + + slice_start = 0 + poly_slice = np.empty((len(polygon_list), 2), dtype=np.long) + for i, polygon in enumerate(polygon_list): + slice_end = slice_start + polygon.shape[0] + poly_slice[i][0] = slice_start + poly_slice[i][1] = slice_end + slice_start = slice_end + pos = torch.tensor(pos, dtype=torch.float) + is_endpoint = torch.tensor(is_endpoint, dtype=torch.bool) + poly_slice = torch.tensor(poly_slice, dtype=torch.long) + batch = torch.tensor(batch, dtype=torch.long) + tensorpoly = TensorPoly(pos=pos, poly_slice=poly_slice, batch=batch, batch_size=batch_size, is_endpoint=is_endpoint) + return tensorpoly + + +def _get_to_padded_index(poly_slice, node_count, padding): + """ + Pad each polygon with a cyclic padding scheme on both sides. + Increases length of x by (padding[0] + padding[1])*polygon_count values. + + :param poly_slice: + :param padding: + :return: + """ + assert len(poly_slice.shape) == 2, "poly_slice should have shape (poly_count, 2), not {}".format(poly_slice.shape) + poly_count = poly_slice.shape[0] + range_tensor = torch.arange(node_count, device=poly_slice.device) + to_padded_index = torch.empty((node_count + (padding[0] + padding[1])*poly_count, ), dtype=torch.long, device=poly_slice.device) + to_unpadded_poly_slice = torch.empty_like(poly_slice) + start = 0 + for poly_i in range(poly_count): + poly_indices = range_tensor[poly_slice[poly_i, 0]:poly_slice[poly_i, 1]] + + # Repeat poly_indices if necessary when padding exceeds polygon length + vertex_count = poly_indices.shape[0] + left_repeats = padding[0] // vertex_count + left_padding_remaining = padding[0] % vertex_count + right_repeats = padding[1] // vertex_count + right_padding_remaining = padding[1] % vertex_count + total_repeats = left_repeats + right_repeats + poly_indices = poly_indices.repeat(total_repeats + 1) # +1 includes the original polygon + + if left_padding_remaining: + poly_indices = torch.cat([poly_indices[-left_padding_remaining:], + poly_indices, + poly_indices[:right_padding_remaining]]) + else: + poly_indices = torch.cat([poly_indices, poly_indices[:right_padding_remaining]]) + + end = start + poly_indices.shape[0] + to_padded_index[start:end] = poly_indices + to_unpadded_poly_slice[poly_i, 0] = start # Init value + to_unpadded_poly_slice[poly_i, 1] = end # Init value + start = end + to_unpadded_poly_slice[:, 0] += padding[0] # Shift all inited values to the right one + to_unpadded_poly_slice[:, 1] -= padding[1] # Shift all inited values to the right one + return to_padded_index, to_unpadded_poly_slice + + +def tensorpoly_pad(tensorpoly, padding): + to_padded_index, to_unpadded_poly_slice = _get_to_padded_index(tensorpoly.poly_slice, tensorpoly.num_nodes, padding) + tensorpoly.to_padded_index = to_padded_index + tensorpoly.to_unpadded_poly_slice = to_unpadded_poly_slice + return tensorpoly + + +def main(): + device = "cuda" + + np.random.seed(0) + torch.manual_seed(0) + padding = (0, 1) + + batch_size = 2 + poly_count = 3 + vertex_min_count = 4 + vertex_max_count = 5 + + polygons_batch = [] + for batch_i in range(batch_size): + polygons = [] + for poly_i in range(poly_count): + vertex_count = np.random.randint(vertex_min_count, vertex_max_count) + polygon = np.random.uniform(0, 1, (vertex_count, 2)) + polygons.append(polygon) + polygons_batch.append(polygons) + print("polygons_batch:") + print(polygons_batch) + tensorpoly = polygons_to_tensorpoly(polygons_batch) + print("batch:") + print(tensorpoly.batch) + print("pos:") + print(tensorpoly.pos.shape) + print("poly_slice:") + print(tensorpoly.poly_slice.shape) + print(tensorpoly.poly_slice) + + tensorpoly.to(device) + + tensorpoly = tensorpoly_pad(tensorpoly, padding) + to_padded_index = tensorpoly.to_padded_index + print("to_padded_index:") + print(to_padded_index.shape) + print(to_padded_index) + + +if __name__ == "__main__": + main() diff --git a/torch_lydorn/torchvision/transforms/tensorskeleton.py b/torch_lydorn/torchvision/transforms/tensorskeleton.py new file mode 100644 index 0000000000000000000000000000000000000000..a5d0df101a199c342371fd5cb569b2ffd440129c --- /dev/null +++ b/torch_lydorn/torchvision/transforms/tensorskeleton.py @@ -0,0 +1,265 @@ +import sys +import time +from typing import List + +import matplotlib.pyplot as plt + +import numpy as np +import skan +import scipy.sparse + +import torch + +# GET_POS_INDEX_TIME = 0 + + +class Skeleton: + def __init__(self, coordinates=None, paths=None, degrees=None): + if coordinates is None: + self.coordinates = np.empty((0, 2), dtype=np.float) + else: + self.coordinates = coordinates + if paths is None: + self.paths = Paths() + else: + self.paths = paths + if degrees is None: + self.degrees = np.empty(0, dtype=np.long) + else: + self.degrees = degrees + + +class Paths: + def __init__(self, indices=None, indptr=None): + if indices is None: + self.indices = np.empty(0, dtype=np.long) + else: + self.indices = indices + if indptr is None: + self.indptr = np.empty(0, dtype=np.long) + else: + self.indptr = indptr + + +class TensorSkeleton(object): + def __init__(self, pos, degrees, path_index, path_delim, batch, batch_delim, batch_size): + """ + In the text below, we use the following notation: + - B: batch size + - N: the number of points in all skeletons, + - P: the number of paths in the skeletons + - J: the number of junction nodes + - Sd: the sum of the degrees of all the junction nodes + + :param pos (N, 2): union of skeleton points in ij format + :param degrees (N,): Degrees of each node in the graph + :param path_index (N - J + Sd,): Indices in pos of all paths (equivalent to 'indices' in the paths crs matrix) + :param path_delim (P + 1,): Indices in path_index delimiting each path (equivalent to 'indptr' in the paths crs matrix) + :param batch (N,): batch index of each point + :param batch_delim (B + 1,): Indices in path_delim delimiting each batch + """ + assert pos.shape[0] == batch.shape[0] + self.pos = pos + self.degrees = degrees + self.path_index = path_index + self.path_delim = path_delim + self.batch = batch + self.batch_delim = batch_delim + self.batch_size = batch_size + + @property + def num_nodes(self): + return self.pos.shape[0] + + @property + def num_paths(self): + return max(0, self.path_delim.shape[0] - 1) + + def to(self, device): + self.pos = self.pos.to(device) + self.degrees = self.degrees.to(device) + self.path_index = self.path_index.to(device) + self.path_delim = self.path_delim.to(device) + self.batch = self.batch.to(device) + self.batch_delim = self.batch_delim.to(device) + + +def skeletons_to_tensorskeleton(skeletons_batch: List[Skeleton], device: str=None) -> TensorSkeleton: + """ + In the text below, we use the following notation: + - B: batch size + - N: the number of points in all skeletons, + - P: the number of paths in the skeletons + - J: the number of junction nodes + - Sd: the sum of the degrees of all the junction nodes + + Parametrizes B skeletons into PyTorch tensors: + - pos (N, 2): union of skeleton points in ij format + - path_index (N - J + Sd,): Indices in pos of all paths (equivalent to 'indices' in the paths crs matrix) + - path_delim (P + 1,): Indices in path_index delimiting each path (equivalent to 'indptr' in the paths crs matrix) + - batch (N,): batch index of each point + + :param skeletons_batch: Batch of coordinates of skeletons [Skeleton(coordinates, paths(indices, indptr), degrees), ...] + :return: TensorSkeleton(pos, path_index, path_delim, batch, batch_size) + """ + batch_size = len(skeletons_batch) + pos_list = [] + degrees_list = [] + path_index_offset = 0 + path_index_list = [] + path_delim_offset = 0 + path_delim_list = [] + batch_list = [] + batch_delim_offset = 0 + batch_delim_list = [] + if 0 < batch_size: + batch_delim_list.append(0) + for batch_i, skeleton in enumerate(skeletons_batch): + n_points = skeleton.coordinates.shape[0] + paths_length = skeleton.paths.indices.shape[0] + n_paths = max(0, skeleton.paths.indptr.shape[0] - 1) + pos_list.append(skeleton.coordinates) + degrees_list.append(skeleton.degrees) + path_index = skeleton.paths.indices + path_index_offset + path_index_list.append(path_index) + if batch_i < batch_size - 1: + # Remove last item of indptr because it will be repeated by the first item of the next indptr + path_delim = skeleton.paths.indptr[:-1] + else: + path_delim = skeleton.paths.indptr + path_delim += path_delim_offset + path_delim_list.append(path_delim) + batch_list.append(batch_i * np.ones(n_points, dtype=np.long)) + + # Setup next batch: + path_index_offset += n_points + path_delim_offset += paths_length + batch_delim_offset += n_paths + + batch_delim_list.append(batch_delim_offset) + + pos = np.concatenate(pos_list, axis=0) + degrees = np.concatenate(degrees_list, axis=0) + path_index = np.concatenate(path_index_list, axis=0) + path_delim = np.concatenate(path_delim_list, axis=0) + batch = np.concatenate(batch_list, axis=0) + + pos = torch.tensor(pos, dtype=torch.float, device=device) + degrees = torch.tensor(degrees, dtype=torch.long, device=device) + path_index = torch.tensor(path_index, dtype=torch.long, device=device) + path_delim = torch.tensor(path_delim, dtype=torch.long, device=device) + batch = torch.tensor(batch, dtype=torch.long, device=device) + batch_delim = torch.tensor(batch_delim_list, dtype=torch.long, device=device) + tensorpoly = TensorSkeleton(pos=pos, degrees=degrees, path_index=path_index, path_delim=path_delim, batch=batch, batch_delim=batch_delim, batch_size=batch_size) + + # toc = time.time() + # print(f"polylines_to_tensorskeleton: {toc - tic}s") + # print(f"Get pos index total: {GET_POS_INDEX_TIME}s") + + return tensorpoly + + +def tensorskeleton_to_skeletons(tensorskeleton: TensorSkeleton) -> List[Skeleton]: + skeletons_list = [] + path_index_offset = 0 + path_delim_offset = 0 + for batch_i in range(tensorskeleton.batch_size): + batch_slice = tensorskeleton.batch_delim[batch_i:batch_i+2] + # print("batch_slice:", batch_slice) + # print("path_delim:", tensorskeleton.path_delim.shape) + indptr = tensorskeleton.path_delim[batch_slice[0]:batch_slice[1] + 1].cpu().numpy() + # print("indptr:", indptr) + # print("path_index:", tensorskeleton.path_index.shape) + indices = tensorskeleton.path_index[indptr[0]:indptr[-1]].cpu().numpy() + # print("indices:", indices) + if 2 <= indptr.shape[0]: + coordinates = tensorskeleton.pos[tensorskeleton.batch == batch_i].detach().cpu().numpy() + + indices = indices - path_index_offset + indptr = indptr - path_delim_offset + skeleton = Skeleton(coordinates, Paths(indices, indptr)) + skeletons_list.append(skeleton) + + n_points = coordinates.shape[0] + paths_length = indices.shape[0] + path_index_offset += n_points + path_delim_offset += paths_length + else: + # Empty skeleton + skeleton = Skeleton() + skeletons_list.append(skeleton) + + return skeletons_list + + +def plot_skeleton(skeleton: Skeleton): + for path_i in range(skeleton.paths.indptr.shape[0] - 1): + start, stop = skeleton.paths.indptr[path_i:path_i + 2] + path_indices = skeleton.paths.indices[start:stop] + path_coordinates = skeleton.coordinates[path_indices] + plt.plot(path_coordinates[:, 1], path_coordinates[:, 0]) + + +def main(): + + device = "cuda" + + np.random.seed(0) + torch.manual_seed(0) + + spatial_shape = (10, 10) + + skan_skeletons_batch = [] + skeleton_image = np.zeros(spatial_shape, dtype=np.bool) + skeleton_image[2, :] = True + skeleton_image[:, 2] = True + skeleton_image[7, :] = True + skeleton_image[:, 7] = True + skan_skeleton = skan.Skeleton(skeleton_image, keep_images=False) + skan_skeletons_batch.append(skan_skeleton) + # plt.imshow(skeleton_image) + plot_skeleton(skan_skeleton) + plt.show() + + skeleton_image = np.zeros(spatial_shape, dtype=np.bool) + skeleton_image[5, :] = True + skeleton_image[:, 5] = True + skan_skeleton = skan.Skeleton(skeleton_image, keep_images=False) + skan_skeletons_batch.append(skan_skeleton) + # plt.imshow(skeleton_image) + plot_skeleton(skan_skeleton) + plt.show() + + skeletons_batch = [Skeleton(skan_skeleton.coordinates, Paths(skan_skeleton.paths.indices, skan_skeleton.paths.indptr)) for skan_skeleton in skan_skeletons_batch] + + print("# --- skeletons_to_tensorskeleton() --- #") + tensorskeleton = skeletons_to_tensorskeleton(skeletons_batch, device=device) + print("# --- --- #") + # print("batch:") + # print(tensorskeleton.batch) + # print("pos:") + # print(tensorskeleton.pos.shape) + # print(tensorskeleton.pos) + print("path_index:") + print(tensorskeleton.path_index.shape) + print(tensorskeleton.path_index) + print("path_delim:") + print(tensorskeleton.path_delim.shape) + print(tensorskeleton.path_delim) + print("batch_delim:") + print(tensorskeleton.batch_delim.shape) + print(tensorskeleton.batch_delim) + + print("# --- tensorskeleton_to_skeletons() --- #") + skeletons_batch = tensorskeleton_to_skeletons(tensorskeleton) + + # Plot + for skeleton in skeletons_batch: + plot_skeleton(skeleton) + plt.show() + + print("# --- --- #") + + +if __name__ == "__main__": + main() diff --git a/torch_lydorn/torchvision/transforms/to_patches.py b/torch_lydorn/torchvision/transforms/to_patches.py new file mode 100644 index 0000000000000000000000000000000000000000..9b7d3fe2a79b1bdeffd8073bcc9fd4c6bc4fbdad --- /dev/null +++ b/torch_lydorn/torchvision/transforms/to_patches.py @@ -0,0 +1,42 @@ +from tqdm import tqdm + +import torch + +from lydorn_utils import image_utils +from lydorn_utils import polygon_utils + + +class ToPatches(object): + """Splits sample into patches""" + + def __init__(self, stride, size): + self.stride = stride + self.size = size + + def _to_patch(self, sample): + image = sample["image"] + gt_polygons = sample["gt_polygons"] + + patch_boundingboxes = image_utils.compute_patch_boundingboxes(image.shape[0:2], + stride=self.stride, + patch_res=self.size) + + patches = [] + for patch_boundingbox in tqdm(patch_boundingboxes, desc="Patching", leave=False): + # Crop image + patch_image = image[patch_boundingbox[0]:patch_boundingbox[2], patch_boundingbox[1]:patch_boundingbox[3], :] + patch_gt_polygon = polygon_utils.crop_polygons_to_patch_if_touch(gt_polygons, patch_boundingbox) + if len(patch_gt_polygon) == 0: + patch_gt_polygon = None + sample["image"] = patch_image + sample["gt_polygons"] = patch_gt_polygon + sample["patch_bbox"] = torch.tensor(patch_boundingbox) + patches.append(sample.copy()) + return patches + + def __call__(self, data_list): + patch_list = [] + for data in data_list: + patches = self._to_patch(data) + patch_list.extend(patches) + return patch_list diff --git a/torch_lydorn/torchvision/transforms/transform_by_key.py b/torch_lydorn/torchvision/transforms/transform_by_key.py new file mode 100644 index 0000000000000000000000000000000000000000..e175b08395aabd1762ec7965dec7c58ec0f3da6f --- /dev/null +++ b/torch_lydorn/torchvision/transforms/transform_by_key.py @@ -0,0 +1,47 @@ +def format_key(key): + if type(key) == list: + for k in key: + assert type(k) == str, "keys should be strings" + else: + assert type(key) == str + return key + + +class TransformByKey(object): + """Performs data[outkey[0]], data[outkey[1]], ... = transform(data[key[0]], data[key[1]], ..)""" + + def __init__(self, transform, key=None, outkey=None, ignore_key_error=False, **kwargs): + self.transform = transform + if key is None: + self.key = None + else: + self.key = format_key(key) + if outkey is None: + self.outkey = self.key + else: + self.outkey = format_key(outkey) + self.ignore_key_error = ignore_key_error + self.kwargs = kwargs + + def __call__(self, data): + assert type(data) == dict, "Input data should be a dictionary, not a {}".format(type(data)) + try: + if self.key is None: + output = self.transform(**self.kwargs) + elif type(self.key) == str: + output = self.transform(data[self.key], **self.kwargs) + else: + inputs = [data[k] for k in self.key] + output = self.transform(*inputs, **self.kwargs) + + if type(self.outkey) == str: + data[self.outkey] = output + else: + assert type(output) == tuple, "Output should be tuple, not {} because outkey is {}".format(type(output), type(self.outkey)) + assert len(self.outkey) == len(output), "len(outkey) and len(output) should be the same for a 1-to-1 matching." + for k, o in zip(self.outkey, output): + data[k] = o + except KeyError as e: + if not self.ignore_key_error: + raise e + return data \ No newline at end of file diff --git a/torch_lydorn/torchvision/transforms/transforms.py b/torch_lydorn/torchvision/transforms/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..eabc5a383823ce2e32ee4be514cff8d41b255c9c --- /dev/null +++ b/torch_lydorn/torchvision/transforms/transforms.py @@ -0,0 +1,79 @@ +import random + +from .functional import to_tensor, center_crop + +__all__ = ["CenterCrop", "ToTensor", "RandomBool", "ConditionApply"] + + +class CenterCrop(object): + def __init__(self, output_size): + self.output_size = output_size + + def __call__(self, tensor): + return center_crop(tensor, self.output_size) + + def __repr__(self): + return self.__class__.__name__ + '()' + + +class ToTensor(object): + """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor without typecasting and rescaling. + Converts a PIL Image or numpy.ndarray (H x W x C) in the range + [0, 255] to a torch.ByteTensor of shape (C x H x W) in the range [0, 255] + if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1) + or if the numpy.ndarray has dtype = np.uint8 + """ + + def __call__(self, pic): + """ + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + Returns: + Tensor: Converted image. + """ + return to_tensor(pic) + + def __repr__(self): + return self.__class__.__name__ + '()' + + +class RandomBool(object): + """Produce a random boolean with p probability for it to be True + + Args: + p (float): probability + """ + def __init__(self, p=0.5): + self.p = p + + def __call__(self): + return random.random() < self.p + + +class ConditionApply(object): + """Apply a transformation if condition is met + + Args: + transform: + """ + + def __init__(self, transform): + self.transform = transform + + def __call__(self, tensor, condition): + """ + + :param tensor: + :param condition (bool): True: apply, False: do not apply + :return: + """ + if condition: + tensor = self.transform(tensor) + return tensor + + def __repr__(self): + format_string = self.__class__.__name__ + '(' + format_string += '\n' + format_string += ' {0}'.format(self.transform) + format_string += '\n)' + return format_string