def __init__(self, device='cpu', img_size=608, thresh=0.7):
self.detector = YOLOv3(device=device, img_size=img_size,
person_detector=True, video=True,
return_dict=True)
self.tracker = Sort()
self.tresh = thresh
self.device = device
def __init__(
self,
device=None,
batch_size=12,
display=False,
detection_threshold=0.7,
detector_type='yolo',
yolo_img_size=608,
output_format='list',
detector_checkpoint=None,
detector_config=None
):
'''
Multi Person Tracker
:param device (str, 'cuda' or 'cpu'): torch device for model and inputs
:param batch_size (int): batch size for detection model
:param display (bool): display the results of multi person tracking
:param detection_threshold (float): threshold to filter detector predictions
:param detector_type (str, 'maskrcnn' or 'yolo'): detector architecture
:param yolo_img_size (int): yolo detector input image size
:param output_format (str, 'dict' or 'list'): result output format
'''
if device is not None:
self.device = device
else:
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.batch_size = batch_size
self.display = display
self.detection_threshold = detection_threshold
self.output_format = output_format
self.detector_type = detector_type
self.detector_checkpoint = detector_checkpoint[0] if type(detector_checkpoint) == tuple else detector_checkpoint,
self.detector_config = detector_config[0] if type(detector_config) == tuple else detector_config
if self.detector_type == 'maskrcnn':
self.detector = keypointrcnn_resnet50_fpn(pretrained=True).to(self.device).eval()
elif self.detector_type == 'yolo':
self.detector = YOLOv3(
device=self.device, img_size=yolo_img_size, person_detector=True, video=True, return_dict=True
)
# output [{'boxes': tensor([], size=(0, 4)),
# 'scores': tensor([]),
# 'classes': tensor([])}]
# x = torch.Tensor([np.random.rand(3, 300, 400), np.random.rand(3, 300, 400)])
# print(self.detector(x))
elif self.detector_type == 'retina':
self.detector = init_detector(
self.detector_config, self.detector_checkpoint[0], device='cuda:0'
)
else:
raise ModuleNotFoundError
self.tracker = Sort()
def __init__(
self,
device=None,
batch_size=12,
display=False,
detection_threshold=0.7,
detector_type='yolo',
yolo_img_size=608,
output_format='list',
):
'''
Multi Person Tracker
:param device (str, 'cuda' or 'cpu'): torch device for model and inputs
:param batch_size (int): batch size for detection model
:param display (bool): display the results of multi person tracking
:param detection_threshold (float): threshold to filter detector predictions
:param detector_type (str, 'maskrcnn' or 'yolo'): detector architecture
:param yolo_img_size (int): yolo detector input image size
:param output_format (str, 'dict' or 'list'): result output format
'''
if device is not None:
self.device = device
else:
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.batch_size = batch_size
self.display = display
self.detection_threshold = detection_threshold
self.output_format = output_format
if detector_type == 'maskrcnn':
self.detector = keypointrcnn_resnet50_fpn(pretrained=True).to(
self.device).eval()
elif detector_type == 'yolo':
self.detector = YOLOv3(device=self.device,
img_size=yolo_img_size,
person_detector=True,
video=True,
return_dict=True)
else:
raise ModuleNotFoundError
self.tracker = Sort()
help="number of cpu threads to use during batch generation")
parser.add_argument("--img_size",
type=int,
default=416,
help="size of each image dimension")
parser.add_argument("--checkpoint_model",
type=str,
help="path to checkpoint model")
opt = parser.parse_args()
print(opt)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs("output", exist_ok=True)
model = YOLOv3(device=device, img_size=opt.img_size, person_detector=False)
dataloader = DataLoader(
ImageFolder(opt.image_folder, img_size=opt.img_size),
batch_size=opt.batch_size,
shuffle=False,
num_workers=opt.n_cpu,
)
classes = load_classes(opt.class_path) # Extracts class labels from file
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available(
) else torch.FloatTensor
imgs = [] # Stores image paths
img_detections = [] # Stores detections for each image index