|
|
import detectron2.data.transforms as T |
|
|
import torch |
|
|
from detectron2.checkpoint import DetectionCheckpointer |
|
|
from detectron2.config import CfgNode, instantiate |
|
|
from detectron2.data import MetadataCatalog |
|
|
from omegaconf import OmegaConf |
|
|
|
|
|
|
|
|
class DefaultPredictor_Lazy: |
|
|
"""Create a simple end-to-end predictor with the given config that runs on single device for a |
|
|
single input image. |
|
|
|
|
|
Compared to using the model directly, this class does the following additions: |
|
|
|
|
|
1. Load checkpoint from the weights specified in config (cfg.MODEL.WEIGHTS). |
|
|
2. Always take BGR image as the input and apply format conversion internally. |
|
|
3. Apply resizing defined by the config (`cfg.INPUT.{MIN,MAX}_SIZE_TEST`). |
|
|
4. Take one input image and produce a single output, instead of a batch. |
|
|
|
|
|
This is meant for simple demo purposes, so it does the above steps automatically. |
|
|
This is not meant for benchmarks or running complicated inference logic. |
|
|
If you'd like to do anything more complicated, please refer to its source code as |
|
|
examples to build and use the model manually. |
|
|
|
|
|
Attributes: |
|
|
metadata (Metadata): the metadata of the underlying dataset, obtained from |
|
|
test dataset name in the config. |
|
|
|
|
|
|
|
|
Examples: |
|
|
:: |
|
|
pred = DefaultPredictor(cfg) |
|
|
inputs = cv2.imread("input.jpg") |
|
|
outputs = pred(inputs) |
|
|
""" |
|
|
|
|
|
def __init__(self, cfg): |
|
|
""" |
|
|
Args: |
|
|
cfg: a yacs CfgNode or a omegaconf dict object. |
|
|
""" |
|
|
if isinstance(cfg, CfgNode): |
|
|
self.cfg = cfg.clone() |
|
|
self.model = build_model(self.cfg) |
|
|
if len(cfg.DATASETS.TEST): |
|
|
test_dataset = cfg.DATASETS.TEST[0] |
|
|
|
|
|
checkpointer = DetectionCheckpointer(self.model) |
|
|
checkpointer.load(cfg.MODEL.WEIGHTS) |
|
|
|
|
|
self.aug = T.ResizeShortestEdge( |
|
|
[cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST |
|
|
) |
|
|
|
|
|
self.input_format = cfg.INPUT.FORMAT |
|
|
else: |
|
|
self.cfg = cfg |
|
|
self.model = instantiate(cfg.model) |
|
|
test_dataset = OmegaConf.select(cfg, "dataloader.test.dataset.names", default=None) |
|
|
if isinstance(test_dataset, (list, tuple)): |
|
|
test_dataset = test_dataset[0] |
|
|
|
|
|
checkpointer = DetectionCheckpointer(self.model) |
|
|
checkpointer.load(OmegaConf.select(cfg, "train.init_checkpoint", default="")) |
|
|
|
|
|
mapper = instantiate(cfg.dataloader.test.mapper) |
|
|
self.aug = mapper.augmentations |
|
|
self.input_format = mapper.image_format |
|
|
|
|
|
self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') |
|
|
self.model.eval().to(self.device) |
|
|
if test_dataset: |
|
|
self.metadata = MetadataCatalog.get(test_dataset) |
|
|
assert self.input_format in ["RGB", "BGR"], self.input_format |
|
|
|
|
|
def __call__(self, original_image): |
|
|
""" |
|
|
Args: |
|
|
original_image (np.ndarray): an image of shape (H, W, C) (in BGR order). |
|
|
|
|
|
Returns: |
|
|
predictions (dict): |
|
|
the output of the model for one image only. |
|
|
See :doc:`/tutorials/models` for details about the format. |
|
|
""" |
|
|
with torch.no_grad(): |
|
|
if self.input_format == "RGB": |
|
|
original_image = original_image[:, :, ::-1] |
|
|
height, width = original_image.shape[:2] |
|
|
image = self.aug(T.AugInput(original_image)).apply_image(original_image) |
|
|
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1)) |
|
|
inputs = {"image": image, "height": height, "width": width} |
|
|
predictions = self.model([inputs])[0] |
|
|
return predictions |
|
|
|
