ultralytics 8.0.143 add Model base class (#3934)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

ultralytics/models/sam/model.py

@@ -3,51 +3,38 @@
 SAM model interface
 """
 
-from ultralytics.cfg import get_cfg
+from ultralytics.engine.model import Model
 from ultralytics.utils.torch_utils import model_info
 
 from .build import build_sam
 from .predict import Predictor
 
 
-class SAM:
+class SAM(Model):
+    """
+    SAM model interface.
+    """
 
     def __init__(self, model='sam_b.pt') -> None:
         if model and not model.endswith('.pt') and not model.endswith('.pth'):
             # Should raise AssertionError instead?
             raise NotImplementedError('Segment anything prediction requires pre-trained checkpoint')
-        self.model = build_sam(model)
-        self.task = 'segment'  # required
-        self.predictor = None  # reuse predictor
+        super().__init__(model=model, task='segment')
+
+    def _load(self, weights: str, task=None):
+        self.model = build_sam(weights)
 
     def predict(self, source, stream=False, bboxes=None, points=None, labels=None, **kwargs):
         """Predicts and returns segmentation masks for given image or video source."""
         overrides = dict(conf=0.25, task='segment', mode='predict', imgsz=1024)
-        overrides.update(kwargs)  # prefer kwargs
-        if not self.predictor:
-            self.predictor = Predictor(overrides=overrides)
-            self.predictor.setup_model(model=self.model)
-        else:  # only update args if predictor is already setup
-            self.predictor.args = get_cfg(self.predictor.args, overrides)
-        return self.predictor(source, stream=stream, bboxes=bboxes, points=points, labels=labels)
-
-    def train(self, **kwargs):
-        """Function trains models but raises an error as SAM models do not support training."""
-        raise NotImplementedError("SAM models don't support training")
-
-    def val(self, **kwargs):
-        """Run validation given dataset."""
-        raise NotImplementedError("SAM models don't support validation")
+        kwargs.update(overrides)
+        prompts = dict(bboxes=bboxes, points=points, labels=labels)
+        super().predict(source, stream, prompts=prompts, **kwargs)
 
     def __call__(self, source=None, stream=False, bboxes=None, points=None, labels=None, **kwargs):
         """Calls the 'predict' function with given arguments to perform object detection."""
         return self.predict(source, stream, bboxes, points, labels, **kwargs)
 
-    def __getattr__(self, attr):
-        """Raises error if object has no requested attribute."""
-        name = self.__class__.__name__
-        raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.__doc__}")
-
     def info(self, detailed=False, verbose=True):
         """
         Logs model info.
@@ -57,3 +44,7 @@ class SAM:
             verbose (bool): Controls verbosity.
         """
         return model_info(self.model, detailed=detailed, verbose=verbose)
+
+    @property
+    def task_map(self):
+        return {'segment': {'predictor': Predictor}}

ultralytics/models/sam/predict.py

@@ -28,6 +28,8 @@ class Predictor(BasePredictor):
         # Args for set_image
         self.im = None
         self.features = None
+        # Args for set_prompts
+        self.prompts = {}
         # Args for segment everything
         self.segment_all = False
 
@@ -92,6 +94,10 @@ class Predictor(BasePredictor):
                 of masks and H=W=256. These low resolution logits can be passed to
                 a subsequent iteration as mask input.
         """
+        # Get prompts from self.prompts first
+        bboxes = self.prompts.pop('bboxes', bboxes)
+        points = self.prompts.pop('points', points)
+        masks = self.prompts.pop('masks', masks)
         if all(i is None for i in [bboxes, points, masks]):
             return self.generate(im, *args, **kwargs)
         return self.prompt_inference(im, bboxes, points, labels, masks, multimask_output)
@@ -348,6 +354,10 @@ class Predictor(BasePredictor):
             self.im = im
             break
 
+    def set_prompts(self, prompts):
+        """Set prompts in advance."""
+        self.prompts = prompts
+
     def reset_image(self):
         self.im = None
         self.features = None