| import os |
|
|
| os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" |
|
|
| from .ct.detection.utils import grab_all_frames, get_valid_faces, sample_chunks |
| from .ct.operations import multiple_tracking |
| import numpy as np |
| from .ct.face_alignment import LandmarkPredictor |
| from .ct.detection import FaceDetector |
| import cv2 |
| from .utils import flatten,partition |
|
|
|
|
| detector = FaceDetector(0) |
| predictor = LandmarkPredictor(0) |
|
|
|
|
| def get_five(ldm68): |
| groups = [range(36, 42), range(42, 48), [30], [48], [54]] |
| points = [] |
| for group in groups: |
| points.append(ldm68[group].mean(0)) |
| return np.array(points) |
|
|
|
|
| def get_bbox(mask): |
| try: |
| y, x = np.nonzero(mask[..., 0]) |
| return x.min() - 1, y.min() - 1, x.max() + 1, y.max() + 1 |
| except: |
| return None |
|
|
|
|
| def get_bigger_box(image, box, scale=0.5): |
| height, width = image.shape[:2] |
| box = np.rint(box).astype(np.int) |
| new_box = box.reshape(2, 2) |
| size = new_box[1] - new_box[0] |
| diff = scale * size |
| diff = diff[None, :] * np.array([-1, 1])[:, None] |
| new_box = new_box + diff |
| new_box[:, 0] = np.clip(new_box[:, 0], 0, width - 1) |
| new_box[:, 1] = np.clip(new_box[:, 1], 0, height - 1) |
| new_box = np.rint(new_box).astype(np.int) |
| return new_box.reshape(-1) |
|
|
|
|
| def process_bigger_clips(clips, dete_res, clip_size, step, scale=0.5): |
| assert len(clips) % clip_size == 0 |
| detect_results = sample_chunks(dete_res, clip_size, step) |
| clips = sample_chunks(clips, clip_size, step) |
| new_clips = [] |
| for i, (frame_clip, record_clip) in enumerate(zip(clips, detect_results)): |
| tracks = multiple_tracking(record_clip) |
| for j, track in enumerate(tracks): |
| new_images = [] |
| for (box, ldm, _), frame in zip(track, frame_clip): |
| big_box = get_bigger_box(frame, box, scale) |
| x1, y1, x2, y2 = big_box |
| top_left = big_box[:2][None, :] |
| new_ldm5 = ldm - top_left |
| box = np.rint(box).astype(np.int) |
| new_box = (box.reshape(2, 2) - top_left).reshape(-1) |
| feed = LandmarkPredictor.prepare_feed(frame, box) |
| ldm68 = predictor(feed) - top_left |
| new_images.append( |
| (frame[y1:y2, x1:x2], big_box, new_box, new_ldm5, ldm68) |
| ) |
| new_clips.append(new_images) |
| return new_clips |
|
|
|
|
| def post(detected_faces): |
| return [[face[:4], None, face[-1]] for face in detected_faces] |
|
|
|
|
| def check(detect_res): |
| return min([len(faces) for faces in detect_res]) != 0 |
|
|
|
|
| def detect_all(file, sfd_only=False, return_frames=False, max_size=None): |
| frames = grab_all_frames(file, max_size=max_size, cvt=True) |
| if not sfd_only: |
| detect_res = flatten( |
| [detector.detect(item) for item in partition(frames, 50)] |
| ) |
| detect_res = get_valid_faces(detect_res, thres=0.5) |
| else: |
| raise NotImplementedError |
|
|
| all_68 = get_lm68(frames, detect_res) |
| if not return_frames: |
| return detect_res, all_68 |
| else: |
| return detect_res, all_68, frames |
|
|
|
|
| def get_lm68(frames, detect_res): |
| assert len(frames) == len(detect_res) |
| frame_count = len(frames) |
| all_68 = [] |
| for i in range(frame_count): |
| frame = frames[i] |
| faces = detect_res[i] |
| if len(faces) == 0: |
| res_68 = [] |
| else: |
| feeds = [] |
| for face in faces: |
| assert len(face) == 3 |
| box = face[0] |
| feed = LandmarkPredictor.prepare_feed(frame, box) |
| feeds.append(feed) |
| res_68 = predictor(feeds) |
| assert len(res_68) == len(faces) |
| for face, l_68 in zip(faces, res_68): |
| if face[1] is None: |
| face[1] = get_five(l_68) |
| all_68.append(res_68) |
|
|
| assert len(all_68) == len(detect_res) |
| return all_68 |
|
|
