| import cv2 |
| |
| import numpy as np |
| def flatten(l): |
| return [item for sublist in l for item in sublist] |
|
|
| def chunks(l, n, step=None): |
| if step is None: |
| step = n |
| return [l[i : i + n] for i in range(0, len(l), step)] |
|
|
|
|
| def sample_chunks(l, n, step=None): |
| return [l[i : i + n] for i in range(0, len(l), step) if i + n <= len(l)] |
|
|
|
|
| def grab_all_frames(path, max_size, cvt=False): |
| capture = cv2.VideoCapture(path) |
| ret = True |
| frames = [] |
| while ret: |
| ret, frame = capture.read() |
| if ret: |
| if cvt: |
| frame = frame[..., ::-1] |
| frames.append(frame) |
| if len(frames) == max_size: |
| break |
| capture.release() |
| return frames |
|
|
|
|
| def get_clips_uniform(path, count, clip_size): |
| capture = cv2.VideoCapture(path) |
| n_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) |
| max_clip_available = n_frames + 1 - clip_size |
| if count > max_clip_available: |
| count = max_clip_available |
| final_start = max_clip_available - 1 |
| start_indices = np.linspace(0, final_start, count, endpoint=True, dtype=np.int) |
| all_clip_idx = [list(range(start, start + clip_size)) for start in start_indices] |
| valid = set(flatten(all_clip_idx)) |
| max_idx = max(valid) |
|
|
| frames = {} |
| for idx in range(max_idx + 1): |
| |
| ret = capture.grab() |
| if not ret: |
| continue |
|
|
| if idx in valid: |
| ret, frame = capture.retrieve() |
| if not ret or frame is None: |
| continue |
| else: |
| |
| frames[idx] = frame |
|
|
| capture.release() |
| clips = [] |
| for clip_idx in all_clip_idx: |
| clip = [] |
| flag = True |
| for idx in clip_idx: |
| if idx not in frames: |
| flag = False |
| break |
| clip.append(frames[idx]) |
| if flag: |
| clips.append(clip) |
| return clips |
|
|
|
|
| def get_valid_faces(detect_results, max_count=10, thres=0.5, at_least=False): |
| new_results = [] |
| for i, faces in enumerate(detect_results): |
| if len(faces) > max_count: |
| faces = faces[:max_count] |
| l = [] |
| for j, face in enumerate(faces): |
| if face[-1] < thres and not (j == 0 and at_least): |
| continue |
| box, lm, score = face |
| box = box.astype(np.float64) |
| lm = lm.astype(np.float64) |
| l.append((box, lm, score)) |
| new_results.append(l) |
| return new_results |
|
|
|
|
| def scale_box(box, scale_h, scale_w, h, w): |
| x1, y1, x2, y2 = box.astype(np.int32) |
| center_x = (x1 + x2) // 2 |
| center_y = (y1 + y2) // 2 |
| box_h = int((y2 - y1) * scale_h) |
| box_w = int((x2 - x1) * scale_w) |
| new_x1 = center_x - box_w // 2 |
| new_x2 = new_x1 + box_w |
| new_y1 = center_y - box_h // 2 |
| new_y2 = new_y1 + box_h |
| new_x1 = max(new_x1, 0) |
| new_y1 = max(new_y1, 0) |
| new_y2 = min(new_y2, h) |
| new_x2 = min(new_x2, w) |
| return new_x1, new_y1, new_x2, new_y2 |
|
|
|
|
| def get_bbox(detect_res): |
| tmp_detect_res = get_valid_faces(detect_res, max_count=4, thres=0.5) |
| all_face_bboxs = [] |
| for faces in tmp_detect_res: |
| all_face_bboxs.extend([face[0] for face in faces]) |
| all_face_bboxs = np.array(all_face_bboxs).astype(np.int) |
| x1 = all_face_bboxs[:, 0].min() |
| x2 = all_face_bboxs[:, 2].max() |
| y1 = all_face_bboxs[:, 1].min() |
| y2 = all_face_bboxs[:, 3].max() |
|
|
| return x1, y1, x2, y2 |
|
|
|
|
| def delta_detect_res(detect_res, x1, y1): |
| diff = np.array([[x1, y1]]) |
| new_detect_res = [] |
| for faces in detect_res: |
| f = [] |
| for face in faces: |
| box, lm, score = face |
| box = box.astype(np.float64) |
| box[[0, 2]] -= x1 |
| box[[1, 3]] -= y1 |
| lm = lm.astype(np.float64) - diff |
| f.append((box, lm, score)) |
| new_detect_res.append(f) |
| return new_detect_res |
|
|
|
|
| def pre_crop(clips, detect_res): |
| box = np.array(get_bbox(detect_res)) |
| w = box[2] - box[0] |
| h = box[3] - box[1] |
| x1, y1, x2, y2 = scale_box( |
| box, 1.5, 1.2 if w > 2 * h else 1.5, clips[0].shape[0], clips[0].shape[1] |
| ) |
| clips = np.array(clips) |
| return clips[:, y1:y2, x1:x2], delta_detect_res(detect_res, x1, y1) |
|
|
