@@ -309,11 +309,16 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
309309 # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
310310 try :
311311 cmd = 'edgetpu_compiler --version'
312- help = 'See https://coral.ai/docs/edgetpu/compiler/'
313- assert platform .system () == 'Linux' , f'export only supported on Linux. { help } '
314- assert subprocess .run (cmd , shell = True ).returncode == 0 , f'export requires edgetpu-compiler. { help } '
315- out = subprocess .run (cmd , shell = True , capture_output = True , check = True )
316- ver = out .stdout .decode ().split ()[- 1 ]
312+ help_url = 'https://coral.ai/docs/edgetpu/compiler/'
313+ assert platform .system () == 'Linux' , f'export only supported on Linux. See { help_url } '
314+ if subprocess .run (cmd , shell = True ).returncode != 0 :
315+ LOGGER .info (f'\n { prefix } export requires Edge TPU compiler. Attempting install from { help_url } ' )
316+ for c in ['curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -' ,
317+ 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list' ,
318+ 'sudo apt-get update' ,
319+ 'sudo apt-get install edgetpu-compiler' ]:
320+ subprocess .run (c , shell = True , check = True )
321+ ver = subprocess .run (cmd , shell = True , capture_output = True , check = True ).stdout .decode ().split ()[- 1 ]
317322
318323 LOGGER .info (f'\n { prefix } starting export with Edge TPU compiler { ver } ...' )
319324 f = str (file ).replace ('.pt' , '-int8_edgetpu.tflite' ) # Edge TPU model
0 commit comments