-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Open
Description
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()
Metadata
Metadata
Assignees
Labels
No labels