probability icon indicating copy to clipboard operation
probability copied to clipboard

Help find dependencies, python 3.9 problem with conversion Converting ONNX -> TensorFlow

Open mrranger opened this issue 9 months ago • 0 comments

Please help me resolve dependency issues.

I'm using Python 3.9 and encountering problems during the conversion step Converting ONNX -> TensorFlow. No matter what I try, I keep running into conflicts or other errors. I cannot find a proper set of dependencies or at least a working requirements.txt.

Some sources say Python 3.9 is required, others say 3.8, and some recommend 3.10. Honestly, I’m exhausted trying to figure this out.

import os
import subprocess
import torch
import onnx
import onnxslim
import sys
import types
from pathlib import Path
import urllib.request
import zipfile
import shutil

def download_yolov5_repo(target_dir):
    zip_url = "https://github.com/ultralytics/yolov5/archive/refs/heads/master.zip"
    zip_path = os.path.join(target_dir, "yolov5.zip")
    print("šŸ“„ Downloading YOLOv5 repository...")
    urllib.request.urlretrieve(zip_url, zip_path)

    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(target_dir)
    os.remove(zip_path)

    extracted_path = os.path.join(target_dir, "yolov5-master")
    yolov5_path = os.path.join(target_dir, "yolov5")
    if os.path.exists(yolov5_path):
        shutil.rmtree(yolov5_path)
    shutil.move(extracted_path, yolov5_path)
    print("āœ… YOLOv5 repository downloaded and prepared.")

def convert_pt_to_onnx(pt_path, onnx_path):
    yolov5_dir = os.path.join(Path(pt_path).parent, "yolov5")
    yolov5_model_path = os.path.join(yolov5_dir, "models", "experimental.py")

    if not os.path.exists(yolov5_model_path):
        print("āš ļø YOLOv5 repo not found, downloading it...")
        download_yolov5_repo(Path(pt_path).parent)

    if yolov5_dir not in sys.path:
        sys.path.insert(0, yolov5_dir)

    try:
        from models.experimental import attempt_load
    except ImportError as e:
        print("āŒ ERROR: Could not import YOLOv5 modules. Ensure YOLOv5 repo is in the correct path.")
        raise e

    model = attempt_load(pt_path, device=torch.device('cpu'))
    model.eval()

    dummy_input = torch.randn(1, 3, 640, 640)

    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,
        opset_version=12,
        do_constant_folding=True,
        input_names=['images'],
        output_names=['output'],
        dynamic_axes={'images': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
    )

    onnxslim.slim(onnx_path, onnx_path)

def convert_onnx_to_tf(onnx_path, tf_path):
    subprocess.run(['onnx-tf', 'convert', '-i', onnx_path, '-o', tf_path], check=True)

def convert_tf_to_tfjs(tf_path, tfjs_path):
    try:
        print("āš™ļø  Importing TensorFlow.js converter inside Python...")
        sys.modules['tensorflow_decision_forests'] = types.SimpleNamespace()
        sys.modules['tensorflow_decision_forests.keras'] = types.SimpleNamespace()
        sys.modules['tensorflow_decision_forests.keras.core'] = types.SimpleNamespace()
        sys.modules['tensorflow_decision_forests.keras.core_inference'] = types.SimpleNamespace()
        sys.modules['tensorflow_decision_forests.tensorflow'] = types.SimpleNamespace()
        sys.modules['tensorflow_decision_forests.tensorflow.ops'] = types.SimpleNamespace()
        sys.modules['tensorflow_decision_forests.tensorflow.ops.inference'] = types.SimpleNamespace()
        sys.modules['tensorflow_decision_forests.tensorflow.ops.inference.api'] = types.SimpleNamespace()
        sys.modules['tensorflow_decision_forests.tensorflow.ops.inference.op'] = types.SimpleNamespace()
        sys.modules['tensorflow_decision_forests.tensorflow.ops.inference.op_dynamic'] = types.SimpleNamespace()

        import tensorflowjs as tfjs

        tfjs.converters.convert_tf_saved_model(
            saved_model_dir=tf_path,
            output_dir=tfjs_path,
            skip_op_check=True,
            strip_debug_ops=True
        )
        print(f"āœ… TensorFlow.js conversion complete! Files saved to: {tfjs_path}")
    except Exception as e:
        print("\nāŒ ERROR: TensorFlow.js conversion failed unexpectedly.")
        raise e

def main():
    folder_path = input("Enter the folder path to search for PyTorch (.pt) models: ")
    pt_files = [f for f in os.listdir(folder_path) if f.endswith(".pt")]

    if not pt_files:
        print("No PyTorch (.pt) models found.")
        return

    print("\nFound PyTorch models:")
    for i, f in enumerate(pt_files, start=1):
        print(f"{i}. {f}")

    idx = int(input("Enter the number of the model you want to convert: ")) - 1
    pt_path = os.path.join(folder_path, pt_files[idx])
    model_name = os.path.splitext(pt_files[idx])[0]

    onnx_path = os.path.join(folder_path, f"{model_name}.onnx")
    tf_path = os.path.join(folder_path, "model_tf")
    tfjs_path = os.path.join(folder_path, f"{model_name}WebModel")

    print("\nConverting PyTorch -> ONNX using YOLOv5...")
    convert_pt_to_onnx(pt_path, onnx_path)

    print("\nConverting ONNX -> TensorFlow...")
    convert_onnx_to_tf(onnx_path, tf_path)

    print("\nConverting TensorFlow -> TensorFlow.js...")
    convert_tf_to_tfjs(tf_path, tfjs_path)

if __name__ == '__main__':
    main()

mrranger avatar May 28 '25 11:05 mrranger