Delete preprocessing

preprocessing/__init__.py

@@ -1 +0,0 @@
-from .face_detector import *

preprocessing/compress_videos.py

@@ -1,45 +0,0 @@
-import argparse
-import os
-import random
-import subprocess
-
-os.environ["MKL_NUM_THREADS"] = "1"
-os.environ["NUMEXPR_NUM_THREADS"] = "1"
-os.environ["OMP_NUM_THREADS"] = "1"
-from functools import partial
-from glob import glob
-from multiprocessing.pool import Pool
-from os import cpu_count
-
-import cv2
-
-cv2.ocl.setUseOpenCL(False)
-cv2.setNumThreads(0)
-from tqdm import tqdm
-
-
-def compress_video(video, root_dir):
-    parent_dir = video.split("/")[-2]
-    out_dir = os.path.join(root_dir, "compressed", parent_dir)
-    os.makedirs(out_dir, exist_ok=True)
-    video_name = video.split("/")[-1]
-    out_path = os.path.join(out_dir, video_name)
-    lvl = random.choice([23, 28, 32])
-    command = "ffmpeg -i {} -c:v libx264 -crf {} -threads 1 {}".format(video, lvl, out_path)
-    try:
-        subprocess.check_output(command, shell=True, stderr=subprocess.STDOUT)
-    except Exception as e:
-        print("Could not process vide", str(e))
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(
-        description="Extracts jpegs from video")
-    parser.add_argument("--root-dir", help="root directory", default="/mnt/sota/datasets/deepfake")
-
-    args = parser.parse_args()
-    videos = [video_path for video_path in glob(os.path.join(args.root_dir, "*/*.mp4"))]
-    with Pool(processes=cpu_count() - 2) as p:
-        with tqdm(total=len(videos)) as pbar:
-            for v in p.imap_unordered(partial(compress_video, root_dir=args.root_dir), videos):
-                pbar.update()

preprocessing/detect_original_faces.py

@@ -1,51 +0,0 @@
-import argparse
-import json
-import os
-from os import cpu_count
-from typing import Type
-
-from torch.utils.data.dataloader import DataLoader
-from tqdm import tqdm
-
-from preprocessing import face_detector, VideoDataset
-from preprocessing.face_detector import VideoFaceDetector
-from preprocessing.utils import get_original_video_paths
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Process a original videos with face detector")
-    parser.add_argument("--root-dir", help="root directory")
-    parser.add_argument("--detector-type", help="type of the detector", default="FacenetDetector",
-                        choices=["FacenetDetector"])
-    args = parser.parse_args()
-    return args
-
-
-def process_videos(videos, root_dir, detector_cls: Type[VideoFaceDetector]):
-    detector = face_detector.__dict__[detector_cls](device="cuda:0")
-    dataset = VideoDataset(videos)
-    loader = DataLoader(dataset, shuffle=False, num_workers=cpu_count() - 2, batch_size=1, collate_fn=lambda x: x)
-    for item in tqdm(loader):
-        result = {}
-        video, indices, frames = item[0]
-        batches = [frames[i:i + detector._batch_size] for i in range(0, len(frames), detector._batch_size)]
-        for j, frames in enumerate(batches):
-            result.update({int(j * detector._batch_size) + i : b for i, b in zip(indices, detector._detect_faces(frames))})
-        id = os.path.splitext(os.path.basename(video))[0]
-        out_dir = os.path.join(root_dir, "boxes")
-        os.makedirs(out_dir, exist_ok=True)
-        with open(os.path.join(out_dir, "{}.json".format(id)), "w") as f:
-            json.dump(result, f)
-
-
-
-
-def main():
-    args = parse_args()
-    originals = get_original_video_paths(args.root_dir)
-    process_videos(originals, args.root_dir, args.detector_type)
-
-
-if __name__ == "__main__":
-    main()

preprocessing/extract_crops.py

