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()
class Manager():
def __init__(self, input, output, max_frames, skip, show, detector):
self.video = VideoManager(input, output, max_frames, skip, show)
self.horses = np.array([], dtype=Horse)
self.detector = detector
self.global_horse_number = 1
def spawnHorse(self, box):
horse = Horse(box, self.global_horse_number)
self.global_horse_number += 1
logger.info(f'spawn horse {horse.number}')
self.horses = np.append(self.horses, horse)
return horse
def removeHorse(self, horse):
logger.info(f'remove horse {horse.number}')
self.horses = self.horses[self.horses != horse]
def detect(self, frame):
boxes, scores = self.detector.detect_boxes(frame)
relevant_boxes = []
for index in range(len(boxes)):
# todo: find appropriate value for low score
if scores[index] > 0.3: relevant_boxes.append(boxes[index])
return np.array(relevant_boxes)
def match(self, frame, detected):
detected_horses = []
for index, box in enumerate(detected):
intersects = False
for horse in detected_horses:
if horse.intersect(box): intersects = True
if intersects: continue
lone_horses = list_diff(self.horses, detected_horses)
horse = find_closest_horse(lone_horses, box)
if horse is None:
horse = self.spawnHorse(box)
detected_horses.append(horse)
else:
horse.detect(box)
detected_horses.append(horse)
lone_horses = list_diff(self.horses, detected_horses)
for horse in lone_horses:
horse.track(frame)
def initialize(self):
raw = self.video.read()
frame = raw.copy()
smooth = raw.copy()
detected = self.detect(frame)
self.match(frame, detected)
for horse in self.horses:
horse.draw(frame)
horse.draw_smooth(smooth)
self.video.write(raw, frame, smooth, self.horses)
def update(self):
raw = self.video.read()
frame = raw.copy()
smooth = raw.copy()
for horse in self.horses:
horse.updated = False
horse.last_detected += 1
detected = self.detect(frame)
self.match(frame, detected)
for horse in self.horses:
if horse.gone():
self.removeHorse(horse)
continue
horse.update(frame, self.horses)
horse.draw(frame)
horse.draw_smooth(smooth)
self.video.write(raw, frame, smooth, self.horses)
def getFrameCount(self):
return self.video.getFrameCount()