| | import torch |
| | import detectron2.data.transforms as T |
| | from detectron2.checkpoint import DetectionCheckpointer |
| | from detectron2.data import ( |
| | MetadataCatalog, |
| | ) |
| | from detectron2.modeling import build_model |
| |
|
| |
|
| | __all__ = [ |
| | "DefaultPredictor", |
| | ] |
| |
|
| |
|
| | class DefaultPredictor: |
| | """ |
| | 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 `cfg.MODEL.WEIGHTS`. |
| | 2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`. |
| | 3. Apply resizing defined by `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 |
| | cfg.DATASETS.TEST. |
| | Examples: |
| | :: |
| | pred = DefaultPredictor(cfg) |
| | inputs = cv2.imread("input.jpg") |
| | outputs = pred(inputs) |
| | """ |
| |
|
| | def __init__(self, cfg): |
| | self.cfg = cfg.clone() |
| | self.model = build_model(self.cfg) |
| | self.model.eval() |
| | if len(cfg.DATASETS.TEST): |
| | self.metadata = MetadataCatalog.get(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 |
| | assert self.input_format in ["RGB", "BGR"], self.input_format |
| |
|
| | def __call__(self, original_image, task): |
| | """ |
| | 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.get_transform(original_image).apply_image(original_image) |
| | image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1)) |
| | |
| | task = f"The task is {task}" |
| |
|
| | inputs = {"image": image, "height": height, "width": width, "task": task} |
| | predictions = self.model([inputs])[0] |
| | return predictions |