@@ -1,86 +0,0 @@
-import argparse
-import json
-import os
-from os import cpu_count
-from pathlib import Path
-
-os.environ["MKL_NUM_THREADS"] = "1"
-os.environ["NUMEXPR_NUM_THREADS"] = "1"
-os.environ["OMP_NUM_THREADS"] = "1"
-from functools import partial
-from glob import glob
-from multiprocessing.pool import Pool
-
-import cv2
-
-cv2.ocl.setUseOpenCL(False)
-cv2.setNumThreads(0)
-from tqdm import tqdm
-
-
-def extract_video(param, root_dir, crops_dir):
-    video, bboxes_path = param
-    with open(bboxes_path, "r") as bbox_f:
-        bboxes_dict = json.load(bbox_f)
-
-    capture = cv2.VideoCapture(video)
-    frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
-
-    for i in range(frames_num):
-        capture.grab()
-        if i % 10 != 0:
-            continue
-        success, frame = capture.retrieve()
-        if not success or str(i) not in bboxes_dict:
-            continue
-        id = os.path.splitext(os.path.basename(video))[0]
-        crops = []
-        bboxes = bboxes_dict[str(i)]
-        if bboxes is None:
-            continue
-        for bbox in bboxes:
-            xmin, ymin, xmax, ymax = [int(b * 2) for b in bbox]
-            w = xmax - xmin
-            h = ymax - ymin
-            p_h = h // 3
-            p_w = w // 3
-            crop = frame[max(ymin - p_h, 0):ymax + p_h, max(xmin - p_w, 0):xmax + p_w]
-            h, w = crop.shape[:2]
-            crops.append(crop)
-        img_dir = os.path.join(root_dir, crops_dir, id)
-        os.makedirs(img_dir, exist_ok=True)
-        for j, crop in enumerate(crops):
-            cv2.imwrite(os.path.join(img_dir, "{}_{}.png".format(i, j)), crop)
-
-
-def get_video_paths(root_dir):
-    paths = []
-    for json_path in glob(os.path.join(root_dir, "*/metadata.json")):
-        dir = Path(json_path).parent
-        with open(json_path, "r") as f:
-            metadata = json.load(f)
-        for k, v in metadata.items():
-            original = v.get("original", None)
-            if not original:
-                original = k
-            bboxes_path = os.path.join(root_dir, "boxes", original[:-4] + ".json")
-            if not os.path.exists(bboxes_path):
-                continue
-            paths.append((os.path.join(dir, k), bboxes_path))
-
-    return paths
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(
-        description="Extracts crops from video")
-    parser.add_argument("--root-dir", help="root directory")
-    parser.add_argument("--crops-dir", help="crops directory")
-
-    args = parser.parse_args()
-    os.makedirs(os.path.join(args.root_dir, args.crops_dir), exist_ok=True)
-    params = get_video_paths(args.root_dir)
-    with Pool(processes=cpu_count()) as p:
-        with tqdm(total=len(params)) as pbar:
-            for v in p.imap_unordered(partial(extract_video, root_dir=args.root_dir, crops_dir=args.crops_dir), params):
-                pbar.update()

preprocessing/extract_images.py

@@ -1,42 +0,0 @@
-import argparse
-import os
-os.environ["MKL_NUM_THREADS"] = "1"
-os.environ["NUMEXPR_NUM_THREADS"] = "1"
-os.environ["OMP_NUM_THREADS"] = "1"
-from functools import partial
-from glob import glob
-from multiprocessing.pool import Pool
-from os import cpu_count
-
-import cv2
-cv2.ocl.setUseOpenCL(False)
-cv2.setNumThreads(0)
-from tqdm import tqdm
-
-
-def extract_video(video, root_dir):
-    capture = cv2.VideoCapture(video)
-    frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
-
-    for i in range(frames_num):
-        capture.grab()
-        success, frame = capture.retrieve()
-        if not success:
-            continue
-        id = os.path.splitext(os.path.basename(video))[0]
-        cv2.imwrite(os.path.join(root_dir, "jpegs", "{}_{}.jpg".format(id, i)), frame, [cv2.IMWRITE_JPEG_QUALITY, 100])
-
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(
-        description="Extracts jpegs from video")
-    parser.add_argument("--root-dir", help="root directory")
-
-    args = parser.parse_args()
-    os.makedirs(os.path.join(args.root_dir, "jpegs"), exist_ok=True)
-    videos = [video_path for video_path in glob(os.path.join(args.root_dir, "*/*.mp4"))]
-    with Pool(processes=cpu_count() - 2) as p:
-        with tqdm(total=len(videos)) as pbar:
-            for v in p.imap_unordered(partial(extract_video, root_dir=args.root_dir), videos):
-                pbar.update()

