pull/17463/head^2
UltralyticsAssistant 4 months ago
parent 85b236c651
commit 9b066bacb2
  1. 6
      ultralytics/engine/exporter.py

@ -1111,8 +1111,8 @@ class Exporter:
raise ValueError("IMX500 export is only supported for YOLOv8 detection models")
check_requirements(("model-compression-toolkit==2.1.1", "sony-custom-layers[torch]"))
import model_compression_toolkit as mct
from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms
import onnx
from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms
def representative_dataset_gen(dataloader=self.get_int8_calibration_dataloader(prefix)):
for batch in dataloader:
@ -1194,7 +1194,7 @@ class Exporter:
iou_threshold=self.args.iou,
max_detections=self.args.max_det,
).to(self.device)
f = Path(str(self.file).replace(self.file.suffix, "_imx500_model.onnx")) # js dir
mct.exporter.pytorch_export_model(model=quant_model, save_model_path=f, repr_dataset=representative_dataset_gen)
@ -1222,7 +1222,7 @@ class Exporter:
output_dir = Path(str(self.file).replace(self.file.suffix, "_imx500_model"))
subprocess.run(["imxconv-pt", "-i", str(f), "-o", str(output_dir)], check=True)
with open(output_dir / "labels.txt", "w") as file:
file.writelines([f"{name}\n" for _, name in self.model.names.items()])

Loading…
Cancel
Save