{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os, sys, time\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"test_dir = \"/kaggle/input/deepfake-detection-challenge/test_videos/\"\n\ntest_videos = sorted([x for x in os.listdir(test_dir) if x[-4:] == \".mp4\"])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gpu = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.insert(0, \"/kaggle/input/xxxxxx/pytorchretinaface\")\nsys.path.insert(0, \"/kaggle/input/zzzzzz/mymodel\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from __future__ import print_function\nimport os\nimport argparse\nimport torch\nimport torch.backends.cudnn as cudnn\nimport numpy as np\nfrom PIL import Image\nfrom skimage import transform as trans\nfrom data import cfg_mnet, cfg_re50\nfrom layers.functions.prior_box import PriorBox\nfrom utils.nms.py_cpu_nms import py_cpu_nms\nimport cv2\nfrom models.retinaface import RetinaFace\nfrom utils.box_utils import decode, decode_landm\nimport time\n\narcface_src = np.array([\n  [122.5, 141.25],\n  [197.5, 141.25],\n  [160., 178.75],\n  [137.5, 225.25],\n  [182.5, 225.25] ], dtype=np.float32 ) # Ziyu\narcface_src = np.expand_dims(arcface_src, axis=0)\n\ndef estimate_norm(lmk, image_size = 112, mode='arcface'):\n    assert lmk.shape==(5,2)\n    tform = trans.SimilarityTransform()\n    lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)\n    min_M = []\n    min_index = []\n    min_error = float('inf') \n    if mode=='arcface':\n        src = arcface_src\n    else:\n        src = src_map[image_size]\n    for i in np.arange(src.shape[0]):\n        tform.estimate(lmk, src[i])\n    M = tform.params[0:2,:]\n    results = np.dot(M, lmk_tran.T)\n    results = results.T\n    error = np.sum(np.sqrt(np.sum((results - src[i]) ** 2,axis=1)))\n#         print(error)\n    if error< min_error:\n        min_error = error\n        min_M = M\n        min_index = i\n    return min_M, min_index\n\ndef norm_crop(img, landmark, image_size=112, mode='arcface'):\n    M, pose_index = estimate_norm(landmark, image_size, mode)\n    warped = cv2.warpAffine(img,M, (image_size, image_size), borderValue = 0.0)\n    return warped\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\n\ndef remove_prefix(state_dict, prefix):\n    ''' Old style model is stored with all names of parameters sharing common prefix 'module.' '''\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\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\ntorch.set_grad_enabled(False)\ncfg = cfg_re50\ncfg['pretrain']=False\n# net and model\n!cp /kaggle/input/xxxxxx/pytorchretinaface/weights ./ -rf\nnet = RetinaFace(cfg=cfg, phase = 'test')\nnet = load_model(net, './weights/Resnet50_Final.pth', False)\nnet.eval()\nprint('Finished loading model!')\ncudnn.benchmark = True\ndevice = torch.device(\"cuda\")\nnet = net.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision import transforms\nfrom model import WSDAN\nfrom util import  batch_augment\ntts=transforms.Compose([\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.4479, 0.3744, 0.3473],std=[0.2537, 0.2502, 0.2424])\n        ])\n\n\ndef extract_frames(data_path, method='cv2'):\n    \"\"\"Method to extract frames, either with ffmpeg or opencv. FFmpeg won't\n    start from 0 so we would have to rename if we want to keep the filenames\n    coherent.\"\"\"\n    data_path=test_dir+data_path\n    if method == 'cv2':\n        reader = cv2.VideoCapture(data_path)\n        outputbuff=[]\n        frames = 0\n        count = 0\n        resize=1\n        while reader.isOpened():\n            success, img = reader.read()\n            img_raw = img\n            if not success:\n                break\n            \n\n            frames +=1\n            if frames==1:\n                im_height, im_width, _ = img.shape\n                scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])\n                scale = scale.to(device)\n                priorbox = PriorBox(cfg, image_size=(im_height, im_width))\n                priors = priorbox.forward()\n                priors = priors.to(device)\n                prior_data = priors.data\n                scale1 = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0],\n                               img.shape[1], img.shape[0], img.shape[1], img.shape[0],\n                               img.shape[1], img.shape[0]])\n                scale1 = scale1.to(device)\n            if frames%10==0:\n                img = img.astype(np.int8)\n                img -= (104, 117, 123)\n                img = img.transpose(2, 0, 1)\n                img = torch.from_numpy(img).unsqueeze(0)\n                img=img.to(device,dtype=torch.float32)\n                loc, conf, landms = net(img)\n                boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])\n                boxes = boxes * scale / resize\n                boxes = boxes.cpu().numpy()\n                scores = conf.squeeze(0).data.cpu().numpy()[:, 1]\n                landms = decode_landm(landms.data.squeeze(0), prior_data, cfg['variance'])\n                landms = landms * scale1 / resize\n                landms = landms.cpu().numpy()\n                inds = np.where(scores > 0.8)[0]\n                if inds.shape[0]==0:\n                    continue\n                boxes = boxes[inds]\n                landms = landms[inds]\n                scores = scores[inds]\n                areas = scores\n                for it in range(areas.shape[0]):\n                    areas[it] = (boxes[it][3]-boxes[it][1])*(boxes[it][2]-boxes[it][0])\n                order = areas.argsort()[::-1][:1]\n                boxes = boxes[order]\n                landms = landms[order]\n                scores = scores[order]\n                landmarks = landms.reshape(5,2).astype(np.int)\n                img=norm_crop(img_raw,landmarks,image_size=320)\n                aligned=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\n                outputbuff.append(aligned)\n                count+=1\n                if count==20:\n                    break\n\n        reader.release()\n        return outputbuff\n\ndef predict_on_video_set(video_paths,model,model2):\n    predictions = []\n    for num in range(len(video_paths)):\n        try:\n            stime=time.time()\n            frames=extract_frames( video_paths[num])\n            frames=torch.cat([tts(i).unsqueeze(0) for i in  frames])\n            images=frames.view(-1,3,320,320).cuda()\n            print(time.time()-stime)\n            y_pred_raw, _,_ = model(images)\n            logits=torch.mean(F.softmax(y_pred_raw,dim=1),dim=0)\n            pred=logits[1].item()\n            images_b=F.interpolate(images,size=300,mode='bilinear')\n            y_pred_raw2, _,_ = model2(images_b)\n            logits2=torch.mean(F.softmax(y_pred_raw2,dim=1),dim=0)\n            pred2=logits2[1].item()\n            pred=(pred*0.7+pred2*0.3)\n            if pred>0.99:\n                pred=0.99\n            if pred<0.01:\n                pred=0.01\n            predictions.append(pred)\n            print(time.time()-stime,pred)\n        except Exception as e:\n            print(e)\n            predictions.append(0.5)            \n    return predictions","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"modelx=WSDAN(num_classes=2, M=8, net='xception', pretrained=False).cuda()\nmodelx.load_state_dict(torch.load('/kaggle/input/zzzzzz/mymodel/ckpt_x.pth')['state_dict'])\nmodely=WSDAN(num_classes=2, M=8, net='efficientnet', pretrained=False).cuda()\nmodely.load_state_dict(torch.load('/kaggle/input/zzzzzz/mymodel/ckpt_e.pth')['state_dict'])\nmodelx.eval()\nmodely.eval()\npredictions = predict_on_video_set(test_videos,modelx,modely)\nsubmission_df = pd.DataFrame({\"filename\": test_videos, \"label\": predictions})\nsubmission_df.to_csv(\"submission.csv\", index=False)","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}