{"cells":[{"metadata":{},"cell_type":"markdown","source":"Efficient Net Solution\n\n- Inference on 30 frames per video at 10 frame intervals"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nos.makedirs(\"/root/.cache/torch/hub\", exist_ok=True)\n\ntry:\n    os.remove(\"/root/.cache/torch/hub/rwightman_gen-efficientnet-pytorch_master\")\nexcept:\n    pass\n\nos.symlink(\"/kaggle/input/rwightman/rwightman_gen-efficientnet-pytorch_master\",\n           \"/root/.cache/torch/hub/rwightman_gen-efficientnet-pytorch_master\")\n\nimport sys\nsys.path.insert(0, \"/kaggle/input/retinaface\")\nsys.path.insert(0, \"/kaggle/input/efnetensembles8\")\n\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport torch\nimport torch.backends.cudnn as cudnn\nimport torchvision\nfrom torch import nn\nfrom torchvision import transforms\nfrom torchvision.ops import nms\nfrom data import cfg_mnet, cfg_re50\nfrom layers.functions.prior_box import PriorBox\nfrom utils.nms.py_cpu_nms import py_cpu_nms\nfrom models.retinaface import RetinaFace\nfrom utils.box_utils import decode, decode_landm\nfrom utils.timer import Timer\nfrom glob import glob\nfrom math import ceil\nfrom os.path import basename\nfrom tqdm.notebook import tqdm\nfrom time import time\nfrom PIL import Image\n\nGPU_ID = 0\n\nVIDEO_GRAB_FRAMES = 60\nVIDEO_MAX_SIZE = 960\n\nFACE_NMS_THRESHOLD = 0.4\nFACE_CONFIDENCE_THRESHOLD = 0.99\nFACE_SCALE = 1.1\n\nNO_FACES_BUT_TOTAL_FRAMES_IN_VIDEO_PROBA = 0.5 # Assume fake\nNO_FACES_BUT_NOT_TOTAL_FRAMES_IN_VIDEO_PROBA = 0.5 # No idea\n\n\nDEFAULT_PROBA = 0.5 # If all else fails\n\nIMAGE_SIZE = 300\nTORCH_DEVICE='cuda:0'\n\nSAMPLE_SIZE=0\n\nVIDEO_FILE_GLOB = \"/kaggle/input/deepfake-detection-challenge/test_videos/*.mp4\"\n\nFACE_MODEL_FILE = '/kaggle/input/retinaface/weights/Resnet50_Final.pth'\nINFERENCE_MODEL_FILES = (\n    \"/kaggle/input/efnetensembles8/effnet_b3_pretrained_on_30_40_validation_plus_augs_dropout_0.2__size_300-2-0.1661-ckpt.pth\",\n)\n\n\ndef process_video(video_file, grab_frames = VIDEO_GRAB_FRAMES, resize = VIDEO_MAX_SIZE, face_scale = FACE_SCALE, debug = False):\n    \n    filename = basename(video_file)\n        \n    if debug: print(f\"Processing {video_file}\")\n\n    frames = []\n        \n    try:\n        cap = cv2.VideoCapture(video_file)\n        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n        if frame_count < 1:\n            frame_count = 300\n            \n        get_frameids = np.linspace(0, frame_count - 1, grab_frames, endpoint=True, dtype=np.int)\n        \n        parsed_frame = 0\n        \n        while True:\n            success = cap.grab()\n            if not success:\n                break\n                \n            if parsed_frame in get_frameids:\n\n                success, frame = cap.retrieve()\n                if not success or frame is None:\n                    break\n\n                if len(frame.shape) == 2:\n                    frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)\n                elif len(frame.shape) == 3:\n                    channels = frame.shape[-1]\n                    if channels == 4:\n                        frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)\n                    elif channels > 4:\n                        frame = cv2.cvtColor(frame[:,:,0:3], cv2.COLOR_RGB2BGR)\n\n                frames.append([frame, parsed_frame, None])\n\n            parsed_frame += 1\n            \n        cap.release()\n    except:\n        pass\n\n    num_frames = len(frames)\n    \n    # Nothing to process\n    #\n    if num_frames < 1:\n        return filename, 0, 0, np.array([]), np.array([])\n    \n    # Resize frames before detection\n    resized_frames = []\n    \n    for i, (frame, frameid, _) in enumerate(frames):\n        resized_frame, downsample = resize_img(frame, resize)\n\n        if debug:\n            print(f\"Resized frame shape is {resized_frame.shape}\")\n\n        upsample = 1/downsample\n        if upsample < 1:\n            upsample = 1\n\n        resized_frames.append(resized_frame)\n        frames[i][2] = upsample # Update upsample ratio\n    \n    # Detect and classify faces\n    confidence_threshold = FACE_CONFIDENCE_THRESHOLD\n    nms_threshold = FACE_NMS_THRESHOLD\n    \n    global face_detector\n    global cfg_re50\n    global image_transform\n    global face_classifiers\n    global device\n    global faces\n        \n    imgs = np.float32(resized_frames)\n    batch_size, im_height, im_width, _ = imgs.shape\n\n    imgs -= (104, 117, 123)\n    imgs = torch.from_numpy(imgs)\n    imgs = imgs.permute(0, 3, 1, 2)\n    imgs = imgs.to(device)\n\n    scale = torch.Tensor([im_width, im_height, im_width, im_height])\n    scale = scale.to(device)\n    scale1 = torch.Tensor([im_width, im_height, im_width, im_height,\n                           im_width, im_height, im_width, im_height,\n                           im_width, im_height])\n\n    scale1 = scale1.to(device)\n\n    b_loc, b_conf, b_landms = face_detector(imgs)  # forward pass\n    \n    priorbox = PriorBox(cfg_re50, image_size=(im_height, im_width))\n    priors = priorbox.forward()\n    priors = priors.to(device)\n    prior_data = priors.data\n    \n    face_records = []\n    faces = []\n    max_faceid_file = 0\n\n    if debug: print(f\"Image shape is {imgs.shape}\")\n            \n    for i in range(len(imgs)):\n\n        frame, frameid, upsample = frames[i]\n        frame_height, frame_width, _ = frame.shape\n            \n        boxes = decode(b_loc[i].data.squeeze(0), prior_data, cfg_re50['variance'])\n        boxes = boxes * scale\n\n        scores = b_conf[i].squeeze(0)\n\n        landms = decode_landm(b_landms[i].data.squeeze(0), prior_data, cfg_re50['variance'])            \n        landms = landms * scale1\n\n        # ignore low scores\n        inds = scores[:,1] > confidence_threshold\n        boxes = boxes[inds]\n        landms = landms[inds]\n        scores = scores[inds][:,1]\n\n        keep = nms(boxes, scores, nms_threshold)\n        boxes = boxes[keep].cpu().numpy()\n        landms = landms[keep].cpu().numpy()\n        scores = scores[keep].cpu().numpy()\n\n        # Upsample\n        boxes = boxes * upsample\n        landms = landms * upsample\n        \n        faceid = 0\n        \n        for box, land, score in zip(boxes, landms, scores):\n            \n            score = round(score, 2)\n            \n            w1,h1,w2,h2 = map(int, box)\n\n            face = crop_and_align_face(w1, w2, h1, h2, frame, scale=face_scale)\n            face = image_transform(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))\n            \n            face_record = [frameid, faceid, score, *land]\n            face_records.append(face_record)\n            faces.append(face)\n\n            faceid += 1            \n            \n        if faceid > max_faceid_file:\n            max_faceid_file = faceid\n            \n    del imgs, scale, scale1\n    \n    if len(faces) < 1:\n        print(\"Found no faces\")\n\n        return filename, num_frames, max_faceid_file, np.array([]), np.array([])\n\n    probas = []\n    with torch.set_grad_enabled(False):\n        faces = torch.stack([i for i in faces]).to(device)\n        for face_classifier in face_classifiers:\n            probas.append(torch.sigmoid(face_classifier(faces).squeeze()).cpu().numpy())\n\n    return filename, num_frames, max_faceid_file, np.array(face_records), np.array(probas)\n\ndef check_keys(model, pretrained_state_dict):\n    ckpt_keys = set(pretrained_state_dict.keys())\n    model_keys = set(model.state_dict().keys())\n    used_pretrained_keys = model_keys & ckpt_keys\n    unused_pretrained_keys = ckpt_keys - model_keys\n    missing_keys = model_keys - ckpt_keys\n    print('Missing keys:{}'.format(len(missing_keys)))\n    print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys)))\n    print('Used keys:{}'.format(len(used_pretrained_keys)))\n    assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'\n    return True\n\ndef remove_prefix(state_dict, prefix):\n    ''' Old style model is stored with all names of parameters sharing common prefix 'module.' '''\n    print('remove prefix \\'{}\\''.format(prefix))\n    f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x\n    return {f(key): value for key, value in state_dict.items()}\n\ndef load_model(model, pretrained_path, load_to_cpu):\n    print('Loading pretrained model from {}'.format(pretrained_path))\n    if load_to_cpu:\n        pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)\n    else:\n        device = torch.cuda.current_device()\n        pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))\n    if \"state_dict\" in pretrained_dict.keys():\n        pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')\n    else:\n        pretrained_dict = remove_prefix(pretrained_dict, 'module.')\n    check_keys(model, pretrained_dict)\n    model.load_state_dict(pretrained_dict, strict=False)\n    return model\n\ndef resize_img(img, max_size):\n\n    im_shape = img.shape\n    im_size_max = np.max(im_shape[0:2])\n    downsample = float(max_size) / float(im_size_max)\n    \n    if downsample < 1:\n        img = cv2.resize(img, None, None, fx=downsample, fy=downsample, interpolation=cv2.INTER_LINEAR)\n        \n    return img, downsample\n\ndef crop_and_align_face(w1, w2, h1, h2, frame, scale=1.1):\n    size_bb = int(max(w2-w1, h2-h1) * scale)\n    height, width = frame.shape[:2]\n    center_w, center_h = (w1 + w2) // 2, (h1 + h2) // 2\n    # Check for out of bounds, x-y top left corner\n    w1 = max(int(center_w - size_bb // 2), 0)\n    h1 = max(int(center_h - size_bb // 2), 0)\n    size_bb = min(width - w1, size_bb)\n    size_bb = min(height - h1, size_bb)\n    square_face = frame[h1:h1+size_bb, w1:w1+size_bb]\n\n    return square_face\n\n## Main\n\nproba_by_filename = {}\n\nimage_transform = transforms.Compose([\n    transforms.ToPILImage(),\n    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE), Image.BICUBIC),\n    transforms.ToTensor(),\n    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))\n])\n\ncfg_re50['pretrain'] = False\n\nface_detector = RetinaFace(cfg=cfg_re50, phase='test')\nface_detector = load_model(face_detector, FACE_MODEL_FILE, False)\nprint('Finished loading model!')\ncudnn.benchmark = True\ndevice = torch.device(\"cuda\")\nface_detector = face_detector.to(device)\nfor param in face_detector.parameters():\n    param.requires_grad = False\nface_detector.eval()\n\ndevice = torch.device(TORCH_DEVICE)\n\nface_classifiers = []\n\nfor i, model_file in enumerate(INFERENCE_MODEL_FILES):\n\n    face_classifiers.append(\n        nn.Sequential(\n            torch.hub.load('rwightman/gen-efficientnet-pytorch', 'tf_efficientnet_b3_ns', pretrained=False),\n            nn.Dropout(0.0),\n            nn.Linear(1000, 1)\n    ))\n    face_classifiers[i].load_state_dict(torch.load(model_file, map_location={'cuda:0':TORCH_DEVICE})['state_dict'])\n    face_classifiers[i] = face_classifiers[i].to(device)\n    face_classifiers[i].eval()\n    print(f\"Loaded model from {model_file}\")\n\nvideo_files = glob(VIDEO_FILE_GLOB)\n\nif SAMPLE_SIZE > 0:\n    video_files = video_files[0:SAMPLE_SIZE]\n    print(f\"Sampled {len(video_files)} video files\")\n\n\nprint(f\"Found {len(video_files)} video files\")\n\n###########################\n# Process Videos\n#\n\nvideo_res = {}\n\npbar = tqdm(enumerate(video_files), total=len(video_files), dynamic_ncols=True)\n\nfor i, video_file in pbar:\n\n    filename, num_frames, max_faceid_file, face_records, probabilities = process_video(video_file)\n    \n    no_faces = len(probabilities) < 1\n    multiple_faces = max_faceid_file > 1\n    \n    video_res[filename] = {'fr': face_records,\n                           'probabilities': probabilities,\n                           'no_faces': no_faces,\n                           'num_frames': num_frames,\n                           'max_faces': max_faceid_file,\n                           'multiple_faces': multiple_faces}\n\n    video_proba = DEFAULT_PROBA\n    \n    if no_faces:\n        video_proba = NO_FACES_BUT_NOT_TOTAL_FRAMES_IN_VIDEO_PROBA\n        if num_frames == VIDEO_GRAB_FRAMES:\n            video_proba = NO_FACES_BUT_TOTAL_FRAMES_IN_VIDEO_PROBA\n    elif multiple_faces:\n        # Get probabilities by frame, pick max per frame then average all max\n        max_probs = []\n\n        for i in range(len(face_classifiers)):\n            dd = list(zip(video_res[filename]['fr'][:,0], video_res[filename]['probabilities'][i]))\n            myd = {}\n            for frameid, prob in dd:\n                if frameid not in myd:\n                    myd[frameid] = [prob]\n                else:\n                    myd[frameid].append(prob)\n\n            max_probs.append([])\n            for frameid in myd:\n                p = np.array(myd[frameid])\n                res = DEFAULT_PROBA\n                if (p > 0.5).all():\n                    res = np.max(p)\n                elif (p <= 0.5).all():\n                    res = np.min(p)\n                max_probs[i].append(res)\n\n        video_proba = np.mean(max_probs)\n    else:\n        video_proba = video_res[filename]['probabilities'].mean()\n\n    video_res[filename]['prediction'] = video_proba\n\n###########################\n# Submission code\n#\n\nsubmission_videos = []\nsubmission_predictions = []\n\nfor video_file in video_files:\n    basefile = basename(video_file)\n    proba = DEFAULT_PROBA\n    if basefile in video_res:\n        proba = video_res[basefile]['prediction']\n    submission_videos.append(basefile)\n    submission_predictions.append(proba)\n    \nsubmission_df = pd.DataFrame({\"filename\": submission_videos, \"label\": submission_predictions})\nsubmission_df.to_csv(\"submission.csv\", index=False)\n\nprint(\"Submit Results\")","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}