ibaigorordo / onnx-topformer-semantic-segmentation

Python scripts performing semantic segmentation using the TopFormer model in ONNX.

ONNX TopFormer Semantic Segmentation

Python scripts performing semantic segmentation using the TopFormer model in ONNX.

Original image: https://en.wikipedia.org/wiki/File:Beatles_-_Abbey_Road.jpg

Requirements

• Check the requirements.txt file.
• For ONNX, if you have a NVIDIA GPU, then install the onnxruntime-gpu, otherwise use the onnxruntime library.
• Additionally, pafy and youtube-dl are required for youtube video inference.

Installation

git clone https://github.com/ibaiGorordo/ONNX-TopFormer-Semantic-Segmentation.git
cd ONNX-TopFormer-Semantic-Segmentation
pip install -r requirements.txt

ONNX Runtime

For Nvidia GPU computers: pip install onnxruntime-gpu

Otherwise: pip install onnxruntime

For youtube video inference

pip install youtube_dl
pip install git+https://github.com/zizo-pro/pafy@b8976f22c19e4ab5515cacbfae0a3970370c102b

ONNX model

The model was converted from the Pytorch implementation using the code in the original repository. Download the converted ONNX model from the drive file and save them into the models folder.

• The License of the models is Apache-2.0 License: https://github.com/hustvl/TopFormer/blob/main/LICENSE

Pytorch model

The original Pytorch model can be found in this repository: https://github.com/hustvl/TopFormer

Examples

• Image inference:

python image_semantic_segmentation.py

• Webcam inference:

python webcam_semantic_segmentation.py

• Video inference: https://youtu.be/JkOSbtKfIFo

python video_semantic_segmentation.py

Original video: https://youtu.be/yWHdkK5j4yk

References:

• TopFormer model: https://github.com/hustvl/TopFormer
• PINTO0309's model zoo: https://github.com/PINTO0309/PINTO_model_zoo
• Original paper: https://arxiv.org/abs/2204.05525

onnx-topformer-semantic-segmentation's Issues

Can you explain how the segmentation is added to the image?

Your code is amazingly simple and good readable. Thank you.

I am trying to see how to get my model working in ONNX (WongKinYiu/yolov7#797 (comment)). I am working on Yolo segmentation and running through a bunch of examples.

Your output is float32[1,150,64,64], how is that converted to the masks you put on the image?

`
def process_output(self, outputs):

	return np.squeeze(np.argmax(outputs[0], axis=1))

def draw_segmentation(self, image, alpha = 0.5):

	return util_draw_seg(self.seg_map, image, alpha)

`