Spaces:
Running
Running
File size: 3,866 Bytes
adbf50c |
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 |
import numpy as np
import cv2
def extract_frames(filename, num_frames, model, image_size=(380, 380)):
cap = cv2.VideoCapture(filename)
if not cap.isOpened():
return [], []
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
frame_idxs = np.linspace(0, frame_count - 1, min(num_frames, frame_count), dtype=int)
croppedfaces, idx_list = [], []
for idx in frame_idxs:
ret = cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
if not ret:
continue
ret, frame = cap.read()
if not ret:
continue
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
faces = extract_face(frame, model, image_size)
if faces:
croppedfaces.extend(faces)
idx_list.extend([idx] * len(faces))
cap.release()
return croppedfaces, idx_list
def extract_face(frame, model, image_size=(380, 380)):
faces = model.predict_jsons(frame)
if len(faces[0]['bbox']) == 0:
scale_factor = 0.5
small_frame = cv2.resize(frame, None, fx=scale_factor, fy=scale_factor)
faces = model.predict_jsons(small_frame)
if len(faces[0]['bbox']) == 0:
return []
for face in faces:
face['bbox'] = [coord/scale_factor for coord in face['bbox']]
croppedfaces = []
for face_idx in range(len(faces)):
x0, y0, x1, y1 = faces[face_idx]['bbox']
bbox = np.array([[x0, y0], [x1, y1]])
face = crop_face(frame, None, bbox, False, crop_by_bbox=True, only_img=True, phase='test')
croppedfaces.append(cv2.resize(face, dsize=image_size).transpose((2, 0, 1)))
return croppedfaces
def crop_face(img, landmark=None, bbox=None, margin=False, crop_by_bbox=True, abs_coord=False, only_img=False, phase='test'):
assert phase in ['train', 'val', 'test']
assert landmark is not None or bbox is not None
H, W = len(img), len(img[0])
if crop_by_bbox:
x0, y0 = bbox[0]
x1, y1 = bbox[1]
w = x1 - x0
h = y1 - y0
w0_margin = w/4
w1_margin = w/4
h0_margin = h/4
h1_margin = h/4
else:
x0, y0 = landmark[:68,0].min(), landmark[:68,1].min()
x1, y1 = landmark[:68,0].max(), landmark[:68,1].max()
w = x1 - x0
h = y1 - y0
w0_margin = w/8
w1_margin = w/8
h0_margin = h/2
h1_margin = h/5
if margin:
w0_margin *= 4
w1_margin *= 4
h0_margin *= 2
h1_margin *= 2
elif phase == 'train':
rand_factor = np.random.rand(4) * 0.6 + 0.2
w0_margin *= rand_factor[0]
w1_margin *= rand_factor[1]
h0_margin *= rand_factor[2]
h1_margin *= rand_factor[3]
else:
w0_margin *= 0.5
w1_margin *= 0.5
h0_margin *= 0.5
h1_margin *= 0.5
y0_new = max(0, int(y0-h0_margin))
y1_new = min(H, int(y1+h1_margin)+1)
x0_new = max(0, int(x0-w0_margin))
x1_new = min(W, int(x1+w1_margin)+1)
img_cropped = img[y0_new:y1_new, x0_new:x1_new]
if only_img:
return img_cropped
if landmark is not None:
landmark_cropped = np.zeros_like(landmark)
for i, (p,q) in enumerate(landmark):
landmark_cropped[i] = [p-x0_new, q-y0_new]
else:
landmark_cropped = None
if bbox is not None:
bbox_cropped = np.zeros_like(bbox)
for i, (p,q) in enumerate(bbox):
bbox_cropped[i] = [p-x0_new, q-y0_new]
else:
bbox_cropped = None
if abs_coord:
return img_cropped, landmark_cropped, bbox_cropped, (y0-y0_new, x0-x0_new, y1_new-y1, x1_new-x1), y0_new, y1_new, x0_new, x1_new
else:
return img_cropped, landmark_cropped, bbox_cropped, (y0-y0_new, x0-x0_new, y1_new-y1, x1_new-x1) |