ultralytics 8.0.89 SAM predict and auto-annotate (#2298)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>
Co-authored-by: Paula Derrenger <107626595+pderrenger@users.noreply.github.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
Co-authored-by: Snyk bot <snyk-bot@snyk.io>
Co-authored-by: Laughing-q <1185102784@qq.com>

ultralytics/yolo/v8/classify/predict.py

@@ -9,8 +9,14 @@ from ultralytics.yolo.utils import DEFAULT_CFG, ROOT
 
 class ClassificationPredictor(BasePredictor):
 
+    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
+        super().__init__(cfg, overrides, _callbacks)
+        self.args.task = 'classify'
+
     def preprocess(self, img):
         """Converts input image to model-compatible data type."""
+        if not isinstance(img, torch.Tensor):
+            img = torch.stack([self.transforms(im) for im in img], dim=0)
         img = (img if isinstance(img, torch.Tensor) else torch.from_numpy(img)).to(self.model.device)
         return img.half() if self.model.fp16 else img.float()  # uint8 to fp16/32
 
@@ -19,7 +25,7 @@ class ClassificationPredictor(BasePredictor):
         results = []
         for i, pred in enumerate(preds):
             orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
-            path, _, _, _, _ = self.batch
+            path = self.batch[0]
             img_path = path[i] if isinstance(path, list) else path
             results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, probs=pred))
 

ultralytics/yolo/v8/classify/train.py

@@ -5,10 +5,11 @@ import torchvision
 
 from ultralytics.nn.tasks import ClassificationModel, attempt_load_one_weight
 from ultralytics.yolo import v8
-from ultralytics.yolo.data import build_classification_dataloader
+from ultralytics.yolo.data import ClassificationDataset, build_dataloader
 from ultralytics.yolo.engine.trainer import BaseTrainer
 from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, RANK, colorstr
-from ultralytics.yolo.utils.torch_utils import is_parallel, strip_optimizer
+from ultralytics.yolo.utils.plotting import plot_images, plot_results
+from ultralytics.yolo.utils.torch_utils import is_parallel, strip_optimizer, torch_distributed_zero_first
 
 
 class ClassificationTrainer(BaseTrainer):
@@ -71,14 +72,16 @@ class ClassificationTrainer(BaseTrainer):
 
         return  # dont return ckpt. Classification doesn't support resume
 
+    def build_dataset(self, img_path, mode='train'):
+        dataset = ClassificationDataset(root=img_path, imgsz=self.args.imgsz, augment=mode == 'train')
+        return dataset
+
     def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode='train'):
         """Returns PyTorch DataLoader with transforms to preprocess images for inference."""
-        loader = build_classification_dataloader(path=dataset_path,
-                                                 imgsz=self.args.imgsz,
-                                                 batch_size=batch_size if mode == 'train' else (batch_size * 2),
-                                                 augment=mode == 'train',
-                                                 rank=rank,
-                                                 workers=self.args.workers)
+        with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
+            dataset = self.build_dataset(dataset_path, mode)
+
+        loader = build_dataloader(dataset, batch_size, self.args.workers, rank=rank)
         # Attach inference transforms
         if mode != 'train':
             if is_parallel(self.model):
@@ -124,6 +127,10 @@ class ClassificationTrainer(BaseTrainer):
         """Resumes training from a given checkpoint."""
         pass
 
+    def plot_metrics(self):
+        """Plots metrics from a CSV file."""
+        plot_results(file=self.csv, classify=True)  # save results.png
+
     def final_eval(self):
         """Evaluate trained model and save validation results."""
         for f in self.last, self.best:
@@ -138,6 +145,13 @@ class ClassificationTrainer(BaseTrainer):
                 #     self.run_callbacks('on_fit_epoch_end')
         LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}")
 
+    def plot_training_samples(self, batch, ni):
+        """Plots training samples with their annotations."""
+        plot_images(images=batch['img'],
+                    batch_idx=torch.arange(len(batch['img'])),
+                    cls=batch['cls'].squeeze(-1),
+                    fname=self.save_dir / f'train_batch{ni}.jpg')
+
 
 def train(cfg=DEFAULT_CFG, use_python=False):
     """Train the YOLO classification model."""

ultralytics/yolo/v8/classify/val.py

@@ -1,9 +1,12 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-from ultralytics.yolo.data import build_classification_dataloader
+import torch
+
+from ultralytics.yolo.data import ClassificationDataset, build_dataloader
 from ultralytics.yolo.engine.validator import BaseValidator
 from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER
 from ultralytics.yolo.utils.metrics import ClassifyMetrics, ConfusionMatrix
+from ultralytics.yolo.utils.plotting import plot_images
 
 
 class ClassificationValidator(BaseValidator):
@@ -52,20 +55,36 @@ class ClassificationValidator(BaseValidator):
         self.metrics.process(self.targets, self.pred)
         return self.metrics.results_dict
 
+    def build_dataset(self, img_path):
+        dataset = ClassificationDataset(root=img_path, imgsz=self.args.imgsz, augment=False)
+        return dataset
+
     def get_dataloader(self, dataset_path, batch_size):
         """Builds and returns a data loader for classification tasks with given parameters."""
-        return build_classification_dataloader(path=dataset_path,
-                                               imgsz=self.args.imgsz,
-                                               batch_size=batch_size,
-                                               augment=False,
-                                               shuffle=False,
-                                               workers=self.args.workers)
+        dataset = self.build_dataset(dataset_path)
+        return build_dataloader(dataset, batch_size, self.args.workers, rank=-1)
 
     def print_results(self):
         """Prints evaluation metrics for YOLO object detection model."""
         pf = '%22s' + '%11.3g' * len(self.metrics.keys)  # print format
         LOGGER.info(pf % ('all', self.metrics.top1, self.metrics.top5))
 
+    def plot_val_samples(self, batch, ni):
+        """Plot validation image samples."""
+        plot_images(images=batch['img'],
+                    batch_idx=torch.arange(len(batch['img'])),
+                    cls=batch['cls'].squeeze(-1),
+                    fname=self.save_dir / f'val_batch{ni}_labels.jpg',
+                    names=self.names)
+
+    def plot_predictions(self, batch, preds, ni):
+        """Plots predicted bounding boxes on input images and saves the result."""
+        plot_images(batch['img'],
+                    batch_idx=torch.arange(len(batch['img'])),
+                    cls=torch.argmax(preds, dim=1),
+                    fname=self.save_dir / f'val_batch{ni}_pred.jpg',
+                    names=self.names)  # pred
+
 
 def val(cfg=DEFAULT_CFG, use_python=False):
     """Validate YOLO model using custom data."""