| | 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 |
| |
|
