ultralytics 8.0.81 single-line docstring updates (#2061)

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

ultralytics/yolo/v8/segment/predict.py

@@ -41,6 +41,7 @@ class SegmentationPredictor(DetectionPredictor):
 
 
 def predict(cfg=DEFAULT_CFG, use_python=False):
+    """Runs YOLO object detection on an image or video source."""
     model = cfg.model or 'yolov8n-seg.pt'
     source = cfg.source if cfg.source is not None else ROOT / 'assets' if (ROOT / 'assets').exists() \
         else 'https://ultralytics.com/images/bus.jpg'

ultralytics/yolo/v8/segment/train.py

@@ -18,12 +18,14 @@ from ultralytics.yolo.v8.detect.train import Loss
 class SegmentationTrainer(v8.detect.DetectionTrainer):
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
+        """Initialize a SegmentationTrainer object with given arguments."""
         if overrides is None:
             overrides = {}
         overrides['task'] = 'segment'
         super().__init__(cfg, overrides, _callbacks)
 
     def get_model(self, cfg=None, weights=None, verbose=True):
+        """Return SegmentationModel initialized with specified config and weights."""
         model = SegmentationModel(cfg, ch=3, nc=self.data['nc'], verbose=verbose and RANK == -1)
         if weights:
             model.load(weights)
@@ -31,15 +33,18 @@ class SegmentationTrainer(v8.detect.DetectionTrainer):
         return model
 
     def get_validator(self):
+        """Return an instance of SegmentationValidator for validation of YOLO model."""
         self.loss_names = 'box_loss', 'seg_loss', 'cls_loss', 'dfl_loss'
         return v8.segment.SegmentationValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args))
 
     def criterion(self, preds, batch):
+        """Returns the computed loss using the SegLoss class on the given predictions and batch."""
         if not hasattr(self, 'compute_loss'):
             self.compute_loss = SegLoss(de_parallel(self.model), overlap=self.args.overlap_mask)
         return self.compute_loss(preds, batch)
 
     def plot_training_samples(self, batch, ni):
+        """Creates a plot of training sample images with labels and box coordinates."""
         images = batch['img']
         masks = batch['masks']
         cls = batch['cls'].squeeze(-1)
@@ -49,6 +54,7 @@ class SegmentationTrainer(v8.detect.DetectionTrainer):
         plot_images(images, batch_idx, cls, bboxes, masks, paths=paths, fname=self.save_dir / f'train_batch{ni}.jpg')
 
     def plot_metrics(self):
+        """Plots training/val metrics."""
         plot_results(file=self.csv, segment=True)  # save results.png
 
 
@@ -61,6 +67,7 @@ class SegLoss(Loss):
         self.overlap = overlap
 
     def __call__(self, preds, batch):
+        """Calculate and return the loss for the YOLO model."""
         loss = torch.zeros(4, device=self.device)  # box, cls, dfl
         feats, pred_masks, proto = preds if len(preds) == 3 else preds[1]
         batch_size, _, mask_h, mask_w = proto.shape  # batch size, number of masks, mask height, mask width
@@ -147,6 +154,7 @@ class SegLoss(Loss):
 
 
 def train(cfg=DEFAULT_CFG, use_python=False):
+    """Train a YOLO segmentation model based on passed arguments."""
     model = cfg.model or 'yolov8n-seg.pt'
     data = cfg.data or 'coco128-seg.yaml'  # or yolo.ClassificationDataset("mnist")
     device = cfg.device if cfg.device is not None else ''

ultralytics/yolo/v8/segment/val.py

@@ -17,16 +17,19 @@ from ultralytics.yolo.v8.detect import DetectionValidator
 class SegmentationValidator(DetectionValidator):
 
     def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
+        """Initialize SegmentationValidator and set task to 'segment', metrics to SegmentMetrics."""
         super().__init__(dataloader, save_dir, pbar, args, _callbacks)
         self.args.task = 'segment'
         self.metrics = SegmentMetrics(save_dir=self.save_dir)
 
     def preprocess(self, batch):
+        """Preprocesses batch by converting masks to float and sending to device."""
         batch = super().preprocess(batch)
         batch['masks'] = batch['masks'].to(self.device).float()
         return batch
 
     def init_metrics(self, model):
+        """Initialize metrics and select mask processing function based on save_json flag."""
         super().init_metrics(model)
         self.plot_masks = []
         if self.args.save_json:
@@ -36,10 +39,12 @@ class SegmentationValidator(DetectionValidator):
             self.process = ops.process_mask  # faster
 
     def get_desc(self):
+        """Return a formatted description of evaluation metrics."""
         return ('%22s' + '%11s' * 10) % ('Class', 'Images', 'Instances', 'Box(P', 'R', 'mAP50', 'mAP50-95)', 'Mask(P',
                                          'R', 'mAP50', 'mAP50-95)')
 
     def postprocess(self, preds):
+        """Postprocesses YOLO predictions and returns output detections with proto."""
         p = ops.non_max_suppression(preds[0],
                                     self.args.conf,
                                     self.args.iou,
@@ -119,6 +124,7 @@ class SegmentationValidator(DetectionValidator):
             #    save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt')
 
     def finalize_metrics(self, *args, **kwargs):
+        """Sets speed and confusion matrix for evaluation metrics."""
         self.metrics.speed = self.speed
         self.metrics.confusion_matrix = self.confusion_matrix
 
@@ -160,6 +166,7 @@ class SegmentationValidator(DetectionValidator):
         return torch.tensor(correct, dtype=torch.bool, device=detections.device)
 
     def plot_val_samples(self, batch, ni):
+        """Plots validation samples with bounding box labels."""
         plot_images(batch['img'],
                     batch['batch_idx'],
                     batch['cls'].squeeze(-1),
@@ -170,6 +177,7 @@ class SegmentationValidator(DetectionValidator):
                     names=self.names)
 
     def plot_predictions(self, batch, preds, ni):
+        """Plots batch predictions with masks and bounding boxes."""
         plot_images(batch['img'],
                     *output_to_target(preds[0], max_det=15),
                     torch.cat(self.plot_masks, dim=0) if len(self.plot_masks) else self.plot_masks,
@@ -184,6 +192,7 @@ class SegmentationValidator(DetectionValidator):
         from pycocotools.mask import encode  # noqa
 
         def single_encode(x):
+            """Encode predicted masks as RLE and append results to jdict."""
             rle = encode(np.asarray(x[:, :, None], order='F', dtype='uint8'))[0]
             rle['counts'] = rle['counts'].decode('utf-8')
             return rle
@@ -204,6 +213,7 @@ class SegmentationValidator(DetectionValidator):
                 'segmentation': rles[i]})
 
     def eval_json(self, stats):
+        """Return COCO-style object detection evaluation metrics."""
         if self.args.save_json and self.is_coco and len(self.jdict):
             anno_json = self.data['path'] / 'annotations/instances_val2017.json'  # annotations
             pred_json = self.save_dir / 'predictions.json'  # predictions
@@ -232,6 +242,7 @@ class SegmentationValidator(DetectionValidator):
 
 
 def val(cfg=DEFAULT_CFG, use_python=False):
+    """Validate trained YOLO model on validation data."""
     model = cfg.model or 'yolov8n-seg.pt'
     data = cfg.data or 'coco128-seg.yaml'