{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Information\n\nBasically, we ensembled three models:\n* a WS-DAN w/ Xception feature extractor (70% weight)\n  - This is the best single model we have. Maybe only use this one is enough.\n* a WS-DAN w/ EfficientNet-b3 feature extractor (10% weight)\n  - This one is semi-finished. We have no time to train it well.\n* a very common Xception classifier (20% weight)\n  - We kept it just because we have no other good model to ensemble.\n  \nWS-DAN is the core part of our method.\n\nTraining code and guideline for reproducing can be found [on GitHub](https://github.com/cuihaoleo/kaggle-dfdc)."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport sys\nimport time\n\nimport cv2\nimport numpy as np\nfrom collections import defaultdict\n\nfrom PIL import Image\nimport torch\nimport torch.nn.functional as F\nfrom torchvision import transforms as T\n\nsys.path.insert(0, \"../input/dfdc-pretrained-2\")\nsys.path.insert(0, \"../input/dfdc-pretrained-2/Pytorch_Retinaface\")\nfrom face_utils import norm_crop, FaceDetector\nfrom model_def import WSDAN, xception","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class DFDCLoader:\n    def __init__(self, video_dir, face_detector, transform=None,\n                 batch_size=25, frame_skip=9, face_limit=25):\n        self.video_dir = video_dir\n        self.file_list = sorted(f for f in os.listdir(video_dir) if f.endswith(\".mp4\"))\n\n        self.transform = transform\n        self.face_detector = face_detector\n\n        self.batch_size = batch_size\n        self.frame_skip = frame_skip\n        self.face_limit = face_limit\n\n        self.record = defaultdict(list)\n        self.score = defaultdict(lambda: 0.5)\n        self.feedback_queue = []\n\n    def iter_one_face(self):\n        for fname in self.file_list:\n            path = os.path.join(self.video_dir, fname)\n            reader = cv2.VideoCapture(path)\n            face_count = 0\n\n            while True:\n                for _ in range(self.frame_skip):\n                    reader.grab()\n\n                success, img = reader.read()\n                if not success:\n                    break\n\n                boxes, landms = self.face_detector.detect(img)\n                if boxes.shape[0] == 0:\n                    continue\n\n                areas = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n                order = areas.argmax()\n\n                boxes = boxes[order]\n                landms = landms[order]\n\n                # Crop faces\n                landmarks = landms.numpy().reshape(5, 2).astype(np.int)\n                img = norm_crop(img, landmarks, image_size=320)\n                aligned = Image.fromarray(img[:, :, ::-1])\n\n                if self.transform:\n                    aligned = self.transform(aligned)\n\n                yield fname, aligned\n\n                # Early stop\n                face_count += 1\n                if face_count == self.face_limit:\n                    break\n\n            reader.release()\n\n    def __iter__(self):\n        self.record.clear()\n        self.feedback_queue.clear()\n\n        batch_buf = []\n        t0 = time.time()\n        batch_count = 0\n\n        for fname, face in self.iter_one_face():\n            self.feedback_queue.append(fname)\n            batch_buf.append(face)\n\n            if len(batch_buf) == self.batch_size:\n                yield torch.stack(batch_buf)\n\n                batch_count += 1\n                batch_buf.clear()\n\n                if batch_count % 10 == 0:\n                    elapsed = 1000 * (time.time() - t0)\n                    print(\"T: %.2f ms / batch\" % (elapsed / batch_count))\n\n        if len(batch_buf) > 0:\n            yield torch.stack(batch_buf)\n\n    def feedback(self, pred):\n        accessed = set()\n\n        for score in pred:\n            fname = self.feedback_queue.pop(0)\n            accessed.add(fname)\n            self.record[fname].append(score)\n\n        for fname in sorted(accessed):\n            self.score[fname] = np.mean(self.record[fname])\n            print(\"[%s] %.6f\" % (fname, self.score[fname]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def main():\n    torch.set_grad_enabled(False)\n    torch.backends.cudnn.benchmark = True\n\n    test_dir = \"../input/deepfake-detection-challenge/test_videos\"\n    csv_path = \"../input/deepfake-detection-challenge/sample_submission.csv\"\n\n    face_detector = FaceDetector()\n    face_detector.load_checkpoint(\"../input/dfdc-pretrained-2/RetinaFace-Resnet50-fixed.pth\")\n    loader = DFDCLoader(test_dir, face_detector, T.ToTensor())\n\n    model1 = xception(num_classes=2, pretrained=False)\n    ckpt = torch.load(\"../input/dfdc-pretrained-2/xception-hg-2.pth\")\n    model1.load_state_dict(ckpt[\"state_dict\"])\n    model1 = model1.cuda()\n    model1.eval()\n\n    model2 = WSDAN(num_classes=2, M=8, net=\"xception\", pretrained=False).cuda()\n    ckpt = torch.load(\"../input/dfdc-pretrained-2/ckpt_x.pth\")\n    model2.load_state_dict(ckpt[\"state_dict\"])\n    model2.eval()\n\n    model3 = WSDAN(num_classes=2, M=8, net=\"efficientnet\", pretrained=False).cuda()\n    ckpt = torch.load(\"../input/dfdc-pretrained-2/ckpt_e.pth\")\n    model3.load_state_dict(ckpt[\"state_dict\"])\n    model3.eval()\n\n    zhq_nm_avg = torch.Tensor([.4479, .3744, .3473]).view(1, 3, 1, 1).cuda()\n    zhq_nm_std = torch.Tensor([.2537, .2502, .2424]).view(1, 3, 1, 1).cuda()\n\n    for batch in loader:\n        batch = batch.cuda(non_blocking=True)\n\n        i1 = F.interpolate(batch, size=299, mode=\"bilinear\")\n        i1.sub_(0.5).mul_(2.0)\n        o1 = model1(i1).softmax(-1)[:, 1].cpu().numpy()\n\n        i2 = (batch - zhq_nm_avg) / zhq_nm_std\n        o2, _, _ = model2(i2)\n        o2 = o2.softmax(-1)[:, 1].cpu().numpy()\n\n        i3 = F.interpolate(i2, size=300, mode=\"bilinear\")\n        o3, _, _ = model3(i3)\n        o3 = o3.softmax(-1)[:, 1].cpu().numpy()\n\n        out = 0.2 * o1 + 0.7 * o2 + 0.1 * o3\n        loader.feedback(out)\n\n    with open(csv_path) as fin, open(\"submission.csv\", \"w\") as fout:\n        fout.write(next(fin))\n        for line in fin:\n            fname = line.split(\",\", 1)[0]\n            pred = loader.score[fname]\n            print(\"%s,%.6f\" % (fname, pred), file=fout)\n\nmain()","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}