Predictor support (#65)

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

ultralytics/yolo/v8/detect/__init__.py

@@ -1,2 +1,3 @@
+from ultralytics.yolo.v8.detect.predict import DetectionPredictor, predict
 from ultralytics.yolo.v8.detect.train import DetectionTrainer, train
 from ultralytics.yolo.v8.detect.val import DetectionValidator, val

ultralytics/yolo/v8/detect/predict.py

@@ -0,0 +1,97 @@
+import hydra
+import torch
+
+from ultralytics.yolo.engine.predictor import BasePredictor
+from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG
+from ultralytics.yolo.utils import ops
+from ultralytics.yolo.utils.plotting import Annotator, colors, save_one_box
+
+
+class DetectionPredictor(BasePredictor):
+
+    def get_annotator(self, img):
+        return Annotator(img, line_width=self.args.line_thickness, example=str(self.model.names))
+
+    def preprocess(self, img):
+        img = torch.from_numpy(img).to(self.model.device)
+        img = img.half() if self.model.fp16 else img.float()  # uint8 to fp16/32
+        img /= 255  # 0 - 255 to 0.0 - 1.0
+        return img
+
+    def postprocess(self, preds, img, orig_img):
+        preds = ops.non_max_suppression(preds,
+                                        self.args.conf_thres,
+                                        self.args.iou_thres,
+                                        agnostic=self.args.agnostic_nms,
+                                        max_det=self.args.max_det)
+
+        for i, pred in enumerate(preds):
+            shape = orig_img[i].shape if self.webcam else orig_img.shape
+            pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round()
+
+        return preds
+
+    def write_results(self, idx, preds, batch):
+        p, im, im0 = batch
+        log_string = ""
+        if len(im.shape) == 3:
+            im = im[None]  # expand for batch dim
+        self.seen += 1
+        im0 = im0.copy()
+        if self.webcam:  # batch_size >= 1
+            log_string += f'{idx}: '
+            frame = self.dataset.count
+        else:
+            frame = getattr(self.dataset, 'frame', 0)
+
+        self.data_path = p
+        # save_path = str(self.save_dir / p.name)  # im.jpg
+        self.txt_path = str(self.save_dir / 'labels' / p.stem) + ('' if self.dataset.mode == 'image' else f'_{frame}')
+        log_string += '%gx%g ' % im.shape[2:]  # print string
+        self.annotator = self.get_annotator(im0)
+
+        det = preds[idx]
+        if len(det) == 0:
+            return log_string
+        for c in det[:, 5].unique():
+            n = (det[:, 5] == c).sum()  # detections per class
+            log_string += f"{n} {self.model.names[int(c)]}{'s' * (n > 1)}, "
+
+        # write
+        gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
+        for *xyxy, conf, cls in reversed(det):
+            if self.args.save_txt:  # Write to file
+                xywh = (ops.xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
+                line = (cls, *xywh, conf) if self.args.save_conf else (cls, *xywh)  # label format
+                with open(f'{self.txt_path}.txt', 'a') as f:
+                    f.write(('%g ' * len(line)).rstrip() % line + '\n')
+
+            if self.save_img or self.args.save_crop or self.args.view_img:  # Add bbox to image
+                c = int(cls)  # integer class
+                label = None if self.args.hide_labels else (
+                    self.model.names[c] if self.args.hide_conf else f'{self.model.names[c]} {conf:.2f}')
+                self.annotator.box_label(xyxy, label, color=colors(c, True))
+            if self.args.save_crop:
+                imc = im0.copy()
+                save_one_box(xyxy,
+                             imc,
+                             file=self.save_dir / 'crops' / self.model.model.names[c] / f'{self.data_path.stem}.jpg',
+                             BGR=True)
+
+        return log_string
+
+
+@hydra.main(version_base=None, config_path=DEFAULT_CONFIG.parent, config_name=DEFAULT_CONFIG.name)
+def predict(cfg):
+    cfg.model = cfg.model or "n.pt"
+    sz = cfg.img_size
+    if type(sz) != int:  # recieved listConfig
+        cfg.img_size = [sz[0], sz[0]] if len(cfg.img_size) == 1 else [sz[0], sz[1]]  # expand
+    else:
+        cfg.img_size = [sz, sz]
+    predictor = DetectionPredictor(cfg)
+    predictor()
+
+
+if __name__ == "__main__":
+    predict()

ultralytics/yolo/v8/detect/val.py

@@ -63,7 +63,7 @@ class DetectionValidator(BaseValidator):
         self.seen = 0
         self.confusion_matrix = ConfusionMatrix(nc=self.nc)
         self.metrics = Metric()
-        self.loss = torch.zeros(4, device=self.device)
+        self.loss = torch.zeros(3, device=self.device)
         self.jdict = []
         self.stats = []