Add SegmentationModel unsupported warning (#9632)

* Add SegmentationModel unsupported warning Signed-off-by: Glenn Jocher <[email protected]> * Update hubconf.py Signed-off-by: Glenn Jocher <[email protected]> * Update hubconf.py Signed-off-by: Glenn Jocher <[email protected]> Signed-off-by: Glenn Jocher <[email protected]>
ultralytics · Sep 28, 2022 · 799e3d0 · 799e3d0
1 parent 2373d54
commit 799e3d0
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/hubconf.py b/hubconf.py
@@ -30,7 +30,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
 
     from models.common import AutoShape, DetectMultiBackend
     from models.experimental import attempt_load
-    from models.yolo import ClassificationModel, DetectionModel
+    from models.yolo import ClassificationModel, DetectionModel, SegmentationModel
     from utils.downloads import attempt_download
     from utils.general import LOGGER, check_requirements, intersect_dicts, logging
     from utils.torch_utils import select_device
@@ -47,8 +47,11 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
                 model = DetectMultiBackend(path, device=device, fuse=autoshape)  # detection model
                 if autoshape:
                     if model.pt and isinstance(model.model, ClassificationModel):
-                        LOGGER.warning('WARNING ⚠️ YOLOv5 v6.2 ClassificationModel is not yet AutoShape compatible. '
+                        LOGGER.warning('WARNING ⚠️ YOLOv5 ClassificationModel is not yet AutoShape compatible. '
                                        'You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224).')
+                    elif model.pt and isinstance(model.model, SegmentationModel):
+                        LOGGER.warning('WARNING ⚠️ YOLOv5 SegmentationModel is not yet AutoShape compatible. '
+                                       'You will not be able to run inference with this model.')
                     else:
                         model = AutoShape(model)  # for file/URI/PIL/cv2/np inputs and NMS
             except Exception: