def track_faces(
self,
clip_dir: str,
out_base_dir: str,
draw_on_dir: str = None,
detect_only: bool = False,
):
"""
This is NOT recognition
Tracking should be based on smooth object motion, not face recognition
Steps Every frame:
detect faces
for old-face in tracked-faces:
for new-face in detected-faces:
if new-face in old-face-region and old-face in new-face-region:
match new-face and old-face
break
if old-face not in matches and tracker.update(img) > thresh:
match to tracked location
for new-face not in matches:
create new tracked-face
"""
# Setup
# load image paths
frames: List[os.DirEntry] = load_and_sort_dir(clip_dir)
draw_on_frames: List[os.DirEntry] = load_and_sort_dir(draw_on_dir)
assert len(draw_on_frames) in (0, len(frames))
# create output directory
out_dir: str = create_output_dir(out_base_dir)
# initialize variables required for object tracking
new_face_id: Iterator[int] = count(start=1)
tracked_faces: Dict[int, TrackedFace] = {}
prev_img = None
# Iterate Through Video Frames
for frame, draw_on_frame in zip_longest(frames, draw_on_frames):
# Read Images
# read image to process
img: np.ndarray = cv.imread(frame.path)
# read image to draw on (if different)
out_img = img.copy() if draw_on_frame is None else cv.imread(
draw_on_frame.path)
# ensure out_img is at least as large as img
assert len(img.shape) == len(out_img.shape) and all(
out_dim >= in_dim
for in_dim, out_dim in zip(img.shape, out_img.shape))
detected_face_boxes: List[Box] = self.detect_faces(img)
# If tracking is disabled, draw the boxes and move to next frame
if detect_only:
write_boxes(
out_path=os.path.join(out_dir, frame.name),
out_img=out_img,
boxes=detected_face_boxes,
)
continue
current_ids_to_detection_idx: Dict[int, Optional[int]] = {}
lost_tracked_face_ids: List[int] = []
# Iterate over the known (tracked) faces
for tracked_face in tracked_faces.values():
# Update the tracker with the new image
# Tracker generates new predicted_rect from previous predicted_rect
# Tracker returns its confidence that the face is inside new predicted_rect
predicted_rect_confidence: float = tracked_face.tracker.update(img)
if predicted_rect_confidence < self.tracking_threshold:
# We've lost te object. Maybe due to a cut. Can't simply look for closest faces.
# We assume the face is no longer present in img and stop tracking it
print(
f"Too low: id={tracked_face.id_}, conf={predicted_rect_confidence}, frame={frame.name}"
)
lost_tracked_face_ids.append(tracked_face.id_)
# TODO: In this case, maybe matchTemplate with found faces to see if one is above thresh
continue
predicted_rect: dlib.rectangle = tracked_face.tracker.get_position(
)
tracked_last_rect: dlib.rectangle = tracked_face.box.to_dlib_rect()
# Iterate over newly detected faces
for detected_i, detected_face_box in enumerate(
detected_face_boxes):
# TODO Maybe just do distance based
# add confidence here?
# I think track motion and distance
detected_rect = detected_face_box.to_dlib_rect()
if (
# TODO: verify these are good checks. Maybe check that the l2 dist is minimal instead
# need to make sure not modifying tracked faces as we go if we start computing minimums
# THEY ARENT
# sanity check: face hasn't moved too much
tracked_last_rect.contains(detected_rect.center())
and detected_rect.contains(tracked_last_rect.center())
# sanity check: tracker prediction isn't too far from detection
and detected_rect.contains(predicted_rect.center()) and
predicted_rect.contains(detected_rect.center())):
# detected_face_box and tracked_face are the same face
# tracker was already update to this location
if tracked_face.id_ in current_ids_to_detection_idx:
print(
f'[ERROR] {tracked_face.id_} found multiple times. Keeping first match'
)
else:
tracked_face.box = detected_face_box
current_ids_to_detection_idx[
tracked_face.id_] = detected_i
new_tracker = dlib.correlation_tracker()
new_tracker.start_track(image=img,
bounding_box=detected_rect)
tracked_face.tracker = new_tracker
if tracked_face.id_ not in current_ids_to_detection_idx:
assert predicted_rect_confidence >= self.tracking_threshold
# Didn't detect this face, but tracker is confident it is at the predicted location.
# We assume detector gave false negative
tracked_face.box = Box.from_dlib_rect(predicted_rect)
# tracker was updated to predicted_rect in update() call in condition
current_ids_to_detection_idx[tracked_face.id_] = None
# Remove lost face ids
for lost_tracked_face_id in lost_tracked_face_ids:
del tracked_faces[lost_tracked_face_id]
tracked_detection_idxs = current_ids_to_detection_idx.values()
# Track new faces
for detected_i, detected_face_box in enumerate(detected_face_boxes):
if detected_i not in tracked_detection_idxs:
# Assume new face has entered frame and start tracking it
id_ = next(new_face_id)
tracker: dlib.correlation_tracker = dlib.correlation_tracker()
tracker.start_track(
image=img, bounding_box=detected_face_box.to_dlib_rect())
tracked_faces[id_] = TrackedFace(id_=id_,
box=detected_face_box,
tracker=tracker)
current_ids_to_detection_idx[id_] = detected_i
tracked_detection_idxs = current_ids_to_detection_idx.values()
assert all(i in tracked_detection_idxs
for i in range(len(detected_face_boxes)))
assert len(current_ids_to_detection_idx) == len(tracked_faces)
write_boxes(
out_path=os.path.join(out_dir, frame.name),
out_img=out_img,
boxes=[face.box for face in tracked_faces.values()],
labelss=[[(f'Person {face.id_}', Point(1, -9))]
for face in tracked_faces.values()],
)
def track_faces(
self,
clip_dir: str,
out_base_dir: str,
draw_on_dir: str = None,
detect_only: bool = False,
):
# Setup
# load image paths
frames: List[os.DirEntry] = load_and_sort_dir(clip_dir)
draw_on_frames: List[os.DirEntry] = load_and_sort_dir(draw_on_dir)
assert len(draw_on_frames) in (0, len(frames))
# create output directory
out_dir: str = create_output_dir(out_base_dir)
# initialize variables required for object tracking
new_face_id: Iterator[int] = count(start=1)
tracked_faces: Dict[int, TrackedFace] = {}
# Iterate Through Video Frames
for frame, draw_on_frame in zip_longest(frames, draw_on_frames):
# load new frame
img = cv.imread(frame.path)
# load out_img
out_img: np.ndarray = (img.copy() if draw_on_frame is None else
cv.imread(draw_on_frame.path))
# ensure out_img is at least as large as img
assert len(img.shape) == len(out_img.shape) and all(
out_dim >= in_dim
for in_dim, out_dim in zip(img.shape, out_img.shape))
detected_face_boxes: List[Box] = self.detect_face_boxes(img)
# If tracking is disabled, draw the boxes and move to next frame
if detect_only:
write_boxes(
out_path=os.path.join(out_dir, frame.name),
out_img=out_img,
boxes=detected_face_boxes,
)
continue
detected_faces: List[GenderedFace] = gender_faces(
img=img,
faces=[
self.recognize_face(img, detected_face_box)
for detected_face_box in detected_face_boxes
],
)
current_face_ids: Set[int] = set()
lost_face_ids: Set[int] = set()
# Iterate over the known (tracked) faces
for tracked_face in tracked_faces.values():
matched_detected_faces: List[GenderedFace] = [
detected_face for detected_face in detected_faces
if self.faces_match(tracked_face, detected_face)
]
if not matched_detected_faces:
# Tracked face was not matched to and detected face
# Increment staleness since we didn't detect this face
tracked_face.staleness += 1
# Update tracker with img and get confidence
tracked_confidence: float = tracked_face.tracker.update(
img)
if (tracked_face.staleness < self.tracking_expiry
and tracked_confidence >= self.tracking_threshold):
# Assume face is still in frame but we failed to detect
# Update box with predicted location box
predicted_box: Box = Box.from_dlib_rect(
tracked_face.tracker.get_position())
tracked_face.box = predicted_box
current_face_ids.add(tracked_face.id_)
else:
# Assume face has left frame because either it is too stale or confidence is too low
if self.remember_identities:
# Set effectively infinite staleness to force tracker reset if face is found again later
tracked_face.staleness = sys.maxsize
else:
lost_face_ids.add(tracked_face.id_)
continue
# Tracked face was matched to one or more detected faces
# Multiple matches should rarely happen if faces in frame are distinct. We take closest to prev location
# TODO: Handle same person multiple times in frame
matched_detected_face = min(
matched_detected_faces,
key=lambda face: tracked_face.box.distance_to(face.box),
)
# Update tracked_face
tracked_face.descriptor = matched_detected_face.descriptor
tracked_face.shape = matched_detected_face.descriptor
tracked_face.box = matched_detected_face.box
if tracked_face.staleness >= self.tracking_expiry:
# Face was not present in last frame so reset tracker
tracked_face.tracker = dlib.correlation_tracker()
tracked_face.tracker.start_track(
image=img,
bounding_box=tracked_face.box.to_dlib_rect())
else:
# Face was present in last frame so just update guess
tracked_face.tracker.update(
image=img, guess=tracked_face.box.to_dlib_rect())
tracked_face.staleness = 0
tracked_face.gender = matched_detected_face.gender
tracked_face.gender_confidence = matched_detected_face.gender_confidence
# Add tracked_face to current_ids to reflect that it is in the frame
current_face_ids.add(tracked_face.id_)
# remove matched_detected_face from detected_faces
detected_faces.remove(matched_detected_face)
# Delete all faces that were being tracked but are now lost
# lost_face_ids will always be empty if self.remember_identities is True
for id_ in lost_face_ids:
del tracked_faces[id_]
for new_face in detected_faces:
# This is a new face (previously unseen)
id_ = next(new_face_id)
tracker: dlib.correlation_tracker = dlib.correlation_tracker()
tracker.start_track(image=img,
bounding_box=new_face.box.to_dlib_rect())
tracked_faces[id_] = TrackedFace(
box=new_face.box,
descriptor=new_face.descriptor,
shape=new_face.shape,
id_=id_,
tracker=tracker,
gender=new_face.gender,
gender_confidence=new_face.gender_confidence,
)
current_face_ids.add(id_)
write_boxes(
out_path=os.path.join(out_dir, frame.name),
out_img=out_img,
boxes=[tracked_faces[id_].box for id_ in current_face_ids],
labelss=[[
(
f'Person {id_}',
Point(3, 14),
),
(
f'{tracked_faces[id_].gender.name[0].upper()}: {round(100 * tracked_faces[id_].gender_confidence, 1)}%',
Point(3, 30),
),
] for id_ in current_face_ids],
color=Color.yellow(),
)
print(
f"Processed {frame.name}. Currently tracking {len(tracked_faces)} faces"
)
return out_dir