Intae/deepfake
1
1import os2 3import cv24import numpy as np5import torch6from PIL import Image7from albumentations.augmentations.functional import image_compression8from facenet_pytorch.models.mtcnn import MTCNN9from concurrent.futures import ThreadPoolExecutor10 11from torchvision.transforms import Normalize12 13mean = [0.485, 0.456, 0.406]14std = [0.229, 0.224, 0.225]15normalize_transform = Normalize(mean, std)16 17 18class VideoReader:19 """Helper class for reading one or more frames from a video file."""20 21 def __init__(self, verbose=True, insets=(0, 0)):22 """Creates a new VideoReader.23 24 Arguments:25 verbose: whether to print warnings and error messages26 insets: amount to inset the image by, as a percentage of27 (width, height). This lets you "zoom in" to an image28 to remove unimportant content around the borders.29 Useful for face detection, which may not work if the30 faces are too small.31 """32 self.verbose = verbose33 self.insets = insets34 35 def read_frames(self, path, num_frames, jitter=0, seed=None):36 """Reads frames that are always evenly spaced throughout the video.37 38 Arguments:39 path: the video file40 num_frames: how many frames to read, -1 means the entire video41 (warning: this will take up a lot of memory!)42 jitter: if not 0, adds small random offsets to the frame indices;43 this is useful so we don't always land on even or odd frames44 seed: random seed for jittering; if you set this to a fixed value,45 you probably want to set it only on the first video46 """47 assert num_frames > 048 49 capture = cv2.VideoCapture(path)50 frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))51 if frame_count <= 0: return None52 53 frame_idxs = np.linspace(0, frame_count - 1, num_frames, endpoint=True, dtype=np.int)54 if jitter > 0:55 np.random.seed(seed)56 jitter_offsets = np.random.randint(-jitter, jitter, len(frame_idxs))57 frame_idxs = np.clip(frame_idxs + jitter_offsets, 0, frame_count - 1)58 59 result = self._read_frames_at_indices(path, capture, frame_idxs)60 capture.release()61 return result62 63 def read_random_frames(self, path, num_frames, seed=None):64 """Picks the frame indices at random.65 66 Arguments:67 path: the video file68 num_frames: how many frames to read, -1 means the entire video69 (warning: this will take up a lot of memory!)70 """71 assert num_frames > 072 np.random.seed(seed)73 74 capture = cv2.VideoCapture(path)75 frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))76 if frame_count <= 0: return None77 78 frame_idxs = sorted(np.random.choice(np.arange(0, frame_count), num_frames))79 result = self._read_frames_at_indices(path, capture, frame_idxs)80 81 capture.release()82 return result83 84 def read_frames_at_indices(self, path, frame_idxs):85 """Reads frames from a video and puts them into a NumPy array.86 87 Arguments:88 path: the video file89 frame_idxs: a list of frame indices. Important: should be90 sorted from low-to-high! If an index appears multiple91 times, the frame is still read only once.92 93 Returns:94 - a NumPy array of shape (num_frames, height, width, 3)95 - a list of the frame indices that were read96 97 Reading stops if loading a frame fails, in which case the first98 dimension returned may actually be less than num_frames.99 100 Returns None if an exception is thrown for any reason, or if no101 frames were read.102 """103 assert len(frame_idxs) > 0104 capture = cv2.VideoCapture(path)105 result = self._read_frames_at_indices(path, capture, frame_idxs)106 capture.release()107 return result108 109 def _read_frames_at_indices(self, path, capture, frame_idxs):110 try:111 frames = []112 idxs_read = []113 for frame_idx in range(frame_idxs[0], frame_idxs[-1] + 1):114 # Get the next frame, but don't decode if we're not using it.115 ret = capture.grab()116 if not ret:117 if self.verbose:118 print("Error grabbing frame %d from movie %s" % (frame_idx, path))119 break120 121 # Need to look at this frame?122 current = len(idxs_read)123 if frame_idx == frame_idxs[current]:124 ret, frame = capture.retrieve()125 if not ret or frame is None:126 if self.verbose:127 print("Error retrieving frame %d from movie %s" % (frame_idx, path))128 break129 130 frame = self._postprocess_frame(frame)131 frames.append(frame)132 idxs_read.append(frame_idx)133 134 if len(frames) > 0:135 return np.stack(frames), idxs_read136 if self.verbose:137 print("No frames read from movie %s" % path)138 return None139 except:140 if self.verbose:141 print("Exception while reading movie %s" % path)142 return None143 144 def read_middle_frame(self, path):145 """Reads the frame from the middle of the video."""146 capture = cv2.VideoCapture(path)147 frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))148 result = self._read_frame_at_index(path, capture, frame_count // 2)149 capture.release()150 return result151 152 def read_frame_at_index(self, path, frame_idx):153 """Reads a single frame from a video.154 155 If you just want to read a single frame from the video, this is more156 efficient than scanning through the video to find the frame. However,157 for reading multiple frames it's not efficient.158 159 My guess is that a "streaming" approach is more efficient than a160 "random access" approach because, unless you happen to grab a keyframe,161 the decoder still needs to read all the previous frames in order to162 reconstruct the one you're asking for.163 164 Returns a NumPy array of shape (1, H, W, 3) and the index of the frame,165 or None if reading failed.166 """167 capture = cv2.VideoCapture(path)168 result = self._read_frame_at_index(path, capture, frame_idx)169 capture.release()170 return result171 172 def _read_frame_at_index(self, path, capture, frame_idx):173 capture.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)174 ret, frame = capture.read()175 if not ret or frame is None:176 if self.verbose:177 print("Error retrieving frame %d from movie %s" % (frame_idx, path))178 return None179 else:180 frame = self._postprocess_frame(frame)181 return np.expand_dims(frame, axis=0), [frame_idx]182 183 def _postprocess_frame(self, frame):184 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)185 186 if self.insets[0] > 0:187 W = frame.shape[1]188 p = int(W * self.insets[0])189 frame = frame[:, p:-p, :]190 191 if self.insets[1] > 0:192 H = frame.shape[1]193 q = int(H * self.insets[1])194 frame = frame[q:-q, :, :]195 196 return frame197 198 199class FaceExtractor:200 def __init__(self, video_read_fn):201 self.video_read_fn = video_read_fn202 self.detector = MTCNN(margin=0, thresholds=[0.7, 0.8, 0.8])203 204 def process_videos(self, input_dir, filenames, video_idxs):205 videos_read = []206 frames_read = []207 frames = []208 results = []209 for video_idx in video_idxs:210 # Read the full-size frames from this video.211 filename = filenames[video_idx]212 video_path = os.path.join(input_dir, filename)213 result = self.video_read_fn(video_path)214 # Error? Then skip this video.215 if result is None: continue216 217 videos_read.append(video_idx)218 219 # Keep track of the original frames (need them later).220 my_frames, my_idxs = result221 222 frames.append(my_frames)223 frames_read.append(my_idxs)224 for i, frame in enumerate(my_frames):225 h, w = frame.shape[:2]226 img = Image.fromarray(frame.astype(np.uint8))227 img = img.resize(size=[s // 2 for s in img.size])228 229 batch_boxes, probs = self.detector.detect(img, landmarks=False)230 231 faces = []232 scores = []233 if batch_boxes is None:234 continue235 for bbox, score in zip(batch_boxes, probs):236 if bbox is not None:237 xmin, ymin, xmax, ymax = [int(b * 2) for b in bbox]238 w = xmax - xmin239 h = ymax - ymin240 p_h = h // 3241 p_w = w // 3242 crop = frame[max(ymin - p_h, 0):ymax + p_h, max(xmin - p_w, 0):xmax + p_w]243 faces.append(crop)244 scores.append(score)245 246 frame_dict = {"video_idx": video_idx,247 "frame_idx": my_idxs[i],248 "frame_w": w,249 "frame_h": h,250 "faces": faces,251 "scores": scores}252 results.append(frame_dict)253 254 return results255 256 def process_video(self, video_path):257 """Convenience method for doing face extraction on a single video."""258 input_dir = os.path.dirname(video_path)259 filenames = [os.path.basename(video_path)]260 return self.process_videos(input_dir, filenames, [0])261 262 263 264def confident_strategy(pred, t=0.8):265 pred = np.array(pred)266 sz = len(pred)267 fakes = np.count_nonzero(pred > t)268 # 11 frames are detected as fakes with high probability269 if fakes > sz // 2.5 and fakes > 11:270 return np.mean(pred[pred > t])271 elif np.count_nonzero(pred < 0.2) > 0.9 * sz:272 return np.mean(pred[pred < 0.2])273 else:274 return np.mean(pred)275 276strategy = confident_strategy277 278 279def put_to_center(img, input_size):280 img = img[:input_size, :input_size]281 image = np.zeros((input_size, input_size, 3), dtype=np.uint8)282 start_w = (input_size - img.shape[1]) // 2283 start_h = (input_size - img.shape[0]) // 2284 image[start_h:start_h + img.shape[0], start_w: start_w + img.shape[1], :] = img285 return image286 287 288def isotropically_resize_image(img, size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC):289 h, w = img.shape[:2]290 if max(w, h) == size:291 return img292 if w > h:293 scale = size / w294 h = h * scale295 w = size296 else:297 scale = size / h298 w = w * scale299 h = size300 interpolation = interpolation_up if scale > 1 else interpolation_down301 resized = cv2.resize(img, (int(w), int(h)), interpolation=interpolation)302 return resized303 304 305def predict_on_video(face_extractor, video_path, batch_size, input_size, models, strategy=np.mean,306 apply_compression=False):307 batch_size *= 4308 try:309 faces = face_extractor.process_video(video_path)310 if len(faces) > 0:311 x = np.zeros((batch_size, input_size, input_size, 3), dtype=np.uint8)312 n = 0313 for frame_data in faces:314 for face in frame_data["faces"]:315 resized_face = isotropically_resize_image(face, input_size)316 resized_face = put_to_center(resized_face, input_size)317 if apply_compression:318 resized_face = image_compression(resized_face, quality=90, image_type=".jpg")319 if n + 1 < batch_size:320 x[n] = resized_face321 n += 1322 else:323 pass324 if n > 0:325 x = torch.tensor(x).float()326 # Preprocess the images.327 x = x.permute((0, 3, 1, 2))328 for i in range(len(x)):329 x[i] = normalize_transform(x[i] / 255.)330 # Make a prediction, then take the average.331 with torch.no_grad():332 preds = []333 for model in models:334 y_pred = model(x[:n])335 y_pred = torch.sigmoid(y_pred.squeeze())336 bpred = y_pred[:n].cpu().numpy()337 preds.append(strategy(bpred))338 return np.mean(preds)339 except Exception as e:340 print("Prediction error on video %s: %s" % (video_path, str(e)))341 342 return 0.5343 344 345def predict_on_video_set(face_extractor, videos, input_size, num_workers, test_dir, frames_per_video, models,346 strategy=np.mean,347 apply_compression=False):348 def process_file(i):349 filename = videos[i]350 y_pred = predict_on_video(face_extractor=face_extractor, video_path=os.path.join(test_dir, filename),351 input_size=input_size,352 batch_size=frames_per_video,353 models=models, strategy=strategy, apply_compression=apply_compression)354 return y_pred355 356 with ThreadPoolExecutor(max_workers=num_workers) as ex:357 predictions = ex.map(process_file, range(len(videos)))358 return list(predictions)359 360