](https://colab.research.google.com/github/zcemycl/TF2DeepFloorplan/blob/master/deepfloorplan.ipynb)  [](https://coveralls.io/github/zcemycl/TF2DeepFloorplan?branch=main)[](https://hits.seeyoufarm.com)
This repo contains a basic procedure to train and deploy the DNN model suggested by the paper ['Deep Floor Plan Recognition using a Multi-task Network with Room-boundary-Guided Attention'](https://arxiv.org/abs/1908.11025). It rewrites the original codes from [zlzeng/DeepFloorplan](https://github.com/zlzeng/DeepFloorplan) into newer versions of Tensorflow and Python.
Network Architectures from the paper,
### Additional feature (pygame)
## Requirements
Depends on different applications, the following installation methods can
|OS|Hardware|Application|Command|
|---|---|---|---|
|Ubuntu|CPU|Model Development|`pip install -e .[tfcpu,dev,testing,linting]`|
|Ubuntu|GPU|Model Development|`pip install -e .[tfgpu,dev,testing,linting]`|
|MacOS|M1 Chip|Model Development|`pip install -e .[tfmacm1,dev,testing,linting]`|
|Ubuntu|GPU|Model Deployment API|`pip install -e .[tfgpu,api]`|
|Ubuntu|GPU|Everything|`pip install -e .[tfgpu,api,dev,testing,linting,game]`|
|Agnostic|...|Docker|(to be updated)|
|Ubuntu|GPU|Notebook|`pip install -e .[tfgpu,jupyter]`|
|Ubuntu|GPU|Game|`pip install -e .[tfgpu,game]`|
## How to run?
1. Install packages.
```
# Option 1
python -m venv venv
source venv/bin/activate
pip install --upgrade pip setuptools wheel
# Option 2 (Preferred)
conda create -n venv python=3.8 cudatoolkit=10.1 cudnn=7.6.5
conda activate venv
# common install
pip install -e .[tfgpu,api,dev,testing,linting]
```
2. According to the original repo, please download r3d dataset and transform it to tfrecords `r3d.tfrecords`. Friendly reminder: there is another dataset r2v used to train their original repo's model, I did not use it here cos of limited access. Please see the link here [https://github.com/zlzeng/DeepFloorplan/issues/17](https://github.com/zlzeng/DeepFloorplan/issues/17).
3. Run the `train.py` file to initiate the training, model checkpoint is stored as `log/store/G` and weight is in `model/store`,
```
python -m dfp.train [--batchsize 2][--lr 1e-4][--epochs 1000]
[--logdir 'log/store'][--modeldir 'model/store']
[--save-tensor-interval 10][--save-model-interval 20]
[--tfmodel 'subclass'/'func'][--feature-channels 256 128 64 32]
[--backbone 'vgg16'/'mobilenetv1'/'mobilenetv2'/'resnet50']
[--feature-names block1_pool block2_pool block3_pool block4_pool block5_pool]
```
- for example,
```
python -m dfp.train --batchsize=4 --lr=5e-4 --epochs=100
--logdir=log/store --modeldir=model/store
```
4. Run Tensorboard to view the progress of loss and images via,
```
tensorboard --logdir=log/store
```
5. Convert model to tflite via `convert2tflite.py`.
```
python -m dfp.convert2tflite [--modeldir model/store]
[--tflitedir model/store/model.tflite]
[--loadmethod 'log'/'none'/'pb']
[--quantize][--tfmodel 'subclass'/'func']
[--feature-channels 256 128 64 32]
[--backbone 'vgg16'/'mobilenetv1'/'mobilenetv2'/'resnet50']
[--feature-names block1_pool block2_pool block3_pool block4_pool block5_pool]
```
6. Download and unzip model from google drive,
```
gdown https://drive.google.com/uc?id=1czUSFvk6Z49H-zRikTc67g2HUUz4imON # log files 112.5mb
unzip log.zip
gdown https://drive.google.com/uc?id=1tuqUPbiZnuubPFHMQqCo1_kFNKq4hU8i # pb files 107.3mb
unzip model.zip
gdown https://drive.google.com/uc?id=1B-Fw-zgufEqiLm00ec2WCMUo5E6RY2eO # tfilte file 37.1mb
unzip tflite.zip
```
7. Deploy the model via `deploy.py`, please be aware that load method parameter should match with weight input.
```
python -m dfp.deploy [--image 'path/to/image']
[--postprocess][--colorize][--save 'path/to/output_image']
[--loadmethod 'log'/'pb'/'tflite']
[--weight 'log/store/G'/'model/store'/'model/store/model.tflite']
[--tfmodel 'subclass'/'func']
[--feature-channels 256 128 64 32]
[--backbone 'vgg16'/'mobilenetv1'/'mobilenetv2'/'resnet50']
[--feature-names block1_pool block2_pool block3_pool block4_pool block5_pool]
```
- for example,
```
python -m dfp.deploy --image floorplan.jpg --weight log/store/G
--postprocess --colorize --save output.jpg --loadmethod log
```
8. Play with pygame.
```
python -m dfp.game
```
## Docker for API
1. Build and run docker container. (Please train your weight, google drive does not work currently due to its update.)
```
docker build -t tf_docker -f Dockerfile .
docker run -d -p 1111:1111 tf_docker:latest
docker run --gpus all -d -p 1111:1111 tf_docker:latest
# special for hot reloading flask
docker run -v ${PWD}/src/dfp/app.py:/src/dfp/app.py -v ${PWD}/src/dfp/deploy.py:/src/dfp/deploy.py -d -p 1111:1111 tf_docker:latest
docker logs `docker ps | grep "tf_docker:latest" | awk '{ print $1 }'` --follow
```
2. Call the api for output.
```
curl -H "Content-Type: application/json" --request POST \
-d '{"uri":"https://cdn.cnn.com/cnnnext/dam/assets/200212132008-04-london-rental-market-intl-exlarge-169.jpg","colorize":1,"postprocess":0}' \
http://0.0.0.0:1111/uri --output /tmp/tmp.jpg
curl --request POST -F "file=@resources/30939153.jpg" \
-F "postprocess=0" -F "colorize=0" http://0.0.0.0:1111/upload --output out.jpg
```
3. If you run `app.py` without docker, the second curl for file upload will not work.
## Google Colab
1. Click on [](https://colab.research.google.com/github/zcemycl/TF2DeepFloorplan/blob/master/deepfloorplan.ipynb) and authorize access.
2. Run the first 2 code cells for installation.
3. Go to Runtime Tab, click on Restart runtime. This ensures the packages installed are enabled.
4. Run the rest of the notebook.
## How to Contribute?
1. Git clone this repo.
2. Install required packages and pre-commit-hooks.
```
pip install -e .[tfgpu,api,dev,testing,linting]
pre-commit install
pre-commit run
pre-commit run --all-files
# pre-commit uninstall/ pip uninstall pre-commit
```
3. Create issues. Maintainer will decide if it requires branch. If so,
```
git fetch origin
git checkout xx-features
```
4. Stage your files, Commit and Push to branch.
5. After pull and merge requests, the issue is solved and the branch is deleted. You can,
```
git checkout main
git pull
git remote prune origin
git branch -d xx-features
```
## Results
- From `train.py` and `tensorboard`.
|Compare Ground Truth (top)against Outputs (bottom)|Total Loss| |:-------------------------:|:-------------------------:| |
|
|
|Boundary Loss|Room Loss|
|
|
|
- From `deploy.py` and `utils/legend.py`.
|Input|Legend|Output|
|:-------------------------:|:-------------------------:|:-------------------------:|
|
|
|
|
|`--colorize`|`--postprocess`|`--colorize``--postprocess`| |
|
|
|
## Optimization
- Backbone Comparison in Size
|Backbone|log|pb|tflite|toml|
|---|---|---|---|---|
|VGG16|130.5Mb|119Mb|45.3Mb|[link](docs/experiments/vgg16/exp1)|
|MobileNetV1|102.1Mb|86.7Mb|50.2Mb|[link](docs/experiments/mobilenetv1/exp1)|
|MobileNetV2|129.3Mb|94.4Mb|57.9Mb|[link](docs/experiments/mobilenetv2/exp1)|
|ResNet50|214Mb|216Mb|107.2Mb|[link](docs/experiments/resnet50/exp1)|
- Feature Selection Comparison in Size
|Backbone|Feature Names|log|pb|tflite|toml|
|---|---|---|---|---|---|
|MobileNetV1|"conv_pw_1_relu", "conv_pw_3_relu",
"conv_pw_5_relu",
"conv_pw_7_relu",
"conv_pw_13_relu"|102.1Mb|86.7Mb|50.2Mb|[link](docs/experiments/mobilenetv1/exp1)| |MobileNetV1|"conv_pw_1_relu",
"conv_pw_3_relu",
"conv_pw_5_relu",
"conv_pw_7_relu",
"conv_pw_12_relu"|84.5Mb|82.3Mb|49.2Mb|[link](docs/experiments/mobilenetv1/exp2)| - Feature Channels Comparison in Size |Backbone|Channels|log|pb|tflite|toml| |---|---|---|---|---|---| |VGG16|[256,128,64,32]|130.5Mb|119Mb|45.3Mb|[link](docs/experiments/vgg16/exp1)| |VGG16|[128,64,32,16]|82.4Mb|81.6Mb|27.3Mb|| |VGG16|[32,32,32,32]|73.2Mb|67.5Mb|18.1Mb|[link](docs/experiments/vgg16/exp2)| - tfmot - Pruning (not working) - Clustering (not working) - Post training Quantization (work the best) - Training aware Quantization (not supported by the version) ================================================ FILE: deepfloorplan.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "### Installation\n", "1. Run the first 2 cells\n", "2. Restart runtime\n", "3. Run the rest of the jupyter notebooks (do not run the first 2 cells again)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "YPHHCUKZn89j", "outputId": "0748f66f-8128-4e89-a845-482efb2d0c8c" }, "outputs": [], "source": [ "!git clone -b main https://github.com/zcemycl/TF2DeepFloorplan.git\n", "!pip install gdown\n", "!pip install --upgrade --no-cache-dir gdown\n", "!gdown https://drive.google.com/uc?id=1czUSFvk6Z49H-zRikTc67g2HUUz4imON\n", "!unzip log.zip\n", "!rm log.zip" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# gpu\n", "# !cd TF2DeepFloorplan && pip install -e .[tfgpu]\n", "# cpu\n", "!cd TF2DeepFloorplan && pip install -e .[tfcpu]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Main Script" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "C4VRom9mqBPT", "outputId": "74d58dfd-60cb-44a2-b992-94fff7cc83f6" }, "outputs": [], "source": [ "import tensorflow as tf\n", "import sys\n", "from dfp.net import *\n", "from dfp.data import *\n", "import matplotlib.image as mpimg\n", "import matplotlib.pyplot as plt\n", "from argparse import Namespace\n", "import os\n", "import gc\n", "os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'\n", "from dfp.utils.rgb_ind_convertor import *\n", "from dfp.utils.util import *\n", "from dfp.utils.legend import *\n", "from dfp.utils.settings import *\n", "from dfp.deploy import *\n", "print(tf.test.is_gpu_available())\n", "print(tf.config.list_physical_devices('GPU'))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "30DTDbxbwm3O" }, "outputs": [], "source": [ "img_path = './TF2DeepFloorplan/resources/30939153.jpg'\n", "inp = mpimg.imread(img_path)\n", "args = parse_args(\"--tomlfile ./TF2DeepFloorplan/docs/notebook.toml\".split())\n", "args = overwrite_args_with_toml(args)\n", "args.image = img_path" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "xzqbdPC0uJNc", "outputId": "e57d885d-2f31-4077-cfbe-8ef738c5466c" }, "outputs": [], "source": [ "result = main(args)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 252 }, "id": "2xVIt5LEusqf", "outputId": "cb197ad8-6971-420e-aae7-3d4a9142cf8c" }, "outputs": [], "source": [ "plt.subplot(1,2,1)\n", "plt.imshow(inp); plt.xticks([]); plt.yticks([]);\n", "plt.subplot(1,2,2)\n", "plt.imshow(result); plt.xticks([]); plt.yticks([]);" ] }, { "cell_type": "markdown", "metadata": { "id": "Jto5H5cXypOD" }, "source": [ "## Breakdown of postprocessing (step by step)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "57rg5h7XywwU", "outputId": "91f8d2d0-e32d-466d-e830-010607016fec" }, "outputs": [], "source": [ "model,img,shp = init(args)\n", "logits_cw,logits_r = predict(model,img,shp)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 252 }, "id": "2aUCqpT6zPmv", "outputId": "1b3901fa-a7a0-4d68-d9e2-21455c4dc26f" }, "outputs": [], "source": [ "logits_r = tf.image.resize(logits_r,shp[:2])\n", "logits_cw = tf.image.resize(logits_cw,shp[:2])\n", "r = convert_one_hot_to_image(logits_r)[0].numpy()\n", "cw = convert_one_hot_to_image(logits_cw)[0].numpy()\n", "plt.subplot(1,2,1)\n", "plt.imshow(r.squeeze()); plt.xticks([]); plt.yticks([]);\n", "plt.subplot(1,2,2)\n", "plt.imshow(cw.squeeze()); plt.xticks([]); plt.yticks([]);" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 252 }, "id": "UYf4WVVCzgqj", "outputId": "423422f1-b292-4b0e-cfc7-e639db1115ca" }, "outputs": [], "source": [ "r_color,cw_color = colorize(r.squeeze(),cw.squeeze())\n", "plt.subplot(1,2,1)\n", "plt.imshow(r_color); plt.xticks([]); plt.yticks([]);\n", "plt.subplot(1,2,2)\n", "plt.imshow(cw_color); plt.xticks([]); plt.yticks([]);" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 287 }, "id": "LTm_qYqa0HGc", "outputId": "83a183c7-2a6a-4144-8462-2a95aabdaed3" }, "outputs": [], "source": [ "newr,newcw = post_process(r,cw,shp)\n", "plt.subplot(1,2,1)\n", "plt.imshow(newr.squeeze()); plt.xticks([]); plt.yticks([]);\n", "plt.subplot(1,2,2)\n", "plt.imshow(newcw.squeeze()); plt.xticks([]); plt.yticks([]);" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 252 }, "id": "S5MCeHF30ygl", "outputId": "bb88248d-331a-496e-8b08-78c7a0306fa4" }, "outputs": [], "source": [ "newr_color,newcw_color = colorize(newr.squeeze(),newcw.squeeze())\n", "plt.subplot(1,2,1)\n", "plt.imshow(newr_color); plt.xticks([]); plt.yticks([]);\n", "plt.subplot(1,2,2)\n", "plt.imshow(newcw_color); plt.xticks([]); plt.yticks([]);" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 252 }, "id": "18UYo3rz0918", "outputId": "2a2319a4-668e-40b9-837d-964421f87c14" }, "outputs": [], "source": [ "plt.imshow(newr_color+newcw_color); plt.xticks([]); plt.yticks([]);" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 252 }, "id": "dydb1kWl13hL", "outputId": "16e9a5bc-3809-41bb-9a1a-f49c7deb5aa4" }, "outputs": [], "source": [ "over255 = lambda x: [p/255 for p in x]\n", "colors2 = [over255(rgb) for rgb in list(floorplan_fuse_map.values())]\n", "colors = [\"background\", \"closet\", \"bathroom\",\n", " \"living room\\nkitchen\\ndining room\",\n", " \"bedroom\",\"hall\",\"balcony\",\"not used\",\"not used\",\n", " \"door/window\",\"wall\"]\n", "f = lambda m,c: plt.plot([],[],marker=m, color=c, ls=\"none\")[0]\n", "handles = [f(\"s\", colors2[i]) for i in range(len(colors))]\n", "labels = colors\n", "legend = plt.legend(handles, labels, loc=3,framealpha=1, frameon=True)\n", "\n", "fig = legend.figure\n", "fig.canvas.draw()\n", "plt.xticks([]); plt.yticks([]);\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "iy8nx0WZ2QGS" }, "outputs": [], "source": [] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [], "name": "deepfloorplan.ipynb", "provenance": [] }, "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.16" } }, "nbformat": 4, "nbformat_minor": 1 } ================================================ FILE: docs/app.toml ================================================ tfmodel = 'subclass' image = '' postprocess = 1 colorize = 1 save = '' weight = 'log/store/G' loadmethod = 'log' feature_channels = [256, 128, 64, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/cfg_test.md ================================================ ## Some coding tests - black (pyproject.toml) - isort (pyproject.toml, isort.cfg) - flake8 - pytest (setup.cfg) - pytest-cov (setup.cfg) - mypy (mypy.ini) - Other Configuration files: tox.ini, .editorconfig ### References 1. https://pycqa.github.io/isort/docs/configuration/config_files.html 2. https://stackoverflow.com/questions/14399534/reference-requirements-txt-for-the-install-requires-kwarg-in-setuptools-setup-py 3. https://ianhopkinson.org.uk/2022/02/understanding-setup-py-setup-cfg-and-pyproject-toml-in-python/ 4. https://towardsdatascience.com/pre-commit-hooks-you-must-know-ff247f5feb7e 5. https://mypy.readthedocs.io/en/stable/running_mypy.html 6. https://mypy.readthedocs.io/en/stable/config_file.html#config-file 7. https://pre-commit.com/ ================================================ FILE: docs/experiments/mobilenetv1/exp1/compress.toml ================================================ tfmodel = 'func' modeldir = 'model/store_mobilenetv1_exp1' tflitedir = 'model/store_mobilenetv1_exp1_model.tflite' quantize = 1 compress_mode = 'quantization' loadmethod = 'pb' feature_channels = [256, 128, 64, 32] backbone = 'mobilenetv1' feature_names = ["conv_pw_1_relu", "conv_pw_3_relu", "conv_pw_5_relu", "conv_pw_7_relu", "conv_pw_13_relu",] ================================================ FILE: docs/experiments/mobilenetv1/exp1/deploy.toml ================================================ tfmodel = 'func' image = 'resources/30939153.jpg' postprocess = 1 colorize = 1 save = '' weight = 'log/store_mobilenetv1_exp1/G' loadmethod = 'log' # weight = 'model/store_mobilenetv1_exp1' # loadmethod = 'pb' # weight = 'model/store_mobilenetv1_exp1_model.tflite' # loadmethod = 'tflite' feature_channels = [256, 128, 64, 32] backbone = 'mobilenetv1' feature_names = ["conv_pw_1_relu", "conv_pw_3_relu", "conv_pw_5_relu", "conv_pw_7_relu", "conv_pw_13_relu",] ================================================ FILE: docs/experiments/mobilenetv1/exp1/train.toml ================================================ tfmodel = 'func' batchsize = 8 lr = 1e-4 wd = 1e-5 epochs = 100 logdir = 'log/store_mobilenetv1_exp1' modeldir = 'model/store_mobilenetv1_exp1' weight = '' save_tensor_interval = 10 save_model_interval = 20 feature_channels = [256, 128, 64, 32] backbone = 'mobilenetv1' feature_names = ["conv_pw_1_relu", "conv_pw_3_relu", "conv_pw_5_relu", "conv_pw_7_relu", "conv_pw_13_relu",] ================================================ FILE: docs/experiments/mobilenetv1/exp2/compress.toml ================================================ tfmodel = 'func' modeldir = 'model/store_mobilenetv1_exp2' tflitedir = 'model/store_mobilenetv1_exp2_model.tflite' quantize = 1 compress_mode = 'quantization' loadmethod = 'pb' feature_channels = [256, 128, 64, 32] backbone = 'mobilenetv1' feature_names = ["conv_pw_1_relu", "conv_pw_3_relu", "conv_pw_5_relu", "conv_pw_7_relu", "conv_pw_12_relu",] ================================================ FILE: docs/experiments/mobilenetv1/exp2/deploy.toml ================================================ tfmodel = 'func' image = 'resources/30939153.jpg' postprocess = 1 colorize = 1 save = '' weight = 'log/store_mobilenetv1_exp2/G' loadmethod = 'log' # weight = 'model/store_mobilenetv1_exp2' # loadmethod = 'pb' # weight = 'model/store_mobilenetv1_exp2_model.tflite' # loadmethod = 'tflite' feature_channels = [256, 128, 64, 32] backbone = 'mobilenetv1' feature_names = ["conv_pw_1_relu", "conv_pw_3_relu", "conv_pw_5_relu", "conv_pw_7_relu", "conv_pw_12_relu",] ================================================ FILE: docs/experiments/mobilenetv1/exp2/train.toml ================================================ tfmodel = 'func' batchsize = 8 lr = 1e-4 wd = 1e-5 epochs = 100 logdir = 'log/store_mobilenetv1_exp2' modeldir = 'model/store_mobilenetv1_exp2' weight = '' save_tensor_interval = 10 save_model_interval = 20 feature_channels = [256, 128, 64, 32] backbone = 'mobilenetv1' feature_names = ["conv_pw_1_relu", "conv_pw_3_relu", "conv_pw_5_relu", "conv_pw_7_relu", "conv_pw_12_relu",] ================================================ FILE: docs/experiments/mobilenetv2/exp1/train.toml ================================================ tfmodel = 'func' batchsize = 8 lr = 1e-4 wd = 1e-5 epochs = 100 logdir = 'log/store_mobilenetv2_exp1' modeldir = 'model/store_mobilenetv2_exp1' weight = '' save_tensor_interval = 10 save_model_interval = 20 feature_channels = [256, 128, 64, 32] backbone = 'mobilenetv2' feature_names = ["block_1_expand_relu", "block_3_expand_relu", "block_5_expand_relu", "block_13_expand_relu", "out_relu",] ================================================ FILE: docs/experiments/resnet50/exp1/train.toml ================================================ tfmodel = 'func' batchsize = 8 lr = 1e-4 wd = 1e-5 epochs = 100 logdir = 'log/store_resnet50_exp1' modeldir = 'model/store_resnet50_exp1' weight = '' save_tensor_interval = 10 save_model_interval = 20 feature_channels = [256, 128, 64, 32] backbone = 'resnet50' feature_names = ["conv1_relu", "conv2_block3_out", "conv3_block4_out", "conv4_block6_out", "conv5_block3_out",] ================================================ FILE: docs/experiments/vgg16/exp1/compress.toml ================================================ tfmodel = 'subclass' modeldir = 'model/store_vgg16_exp1' tflitedir = 'model/store_vgg16_exp1_model.tflite' quantize = 1 compress_mode = 'quantization' loadmethod = 'pb' feature_channels = [256, 128, 64, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/experiments/vgg16/exp1/compress_log.toml ================================================ tfmodel = 'subclass' modeldir = 'log/store_vgg16_exp1/G' tflitedir = 'model/store_vgg16_exp1_model_log.tflite' quantize = 1 compress_mode = 'quantization' loadmethod = 'log' feature_channels = [256, 128, 64, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/experiments/vgg16/exp1/deploy.toml ================================================ tfmodel = 'subclass' image = 'resources/30939153.jpg' postprocess = 1 colorize = 1 save = '' # weight = 'log/store_vgg16_exp1/G' # loadmethod = 'log' # weight = 'model/store_vgg16_exp1' # loadmethod = 'pb' weight = 'model/store_vgg16_exp1_model.tflite' loadmethod = 'tflite' feature_channels = [256, 128, 64, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/experiments/vgg16/exp1/train.toml ================================================ tfmodel = 'subclass' batchsize = 8 lr = 1e-4 wd = 1e-5 epochs = 100 logdir = 'log/store_vgg16_exp1' modeldir = 'model/store_vgg16_exp1' weight = '' save_tensor_interval = 10 save_model_interval = 20 feature_channels = [256, 128, 64, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/experiments/vgg16/exp2/compress.toml ================================================ tfmodel = 'subclass' modeldir = 'model/store_vgg16_exp2' tflitedir = 'model/store_vgg16_exp2_model.tflite' quantize = 1 compress_mode = 'quantization' loadmethod = 'pb' feature_channels = [32, 32, 32, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/experiments/vgg16/exp2/deploy.toml ================================================ tfmodel = 'subclass' image = 'resources/30939153.jpg' postprocess = 1 colorize = 1 save = '' weight = 'log/store_vgg16_exp2/G' loadmethod = 'log' # weight = 'model/store_vgg16_exp2' # loadmethod = 'pb' # weight = 'model/store_vgg16_exp2_model.tflite' # loadmethod = 'tflite' feature_channels = [32, 32, 32, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/experiments/vgg16/exp2/train.toml ================================================ tfmodel = 'subclass' batchsize = 8 lr = 1e-4 wd = 1e-5 epochs = 100 logdir = 'log/store_vgg16_exp2' modeldir = 'model/store_vgg16_exp2' weight = '' save_tensor_interval = 10 save_model_interval = 20 feature_channels = [32, 32, 32, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/experiments/vgg16/exp3/compress.toml ================================================ tfmodel = 'func' modeldir = 'model/store_vgg16_exp3' tflitedir = 'model/store_vgg16_exp3_model.tflite' quantize = 1 compress_mode = 'quantization' loadmethod = 'pb' feature_channels = [256, 128, 64, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/experiments/vgg16/exp3/deploy.toml ================================================ tfmodel = 'func' image = 'resources/30939153.jpg' postprocess = 1 colorize = 1 save = '' # weight = 'log/store_vgg16_exp3/G' # loadmethod = 'log' weight = 'model/store_vgg16_exp3' loadmethod = 'pb' # weight = 'model/store_vgg16_exp3_model.tflite' # loadmethod = 'tflite' feature_channels = [256, 128, 64, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/experiments/vgg16/exp3/train.toml ================================================ tfmodel = 'func' batchsize = 8 lr = 1e-4 wd = 1e-5 epochs = 100 logdir = 'log/store_vgg16_exp3' modeldir = 'model/store_vgg16_exp3' weight = '' save_tensor_interval = 10 save_model_interval = 20 feature_channels = [256, 128, 64, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/game.toml ================================================ tfmodel = 'subclass' # image = 'resources/123.jpg' # image = 'resources/example4.png' image = 'resources/example5.jpg' CASTED_RAYS = 150 postprocess = 1 colorize = 0 save = '' weight = 'model/store_vgg16_exp2_model.tflite' loadmethod = 'tflite' feature_channels = [32, 32, 32, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/notebook.toml ================================================ tfmodel = 'subclass' image = './TF2DeepFloorplan/resources/30939153.jpg' postprocess = 1 colorize = 1 save = '' weight = './log/store/G' loadmethod = 'log' feature_channels = [256, 128, 64, 32] backbone = 'vgg16' feature_names = ["block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool",] ================================================ FILE: docs/pytest.md ================================================ ## Some coding tests - pytest (setup.cfg) - pytest-cov (setup.cfg) - pytest-mock ## Keywords - Fixture (share data across test cases) - Unit Test (test for indepedent component) - Mock Test (test for depedent component) - Parametrizing (various test cases for same component) - Integration Test (end-to-end) ### References 1. https://blogs.sap.com/2022/02/16/how-to-write-independent-unit-test-with-pytest-and-mock-techniques/ ================================================ FILE: install/environment.yml ================================================ channels: - defaults dependencies: - python=3.8 - cudatoolkit=10.1 - cudnn=7.6.5 - pip: - -e .[tfgpu,api,dev,testing,linting] ================================================ FILE: mypy.ini ================================================ [mypy] ignore_missing_imports = True files = src/, tests/ exclude = (?x)( ^app\.py$ # files named "one.py" | app\.py$ # or files ending with "two.pyi" | ^three\. # or files starting with "three." | venv | resources | model | log | src/app.py ) ================================================ FILE: pyproject.toml ================================================ [tool.black] line-length = 79 include = ''' /( src | tests )/ ''' exclude = ''' /( \.git | \.mypy_cache | \.nox | __pycache__ | \.pyc$ | \.ipynb$ | \.md$ | build | venv | dist | \.eggs )/ ''' [tool.isort] profile = "black" src_paths = [".","src", "tests"] skip=[".tox",".nox","venv","build","dist","resources","model","log"] skip_glob=["venv/*","build/*","dist/*","resources/*"] sections="FUTURE,STDLIB,SETUPTOOLS,TEST,THIRDPARTY,FIRSTPARTY,LOCALFOLDER" known_first_party = "src" ================================================ FILE: requirements.txt ================================================ matplotlib numpy opencv-python pdbpp scipy Pillow gdown protobuf==3.20.0 chardet types-requests pytype dynaconf ================================================ FILE: setup.cfg ================================================ [metadata] name = dfp description = Deep Floorplan (dfp) author = Leo Leung [options] zip_safe = False packages = find: include_package_data = True setup_requires = setuptools_scm package_dir = = src [options.packages.find] where = src exclude = tests [options.extras_require] tfcpu = tensorflow-cpu tensorboard tensorflow-model-optimization tfgpu = tensorflow-gpu==2.3.0 tensorboard==2.3.0 tensorflow-model-optimization==0.5.0 tfmacm1 = tensorflow-macos tensorflow-metal tensorboard tensorflow-model-optimization api = Flask Flask-Bcrypt Flask-Classful Flask-Cors Flask-Ext Flask-Jsonpify Flask-Markdown Flask-WTF dev = pre-commit testing = pytest pytest-cov pytest-mock pytest-flask linting = black==22.3.0 isort==5.10.1 flake8==4.0.1 mypy==0.961 jupyter = jupyterlab notebook game = pygame numba [tool:pytest] testpaths = tests addopts = --cov src --cov-report term-missing --disable-warnings --verbose norecursedirs = dist build .tox resources log model venv [tool.setuptools_scm] version_scheme = guess-next-dev [bdist_wheel] universal = 1 [flake8] ignore = E203 W503 W291 W293 max-line-length = 79 exclude = .tox dist .eggs venv log model resources ================================================ FILE: setup.py ================================================ import os from setuptools import setup with open("requirements.txt") as f: required = f.read().splitlines() if __name__ == "__main__": use_scm_version = not os.environ.get("AM_I_IN_A_DOCKER_CONTAINER", False) setup( use_scm_version=use_scm_version, install_requires=required, test_suite="tests", ) ================================================ FILE: src/dfp/__init__.py ================================================ ================================================ FILE: src/dfp/app.py ================================================ import multiprocessing as mp import os import random from argparse import Namespace import matplotlib.image as mpimg import numpy as np import requests from flask import Flask, request, send_file from werkzeug.datastructures import FileStorage from .deploy import main from .utils.settings import overwrite_args_with_toml app = Flask(__name__) app.config["UPLOAD_EXTENSIONS"] = [".jpg", ".png", ".jpeg"] args = Namespace(tomlfile="docs/app.toml") args = overwrite_args_with_toml(args) finname = "resources/30939153.jpg" output = "/tmp" def saveStreamFile(stream: FileStorage, fnum: str): stream.save(fnum + ".jpg") def saveStreamURI(stream: bytes, fnum: str): with open(fnum + ".jpg", "wb") as handler: handler.write(stream) @app.route("/") def home(): return {"message": "Hello Flask!"} @app.route("/upload", methods=["POST"]) def dummy(): finname = "resources/30939153.jpg" fnum = str(random.randint(0, 10000)) foutname = fnum + "-out.jpg" if "file" in request.files: saveStreamFile(request.files["file"], fnum) finname = fnum + ".jpg" postprocess = ( False if "postprocess" not in request.form.keys() else bool(int(request.form.getlist("postprocess")[0])) ) colorize = ( False if "colorize" not in request.form.keys() else bool(int(request.form.getlist("colorize")[0])) ) args.image = finname args.postprocess = postprocess args.colorize = colorize args.save = os.path.join(output, foutname) app.logger.info(args) with mp.Pool() as pool: result = pool.map(main, [args])[0] app.logger.info(f"Output Image shape: {np.array(result).shape}") if args.save: mpimg.imsave(args.save, np.array(result).astype(np.uint8)) try: callback = send_file( os.path.join(output, foutname), mimetype="image/jpg" ) return callback, 200 except Exception: return {"message": "send error"}, 400 finally: os.system("rm " + os.path.join(output, foutname)) if finname != "resources/30939153.jpg": os.system("rm " + finname) return {"message": "hello"} @app.route("/uri", methods=["POST"]) def process_image(): fnum = str(random.randint(0, 10000)) finname = "resources/30939153.jpg" foutname = fnum + "-out.jpg" postprocess = ( bool(request.json["postprocess"]) if request.json and "postprocess" in request.json.keys() else False ) colorize = ( bool(request.json["colorize"]) if request.json and "colorize" in request.json.keys() else False ) # input image: uri if request.json and "uri" in request.json.keys(): app.logger.info("URI mode...") uri = request.json["uri"] try: data = requests.get(uri).content saveStreamURI(data, fnum) finname = fnum + ".jpg" except Exception: return {"message": "input error"}, 400 args.image = finname args.postprocess = postprocess args.colorize = colorize args.save = os.path.join(output, foutname) app.logger.info(args) with mp.Pool() as pool: result = pool.map(main, [args])[0] app.logger.info(f"Output Image shape: {np.array(result).shape}") if args.save: mpimg.imsave(args.save, np.array(result).astype(np.uint8)) try: callback = send_file( os.path.join(output, foutname), mimetype="image/jpg" ) return callback, 200 except Exception: return {"message": "send error"}, 400 finally: os.system("rm " + os.path.join(output, foutname)) if finname != "resources/30939153.jpg": os.system("rm " + finname) if __name__ == "__main__": app.run(debug=True, host="0.0.0.0", port=1111) ================================================ FILE: src/dfp/convert2tflite.py ================================================ import argparse import os import sys import tempfile from typing import List import tensorflow as tf import tensorflow_model_optimization as tfmot from tqdm import tqdm from .data import decodeAllRaw, loadDataset, preprocess from .net import deepfloorplanModel from .net_func import deepfloorplanFunc from .train import train_step from .utils.settings import overwrite_args_with_toml from .utils.util import ( print_model_weight_clusters, print_model_weights_sparsity, ) def model_init(config: argparse.Namespace) -> tf.keras.Model: if config.loadmethod == "log": if config.tfmodel == "subclass": base_model = deepfloorplanModel(config=config) base_model.build((1, 512, 512, 3)) assert True, "subclass and log are not convertible to tflite." elif config.tfmodel == "func": base_model = deepfloorplanFunc(config=config) base_model.load_weights(config.modeldir) elif config.loadmethod == "pb": base_model = tf.keras.models.load_model(config.modeldir) # need changes later to frozen backbone and constant kernel for layer in base_model.layers: layer.trainable = False return base_model def converter(config: argparse.Namespace): # model = tf.keras.models.load_model(config.modeldir) model = model_init(config) converter = tf.lite.TFLiteConverter.from_keras_model(model) if config.quantize: converter.optimizations = [tf.lite.Optimize.DEFAULT] # converter.target_spec.supported_ops = [ # tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops. # tf.lite.OpsSet.SELECT_TF_OPS, # enable TensorFlow ops. # tf.float16 # ] converter.experimental_new_converter = True tflite_model = converter.convert() with open(config.tflitedir, "wb") as f: f.write(tflite_model) def parse_args(args: List[str]) -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument( "--tfmodel", type=str, default="subclass", choices=["subclass", "func"] ) p.add_argument("--modeldir", type=str, default="model/store") p.add_argument("--tflitedir", type=str, default="model/store/model.tflite") p.add_argument("--quantize", action="store_true") p.add_argument( "--compress-mode", type=str, default="quantization", choices=["quantization", "prune", "cluster"], ) p.add_argument( "--loadmethod", type=str, default="log", choices=["log", "pb", "none"], ) # log,tflite,pb p.add_argument( "--feature-channels", type=int, action="store", default=[256, 128, 64, 32], nargs=4, ) p.add_argument( "--backbone", type=str, default="vgg16", choices=["vgg16", "resnet50", "mobilenetv1", "mobilenetv2"], ) p.add_argument( "--feature-names", type=str, action="store", nargs=5, default=[ "block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool", ], ) p.add_argument("--tomlfile", type=str, default=None) return p.parse_args(args) def prune(config: argparse.Namespace): base_model = model_init(config) print(base_model.summary()) model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model) print(model_for_pruning.summary()) dataset = loadDataset() optimizer = tf.keras.optimizers.Adam() log_dir = tempfile.mkdtemp() unused_arg = -1 epochs = 4 batches = 8 model_for_pruning.optimizer = optimizer step_callback = tfmot.sparsity.keras.UpdatePruningStep() step_callback.set_model(model_for_pruning) log_callback = tfmot.sparsity.keras.PruningSummaries( log_dir=log_dir ) # Log sparsity and other metrics in Tensorboard. log_callback.set_model(model_for_pruning) step_callback.on_train_begin() # run pruning callback for _ in range(epochs): log_callback.on_epoch_begin(epoch=unused_arg) # run pruning callback for data in tqdm(list(dataset.batch(batches))): step_callback.on_train_batch_begin( batch=unused_arg ) # run pruning callback img, bound, room = decodeAllRaw(data) img, bound, room, hb, hr = preprocess(img, bound, room) _, _, loss_value, _, _ = train_step( model_for_pruning, optimizer, img, hr, hb ) step_callback.on_epoch_end(batch=unused_arg) # run pruning callback # model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model) print(f"log directory: {log_dir}...") model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning) print_model_weights_sparsity(model_for_export) model_for_export.save(log_dir + "/prune") converter = tf.lite.TFLiteConverter.from_keras_model(model_for_export) converter.optimizations = [tf.lite.Optimize.DEFAULT] pruned_tflite_model = converter.convert() os.system(f"mkdir -p {log_dir}/tflite") with open(log_dir + "/tflite/model.tflite", "wb") as f: f.write(pruned_tflite_model) def cluster(config: argparse.Namespace): base_model = model_init(config) cluster_weights = tfmot.clustering.keras.cluster_weights CentroidInitialization = tfmot.clustering.keras.CentroidInitialization clustering_params = { "number_of_clusters": 8, "cluster_centroids_init": CentroidInitialization.DENSITY_BASED, } def apply_clustering_to_conv2d(layer): if isinstance(layer, tf.keras.layers.Conv2D) or isinstance( layer, tf.keras.layers.Conv2DTranspose ): return cluster_weights(layer, **clustering_params) return layer clustered_model = tf.keras.models.clone_model( base_model, clone_function=apply_clustering_to_conv2d, ) print(clustered_model.summary()) dataset = loadDataset() optimizer = tf.keras.optimizers.Adam() for epoch in range(4): print("[INFO] Epoch {}".format(epoch)) for data in tqdm(list(dataset.batch(8))): img, bound, room = decodeAllRaw(data) img, bound, room, hb, hr = preprocess(img, bound, room) _, _, loss_value, _, _ = train_step( clustered_model, optimizer, img, hr, hb ) log_dir = tempfile.mkdtemp() print("log directory: " + log_dir) final_model = tfmot.clustering.keras.strip_clustering(clustered_model) print_model_weight_clusters(final_model) final_model.save(log_dir + "/cluster") converter = tf.lite.TFLiteConverter.from_keras_model(final_model) converter.optimizations = [tf.lite.Optimize.DEFAULT] clustered_tflite_model = converter.convert() os.system(f"mkdir -p {log_dir}/tflite") with open(log_dir + "/tflite/model.tflite", "wb") as f: f.write(clustered_tflite_model) def quantization_aware_training(config: argparse.Namespace): base_model = model_init(config) def apply_quantization_to_conv2D(layer): if isinstance(layer, tf.keras.layers.Conv2D): return tfmot.quantization.keras.quantize_annotate_layer(layer) return layer raise Exception("Not supported") annotated_model = tf.keras.models.clone_model( base_model, clone_function=apply_quantization_to_conv2D, ) quant_aware_model = tfmot.quantization.keras.quantize_apply( annotated_model ) # quant_aware_model = tfmot.quantization.keras.quantize_model(base_model) print(quant_aware_model.summary()) dataset = loadDataset() optimizer = tf.keras.optimizers.Adam() for epoch in range(4): print("[INFO] Epoch {}".format(epoch)) for data in tqdm(list(dataset.batch(8))): img, bound, room = decodeAllRaw(data) img, bound, room, hb, hr = preprocess(img, bound, room) _, _, loss_value, _, _ = train_step( quant_aware_model, optimizer, img, hr, hb ) log_dir = tempfile.mkdtemp() print("log directory: " + log_dir) converter = tf.lite.TFLiteConverter.from_keras_model(quant_aware_model) converter.optimizations = [tf.lite.Optimize.DEFAULT] quantized_tflite_model = converter.convert() os.system(f"mkdir -p {log_dir}/tflite") with open(log_dir + "/tflite/model.tflite", "wb") as f: f.write(quantized_tflite_model) if __name__ == "__main__": args = parse_args(sys.argv[1:]) args = overwrite_args_with_toml(args) print(args) if args.compress_mode == "quantization" or args.quantize: # quantization_aware_training(args) converter(args) elif args.tfmodel == "func": if args.compress_mode == "prune": prune(args) elif args.compress_mode == "cluster": cluster(args) elif args.tfmodel == "subclass": raise Exception( "Pruning or Clustering for Subclass Model are not available." ) ================================================ FILE: src/dfp/data.py ================================================ from typing import Dict, Tuple import matplotlib.pyplot as plt import tensorflow as tf def convert_one_hot_to_image( one_hot: tf.Tensor, dtype: str = "float", act: str = None ) -> tf.Tensor: if act == "softmax": one_hot = tf.keras.activations.softmax(one_hot) [n, h, w, c] = one_hot.shape.as_list() im = tf.reshape(tf.keras.backend.argmax(one_hot, axis=-1), [n, h, w, 1]) if dtype == "int": im = tf.cast(im, dtype=tf.uint8) else: im = tf.cast(im, dtype=tf.float32) return im def _parse_function(example_proto: bytes) -> Dict[str, tf.Tensor]: feature = { "image": tf.io.FixedLenFeature([], tf.string), "boundary": tf.io.FixedLenFeature([], tf.string), "room": tf.io.FixedLenFeature([], tf.string), "door": tf.io.FixedLenFeature([], tf.string), } return tf.io.parse_single_example(example_proto, feature) def decodeAllRaw( x: Dict[str, tf.Tensor] ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: image = tf.io.decode_raw(x["image"], tf.uint8) boundary = tf.io.decode_raw(x["boundary"], tf.uint8) room = tf.io.decode_raw(x["room"], tf.uint8) return image, boundary, room def preprocess( img: tf.Tensor, bound: tf.Tensor, room: tf.Tensor, size: int = 512 ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]: img = tf.cast(img, dtype=tf.float32) img = tf.reshape(img, [-1, size, size, 3]) / 255 bound = tf.reshape(bound, [-1, size, size]) room = tf.reshape(room, [-1, size, size]) hot_b = tf.one_hot(bound, 3, axis=-1) hot_r = tf.one_hot(room, 9, axis=-1) return img, bound, room, hot_b, hot_r def loadDataset(size: int = 512) -> tf.data.Dataset: raw_dataset = tf.data.TFRecordDataset("r3d.tfrecords") parsed_dataset = raw_dataset.map(_parse_function) return parsed_dataset def plotData(data: Dict[str, tf.Tensor]): img, bound, room = decodeAllRaw(data) img, bound, room, hb, hr = preprocess(img, bound, room) plt.subplot(1, 3, 1) plt.imshow(img[0].numpy()) plt.subplot(1, 3, 2) plt.imshow(bound[0].numpy()) plt.subplot(1, 3, 3) plt.imshow(convert_one_hot_to_image(hb)[0].numpy()) def main(dataset: tf.data.Dataset): for ite in range(2): for data in list(dataset.shuffle(400).batch(1)): plotData(data) plt.show() break if __name__ == "__main__": dataset = loadDataset() main(dataset) ================================================ FILE: src/dfp/deploy.py ================================================ import argparse import gc import os import sys from typing import List, Tuple import matplotlib.image as mpimg import matplotlib.pyplot as plt import numpy as np import tensorflow as tf from .data import convert_one_hot_to_image from .net import deepfloorplanModel from .net_func import deepfloorplanFunc from .utils.rgb_ind_convertor import ( floorplan_boundary_map, floorplan_fuse_map, ind2rgb, ) from .utils.settings import overwrite_args_with_toml from .utils.util import fill_break_line, flood_fill, refine_room_region os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" def init( config: argparse.Namespace, ) -> Tuple[tf.keras.Model, tf.Tensor, np.ndarray]: if config.tfmodel == "subclass": model = deepfloorplanModel(config=config) elif config.tfmodel == "func": model = deepfloorplanFunc(config=config) if config.loadmethod == "log": model.load_weights(config.weight) elif config.loadmethod == "pb": model = tf.keras.models.load_model(config.weight) elif config.loadmethod == "tflite": model = tf.lite.Interpreter(model_path=config.weight) model.allocate_tensors() img = mpimg.imread(config.image)[:, :, :3] shp = img.shape img = tf.convert_to_tensor(img, dtype=tf.uint8) img = tf.image.resize(img, [512, 512]) img = tf.cast(img, dtype=tf.float32) img = tf.reshape(img, [-1, 512, 512, 3]) if tf.math.reduce_max(img) > 1.0: img /= 255 if config.loadmethod == "tflite": return model, img, shp model.trainable = False if config.tfmodel == "subclass": model.vgg16.trainable = False return model, img, shp def predict( model: tf.keras.Model, img: tf.Tensor, shp: np.ndarray ) -> Tuple[tf.Tensor, tf.Tensor]: features = [] feature = img for layer in model.vgg16.layers: feature = layer(feature) if layer.name.find("pool") != -1: features.append(feature) x = feature features = features[::-1] del model.vgg16 gc.collect() featuresrbp = [] for i in range(len(model.rbpups)): x = model.rbpups[i](x) + model.rbpcv1[i](features[i + 1]) x = model.rbpcv2[i](x) featuresrbp.append(x) logits_cw = tf.keras.backend.resize_images( model.rbpfinal(x), 2, 2, "channels_last" ) x = features.pop(0) nLays = len(model.rtpups) for i in range(nLays): rs = model.rtpups.pop(0) r1 = model.rtpcv1.pop(0) r2 = model.rtpcv2.pop(0) f = features.pop(0) x = rs(x) + r1(f) x = r2(x) a = featuresrbp.pop(0) x = model.non_local_context(a, x, i) del featuresrbp logits_r = tf.keras.backend.resize_images( model.rtpfinal(x), 2, 2, "channels_last" ) del model.rtpfinal return logits_cw, logits_r def post_process( rm_ind: np.ndarray, bd_ind: np.ndarray, shp: np.ndarray ) -> Tuple[np.ndarray, np.ndarray]: hard_c = (bd_ind > 0).astype(np.uint8) # region from room prediction rm_mask = np.zeros(rm_ind.shape) rm_mask[rm_ind > 0] = 1 # region from close wall line cw_mask = hard_c # regine close wall mask by filling the gap between bright line cw_mask = fill_break_line(cw_mask) cw_mask = np.reshape(cw_mask, (*shp[:2], -1)) fuse_mask = cw_mask + rm_mask fuse_mask[fuse_mask >= 1] = 255 # refine fuse mask by filling the hole fuse_mask = flood_fill(fuse_mask) fuse_mask = fuse_mask // 255 # one room one label new_rm_ind = refine_room_region(cw_mask, rm_ind) # ignore the background mislabeling new_rm_ind = fuse_mask.reshape(*shp[:2], -1) * new_rm_ind new_bd_ind = fill_break_line(bd_ind).squeeze() return new_rm_ind, new_bd_ind def colorize(r: np.ndarray, cw: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: cr = ind2rgb(r, color_map=floorplan_fuse_map) ccw = ind2rgb(cw, color_map=floorplan_boundary_map) return cr, ccw def main(config: argparse.Namespace) -> np.ndarray: model, img, shp = init(config) if config.loadmethod == "tflite": input_details = model.get_input_details() output_details = model.get_output_details() model.set_tensor(input_details[0]["index"], img) model.invoke() ri, cwi = 0, 1 if config.tfmodel == "func": ri, cwi = 1, 0 logits_r = model.get_tensor(output_details[ri]["index"]) logits_cw = model.get_tensor(output_details[cwi]["index"]) logits_cw = tf.convert_to_tensor(logits_cw) logits_r = tf.convert_to_tensor(logits_r) else: if config.tfmodel == "func": logits_r, logits_cw = model.predict(img) elif config.tfmodel == "subclass": if config.loadmethod == "log": logits_cw, logits_r = predict(model, img, shp) elif config.loadmethod == "pb" or config.loadmethod == "none": logits_r, logits_cw = model(img) logits_r = tf.image.resize(logits_r, shp[:2]) logits_cw = tf.image.resize(logits_cw, shp[:2]) r = convert_one_hot_to_image(logits_r)[0].numpy() cw = convert_one_hot_to_image(logits_cw)[0].numpy() if not config.colorize and not config.postprocess: cw[cw == 1] = 9 cw[cw == 2] = 10 r[cw != 0] = 0 return (r + cw).squeeze() elif config.colorize and not config.postprocess: r_color, cw_color = colorize(r.squeeze(), cw.squeeze()) return r_color + cw_color newr, newcw = post_process(r, cw, shp) if not config.colorize and config.postprocess: newcw[newcw == 1] = 9 newcw[newcw == 2] = 10 newr[newcw != 0] = 0 return newr.squeeze() + newcw newr_color, newcw_color = colorize(newr.squeeze(), newcw.squeeze()) result = newr_color + newcw_color print(shp, result.shape) if config.save: mpimg.imsave(config.save, result.astype(np.uint8)) return result def parse_args(args: List[str]) -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument( "--tfmodel", type=str, default="subclass", choices=["subclass", "func"] ) p.add_argument("--image", type=str, default="resources/30939153.jpg") p.add_argument("--weight", type=str, default="log/store/G") p.add_argument("--postprocess", action="store_true") p.add_argument("--colorize", action="store_true") p.add_argument( "--loadmethod", type=str, default="log", choices=["log", "tflite", "pb", "none"], ) # log,tflite,pb p.add_argument("--save", type=str) p.add_argument( "--feature-channels", type=int, action="store", default=[256, 128, 64, 32], nargs=4, ) p.add_argument( "--backbone", type=str, default="vgg16", choices=["vgg16", "resnet50", "mobilenetv1", "mobilenetv2"], ) p.add_argument( "--feature-names", type=str, action="store", nargs=5, default=[ "block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool", ], ) p.add_argument("--tomlfile", type=str, default=None) return p.parse_args(args) def deploy_plot_res(result: np.ndarray): print(result.shape) plt.imshow(result) plt.xticks([]) plt.yticks([]) plt.grid(False) if __name__ == "__main__": args = parse_args(sys.argv[1:]) args = overwrite_args_with_toml(args) result = main(args) deploy_plot_res(result) plt.show() ================================================ FILE: src/dfp/game/__init__.py ================================================ ================================================ FILE: src/dfp/game/__main__.py ================================================ import logging import math import sys import time from typing import Tuple, Union import numpy as np import pygame from numba import jit, njit from .controller import Controller from .model import Model from .view import View logging.basicConfig(level=logging.INFO) @njit def ray_cast_dda( x: float, y: float, angle: float, depth: Union[float, int], w: int, h: int, binary_map: np.ndarray, ) -> Tuple[float, float, bool]: vPlayer = np.array([x, y]) vRayStart = vPlayer.copy() vMouseCell = np.array( [x - math.sin(angle) * depth, y + math.cos(angle) * depth] ) vRayDir = (vMouseCell - vPlayer) / np.linalg.norm((vMouseCell - vPlayer)) vRayUnitStepSize = np.array( [ np.sqrt(1 + (vRayDir[1] / (vRayDir[0] + 1e-10)) ** 2), np.sqrt(1 + (vRayDir[0] / (vRayDir[1] + 1e-10)) ** 2), ] ) vMapCheck = vRayStart.copy() vRayLength1D = np.array([0.0, 0.0]) vStep = np.array([0, 0]) for i in range(2): vStep[i] = -1 if vRayDir[i] < 0 else 1 # vRayLength1D[i] = vRayUnitStepSize[i] # vRayLength1D[i] = ( # (vRayStart[i] - vMapCheck[i]) * vRayUnitStepSize[i] # if vRayDir[i] < 0 # else ((vMapCheck[i] + 1) - vRayStart[i]) * vRayUnitStepSize[i] # ) bTileFound = False fMaxDistance = depth fDistance = 0 while not bTileFound and fDistance < fMaxDistance: if vRayLength1D[0] < vRayLength1D[1]: vMapCheck[0] += vStep[0] fDistance = vRayLength1D[0] vRayLength1D[0] += vRayUnitStepSize[0] elif vRayLength1D[0] > vRayLength1D[1]: vMapCheck[1] += vStep[1] fDistance = vRayLength1D[1] vRayLength1D[1] += vRayUnitStepSize[1] elif vRayLength1D[0] == vRayLength1D[1]: vMapCheck += vStep vRayLength1D += vRayUnitStepSize if 0 <= vMapCheck[0] < w and 0 <= vMapCheck[1] < h: if binary_map[int(vMapCheck[1]), int(vMapCheck[0])] == 1: bTileFound = True else: break return angle, fDistance, bTileFound @jit(nopython=True) def loop_rays(x, y, angle, depth, w, h, binary_map, STEP_ANGLE, CASTED_RAYS): return [ ray_cast_dda(x, y, angle + i * STEP_ANGLE, depth, w, h, binary_map) for i in range(CASTED_RAYS) ] class Game: def __init__(self, tomlfile: str = "docs/game.toml"): start = time.time() self.model = Model(tomlfile=tomlfile) end = time.time() logging.info(f"DFP takes {end-start} s to output map.") self._initialise_game_engine() self.view = View(self.model).register_window(self.win) self.control = Controller(self.model) def _initialise_game_engine(self): pygame.init() self.win = pygame.display.set_mode( (self.model.SCREEN_WIDTH, self.model.SCREEN_HEIGHT) ) pygame.display.set_caption("Deep Floorplan Raycasting") self.clock = pygame.time.Clock() def run(self): while True: for event in pygame.event.get(): if event.type == pygame.QUIT: pygame.quit() sys.exit(0) # 3d background self.view.draw_3d_env() self.control.player_control() self.view.draw_dfp() self.view.show_fps(str(int(self.clock.get_fps()))) self.view.draw_player_loc() # calculate player rays wallangles = loop_rays( self.model.player_x, self.model.player_y, self.model.player_angle - self.model.HALF_FOV, self.model.MAX_DEPTH, self.model.w, self.model.h, self.model.result, self.model.STEP_ANGLE, self.model.args.casted_rays, ) for idx, (angle, fdist, iswall) in enumerate(wallangles): self.view.draw_player_rays(angle, fdist) wall_color = int( 255 * (1 - (3 * fdist / self.model.MAX_DEPTH) ** 2) ) wall_color = max(min(wall_color, 255), 30) fdist *= math.cos(self.model.player_angle - angle) wall_height = 21000 / (fdist + 0.00001) if iswall: self.view.draw_3d_wall(idx, wall_color, wall_height) pygame.display.flip() self.clock.tick() pygame.quit() if __name__ == "__main__": env = Game() env.run() ================================================ FILE: src/dfp/game/controller.py ================================================ import math import sys import pygame from .model import Model class Controller: def __init__(self, model: Model): self.model = model def player_control(self): keys = pygame.key.get_pressed() angle, x, y = ( self.model.player_angle, self.model.player_x, self.model.player_y, ) origx, origy = x, y if keys[pygame.K_LEFT]: angle -= self.model.dangle if keys[pygame.K_RIGHT]: angle += self.model.dangle if keys[pygame.K_w]: x -= math.sin(angle) * self.model.dpos y += math.cos(angle) * self.model.dpos if keys[pygame.K_s]: x += math.sin(angle) * self.model.dpos y -= math.cos(angle) * self.model.dpos if keys[pygame.K_d]: x += math.sin(angle - math.pi / 2) * self.model.dpos y -= math.cos(angle - math.pi / 2) * self.model.dpos if keys[pygame.K_a]: x -= math.sin(angle - math.pi / 2) * self.model.dpos y += math.cos(angle - math.pi / 2) * self.model.dpos if keys[pygame.K_q]: pygame.quit() sys.exit(0) if x < 0: x = 0 if x >= self.model.w: x = self.model.w - 1 if y < 0: y = 0 if y >= self.model.h: y = self.model.h - 1 if self.model.result[int(y), int(x)] == 1: ( self.model.player_angle, self.model.player_x, self.model.player_y, ) = ( angle, origx, origy, ) return self.model.player_angle, self.model.player_x, self.model.player_y = ( angle, x, y, ) ================================================ FILE: src/dfp/game/model.py ================================================ import math import random from argparse import Namespace import numpy as np import pygame from ..deploy import main from ..utils.settings import overwrite_args_with_toml class Model: FOV = math.pi / 3 HALF_FOV = FOV / 2 WALL_COLOR = [100, 100, 100] FLOOR_COLOR = [200, 200, 200] RAY_COLOR = (0, 255, 0) SKY_COLOR = (0, 150, 200) GRD_COLOR = (100, 100, 100) PLAYER_COLOR = (255, 0, 0) TILE_SIZE = 1 dangle = 0.01 dpos = 1 GAME_TEXT_COLOR = pygame.Color("coral") # surf = None def __init__(self, tomlfile: str = "docs/game.toml"): self.tomlfile = tomlfile args = Namespace(tomlfile=tomlfile) self.args = overwrite_args_with_toml(args) self.STEP_ANGLE = self.FOV / self.args.casted_rays self._initialise_map() self._initialise_player_pose() def _initialise_map(self): self.result = main(self.args) self.h, self.w = self.result.shape self.bg = self.rgb_im = np.zeros((self.h, self.w, 3)) self.bg[self.result != 10] = self.FLOOR_COLOR self.bg[self.result == 10] = self.WALL_COLOR self.result[self.result != 10] = 0 self.result[self.result == 10] = 1 self.bg = np.transpose(self.bg, (1, 0, 2)) self.MAX_DEPTH = max(self.h, self.w) self.SCREEN_HEIGHT = self.h self.SCREEN_WIDTH = self.w * 2 self.SCALE = (self.SCREEN_WIDTH / 2) / self.args.casted_rays self.surf = pygame.surfarray.make_surface(self.bg) def _initialise_player_pose(self): posy, posx = np.where(self.result != 1) posidx = random.randint(0, len(posy) - 1) self.player_x, self.player_y = posx[posidx], posy[posidx] self.player_angle = math.pi ================================================ FILE: src/dfp/game/view.py ================================================ from __future__ import annotations import math from typing import Union import pygame from .model import Model class View: def __init__(self, model: Model): self.model = model def register_window(self, win: pygame.Surface) -> View: self.win = win return self def draw_dfp(self): self.win.blit(self.model.surf, (0, 0)) def show_fps(self, fps: str): font = pygame.font.SysFont("Arial", 18) fps_text = font.render(fps, 1, self.model.GAME_TEXT_COLOR) self.win.blit(fps_text, (10, 0)) def draw_player_loc(self): pygame.draw.circle( self.win, self.model.PLAYER_COLOR, (self.model.player_x, self.model.player_y), 8, ) def draw_player_rays( self, angle: float, depth: Union[float, int], ): pygame.draw.line( self.win, self.model.RAY_COLOR, (self.model.player_x, self.model.player_y), ( self.model.player_x - math.sin(angle) * depth, self.model.player_y + math.cos(angle) * depth, ), 3, ) def draw_3d_env(self): pygame.draw.rect( self.win, self.model.GRD_COLOR, (self.model.w, self.model.h / 2, self.model.w, self.model.h), ) pygame.draw.rect( self.win, self.model.SKY_COLOR, (self.model.w, -self.model.h / 2, self.model.w, self.model.h), ) def draw_3d_wall( self, idx: int, wall_color: int, wall_height: Union[int, float] ): pygame.draw.rect( self.win, (wall_color, wall_color, wall_color), ( int(self.model.w + idx * self.model.SCALE), int((self.model.h / 2) - wall_height / 2), self.model.SCALE, wall_height, ), ) ================================================ FILE: src/dfp/loss.py ================================================ from typing import List, Tuple import tensorflow as tf def cross_two_tasks_weight( y1: tf.Tensor, y2: tf.Tensor ) -> Tuple[tf.Tensor, tf.Tensor]: p1, p2 = tf.keras.backend.sum(y1), tf.keras.backend.sum(y2) w1, w2 = p2 / (p1 + p2), p1 / (p1 + p2) return w1, w2 def balanced_entropy(x: tf.Tensor, y: tf.Tensor) -> tf.Tensor: eps = 1e-6 z = tf.keras.activations.softmax(x) cliped_z = tf.keras.backend.clip(z, eps, 1 - eps) log_z = tf.keras.backend.log(cliped_z) num_classes = y.shape.as_list()[-1] ind = tf.keras.backend.argmax(y, axis=-1) total = tf.keras.backend.sum(y) m_c: List[int] = [] n_c: List[int] = [] loss = 0 for c_ in range(num_classes): m_c.append( tf.keras.backend.cast( tf.keras.backend.equal(ind, c_), dtype=tf.int32 ) ) n_c.append( tf.keras.backend.cast( tf.keras.backend.sum(m_c[-1]), dtype=tf.float32 ) ) c: List[int] = [] for i in range(num_classes): c.append(total - n_c[i]) tc = tf.math.add_n(c) for i in range(num_classes): w = c[i] / tc m_c_one_hot = tf.one_hot((i * m_c[i]), num_classes, axis=-1) y_c = m_c_one_hot * y loss += w * tf.keras.backend.mean( -tf.keras.backend.sum(y_c * log_z, axis=1) ) return loss / num_classes ================================================ FILE: src/dfp/net.py ================================================ import argparse import os from typing import Tuple import numpy as np import tensorflow as tf from tensorflow.keras.applications.vgg16 import VGG16 from tensorflow.keras.models import Model os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" # print('Is gpu available: ',tf.test.is_gpu_available()); def conv2d( dim: int, size: int = 3, stride: int = 1, rate: int = 1, pad: str = "same", act: str = "relu", ) -> tf.keras.Sequential: result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2D( dim, size, strides=stride, padding=pad, dilation_rate=rate ) ) if act == "leaky": result.add(tf.keras.layers.LeakyReLU()) elif act == "relu": result.add(tf.keras.layers.ReLU()) return result def max_pool2d( size: int = 2, stride: int = 2, pad: str = "valid" ) -> tf.keras.Sequential: result = tf.keras.Sequential() result.add( tf.keras.layers.MaxPool2D(pool_size=size, strides=stride, padding=pad) ) return result def upconv2d( dim: int, size: int = 4, stride: int = 2, pad: str = "same", act: str = "relu", ) -> tf.keras.Sequential: result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2DTranspose(dim, size, strides=stride, padding=pad) ) if act == "relu": result.add(tf.keras.layers.ReLU()) return result def up_bilinear(dim: int) -> tf.keras.Sequential: result = tf.keras.Sequential() result.add(conv2d(dim, size=1, act="linear")) return result class deepfloorplanModel(Model): def __init__(self, config: argparse.Namespace = None): super(deepfloorplanModel, self).__init__() self.config = config dimlist = [256, 128, 64, 32] self.feature_names = [ "block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool", ] if config is not None: dimlist = config.feature_channels assert ( config.backbone == "vgg16" ), "subclass backbone must be vgg16" self.feature_names = config.feature_names self._vgg16init() # room boundary prediction (rbp) self.rbpups = [upconv2d(dim=d, act="linear") for d in dimlist] self.rbpcv1 = [conv2d(dim=d, act="linear") for d in dimlist] self.rbpcv2 = [conv2d(dim=d) for d in dimlist] self.rbpfinal = up_bilinear(3) # room type prediction (rtp) self.rtpups = [upconv2d(dim=d, act="linear") for d in dimlist] self.rtpcv1 = [conv2d(dim=d, act="linear") for d in dimlist] self.rtpcv2 = [conv2d(dim=d) for d in dimlist] # attention map self.atts1 = [conv2d(dim=dimlist[i]) for i in range(len(dimlist))] self.atts2 = [conv2d(dim=dimlist[i]) for i in range(len(dimlist))] self.atts3 = [ conv2d(dim=1, size=1, act="sigmoid") for i in range(len(dimlist)) ] # reduce the tensor depth self.xs1 = [conv2d(dim=d) for d in dimlist] self.xs2 = [conv2d(dim=1, size=1, act="linear") for d in dimlist] # context conv2d dak = [9, 17, 33, 65] # kernel_shape=[h,v,inc,outc] # horizontal self.hs = [self.constant_kernel((d, 1, 1, 1)) for d in dak] self.hf = [ tf.keras.layers.Conv2D( 1, [dak[i], 1], strides=1, padding="same", trainable=False, use_bias=False, weights=[self.hs[i]], ) for i in range(len(dak)) ] # vertical self.vs = [self.constant_kernel((1, d, 1, 1)) for d in dak] self.vf = [ tf.keras.layers.Conv2D( 1, [1, dak[i]], strides=1, padding="same", trainable=False, use_bias=False, weights=[self.vs[i]], ) for i in range(len(dak)) ] # diagonal self.ds = [self.constant_kernel((d, d, 1, 1), diag=True) for d in dak] self.df = [ tf.keras.layers.Conv2D( 1, dak[i], strides=1, padding="same", trainable=False, use_bias=False, weights=[self.ds[i]], ) for i in range(len(dak)) ] # diagonal flip self.dfs = [ self.constant_kernel((d, d, 1, 1), diag=True, flip=True) for d in dak ] self.dff = [ tf.keras.layers.Conv2D( 1, dak[i], strides=1, padding="same", trainable=False, use_bias=False, weights=[self.dfs[i]], ) for i in range(len(dak)) ] # expand dim self.ed = [conv2d(dim=d, size=1, act="linear") for d in dimlist] # learn rich feature self.lrf = [conv2d(dim=d) for d in dimlist] # final self.rtpfinal = up_bilinear(9) def _vgg16init(self): self.vgg16 = VGG16( weights="imagenet", include_top=False, input_shape=(512, 512, 3) ) for layer in self.vgg16.layers: layer.trainable = False def constant_kernel( self, shape: Tuple[int, int, int, int], val: int = 1, diag: bool = False, flip: bool = False, ) -> np.ndarray: k = np.array([]).astype(int) if not diag: k = val * np.ones(shape) else: w = np.eye(shape[0], shape[1]) if flip: w = w.reshape((shape[0], shape[1], 1)) w = np.flip(w, 1) k = w.reshape(shape) return k def non_local_context( self, t1: tf.Tensor, t2: tf.Tensor, idx: int, stride: int = 4 ) -> tf.Tensor: N, H, W, C = t1.shape.as_list() hs = H // stride if (H // stride) > 1 else (stride - 1) vs = W // stride if (W // stride) > 1 else (stride - 1) hs = hs if (hs % 2 != 0) else hs + 1 vs = hs if (vs % 2 != 0) else vs + 1 a = t1 x = t2 a = self.atts1[idx](a) a = self.atts2[idx](a) a = self.atts3[idx](a) a = tf.keras.activations.sigmoid(a) x = self.xs1[idx](x) x = self.xs2[idx](x) x = a * x h = self.hf[idx](x) v = self.vf[idx](x) d = self.df[idx](x) f = self.dff[idx](x) c1 = a * (h + v + d + f) c1 = self.ed[idx](c1) features = tf.concat([t2, c1], axis=3) out = self.lrf[idx](features) return out def call(self, x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]: features = [] feature = x for layer in self.vgg16.layers: feature = layer(feature) if layer.name in self.feature_names: features.append(feature) x = feature features = features[::-1] featuresrbp = [] for i in range(len(self.rbpups)): x = self.rbpups[i](x) + self.rbpcv1[i](features[i + 1]) x = self.rbpcv2[i](x) featuresrbp.append(x) logits_cw = tf.keras.backend.resize_images( self.rbpfinal(x), 2, 2, "channels_last" ) x = feature for i in range(len(self.rtpups)): x = self.rtpups[i](x) + self.rtpcv1[i](features[i + 1]) x = self.rtpcv2[i](x) x = self.non_local_context(featuresrbp[i], x, i) logits_r = tf.keras.backend.resize_images( self.rtpfinal(x), 2, 2, "channels_last" ) return logits_r, logits_cw ================================================ FILE: src/dfp/net_func.py ================================================ import argparse from typing import Tuple import numpy as np import tensorflow as tf from tensorflow.keras.applications.mobilenet import MobileNet from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2 from tensorflow.keras.applications.resnet50 import ResNet50 from tensorflow.keras.applications.vgg16 import VGG16 from tensorflow.keras.layers import ( Add, Concatenate, Conv2D, Conv2DTranspose, Input, Multiply, ReLU, ) from tensorflow.keras.models import Model def resnet50_backbone(x, feature_names): backbone = ResNet50(weights="imagenet", include_top=False, input_tensor=x) backbone = Model( inputs=x, outputs=backbone.get_layer(feature_names[-1]).output ) for layer in backbone.layers: layer.trainable = False features = [] for layer in feature_names: features.append(backbone.get_layer(layer).output) features = features[::-1] return features def mobilenet_backbone(x, feature_names): backbone = MobileNet(weights="imagenet", include_top=False, input_tensor=x) backbone = Model( inputs=x, outputs=backbone.get_layer(feature_names[-1]).output ) for layer in backbone.layers: layer.trainable = False features = [] for layer in feature_names: features.append(backbone.get_layer(layer).output) features = features[::-1] return features def mobilenetv2_backbone(x, feature_names): backbone = MobileNetV2( weights="imagenet", include_top=False, input_tensor=x ) backbone = Model( inputs=x, outputs=backbone.get_layer(feature_names[-1]).output ) for layer in backbone.layers: layer.trainable = False features = [] for layer in feature_names: features.append(backbone.get_layer(layer).output) features = features[::-1] return features def vgg16_backbone(x, feature_names): backbone = VGG16(weights="imagenet", include_top=False, input_tensor=x) backbone = Model( inputs=x, outputs=backbone.get_layer(feature_names[-1]).output ) for layer in backbone.layers: layer.trainable = False features = [] for layer in feature_names: features.append(backbone.get_layer(layer).output) features = features[::-1] return features def vertical_horizontal_filters( h: int, w: int ) -> Tuple[np.ndarray, np.ndarray]: return np.ones([1, h, 1, 1]), np.ones([w, 1, 1, 1]) def diagonal_filters(h: int, w: int) -> Tuple[np.ndarray, np.ndarray]: d = np.eye(h, w) dr = np.eye(h, w) dr = dr.reshape([h, w, 1]) dr = np.flip(dr, 1) return d.reshape((h, w, 1, 1)), dr.reshape((h, w, 1, 1)) def non_local_context(xf, x_, rbdim, stride=4): config_non_trainable = { "filters": 1, "trainable": False, "padding": "same", "use_bias": False, } _, H, W, _ = xf.shape hs = H // stride if (H // stride) > 1 else (stride - 1) vs = W // stride if (W // stride) > 1 else (stride - 1) hs = hs if (hs % 2 != 0) else hs + 1 vs = hs if (vs % 2 != 0) else vs + 1 vf, hf = vertical_horizontal_filters(vs, hs) df, dfr = diagonal_filters(vs, hs) v = Conv2D( kernel_size=[1, vs], **config_non_trainable, weights=[vf], )(x_) h = Conv2D( kernel_size=[hs, 1], **config_non_trainable, weights=[hf], )(x_) d = Conv2D( kernel_size=[hs, vs], **config_non_trainable, weights=[df], )(x_) dr = Conv2D( kernel_size=[hs, vs], **config_non_trainable, weights=[dfr], )(x_) c = Add()([v, h, d, dr]) c = Multiply()([xf, c]) c = Conv2D(rbdim, 1, strides=1, padding="same", dilation_rate=1)(c) c = Concatenate(axis=3)([x_, c]) x = Conv2D(rbdim, 3, strides=1, padding="same", dilation_rate=1)(c) return x def attention(xf, x_, rbdim): xf = Conv2D(rbdim, 3, strides=1, padding="same", dilation_rate=1)(xf) xf = ReLU()(xf) xf = Conv2D(rbdim, 3, strides=1, padding="same", dilation_rate=1)(xf) xf = ReLU()(xf) xf = Conv2D(1, 1, strides=1, padding="same", dilation_rate=1)(xf) xf = tf.keras.activations.sigmoid(xf) x_ = Conv2D(rbdim, 3, strides=1, padding="same", dilation_rate=1)(x_) x_ = ReLU()(x_) x_ = Conv2D(1, 1, strides=1, padding="same", dilation_rate=1)(x_) x_ = Multiply()([xf, x_]) return non_local_context(xf, x_, rbdim) def deepfloorplanFunc(config: argparse.Namespace = None): inp = Input([512, 512, 3]) if config is None: rbdims = [256, 128, 64, 32] features = vgg16_backbone( inp, [ "block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool", ], ) elif config is not None: rbdims = config.feature_channels if config.backbone == "resnet50": features = resnet50_backbone(inp, config.feature_names) elif config.backbone == "vgg16": features = vgg16_backbone(inp, config.feature_names) elif config.backbone == "mobilenetv1": features = mobilenet_backbone(inp, config.feature_names) elif config.backbone == "mobilenetv2": features = mobilenetv2_backbone(inp, config.feature_names) assert len(features) == 5, "Not enough 5 features..." features_room_boundary = [] x = features[0] for i in range(len(rbdims)): x = Conv2DTranspose(rbdims[i], 4, strides=2, padding="same")(x) xf = Conv2D(rbdims[i], 3, strides=1, padding="same", dilation_rate=1)( features[i + 1] ) x = Add()([x, xf]) x = Conv2D(rbdims[i], 3, strides=1, padding="same", dilation_rate=1)(x) x = ReLU()(x) features_room_boundary.append(x) x = Conv2D(3, 1, strides=1, padding="same", dilation_rate=1)(x) logits_cw = tf.keras.backend.resize_images(x, 2, 2, "channels_last") x = features[0] for i in range(len(rbdims)): x = Conv2DTranspose(rbdims[i], 4, strides=2, padding="same")(x) xf = Conv2D(rbdims[i], 3, strides=1, padding="same", dilation_rate=1)( features[i + 1] ) x = Add()([x, xf]) x = Conv2D(rbdims[i], 3, strides=1, padding="same", dilation_rate=1)(x) x = ReLU()(x) # attention and contexture x = attention(features_room_boundary[i], x, rbdims[i]) x = Conv2D(9, 1, strides=1, padding="same", dilation_rate=1)(x) logits_r = tf.keras.backend.resize_images(x, 2, 2, "channels_last") print(logits_cw.shape, logits_r.shape) return Model(inputs=inp, outputs=[logits_r, logits_cw]) if __name__ == "__main__": model = deepfloorplanFunc() print(model.summary()) for layer in model.layers: print(layer.name, layer.trainable) model.save("/tmp/tmp") ================================================ FILE: src/dfp/train.py ================================================ import argparse import io import os import sys from typing import List, Tuple import matplotlib import matplotlib.pyplot as plt import tensorflow as tf from tqdm import tqdm from .data import ( convert_one_hot_to_image, decodeAllRaw, loadDataset, preprocess, ) from .loss import balanced_entropy, cross_two_tasks_weight from .net import deepfloorplanModel from .net_func import deepfloorplanFunc from .utils.settings import overwrite_args_with_toml os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" def init( config: argparse.Namespace, ) -> Tuple[tf.data.Dataset, tf.keras.Model, tf.keras.optimizers.Optimizer]: dataset = loadDataset() if config.tfmodel == "subclass": model = deepfloorplanModel(config=config) elif config.tfmodel == "func": model = deepfloorplanFunc(config=config) os.system(f"mkdir -p {config.modeldir}") if config.weight: model.load_weights(config.weight) # optim = tf.keras.optimizers.AdamW(learning_rate=config.lr, # weight_decay=config.wd) optim = tf.keras.optimizers.Adam(learning_rate=config.lr) return dataset, model, optim def plot_to_image(figure: matplotlib.figure.Figure) -> tf.Tensor: buf = io.BytesIO() plt.savefig(buf, format="png") plt.close(figure) buf.seek(0) image = tf.image.decode_png(buf.getvalue(), channels=4) image = tf.expand_dims(image, 0) return image def image_grid( img: tf.Tensor, bound: tf.Tensor, room: tf.Tensor, logr: tf.Tensor, logcw: tf.Tensor, ) -> matplotlib.figure.Figure: figure = plt.figure() plt.subplot(2, 3, 1) plt.imshow(img[0].numpy()) plt.xticks([]) plt.yticks([]) plt.grid(False) plt.subplot(2, 3, 2) plt.imshow(bound[0].numpy()) plt.xticks([]) plt.yticks([]) plt.grid(False) plt.subplot(2, 3, 3) plt.imshow(room[0].numpy()) plt.xticks([]) plt.yticks([]) plt.grid(False) plt.subplot(2, 3, 5) plt.imshow(convert_one_hot_to_image(logcw)[0].numpy().squeeze()) plt.xticks([]) plt.yticks([]) plt.grid(False) plt.subplot(2, 3, 6) plt.imshow(convert_one_hot_to_image(logr)[0].numpy().squeeze()) plt.xticks([]) plt.yticks([]) plt.grid(False) return figure @tf.function def train_step( model: tf.keras.Model, optim: tf.keras.optimizers.Optimizer, img: tf.Tensor, hr: tf.Tensor, hb: tf.Tensor, ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]: # forward with tf.GradientTape() as tape: logits_r, logits_cw = model(img) loss1 = balanced_entropy(logits_r, hr) loss2 = balanced_entropy(logits_cw, hb) w1, w2 = cross_two_tasks_weight(hr, hb) loss = w1 * loss1 + w2 * loss2 # backward grads = tape.gradient(loss, model.trainable_weights) optim.apply_gradients(zip(grads, model.trainable_weights)) return logits_r, logits_cw, loss, loss1, loss2 def main(config: argparse.Namespace): # initialization writer = tf.summary.create_file_writer(config.logdir) pltiter = 0 dataset, model, optim = init(config) # training loop for epoch in range(config.epochs): print("[INFO] Epoch {}".format(epoch)) for data in tqdm(list(dataset.shuffle(400).batch(config.batchsize))): img, bound, room = decodeAllRaw(data) img, bound, room, hb, hr = preprocess(img, bound, room) logits_r, logits_cw, loss, loss1, loss2 = train_step( model, optim, img, hr, hb ) # plot progress if pltiter % config.save_tensor_interval == 0: f = image_grid(img, bound, room, logits_r, logits_cw) im = plot_to_image(f) with writer.as_default(): tf.summary.scalar("Loss", loss.numpy(), step=pltiter) tf.summary.scalar("LossR", loss1.numpy(), step=pltiter) tf.summary.scalar("LossB", loss2.numpy(), step=pltiter) tf.summary.image("Data", im, step=pltiter) writer.flush() pltiter += 1 # save model if epoch % config.save_model_interval == 0: model.save_weights(config.logdir + "/G") model.save(config.modeldir) print("[INFO] Saving Model ...") def parse_args(args: List[str]) -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument( "--tfmodel", type=str, default="subclass", choices=["subclass", "func"] ) p.add_argument("--batchsize", type=int, default=2) p.add_argument("--lr", type=float, default=1e-4) p.add_argument("--wd", type=float, default=1e-5) p.add_argument("--epochs", type=int, default=1000) p.add_argument("--logdir", type=str, default="log/store") p.add_argument("--modeldir", type=str, default="model/store") p.add_argument("--weight", type=str) p.add_argument("--save-tensor-interval", type=int, default=10) p.add_argument("--save-model-interval", type=int, default=20) p.add_argument("--tomlfile", type=str, default=None) p.add_argument( "--feature-channels", type=int, action="store", default=[256, 128, 64, 32], nargs=4, ) p.add_argument( "--backbone", type=str, default="vgg16", choices=["vgg16", "resnet50", "mobilenetv1", "mobilenetv2"], ) p.add_argument( "--feature-names", type=str, action="store", nargs=5, default=[ "block1_pool", "block2_pool", "block3_pool", "block4_pool", "block5_pool", ], ) return p.parse_args(args) if __name__ == "__main__": args = parse_args(sys.argv[1:]) args = overwrite_args_with_toml(args) print(args) main(args) ================================================ FILE: src/dfp/utils/__init__.py ================================================ ================================================ FILE: src/dfp/utils/legend.py ================================================ from typing import List import matplotlib import matplotlib.pyplot as plt import numpy as np from .rgb_ind_convertor import floorplan_fuse_map def export_legend( legend: matplotlib.legend.Legend, filename: str = "legend.png", expand: List[int] = [-5, -5, 5, 5], ): fig = legend.figure fig.canvas.draw() bbox = legend.get_window_extent() bbox = bbox.from_extents(*(bbox.extents + np.array(expand))) bbox = bbox.transformed(fig.dpi_scale_trans.inverted()) fig.savefig(filename, dpi="figure", bbox_inches=bbox) def norm255to1(x: List[int]) -> List[float]: return [p / 255 for p in x] def handle(m: str, c: List[float]): return plt.plot([], [], marker=m, color=c, ls="none")[0] def main(): colors = [ "background", "closet", "bathroom", "living room\nkitchen\ndining room", "bedroom", "hall", "balcony", "not used", "not used", "door/window", "wall", ] colors2 = [norm255to1(rgb) for rgb in list(floorplan_fuse_map.values())] handles = [handle("s", colors2[i]) for i in range(len(colors))] labels = colors legend = plt.legend(handles, labels, loc=3, framealpha=1, frameon=True) export_legend(legend) if __name__ == "__main__": main() ================================================ FILE: src/dfp/utils/rgb_ind_convertor.py ================================================ from typing import Dict, List import numpy as np # use for index 2 rgb floorplan_room_map = { 0: [0, 0, 0], # background 1: [192, 192, 224], # closet 2: [192, 255, 255], # bathroom/washroom 3: [224, 255, 192], # livingroom/kitchen/diningroom 4: [255, 224, 128], # bedroom 5: [255, 160, 96], # hall 6: [255, 224, 224], # balcony 7: [224, 224, 224], # not used 8: [224, 224, 128], # not used } # boundary label floorplan_boundary_map = { 0: [0, 0, 0], # background 1: [255, 60, 128], # opening (door&window) 2: [255, 255, 255], # wall line } # boundary label for presentation floorplan_boundary_map_figure = { 0: [255, 255, 255], # background 1: [255, 60, 128], # opening (door&window) 2: [0, 0, 0], # wall line } # merge all label into one multi-class label floorplan_fuse_map = { 0: [0, 0, 0], # background 1: [192, 192, 224], # closet 2: [192, 255, 255], # batchroom/washroom 3: [224, 255, 192], # livingroom/kitchen/dining room 4: [255, 224, 128], # bedroom 5: [255, 160, 96], # hall 6: [255, 224, 224], # balcony 7: [224, 224, 224], # not used 8: [224, 224, 128], # not used 9: [255, 60, 128], # extra label for opening (door&window) 10: [255, 255, 255], # extra label for wall line } # invert the color of wall line and background for presentation floorplan_fuse_map_figure = { 0: [255, 255, 255], # background 1: [192, 192, 224], # closet 2: [192, 255, 255], # batchroom/washroom 3: [224, 255, 192], # livingroom/kitchen/dining room 4: [255, 224, 128], # bedroom 5: [255, 160, 96], # hall 6: [255, 224, 224], # balcony 7: [224, 224, 224], # not used 8: [224, 224, 128], # not used 9: [255, 60, 128], # extra label for opening (door&window) 10: [0, 0, 0], # extra label for wall line } def rgb2ind( im: np.ndarray, color_map: Dict[int, List[int]] = floorplan_room_map ) -> np.ndarray: ind = np.zeros((im.shape[0], im.shape[1])) for i, rgb in color_map.items(): ind[(im == rgb).all(2)] = i # return ind.astype(int) # int => int64 return ind.astype(np.uint8) # force to uint8 def ind2rgb( ind_im: np.ndarray, color_map: Dict[int, List[int]] = floorplan_room_map ) -> np.ndarray: rgb_im = np.zeros((ind_im.shape[0], ind_im.shape[1], 3)) for i, rgb in color_map.items(): rgb_im[(ind_im == i)] = rgb return rgb_im.astype(int) ================================================ FILE: src/dfp/utils/settings.py ================================================ import argparse from argparse import Namespace from dynaconf import Dynaconf def overwrite_args_with_toml(config: argparse.Namespace) -> argparse.Namespace: if config.tomlfile is None: return config settings = Dynaconf( envvar_prefix="DYNACONF", settings_files=[config.tomlfile] ) settings = dict((k.lower(), v) for k, v in settings.as_dict().items()) return Namespace(**settings) ================================================ FILE: src/dfp/utils/util.py ================================================ import cv2 import numpy as np import tensorflow as tf from scipy import ndimage def fast_hist(im: np.ndarray, gt: np.ndarray, n: int = 9) -> np.ndarray: """ n is num_of_classes """ k = (gt >= 0) & (gt < n) return np.bincount( n * gt[k].astype(int) + im[k], minlength=n**2 ).reshape(n, n) def flood_fill(test_array: np.ndarray, h_max: int = 255) -> np.ndarray: """ fill in the hole """ test_array = test_array.squeeze() input_array = np.copy(test_array) el = ndimage.generate_binary_structure(2, 2).astype(int) inside_mask = ndimage.binary_erosion(~np.isnan(input_array), structure=el) output_array = np.copy(input_array) output_array[inside_mask] = h_max output_old_array = np.copy(input_array) output_old_array.fill(0) el = ndimage.generate_binary_structure(2, 1).astype(int) while not np.array_equal(output_old_array, output_array): output_old_array = np.copy(output_array) output_array = np.maximum( input_array, ndimage.grey_erosion(output_array, size=(3, 3), footprint=el), ) return output_array def fill_break_line(cw_mask: np.ndarray) -> np.ndarray: broken_line_h = np.array( [ [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 0, 0, 0, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], ], dtype=np.uint8, ) broken_line_h2 = np.array( [ [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 1, 0, 1, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], ], dtype=np.uint8, ) broken_line_v = np.transpose(broken_line_h) broken_line_v2 = np.transpose(broken_line_h2) cw_mask = cv2.morphologyEx(cw_mask, cv2.MORPH_CLOSE, broken_line_h) cw_mask = cv2.morphologyEx(cw_mask, cv2.MORPH_CLOSE, broken_line_v) cw_mask = cv2.morphologyEx(cw_mask, cv2.MORPH_CLOSE, broken_line_h2) cw_mask = cv2.morphologyEx(cw_mask, cv2.MORPH_CLOSE, broken_line_v2) return cw_mask def refine_room_region(cw_mask: np.ndarray, rm_ind: np.ndarray) -> np.ndarray: label_rm, num_label = ndimage.label((1 - cw_mask)) new_rm_ind = np.zeros(rm_ind.shape) for j in range(1, num_label + 1): mask = (label_rm == j).astype(np.uint8) ys, xs, _ = np.where(mask != 0) area = (np.amax(xs) - np.amin(xs)) * (np.amax(ys) - np.amin(ys)) if area < 100: continue else: room_types, type_counts = np.unique( mask * rm_ind, return_counts=True ) if len(room_types) > 1: room_types = room_types[1:] # ignore background type which is zero type_counts = type_counts[1:] # ignore background count new_rm_ind += mask * room_types[np.argmax(type_counts)] return new_rm_ind def print_model_weights_sparsity(model): for layer in model.layers: if isinstance(layer, tf.keras.layers.Wrapper): weights = layer.trainable_weights else: weights = layer.weights for weight in weights: if "kernel" not in weight.name or "centroid" in weight.name: continue weight_size = weight.numpy().size zero_num = np.count_nonzero(weight == 0) print( f"{weight.name}: {zero_num/weight_size:.2%} sparsity ", f"({zero_num}/{weight_size})", ) def print_model_weight_clusters(model): for layer in model.layers: if isinstance(layer, tf.keras.layers.Wrapper): weights = layer.trainable_weights else: weights = layer.weights for weight in weights: # ignore auxiliary quantization weights if "quantize_layer" in weight.name: continue if "kernel" in weight.name: unique_count = len(np.unique(weight)) print(f"{layer.name}/{weight.name}: {unique_count} clusters ") ================================================ FILE: tests/__init__.py ================================================ ================================================ FILE: tests/test_app.py ================================================ import os from argparse import Namespace from types import TracebackType from typing import Any, Dict, List, Optional, Type import pytest import numpy as np from flask.testing import FlaskClient from pytest_mock import MockFixture from dfp.app import app as create_app class fakeMultiprocessing: def Pool(self): return self def map(self, *args: str, **kwargs: int) -> np.ndarray: return np.zeros([1, 32, 32, 3]) def __enter__(self) -> object: return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_value: Optional[BaseException], traceback: Optional[TracebackType], ): pass def __getitem__(self) -> np.ndarray: return np.zeros([1, 32, 32, 3]) class fakeForm: def __init__(self, data: Dict[str, str]): self.data = data def keys(self): return self.data.keys() def getlist(self, key: str) -> List[str]: return [self.data[key]] def __getitem__(self, key: str) -> str: return self.data[key] class fakeRequest: def __init__(self): self.form = fakeForm( {"postprocess": "0", "colorize": "0", "output": "/tmp"} ) # self.files = [] self.json = fakeForm({}) @pytest.fixture def client(mocker: MockFixture) -> FlaskClient: mp = fakeMultiprocessing() args = Namespace( image="resources/30939153.jpg", weight="", loadmethod="none", postprocess=False, colorize=False, save="/tmp/tmp.jpg", ) content = Namespace(content=None) mocker.patch("dfp.app.mp.Pool", return_value=mp) mocker.patch("dfp.app.Namespace", return_value=args) mocker.patch("dfp.app.saveStreamFile", return_value=None) mocker.patch("dfp.app.saveStreamURI", return_value=None) mocker.patch("dfp.app.requests.get", return_value=content) mocker.patch("dfp.app.os.system", return_value=None) mocker.patch("dfp.app.mpimg.imsave", return_value=None) mocker.patch("dfp.app.send_file", return_value={"message": "success!"}) return create_app.test_client() def test_app_home(client: FlaskClient): resp = client.get("/") assert resp.status_code == 200 assert isinstance(resp.json, dict) assert resp.json.get("message", "Hello Flask!") def test_app_process_image(client: FlaskClient): resp = client.post("/upload") assert resp.status_code == 200 def test_app_mock_process_empty(client: FlaskClient): headers: Dict[Any, Any] = {} data: Dict[Any, Any] = {} resp = client.post("/uri", headers=headers, json=data) assert resp.status_code == 200 assert resp.json.get("message", "success!") def test_app_mock_process_uri(client: FlaskClient): headers: Dict[Any, Any] = {} data = { "uri": "", "postprocess": 1, "colorize": 1, "output": "/tmp", } resp = client.post("/uri", headers=headers, json=data) os.system("rm *.jpg") assert resp.status_code == 200 assert resp.json.get("message", "success!") def test_app_mock_process_file(client: FlaskClient): files = {"file": (open("resources/30939153.jpg", "rb"), "30939153.jpg")} resp = client.post("/upload", data=files) os.system("rm *.jpg") assert resp.status_code == 200 ================================================ FILE: tests/test_convert2tflite.py ================================================ from argparse import Namespace from types import TracebackType from typing import Any, List, Optional, Type from pytest_mock import MockFixture from dfp.convert2tflite import converter, parse_args class fakeConverter: def __init__(self): self.optimizations = [] self.experimental_new_converter = False def convert(self): return None class fakeFile: def write(self, *args: str, **kwargs: int): pass def __enter__(self, *args: str, **kwargs: int): return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_value: Optional[BaseException], traceback: Optional[TracebackType], ): pass class fakeModel: layers: List[Any] = [] def test_parse_args(): args = parse_args(["--quantize"]) assert args.quantize is True def test_converter(mocker: MockFixture): model = fakeModel() con = fakeConverter() f = fakeFile() mocker.patch( "dfp.convert2tflite.tf.keras.models.load_model", return_value=model ) mocker.patch( "dfp.convert2tflite.tf.lite.TFLiteConverter.from_keras_model", return_value=con, ) mocker.patch("dfp.convert2tflite.open", return_value=f) args = Namespace( quantize=True, tflitedir="model/store/model.tflite", modeldir="model/store", compress_mode="quantization", tfmodel="subclass", loadmethod="pb", ) converter(args) ================================================ FILE: tests/test_data.py ================================================ import numpy as np import tensorflow as tf from pytest_mock import MockFixture from dfp.data import ( _parse_function, convert_one_hot_to_image, decodeAllRaw, main, plotData, preprocess, ) def _bytes_feature(value: bytes) -> tf.train.Feature: """Returns a bytes_list from a string / byte.""" if isinstance(value, type(tf.constant(0))): value = ( value.numpy() ) # BytesList won't unpack a string from an EagerTensor. return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) class TestDataCase: def test_convert_one_hot2img(self): hot = tf.random.uniform((1, 16, 16, 3), minval=0, maxval=1) hot = hot / tf.reduce_sum(hot, axis=-1, keepdims=True) res = convert_one_hot_to_image(hot, dtype="int", act="softmax") assert res.numpy().shape == (1, 16, 16, 1) def test_preprocess(self): img = tf.random.normal((26, 26, 3)) b = tf.random.uniform( shape=(26, 26), minval=0, maxval=2, dtype=tf.int32 ) r = tf.random.uniform( shape=(26, 26), minval=0, maxval=8, dtype=tf.int32 ) img_, b_, r_, hb, hr = preprocess(img, b, r, size=26) assert img_.numpy().shape == (1, 26, 26, 3) assert b_.numpy().shape == (1, 26, 26) assert r_.numpy().shape == (1, 26, 26) assert hb.numpy().shape == (1, 26, 26, 3) assert hr.numpy().shape == (1, 26, 26, 9) def test_decodeAllRaw(self): inp = tf.constant("1") x = {"image": inp, "boundary": inp, "room": inp} i, b, r = decodeAllRaw(x) assert i.shape == (1,) assert b.shape == (1,) assert r.shape == (1,) def test_parse_function(self, mocker: MockFixture): encoded_image_string = np.array([[1, 2], [3, 4]]).tobytes() image = tf.compat.as_bytes(encoded_image_string) image = _bytes_feature(image) feature = { "image": image, "boundary": image, "room": image, "door": image, } tf_example = tf.train.Example( features=tf.train.Features(feature=feature) ).SerializeToString() res = _parse_function(tf_example) assert len(list(res.keys())) == 4 def test_plotData(self, mocker: MockFixture): img = tf.random.normal((1, 26, 26, 3)) b = tf.random.uniform( shape=(1, 26, 26), minval=0, maxval=2, dtype=tf.int32 ) r = tf.random.uniform( shape=(1, 26, 26), minval=0, maxval=8, dtype=tf.int32 ) hb = tf.random.uniform(shape=(1, 26, 26, 3), minval=0, maxval=1) hr = tf.random.uniform(shape=(1, 26, 26, 9), minval=0, maxval=1) mocker.patch("dfp.data.decodeAllRaw", return_value=(img, b, r)) mocker.patch("dfp.data.preprocess", return_value=(img, b, r, hb, hr)) eg = tf.convert_to_tensor(np.array([[1, 2], [3, 4]]).tobytes()) plotData({"image": eg}) def test_main(self, mocker: MockFixture): mocker.patch("dfp.data.plotData", return_value=None) mocker.patch("dfp.data.plt.show", return_value=None) arr = np.array([[1, 2], [3, 4]]).tobytes() arr = tf.data.Dataset.from_tensor_slices([arr]) ds = tf.data.Dataset.zip((arr, arr)) main(ds) ================================================ FILE: tests/test_deploy.py ================================================ from argparse import Namespace from typing import Dict, List, Tuple import pytest import numpy as np import tensorflow as tf from pytest_mock import MockFixture from dfp.deploy import ( colorize, deploy_plot_res, init, main, parse_args, post_process, predict, ) class fakeLayer: def __init__(self): self.name = "pool" def __call__(self, x: tf.Tensor) -> tf.Tensor: return x class fakeVGG16: def __init__(self): self.trainable = True self.layer = fakeLayer() self.layers = [self.layer, self.layer] class fakeModel: def __init__(self): self.trainable = True self.vgg16 = fakeVGG16() self.rtpfinal = fakeLayer() self.rbpups = [self.rbpfinal] self.rbpcv1 = [self.rbpfinal] self.rbpcv2 = [self.rbpfinal] self.rtpups = [self.rtpfinal] self.rtpcv1 = [self.rtpfinal] self.rtpcv2 = [self.rtpfinal] def rbpfinal(self, x: tf.Tensor) -> tf.Tensor: return x def non_local_context( self, x: tf.Tensor, *args: int, **kwargs: str ) -> tf.Tensor: return x def predict(self, x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]: a = tf.random.uniform((1, 16, 16, 3), minval=0, maxval=1) b = tf.random.uniform((1, 16, 16, 9), minval=0, maxval=1) return b / tf.reduce_sum(b, axis=-1, keepdims=True), a / tf.reduce_sum( a, axis=-1, keepdims=True ) def invoke(self): pass def get_input_details(self) -> List[Dict[str, Tuple[int, int, int]]]: return [{"index": (512, 512, 3)}] def get_output_details(self) -> List[Dict[str, Tuple[int, int, int]]]: return [{"index": (512, 512, 9)}, {"index": (512, 512, 3)}] def set_tensor(self, ind: Tuple[int, int, int], img: tf.Tensor): pass def get_tensor(self, ind: Tuple[int, int, int]): if ind[2] == 3: a = tf.random.uniform((1, 16, 16, 3), minval=0, maxval=1) return a / tf.reduce_sum(a, axis=-1, keepdims=True) elif ind[2] == 9: b = tf.random.uniform((1, 16, 16, 9), minval=0, maxval=1) return b / tf.reduce_sum(b, axis=-1, keepdims=True) def load_weights(self, x: str): pass def allocate_tensors(self): pass @pytest.fixture def model_img() -> Tuple[fakeModel, tf.Tensor]: model = fakeModel() img = tf.random.normal((1, 16, 16, 3)) return model, img @pytest.mark.parametrize( "h,w,c", [(8, 8, 3), (16, 16, 3), (32, 32, 3), (64, 64, 3)] ) def test_colorize(h: int, w: int, c: int, mocker: MockFixture): inp = np.random.randint(2, size=(h, w)) out = np.zeros((h, w, c)) m = mocker.patch("dfp.deploy.ind2rgb") m.return_value = out r, cw = colorize(inp, inp) assert r.shape == (h, w, c) and cw.shape == (h, w, c) @pytest.mark.parametrize( "h,w,c", [(8, 8, 3), (16, 16, 3), (32, 32, 3), (64, 64, 3)] ) def test_post_process(h: int, w: int, c: int): inp = np.ones((h, w, 1)) r, cw = post_process(inp, inp, np.array([h, w, c])) assert r.shape == (h, w, 1) and cw.shape == (h, w) def test_parse_args(): args = parse_args( ["--postprocess", "--loadmethod", "tflite", "--tfmodel", "subclass"] ) assert args.postprocess is True assert args.loadmethod == "tflite" def test_init_none(mocker: MockFixture): model = fakeModel() mocker.patch("dfp.deploy.deepfloorplanModel", return_value=model) mocker.patch("dfp.deploy.mpimg.imread", return_value=np.zeros([16, 16, 3])) args = parse_args( '--loadmethod none --image "" --tfmodel subclass'.split() ) model_, img, shp = init(args) assert shp == (16, 16, 3) def test_init_log(mocker: MockFixture): model = fakeModel() mocker.patch("dfp.deploy.deepfloorplanModel", return_value=model) mocker.patch("dfp.deploy.mpimg.imread", return_value=np.zeros([16, 16, 3])) args = parse_args( """--loadmethod log --image "" --weight log/store/G --tfmodel subclass""".split() ) model_, img, shp = init(args) assert shp == (16, 16, 3) def test_init_pb(mocker: MockFixture): model = fakeModel() mocker.patch("dfp.deploy.deepfloorplanModel", return_value=model) mocker.patch("dfp.deploy.mpimg.imread", return_value=np.zeros([16, 16, 3])) mocker.patch("dfp.deploy.tf.keras.models.load_model", return_value=model) args = parse_args( """--loadmethod pb --image "" --weight model/store --tfmodel subclass""".split() ) model_, img, shp = init(args) assert shp == (16, 16, 3) def test_init_tflite(mocker: MockFixture): model = fakeModel() mocker.patch("dfp.deploy.deepfloorplanModel", return_value=model) mocker.patch("dfp.deploy.mpimg.imread", return_value=np.zeros([16, 16, 3])) mocker.patch("dfp.deploy.tf.lite.Interpreter", return_value=model) args = parse_args( """--loadmethod tflite --image \"\" --weight model/store/model.tflite --tfmodel subclass""".split() ) model_, img, shp = init(args) assert shp == (16, 16, 3) @pytest.mark.parametrize( "colorize,postprocess,expected", [ (True, True, (16, 16, 3)), (False, False, (16, 16)), (True, False, (16, 16, 3)), (False, True, (16, 16)), ], ) def test_main( colorize: bool, postprocess: bool, expected: Tuple[int, int, int], model_img: Tuple[fakeModel, tf.Tensor], mocker: MockFixture, ): args = Namespace( loadmethod="none", image="", colorize=colorize, postprocess=postprocess, save=True, tfmodel="subclass", ) model, img = model_img mocker.patch( "dfp.deploy.init", return_value=(model.predict, img, [16, 16, 3]) ) mocker.patch("dfp.deploy.mpimg.imsave", return_value=None) res = main(args) assert res.shape == expected def test_main_tflite( model_img: Tuple[fakeModel, tf.Tensor], mocker: MockFixture ): args = Namespace( loadmethod="tflite", image="", colorize=True, postprocess=True, save=False, tfmodel="subclass", ) model, img = model_img mocker.patch("dfp.deploy.init", return_value=(model, img, [16, 16, 3])) res = main(args) assert res.shape == (16, 16, 3) def test_main_log(model_img: Tuple[fakeModel, tf.Tensor], mocker: MockFixture): args = Namespace( loadmethod="log", image="", colorize=True, postprocess=True, save=False, tfmodel="subclass", ) model, img = model_img mocker.patch("dfp.deploy.init", return_value=(model, img, [16, 16, 3])) mocker.patch( "dfp.deploy.predict", return_value=( tf.random.uniform((1, 16, 16, 3), minval=0, maxval=1), tf.random.uniform((1, 16, 16, 9), minval=0, maxval=1), ), ) res = main(args) assert res.shape == (16, 16, 3) def test_deploy_plot_res(): a = np.zeros((16, 16, 3)) deploy_plot_res(a) def test_predict(model_img: Tuple[fakeModel, tf.Tensor], mocker: MockFixture): model, img = model_img shp = np.array([16, 16, 3]) a, b = predict(model, img, shp) assert a.numpy().shape == (1, 32, 32, 3) ================================================ FILE: tests/test_loss.py ================================================ import unittest import tensorflow as tf from dfp.loss import balanced_entropy, cross_two_tasks_weight class TestLossCase(unittest.TestCase): randx = tf.random.normal((1, 32, 32, 3), 0, 1) randy = tf.random.normal((1, 32, 32, 9), 0, 1) def test_balanced_entropy(self): loss = balanced_entropy(self.__class__.randx, self.__class__.randx) self.assertEqual(loss.dtype, tf.float32) def test_weight(self): w1, w2 = cross_two_tasks_weight( self.__class__.randx, self.__class__.randy ) w = w1 + w2 self.assertLessEqual(w.numpy(), 2.0) if __name__ == "__main__": unittest.main() ================================================ FILE: tests/test_net.py ================================================ import unittest import numpy as np import tensorflow as tf from tensorflow.keras.applications.vgg16 import preprocess_input from dfp.net import ( conv2d, deepfloorplanModel, max_pool2d, up_bilinear, upconv2d, ) class TestNetCase(unittest.TestCase): model = deepfloorplanModel() randx = np.random.randn(1, 512, 512, 3) def test_deepfloorplan_forward(self): with tf.device("/cpu:0"): x = preprocess_input(self.__class__.randx) logits_r, logits_cw = self.__class__.model(x) shpr, shpcw = logits_r.numpy().shape, logits_cw.numpy().shape self.assertEqual((shpr, shpcw), ((1, 512, 512, 9), (1, 512, 512, 3))) def test_vgg16(self): gt = [ (1, 256, 256, 64), (1, 128, 128, 128), (1, 64, 64, 256), (1, 32, 32, 512), (1, 16, 16, 512), ] vgg16 = self.__class__.model.vgg16 out = [] x = self.__class__.randx for lay in vgg16.layers: x = lay(x) if lay.name.find("pool") != -1: out.append(x.shape) self.assertEqual(out, gt) def test_conv2d(self): lay = conv2d(1, act="leaky") x = np.random.randn(1, 32, 32, 2) x = lay(x) self.assertEqual(x.numpy().shape, (1, 32, 32, 1)) def test_upconv2d(self): lay = upconv2d(1, act="relu") x = np.random.randn(1, 16, 16, 2) x = lay(x) self.assertEqual(x.numpy().shape, (1, 32, 32, 1)) def test_maxpool2d(self): lay = max_pool2d() x = np.random.randn(1, 32, 32, 1) x = lay(x) self.assertEqual(x.numpy().shape, (1, 16, 16, 1)) def test_upbilinear(self): lay = up_bilinear(2) x = np.random.randn(1, 16, 16, 1) x = lay(x) self.assertEqual(x.numpy().shape, (1, 16, 16, 2)) if __name__ == "__main__": unittest.main() ================================================ FILE: tests/test_train.py ================================================ from argparse import Namespace from types import TracebackType from typing import Any, List, Optional, Tuple, Type import matplotlib import matplotlib.pyplot as plt import tensorflow as tf from pytest_mock import MockFixture from dfp.train import image_grid, init, parse_args, plot_to_image, train_step class fakeModel: def __init__(self): self.trainable_weights = [] def load_weights(self, weights: str): pass def __call__( self, *args: str, **kwargs: int ) -> Tuple[tf.Tensor, tf.Tensor]: a = tf.random.uniform((1, 16, 16, 3), minval=0, maxval=1) b = tf.random.uniform((1, 16, 16, 9), minval=0, maxval=1) return a / tf.reduce_sum(a, axis=-1, keepdims=True), b / tf.reduce_sum( b, axis=-1, keepdims=True ) class fakeOptim: def apply_gradients(self, *args: str, **kwargs: int): pass class fakeTape: def GradientTape(self) -> object: return self def gradient(self, *args: str, **kwargs: int) -> List[Any]: return [] def __enter__(self) -> object: return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_value: Optional[BaseException], traceback: Optional[TracebackType], ): pass class TestTrainCase: def test_image_grid(self): inp123 = tf.random.normal([1, 32, 32, 3]) inp123 = tf.clip_by_value(inp123, 0, 1) inp45 = inp123 / tf.reduce_sum(inp123, axis=-1, keepdims=True) f = image_grid(inp123, inp123, inp123, inp45, inp45) assert isinstance(f, matplotlib.figure.Figure) def test_plot_to_image(self): f = plt.figure() img = plot_to_image(f) assert img.numpy().ndim == 4 def test_parse_args(self): args = parse_args(["--lr", "1e-3", "--epochs", "300"]) assert args.lr == 1e-3 assert args.epochs == 300 def test_init(self, mocker: MockFixture): model = fakeModel() mocker.patch("dfp.train.loadDataset", return_value=None) mocker.patch("dfp.train.deepfloorplanModel", return_value=model) args = Namespace( weight="fakepath", lr=1e-4, tfmodel="subclass", modeldir=None ) ds, model, opt = init(args) assert isinstance(opt, tf.keras.optimizers.Optimizer) def test_trainstep(self, mocker: MockFixture): model = fakeModel() tape = fakeTape() optim = fakeOptim() img = tf.random.uniform((1, 16, 16, 3), minval=0, maxval=1) hr = tf.random.uniform((1, 16, 16, 3), minval=0, maxval=1) hb = tf.random.uniform((1, 16, 16, 9), minval=0, maxval=1) mocker.patch("dfp.train.tf.GradientTape", return_value=tape) logits_r, logits_cw, loss, loss1, loss2 = train_step.__wrapped__( model, optim, img, hr, hb ) assert logits_r.numpy().shape == (1, 16, 16, 3) assert logits_cw.numpy().shape == (1, 16, 16, 9) ================================================ FILE: tests/utils/__init__.py ================================================ ================================================ FILE: tests/utils/test_legend.py ================================================ import os import unittest import matplotlib import matplotlib.pyplot as plt from dfp.utils.legend import export_legend, handle, main, norm255to1 from dfp.utils.rgb_ind_convertor import floorplan_fuse_map class TestLegendCase(unittest.TestCase): rgbs = list(floorplan_fuse_map.values()) colors = [ "background", "closet", "bathroom", "living room\nkitchen\ndining room", "bedroom", "hall", "balcony", "not used", "not used", "door/window", "wall", ] rgbs01 = [norm255to1(rgb) for rgb in rgbs] def test_norm255to1(self): self.assertEqual(len(self.__class__.rgbs01[0]), 3) def test_handle(self): self.__class__.handles = [ handle("s", self.__class__.rgbs01[i]) for i in range(len(self.__class__.colors)) ] self.assertIsInstance( self.__class__.handles[0], matplotlib.lines.Line2D ) def test_export_legend(self): handles = [handle("s", [0.5, 0.5, 0.5])] legend = plt.legend( handles, "Leo", loc=3, framealpha=1, frameon=True, ) export_legend(legend, filename="tmp.png") ans = os.path.isfile("tmp.png") os.system("rm tmp.png") self.assertTrue(ans) def test_main(self): main() ans = os.path.isfile("legend.png") os.system("rm legend.png") self.assertTrue(ans) if __name__ == "__main__": unittest.main() ================================================ FILE: tests/utils/test_rgb_ind_convertor.py ================================================ import unittest import numpy as np from dfp.utils.rgb_ind_convertor import floorplan_fuse_map, ind2rgb, rgb2ind class TestRgbIndConvertorCase(unittest.TestCase): def test_rgb2ind(self): inp = np.array([[[255, 60, 128], [192, 255, 255]]]) expected = [9, 2] out = rgb2ind(inp, floorplan_fuse_map) out = list(out[0]) self.assertListEqual(out, expected) def test_ind2rgb(self): inp = np.array([[9, 2]]) expected = np.array([[[255, 60, 128], [192, 255, 255]]]) out = ind2rgb(inp, floorplan_fuse_map) self.assertSequenceEqual(out.tolist(), expected.tolist()) if __name__ == "__main__": unittest.main() ================================================ FILE: tests/utils/test_util.py ================================================ from typing import List import pytest import numpy as np from dfp.utils.util import ( fast_hist, fill_break_line, flood_fill, refine_room_region, ) @pytest.mark.parametrize("shape", [[16, 16], [32, 32]]) class TestUtilCase: def test_fill_break_line(self, shape: List[int]): inp = np.ones(shape) out = fill_break_line(inp) assert out.shape == tuple(shape) def test_flood_fill(self, shape: List[int]): inp = np.ones(shape) inp = np.reshape(inp, (*shape, -1)) out = flood_fill(inp) assert out.shape == tuple(shape) def test_refine_room_region(self, shape: List[int]): inp = np.random.randint(10, size=shape) inp = np.reshape(inp, (*shape, -1)) out = refine_room_region(inp, inp) assert out.shape == (*shape, 1) def test_fast_hist(self, shape: List[int]): inp = np.random.randint(2, size=shape) out = fast_hist(inp, inp) assert out.shape == (9, 9)