Document Scanner: Semantic Segmentation

Setup

• To create environment:

conda create -n env python=3.8
conda activate env
pip install -r requirements.txt
pip install nvidia-pyindex
pip install onnx-graphsurgeon

Create dataset

• Download raw scroped dataset and background

• To gen full mask

python3 main.py --scenario generate_doc_set

• To resize

python3 main.py --scenario resizer \
-s DOCUMENTS/CHOSEN/images \
-d DOCUMENTS/CHOSEN/resized_images 
-x 640

python3 main.py --scenraio resizer \
-s DOCUMENTS/CHOSEN/masks \
-d DOCUMENTS/CHOSEN/resized_masks \
-x 640

• To gen data (config factor to change the ratio of image size and background size)

python3 main.py --scenario create_dataset

Train

• To train:

python3 main.py --scenario train \
--backbone_model "mbv3" \
--data_dir "document_dataset_resized" \
--batch_size 32 \
--num_epoch 100

Inference

• To infer:

python3 main.py --scenario infer \
--data_infer "./test_images/IMG_20220721_162757.jpg" \
--image_size 384 \
--checkpoint_path "model_repository/mbv3_averaged.pth"

Export model

• To export model to onnx

python3 export_onnx.py

• To export onnx to tflite

onnx2tf -i path/to/model.onnx -o saved_model

Simple Demo

• To run streamlit app:

streamlit run app.py --server.port 8080