CMU-INF-DIVA / detectors

Object detection models

Detectors

Author: Lijun Yu

Email: [e-mail]

A submodule of object detection models.

Models

• Mask R-CNN from Detectron2
• YOLOv5 official
• EfficientDet in PyTorch

API

import torch
from detectors import get_detector

detector_class = get_detector('Mask R-CNN')  # Or YOLOv5, EfficientDet
detector = detector_class(gpu_id=0)

# images: a list of pytorch tensors as H x W x C[BGR] in [0, 256)
images = [torch.zeros(1080, 1920, 3)]

detections = detector(images)
for detection in detections:
  # Detection attributes: object_types, image_boxes, detection_scores

Dependency

See actev_base.

License

See License.