Spaces:
Runtime error
Runtime error
File size: 3,109 Bytes
97a6728 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import torch
import tops
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.structures import Boxes
from detectron2.data import MetadataCatalog
from detectron2 import model_zoo
from typing import Dict
from detectron2.data.transforms import ResizeShortestEdge
from torchvision.transforms.functional import resize
model_urls = {
"COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml": "https://dl.fbaipublicfiles.com/detectron2/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x/139653917/model_final_2d9806.pkl",
}
class MaskRCNNDetector:
def __init__(
self,
cfg_name: str = "COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml",
score_thres: float = 0.9,
class_filter=["person"], # ["car", "bicycle","truck", "bus", "backpack"]
fp16_inference: bool = False
) -> None:
cfg = model_zoo.get_config(cfg_name)
cfg.MODEL.DEVICE = str(tops.get_device())
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_thres
cfg.freeze()
self.cfg = cfg
with tops.logger.capture_log_stdout():
self.model = build_model(cfg)
DetectionCheckpointer(self.model).load(model_urls[cfg_name])
self.model.eval()
self.input_format = cfg.INPUT.FORMAT
self.class_names = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes
self.class_to_keep = set([self.class_names.index(cls_) for cls_ in class_filter])
self.person_class = self.class_names.index("person")
self.fp16_inference = fp16_inference
tops.logger.log("Mask R-CNN built.")
def __call__(self, *args, **kwargs):
return self.forward(*args, **kwargs)
def resize_im(self, im):
H, W = im.shape[1:]
newH, newW = ResizeShortestEdge.get_output_shape(
H, W, self.cfg.INPUT.MIN_SIZE_TEST, self.cfg.INPUT.MAX_SIZE_TEST)
return resize(
im, (newH, newW), antialias=True)
@torch.no_grad()
def forward(self, im: torch.Tensor):
if self.input_format == "BGR":
im = im.flip(0)
else:
assert self.input_format == "RGB"
H, W = im.shape[-2:]
im = self.resize_im(im)
with torch.cuda.amp.autocast(enabled=self.fp16_inference):
output = self.model([{"image": im, "height": H, "width": W}])[0]["instances"]
scores = output.get("scores")
N = len(scores)
classes = output.get("pred_classes")
idx2keep = [i for i in range(N) if classes[i].tolist() in self.class_to_keep]
classes = classes[idx2keep]
assert isinstance(output.get("pred_boxes"), Boxes)
segmentation = output.get("pred_masks")[idx2keep]
assert segmentation.dtype == torch.bool
is_person = classes == self.person_class
return {
"scores": output.get("scores")[idx2keep],
"segmentation": segmentation,
"classes": output.get("pred_classes")[idx2keep],
"is_person": is_person
}
|