2

Is it possible to convert the YOLOv5 PyTorch model to the Tensorflow.js model?

I am developing an object detection web app. so I have trained the data using Yolov5, but now I am looking for the correct method to convert that model to Tf.js.

desertnaut
  • 57,590
  • 26
  • 140
  • 166
Irugal
  • 31
  • 1
  • 2

2 Answers2

3

I think this is what you are looking for https://github.com/zldrobit/tfjs-yolov5-example

Inside the YoloV5 repo, run the export.py command.

python export.py --weights yolov5s.pt --include tfjs

Then cd into the above linked repo and copy the weights folder to the public:

cp ./yolov5s_web_model public/web_model

Don't forget, you'll have to change the names array in src/index.js to match your custom model.

But unfortunately, it seems painfully slow at about 1-2 seconds. I don't think I was able to get WebGL working.

ShermanL
  • 1,372
  • 1
  • 9
  • 8
  • You could get WebGL working by enabling WebGL in `chrome://settings` and `chrome://flags` as https://superuser.com/a/836833/1038955, if you're using Chrome browser. – robit Dec 01 '21 at 08:30
0

With few interim steps, but most of the times it works:

  1. Export PyTorch to ONNX
  2. Convert ONNX to TF Saved Model
  3. Convert TF Saved Model to TFJS Graph Model

When converting from ONNX to TF, you might need to adjust target version if you run into unsupported ops.

Also, make sure to set input resolutions to a fixed values, any dynamic inputs get messed up in this multi-step conversion.

Vladimir Mandic
  • 813
  • 5
  • 11