preprocessing/face_detector.py

@@ -1,72 +0,0 @@
-import os
-os.environ["MKL_NUM_THREADS"] = "1"
-os.environ["NUMEXPR_NUM_THREADS"] = "1"
-os.environ["OMP_NUM_THREADS"] = "1"
-
-from abc import ABC, abstractmethod
-from collections import OrderedDict
-from typing import List
-
-import cv2
-cv2.ocl.setUseOpenCL(False)
-cv2.setNumThreads(0)
-
-from PIL import Image
-from facenet_pytorch.models.mtcnn import MTCNN
-from torch.utils.data import Dataset
-
-
-class VideoFaceDetector(ABC):
-
-    def __init__(self, **kwargs) -> None:
-        super().__init__()
-
-    @property
-    @abstractmethod
-    def _batch_size(self) -> int:
-        pass
-
-    @abstractmethod
-    def _detect_faces(self, frames) -> List:
-        pass
-
-
-class FacenetDetector(VideoFaceDetector):
-
-    def __init__(self, device="cuda:0") -> None:
-        super().__init__()
-        self.detector = MTCNN(margin=0,thresholds=[0.85, 0.95, 0.95], device=device)
-
-    def _detect_faces(self, frames) -> List:
-        batch_boxes, *_ = self.detector.detect(frames, landmarks=False)
-        return [b.tolist() if b is not None else None for b in batch_boxes]
-
-    @property
-    def _batch_size(self):
-        return 32
-
-
-class VideoDataset(Dataset):
-
-    def __init__(self, videos) -> None:
-        super().__init__()
-        self.videos = videos
-
-    def __getitem__(self, index: int):
-        video = self.videos[index]
-        capture = cv2.VideoCapture(video)
-        frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
-        frames = OrderedDict()
-        for i in range(frames_num):
-            capture.grab()
-            success, frame = capture.retrieve()
-            if not success:
-                continue
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            frame = Image.fromarray(frame)
-            frame = frame.resize(size=[s // 2 for s in frame.size])
-            frames[i] = frame
-        return video, list(frames.keys()), list(frames.values())
-
-    def __len__(self) -> int:
-        return len(self.videos)

preprocessing/face_encodings.py

@@ -1,55 +0,0 @@
-import argparse
-import os
-from functools import partial
-from multiprocessing.pool import Pool
-
-from tqdm import tqdm
-
-from preprocessing.utils import get_original_video_paths
-
-os.environ["MKL_NUM_THREADS"] = "1"
-os.environ["NUMEXPR_NUM_THREADS"] = "1"
-os.environ["OMP_NUM_THREADS"] = "1"
-
-import random
-
-import face_recognition
-import numpy as np
-
-
-def write_face_encodings(video, root_dir):
-    video_id, *_ = os.path.splitext(video)
-    crops_dir = os.path.join(root_dir, "crops", video_id)
-    if not os.path.exists(crops_dir):
-        return
-    crop_files = [f for f in os.listdir(crops_dir) if f.endswith("jpg")]
-    if crop_files:
-        crop_files = random.sample(crop_files, min(10, len(crop_files)))
-        encodings = []
-        for crop_file in crop_files:
-            img = face_recognition.load_image_file(os.path.join(crops_dir, crop_file))
-            encoding = face_recognition.face_encodings(img, num_jitters=10)
-            if encoding:
-                encodings.append(encoding[0])
-        np.save(os.path.join(crops_dir, "encodings"), encodings)
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Extract 10 crops encodings for each video")
-    parser.add_argument("--root-dir", help="root directory", default="/home/selim/datasets/deepfake")
-    args = parser.parse_args()
-    return args
-
-
-def main():
-    args = parse_args()
-    originals = get_original_video_paths(args.root_dir, basename=True)
-    with Pool(processes=os.cpu_count() - 4) as p:
-        with tqdm(total=len(originals)) as pbar:
-            for v in p.imap_unordered(partial(write_face_encodings, root_dir=args.root_dir), originals):
-                pbar.update()
-
-
-if __name__ == '__main__':
-    main()

preprocessing/generate_diffs.py

@@ -1,73 +0,0 @@
-import argparse
-import os
-
-os.environ["MKL_NUM_THREADS"] = "1"
-os.environ["NUMEXPR_NUM_THREADS"] = "1"
-os.environ["OMP_NUM_THREADS"] = "1"
-from skimage.measure import compare_ssim
-
-from functools import partial
-from multiprocessing.pool import Pool
-
-from tqdm import tqdm
-
-from preprocessing.utils import get_original_with_fakes
-
-import cv2
-
-cv2.ocl.setUseOpenCL(False)
-cv2.setNumThreads(0)
-
-import numpy as np
-
-cache = {}
-
-
-def save_diffs(pair, root_dir):
-    ori_id, fake_id = pair
-    ori_dir = os.path.join(root_dir, "crops", ori_id)
-    fake_dir = os.path.join(root_dir, "crops", fake_id)
-    diff_dir = os.path.join(root_dir, "diffs", fake_id)
-    os.makedirs(diff_dir, exist_ok=True)
-    for frame in range(320):
-        if frame % 10 != 0:
-            continue
-        for actor in range(2):
-            image_id = "{}_{}.png".format(frame, actor)
-            diff_image_id = "{}_{}_diff.png".format(frame, actor)
-            ori_path = os.path.join(ori_dir, image_id)
-            fake_path = os.path.join(fake_dir, image_id)
-            diff_path = os.path.join(diff_dir, diff_image_id)
-            if os.path.exists(ori_path) and os.path.exists(fake_path):
-                img1 = cv2.imread(ori_path, cv2.IMREAD_COLOR)
-                img2 = cv2.imread(fake_path, cv2.IMREAD_COLOR)
-                try:
-                    d, a = compare_ssim(img1, img2, multichannel=True, full=True)
-                    a = 1 - a
-                    diff = (a * 255).astype(np.uint8)
-                    diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
-                    cv2.imwrite(diff_path, diff)
-                except:
-                    pass
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Extract image diffs")
-    parser.add_argument("--root-dir", help="root directory", default="/mnt/sota/datasets/deepfake")
-    args = parser.parse_args()
-    return args
-
-
-def main():
-    args = parse_args()
-    pairs = get_original_with_fakes(args.root_dir)
-    os.makedirs(os.path.join(args.root_dir, "diffs"), exist_ok=True)
-    with Pool(processes=os.cpu_count() - 2) as p:
-        with tqdm(total=len(pairs)) as pbar:
-            func = partial(save_diffs, root_dir=args.root_dir)
-            for v in p.imap_unordered(func, pairs):
-                pbar.update()
-
-
-if __name__ == '__main__':
-    main()

preprocessing/generate_folds.py

@@ -1,114 +0,0 @@
-import argparse
-import json
-import os
-import random
-from functools import partial
-from multiprocessing.pool import Pool
-from pathlib import Path
-
-os.environ["MKL_NUM_THREADS"] = "1"
-os.environ["NUMEXPR_NUM_THREADS"] = "1"
-os.environ["OMP_NUM_THREADS"] = "1"
-import pandas as pd
-
-from tqdm import tqdm
-
-from preprocessing.utils import get_original_with_fakes
-
-import cv2
-
-cv2.ocl.setUseOpenCL(False)
-cv2.setNumThreads(0)
-
-
-def get_paths(vid, label, root_dir):
-    ori_vid, fake_vid = vid
-    ori_dir = os.path.join(root_dir, "crops", ori_vid)
-    fake_dir = os.path.join(root_dir, "crops", fake_vid)
-    data = []
-    for frame in range(320):
-        if frame % 10 != 0:
-            continue
-        for actor in range(2):
-            image_id = "{}_{}.png".format(frame, actor)
-            ori_img_path = os.path.join(ori_dir, image_id)
-            fake_img_path = os.path.join(fake_dir, image_id)
-            img_path = ori_img_path if label == 0 else fake_img_path
-            try:
-                # img = cv2.imread(img_path)[..., ::-1]
-                if os.path.exists(img_path):
-                    data.append([img_path, label, ori_vid])
-            except:
-                pass
-    return data
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Generate Folds")
-    parser.add_argument("--root-dir", help="root directory", default="/mnt/sota/datasets/deepfake")
-    parser.add_argument("--out", type=str, default="folds02.csv", help="CSV file to save")
-    parser.add_argument("--seed", type=int, default=777, help="Seed to split, default 777")
-    parser.add_argument("--n_splits", type=int, default=16, help="Num folds, default 10")
-    args = parser.parse_args()
-
-    return args
-
-
-def main():
-    args = parse_args()
-    ori_fakes = get_original_with_fakes(args.root_dir)
-    sz = 50 // args.n_splits
-    folds = []
-    for fold in range(args.n_splits):
-        folds.append(list(range(sz * fold, sz * fold + sz if fold < args.n_splits - 1 else 50)))
-    print(folds)
-    video_fold = {}
-    for d in os.listdir(args.root_dir):
-        if "dfdc" in d:
-            part = int(d.split("_")[-1])
-            for f in os.listdir(os.path.join(args.root_dir, d)):
-                if "metadata.json" in f:
-                    with open(os.path.join(args.root_dir, d, "metadata.json")) as metadata_json:
-                        metadata = json.load(metadata_json)
-
-                    for k, v in metadata.items():
-                        fold = None
-                        for i, fold_dirs in enumerate(folds):
-                            if part in fold_dirs:
-                                fold = i
-                                break
-                        assert fold is not None
-                        video_id = k[:-4]
-                        video_fold[video_id] = fold
-    for fold in range(len(folds)):
-        holdoutset = {k for k, v in video_fold.items() if v == fold}
-        trainset = {k for k, v in video_fold.items() if v != fold}
-        assert holdoutset.isdisjoint(trainset), "Folds have leaks"
-    data = []
-    ori_ori = set([(ori, ori) for ori, fake in ori_fakes])
-    with Pool(processes=os.cpu_count()) as p:
-        with tqdm(total=len(ori_ori)) as pbar:
-            func = partial(get_paths, label=0, root_dir=args.root_dir)
-            for v in p.imap_unordered(func, ori_ori):
-                pbar.update()
-                data.extend(v)
-        with tqdm(total=len(ori_fakes)) as pbar:
-            func = partial(get_paths, label=1, root_dir=args.root_dir)
-            for v in p.imap_unordered(func, ori_fakes):
-                pbar.update()
-                data.extend(v)
-    fold_data = []
-    for img_path, label, ori_vid in data:
-        path = Path(img_path)
-        video = path.parent.name
-        file = path.name
-        assert video_fold[video] == video_fold[ori_vid], "original video and fake have leak  {} {}".format(ori_vid,
-                                                                                                           video)
-        fold_data.append([video, file, label, ori_vid, int(file.split("_")[0]), video_fold[video]])
-    random.shuffle(fold_data)
-    pd.DataFrame(fold_data, columns=["video", "file", "label", "original", "frame", "fold"]).to_csv(args.out, index=False)
-
-
-if __name__ == '__main__':
-    main()

preprocessing/generate_landmarks.py

@@ -1,75 +0,0 @@
-import argparse
-import os
-from functools import partial
-from multiprocessing.pool import Pool
-
-
-
-os.environ["MKL_NUM_THREADS"] = "1"
-os.environ["NUMEXPR_NUM_THREADS"] = "1"
-os.environ["OMP_NUM_THREADS"] = "1"
-
-from tqdm import tqdm
-
-
-import cv2
-
-cv2.ocl.setUseOpenCL(False)
-cv2.setNumThreads(0)
-from preprocessing.utils import get_original_video_paths
-
-from PIL import Image
-from facenet_pytorch.models.mtcnn import MTCNN
-import numpy as np
-
-detector = MTCNN(margin=0, thresholds=[0.65, 0.75, 0.75], device="cpu")
-
-
-def save_landmarks(ori_id, root_dir):
-    ori_id = ori_id[:-4]
-    ori_dir = os.path.join(root_dir, "crops", ori_id)
-    landmark_dir = os.path.join(root_dir, "landmarks", ori_id)
-    os.makedirs(landmark_dir, exist_ok=True)
-    for frame in range(320):
-        if frame % 10 != 0:
-            continue
-        for actor in range(2):
-            image_id = "{}_{}.png".format(frame, actor)
-            landmarks_id = "{}_{}".format(frame, actor)
-            ori_path = os.path.join(ori_dir, image_id)
-            landmark_path = os.path.join(landmark_dir, landmarks_id)
-
-            if os.path.exists(ori_path):
-                try:
-                    image_ori = cv2.imread(ori_path, cv2.IMREAD_COLOR)[...,::-1]
-                    frame_img = Image.fromarray(image_ori)
-                    batch_boxes, conf, landmarks = detector.detect(frame_img, landmarks=True)
-                    if landmarks is not None:
-                        landmarks = np.around(landmarks[0]).astype(np.int16)
-                        np.save(landmark_path, landmarks)
-                except Exception as e:
-                    print(e)
-                    pass
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Extract image landmarks")
-    parser.add_argument("--root-dir", help="root directory", default="/mnt/sota/datasets/deepfake")
-    args = parser.parse_args()
-    return args
-
-
-def main():
-    args = parse_args()
-    ids = get_original_video_paths(args.root_dir, basename=True)
-    os.makedirs(os.path.join(args.root_dir, "landmarks"), exist_ok=True)
-    with Pool(processes=os.cpu_count()) as p:
-        with tqdm(total=len(ids)) as pbar:
-            func = partial(save_landmarks, root_dir=args.root_dir)
-            for v in p.imap_unordered(func, ids):
-                pbar.update()
-
-
-if __name__ == '__main__':
-    main()

preprocessing/utils.py

@@ -1,51 +0,0 @@
-import json
-import os
-from glob import glob
-from pathlib import Path
-
-
-def get_original_video_paths(root_dir, basename=False):
-    originals = set()
-    originals_v = set()
-    for json_path in glob(os.path.join(root_dir, "*/metadata.json")):
-        dir = Path(json_path).parent
-        with open(json_path, "r") as f:
-            metadata = json.load(f)
-        for k, v in metadata.items():
-            original = v.get("original", None)
-            if v["label"] == "REAL":
-                original = k
-                originals_v.add(original)
-                originals.add(os.path.join(dir, original))
-    originals = list(originals)
-    originals_v = list(originals_v)
-    print(len(originals))
-    return originals_v if basename else originals
-
-
-def get_original_with_fakes(root_dir):
-    pairs = []
-    for json_path in glob(os.path.join(root_dir, "*/metadata.json")):
-        with open(json_path, "r") as f:
-            metadata = json.load(f)
-        for k, v in metadata.items():
-            original = v.get("original", None)
-            if v["label"] == "FAKE":
-                pairs.append((original[:-4], k[:-4] ))
-
-    return pairs
-
-
-def get_originals_and_fakes(root_dir):
-    originals = []
-    fakes = []
-    for json_path in glob(os.path.join(root_dir, "*/metadata.json")):
-        with open(json_path, "r") as f:
-            metadata = json.load(f)
-        for k, v in metadata.items():
-            if v["label"] == "FAKE":
-                fakes.append(k[:-4])
-            else:
-                originals.append(k[:-4])
-
-    return originals, fakes