first commit

ultralytics/models/yolo/demask/predict.py

@@ -0,0 +1,143 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+from ultralytics.engine.predictor import BasePredictor
+from ultralytics.engine.results import Results
+from ultralytics.utils import ops, LOGGER
+from ultralytics.utils.torch_utils import select_device, smart_inference_mode
+
+from pathlib import Path
+import numpy as np
+import torch
+import cv2
+
+
+class DetectionPredictor(BasePredictor):
+    """
+    A class extending the BasePredictor class for prediction based on a detection model.
+
+    Example:
+        ```python
+        from ultralytics.utils import ASSETS
+        from ultralytics.models.yolo.detect import DetectionPredictor
+
+        args = dict(model='yolov8n.pt', source=ASSETS)
+        predictor = DetectionPredictor(overrides=args)
+        predictor.predict_cli()
+        ```
+    """
+
+    def preprocess(self, im):
+        """
+        Prepares input image before inference.
+
+        Args:
+            im (torch.Tensor | List(np.ndarray)): BCHW for tensor, [(HWC) x B] for list.
+        """
+        not_tensor = not isinstance(im, torch.Tensor)
+        if not_tensor:
+            im = np.stack(self.pre_transform(im))
+            im_rgb = im.copy()
+            # for i in range(im.shape[3] // 3):
+            #     start_idx = i * 3
+            #     channel_idx = list(range(start_idx, start_idx + 3))
+            #     im_rgb[..., channel_idx] = im[..., channel_idx][:, :, ::-1]  # BGR to RGB
+            im = im_rgb.transpose((0, 3, 1, 2))  # BHWC to BCHW
+            im = np.ascontiguousarray(im)  # contiguous
+            im = torch.from_numpy(im)
+
+        im = im.to(self.device)
+        im = im.half() if self.model.fp16 else im.float()  # uint8 to fp16/32
+        if not_tensor:
+            im /= 255  # 0 - 255 to 0.0 - 1.0
+        return im
+
+    def postprocess(self, preds, img, orig_imgs):
+        """Post-processes predictions and returns a list of Results objects."""
+        pred_mask = preds[0] if isinstance(preds, (tuple, list)) else preds
+        return [pred_mask.squeeze().detach().cpu().numpy()]
+        # return (pred_mask.squeeze().detach().cpu().numpy() * 255).astype(np.uint8)
+    
+    @smart_inference_mode()
+    def stream_inference(self, source=None, model=None, *args, **kwargs):
+        """Streams real-time inference on camera feed and saves results to file."""
+        if self.args.verbose:
+            LOGGER.info("")
+
+        # Setup model
+        if not self.model:
+            self.setup_model(model)
+
+        with self._lock:  # for thread-safe inference
+            # Setup source every time predict is called
+            self.setup_source(source if source is not None else self.args.source)
+
+            # Check if save_dir/ label file exists
+            if self.args.save or self.args.save_txt:
+                (self.save_dir / "labels" if self.args.save_txt else self.save_dir).mkdir(parents=True, exist_ok=True)
+
+            # Warmup model
+            if not self.done_warmup:
+                self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else self.dataset.bs, 6, *self.imgsz))
+                self.done_warmup = True
+
+            self.seen, self.windows, self.batch = 0, [], None
+            profilers = (
+                ops.Profile(device=self.device),
+                ops.Profile(device=self.device),
+                ops.Profile(device=self.device),
+            )
+            self.run_callbacks("on_predict_start")
+            for self.batch in self.dataset:
+                self.run_callbacks("on_predict_batch_start")
+                paths, im0s, s = self.batch
+
+                # Preprocess
+                with profilers[0]:
+                    im = self.preprocess(im0s)
+
+                # Inference
+                with profilers[1]:
+                    preds = self.inference(im, *args, **kwargs)
+                    if self.args.embed:
+                        yield from [preds] if isinstance(preds, torch.Tensor) else preds  # yield embedding tensors
+                        continue
+
+                # Postprocess
+                with profilers[2]:
+                    self.results = self.postprocess(preds, im, im0s)
+                self.run_callbacks("on_predict_postprocess_end")
+
+                # Visualize, save, write results
+                n = len(im0s)
+                for i in range(n):
+                    self.seen += 1
+                    # self.results[i].speed = {
+                    #     "preprocess": profilers[0].dt * 1e3 / n,
+                    #     "inference": profilers[1].dt * 1e3 / n,
+                    #     "postprocess": profilers[2].dt * 1e3 / n,
+                    # }
+                    if self.args.verbose or self.args.save or self.args.save_txt or self.args.show:
+                        mask = cv2.resize(self.results[0], (im0s[i].shape[:2][::-1]), interpolation=cv2.INTER_LINEAR)
+                        mask = (mask * 255).astype(np.uint8)
+                        mask = cv2.flip(mask, 1) #左右翻转
+                        # cv2.imwrite(str(self.save_dir / Path(paths[i]).with_suffix('.png').name), mask)
+                        cv2.imwrite(str(self.save_dir / Path(paths[i]).name), mask)
+
+                self.run_callbacks("on_predict_batch_end")
+                yield from self.results
+
+        # Release assets
+        for v in self.vid_writer.values():
+            if isinstance(v, cv2.VideoWriter):
+                v.release()
+
+        # Print final results
+        if self.args.verbose and self.seen:
+            t = tuple(x.t / self.seen * 1e3 for x in profilers)  # speeds per image
+            LOGGER.info(
+                f"Speed: %.1fms preprocess, %.1fms inference, %.1fms postprocess per image at shape "
+                f"{(min(self.args.batch, self.seen), 3, *im.shape[2:])}" % t
+            )
+
+        self.run_callbacks("on_predict_end")
+