General version of Yolov8 export to ONNX at Colab

Replace the codes at
https://colab.research.google.com/drive/1-yZg6hFg27uCPSycRCRtyezHhq_VAHxQ?usp=sharing with the following. For more info. check https://github.com/ibaiGorordo/ONNX-YOLOv8-Object-Detection?ts=2

!pip install ultralytics
from ultralytics import YOLO
from math import ceil

def compute_wh_yolo(skale, w_ori, h_ori):
  w_32 = int(ceil(w_ori * skale / 32.0))
  w_scaled = int(w_32 * 32)
  h_32 = int(ceil(h_ori * skale / 32.0))
  h_scaled = int(h_32 * 32)
  return w_scaled, h_scaled 

task = 'segmentation' #@param ["classification", "detection", "pose_estimation", "segmentation"]
if 'classification' == task:
  s_task = '-cls'
elif 'pose_estimation' == task:
  s_task = '-pose'
elif 'segmentation' == task:
  s_task = '-seg'
else:
  s_task = ''

skale = 'nano' #@param ["nano", "small", "medium", "large", "xlarge"]
if 'nano' == skale:
  s_scale = 'n'
elif 'small' == skale:
  s_scale = 's'
elif 'medium' == skale:
  s_scale = 'm'
elif 'large' == skale:
  s_scale = 'l'
else:
  s_scale = 'x'

#model_name = 'yolov8n' #@param ["yolov8n", "yolov8s", "yolov8m", "yolov8l", "yolov8x"]
#model_name = 'yolov8n-seg' #@param ["yolov8n-seg", "yolov8s-seg", "yolov8m-seg", "yolov8l-seg", "yolov8x-seg"]
model_name = f'yolov8{s_scale}{s_task}'
opset = 10 #@param {type:"slider", min:8, max:18, step:1}
input_width = 1080 #@param {type:"slider", min:32, max:4096, step:32}
input_height = 1920 #@param {type:"slider", min:32, max:4096, step:32}
downsample_ratio = "0.25" #@param [1.0, 0.5, 0.25, 0.125]

input_width, input_height = compute_wh_yolo(float(downsample_ratio), input_width, input_height)
optimize_cpu = False
model = YOLO(f"{model_name}.pt") 
model.export(format="onnx", opset = opset, imgsz=[input_height,input_width], optimize=optimize_cpu)

%cd /content
from google.colab import files
fn_exported = f'{model_name}.onnx'
fn_download = f'{model_name}_opset_{opset}_{input_width}_{input_height}.onnx'
!cp '{fn_exported}' '{fn_download}'
files.download(fn_download)