class Manager():
def __init__(self, input, output, max_frames, skip, show, detector):
self.video = VideoManager(input, output, max_frames, skip, show)
self.horse = None
self.detector = detector
def match(self, frame, boxes, scores):
# if no horse and no boxes: return
# if no horse and boxes: assign first box
# if horse but no boxes: track
# if horse and boxes: find closest box
logger.info(f'detected boxes: {boxes}')
if self.horse is None and len(boxes) == 0: return
elif self.horse is None and len(boxes) > 0:
max_index = np.argmax(scores)
self.horse = Horse(boxes[max_index], 1)
elif self.horse is not None and len(boxes) == 0:
self.horse.track(frame)
elif self.horse is not None and len(boxes) > 0:
box = find_closest_box(self.horse, boxes)
if box is None:
#self.horse.track() # old
self.horse.track(frame)
else:
self.horse.detect(box)
def initialize(self):
raw = self.video.read()
frame = raw.copy()
smooth = raw.copy()
boxes, scores = self.detector.detect_boxes(frame)
self.match(frame, boxes, scores)
if self.horse is not None:
self.horse.draw(frame)
self.horse.draw_smooth(smooth)
horses = [] if self.horse is None else [self.horse]
self.video.write(raw, frame, smooth, horses)
def update(self):
raw = self.video.read()
frame = raw.copy()
smooth = raw.copy()
if self.horse is not None:
self.horse.updated = False
self.horse.last_detected += 1
boxes, scores = self.detector.detect_boxes(frame)
self.match(frame, boxes, scores)
if self.horse is not None:
self.horse.update(frame, [])
self.horse.draw(frame)
self.horse.draw_smooth(smooth)
if self.horse is not None and self.horse.gone():
logger.info('horse is gone')
self.horse = None
horses = [] if self.horse is None else [self.horse]
self.video.write(raw, frame, smooth, horses)
def getFrameCount(self):
return self.video.getFrameCount()