File size: 4,390 Bytes
9e14838 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | import cv2
# from test_tools.utils import flatten
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):
# Get the next frame, but don't decode if we're not using it.
ret = capture.grab()
if not ret:
continue
if idx in valid:
ret, frame = capture.retrieve()
if not ret or frame is None:
continue
else:
# frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
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)
|