144 lines
5.6 KiB
Python
144 lines
5.6 KiB
Python
# 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")
|
|
|