diff --git a/export.py b/export.py index e3cf392b0101..20c1fbc5c7b8 100644 --- a/export.py +++ b/export.py @@ -560,13 +560,20 @@ def run( # Finish f = [str(x) for x in f if x] # filter out '' and None if any(f): + tp = type(model) + dir = Path('segment' if tp is SegmentationModel else 'classify' if tp is ClassificationModel else '') + predict = 'detect.py' if tp is DetectionModel else 'predict.py' h = '--half' if half else '' # --half FP16 inference arg LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)' f"\nResults saved to {colorstr('bold', file.parent.resolve())}" - f"\nDetect: python detect.py --weights {f[-1]} {h}" - f"\nValidate: python val.py --weights {f[-1]} {h}" + f"\nDetect: python {dir / predict} --weights {f[-1]} {h}" + f"\nValidate: python {dir / 'val.py'} --weights {f[-1]} {h}" f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')" f"\nVisualize: https://netron.app") + if tp is ClassificationModel: + LOGGER.warning("WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference") + if tp is SegmentationModel: + LOGGER.warning("WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference") return f # return list of exported files/dirs