# 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")