Spaces:
Running
on
T4
Running
on
T4
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() # cfg can be modified by model | |
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(): # https://github.com/sphinx-doc/sphinx/issues/4258 | |
# Apply pre-processing to image. | |
if self.input_format == "RGB": | |
# whether the model expects BGR inputs or 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 |