{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install /kaggle/input/deepfake-detection-jph/facenet_pytorch-2.0.1-py3-none-any.whl\n!cp -r /kaggle/input/pretrainedmodels /kaggle/working/; cd /kaggle/working/pretrainedmodels/pretrained-models.pytorch ; pip install .\n#!pip install /kaggle/input/deepfake-detection-jph/efficientnet_pytorch-0.6.1/efficientnet_pytorch-0.6.1","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-05-31T08:17:41.472387Z","iopub.execute_input":"2022-05-31T08:17:41.472695Z","iopub.status.idle":"2022-05-31T08:18:28.10503Z","shell.execute_reply.started":"2022-05-31T08:17:41.472645Z","shell.execute_reply":"2022-05-31T08:18:28.104248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport math\nimport pickle\nfrom functools import partial\nfrom collections import defaultdict\n\nfrom PIL import Image\nfrom glob import glob\n\nimport cv2\nimport numpy as np\nimport skimage.measure\nimport albumentations as A\nfrom tqdm.notebook import tqdm \nfrom albumentations.pytorch import ToTensor \n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.autograd import Variable\nfrom torchvision.models.video import mc3_18, r2plus1d_18\n\nfrom facenet_pytorch import MTCNN","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-05-31T08:18:40.112531Z","iopub.execute_input":"2022-05-31T08:18:40.113111Z","iopub.status.idle":"2022-05-31T08:18:43.016112Z","shell.execute_reply.started":"2022-05-31T08:18:40.112903Z","shell.execute_reply":"2022-05-31T08:18:43.015408Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Input data\nINPUT_DIR = \"/kaggle/input/deepfake-detection-challenge/test_videos\"\nPRETRAINED_MODELS_3D = [{'type':'i3d',\n                         'path':\"/kaggle/input/deepfake-detection-jph/j3d_e1_l0.1374.model\"},\n                        {'type':'res34',\n                         'path':\"/kaggle/input/deepfake-detection-jph/res34_1cy_minaug_nonorm_e4_l0.1794.model\"},\n                        {'type':'mc3_112',\n                         'path':\"/kaggle/input/deepfake-detection-jph/mc3_18_112_1cy_lilaug_nonorm_e9_l0.1905.model\"},\n                        {'type':'mc3_224',\n                         'path':\"/kaggle/input/deepfake-detection-jph/mc3_18_112t224_1cy_lilaug_nonorm_e7_l0.1901.model\"},\n                        {'type':'r2p1_112',\n                         'path':\"/kaggle/input/deepfake-detection-jph/r2p1_18_8_112tr_112te_e12_l0.1741.model\"},\n                        {'type':'i3d',\n                         'path':\"/kaggle/input/deepfake-detection-jph/i3dcutmix_e11_l0.1612.model\"},\n                        {'type':'r2p1_112',\n                         'path':\"/kaggle/input/deepfake-detection-jph/r2plus1dcutmix_112_e10_l0.1608.model\"}]\n\n# Face detection\nMAX_FRAMES_TO_LOAD = 100\nMIN_FRAMES_FOR_FACE = 30\nMAX_FRAMES_FOR_FACE = 100\nFACE_FRAMES = 10\nMAX_FACES_HIGHTHRESH = 5\nMAX_FACES_LOWTHRESH = 1\nFACEDETECTION_DOWNSAMPLE = 0.25\nMTCNN_THRESHOLDS = (0.8, 0.8, 0.9)  # Default [0.6, 0.7, 0.7]\nMTCNN_THRESHOLDS_RETRY = (0.5, 0.5, 0.5)\nMMTNN_FACTOR = 0.71  # Default 0.709 p\nTWO_FRAME_OVERLAP = False\n\n# Inference\nPROB_MIN, PROB_MAX = 0.001, 0.999\nREVERSE_PROBS = True\nDEFAULT_MISSING_PRED = 0.5\nUSE_FACE_FUNCTION = np.mean\n\n# 3D inference\nRATIO_3D = 1\nOUTPUT_FACE_SIZE = (256, 256)\nPRE_INFERENCE_CROP = (224, 224)\n\n# 2D\nRATIO_2D = 1","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:18:43.018027Z","iopub.execute_input":"2022-05-31T08:18:43.018304Z","iopub.status.idle":"2022-05-31T08:18:43.029263Z","shell.execute_reply.started":"2022-05-31T08:18:43.018261Z","shell.execute_reply":"2022-05-31T08:18:43.028525Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_video(filename, every_n_frames=None, specific_frames=None, to_rgb=True, rescale=None, inc_pil=False, max_frames=None):\n    \"\"\"Loads a video.\n    Called by:\n    \n    1) The finding faces algorithm where it pulls a frame every FACE_FRAMES frames up to MAX_FRAMES_TO_LOAD at a scale of FACEDETECTION_DOWNSAMPLE, and then half that if there's a CUDA memory error.\n    \n    2) The inference loop where it pulls EVERY frame up to a certain amount which it the last needed frame for each face for that video\"\"\"\n    \n    assert every_n_frames or specific_frames, \"Must supply either every n_frames or specific_frames\"\n    assert bool(every_n_frames) != bool(specific_frames), \"Supply either 'every_n_frames' or 'specific_frames', not both\"\n    \n    cap = cv2.VideoCapture(filename)\n    n_frames_in = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n    width_in = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n    height_in = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n    \n    if rescale:\n        rescale = rescale * 1920./np.max((width_in, height_in))\n    \n    width_out = int(width_in*rescale) if rescale else width_in\n    height_out = int(height_in*rescale) if rescale else height_in\n    \n    if max_frames:\n        n_frames_in = min(n_frames_in, max_frames)\n    \n    if every_n_frames:\n        specific_frames = list(range(0,n_frames_in,every_n_frames))\n    \n    n_frames_out = len(specific_frames)\n    \n    out_pil = []\n\n    out_video = np.empty((n_frames_out, height_out, width_out, 3), np.dtype('uint8'))\n\n    i_frame_in = 0\n    i_frame_out = 0\n    ret = True\n\n    while (i_frame_in < n_frames_in and ret):\n        \n        try:\n            try:\n        \n                if every_n_frames == 1:\n                    ret, frame_in = cap.read()  # Faster if reading all frames\n                else:\n                    ret = cap.grab()\n\n                    if i_frame_in not in specific_frames:\n                        i_frame_in += 1\n                        continue\n\n                    ret, frame_in = cap.retrieve()\n                    \n#                 print(f\"Reading frame {i_frame_in}\")\n\n                if rescale:\n                    frame_in = cv2.resize(frame_in, (width_out, height_out))\n                if to_rgb:\n                    frame_in = cv2.cvtColor(frame_in, cv2.COLOR_BGR2RGB)\n                \n            except Exception as e:\n                print(f\"Error for frame {i_frame_in} for video {filename}: {e}; using 0s\")\n                frame_in = np.zeros((height_out, width_out, 3))\n\n        \n            out_video[i_frame_out] = frame_in\n            i_frame_out += 1\n\n            if inc_pil:\n                try:  # https://www.kaggle.com/zaharch/public-test-errors\n                    pil_img = Image.fromarray(frame_in)\n                except Exception as e:\n                    print(f\"Using a blank frame for video {filename} frame {i_frame_in} as error {e}\")\n                    pil_img = Image.fromarray(np.zeros((224,224,3), dtype=np.uint8))  # Use a blank frame\n                out_pil.append(pil_img)\n\n            i_frame_in += 1\n            \n        except Exception as e:\n            print(f\"Error for file {filename}: {e}\")\n\n    cap.release()\n    \n    if inc_pil:\n        return out_video, out_pil, rescale\n    else:\n        return out_video, rescale\n\ndef get_roi_for_each_face(faces_by_frame, probs, video_shape, temporal_upsample, upsample=1):\n    # Create boolean face array\n    frames_video, rows_video, cols_video, channels_video = video_shape\n    frames_video = math.ceil(frames_video)\n    boolean_face_3d = np.zeros((frames_video, rows_video, cols_video), dtype=np.bool)  # Remove colour channel\n    proba_face_3d = np.zeros((frames_video, rows_video, cols_video)).astype('float32')\n    for i_frame, faces in enumerate(faces_by_frame):\n        if faces is not None:  # May not be a face in the frame\n            for i_face, face in enumerate(faces):\n                left, top, right, bottom = face\n                boolean_face_3d[i_frame, int(top):int(bottom), int(left):int(right)] = True\n                proba_face_3d[i_frame, int(top):int(bottom), int(left):int(right)] = probs[i_frame][i_face]\n                \n    # Replace blank frames if face(s) in neighbouring frames with overlap\n    for i_frame, frame in enumerate(boolean_face_3d):\n        if i_frame == 0 or i_frame == frames_video-1:  # Can't do this for 1st or last frame\n            continue\n        if True not in frame:\n            if TWO_FRAME_OVERLAP:\n                if i_frame > 1:\n                    pre_overlap = boolean_face_3d[i_frame-1] | boolean_face_3d[i_frame-2]\n                else:\n                    pre_overlap = boolean_face_3d[i_frame-1]\n                if i_frame < frames_video-2:\n                    post_overlap = boolean_face_3d[i_frame+1] | boolean_face_3d[i_frame+2]\n                else:\n                    post_overlap = boolean_face_3d[i_frame+1]\n                neighbour_overlap = pre_overlap & post_overlap\n            else:\n                neighbour_overlap = boolean_face_3d[i_frame-1] & boolean_face_3d[i_frame+1]\n            boolean_face_3d[i_frame] = neighbour_overlap\n\n    # Find faces through time\n    id_face_3d, n_faces = skimage.measure.label(boolean_face_3d, return_num=True)\n    region_labels, counts = np.unique(id_face_3d, return_counts=True)\n    # Get rid of background=0\n    region_labels, counts = region_labels[1:], counts[1:]\n    ###################\n    # DESCENDING SIZE #\n    ###################\n    descending_size = np.argsort(counts)[::-1]\n    labels_by_size = region_labels[descending_size]\n    ####################\n    # DESCENDING PROBS #\n    ####################\n    probs = [np.mean(proba_face_3d[id_face_3d == i_face]) for i_face in region_labels]\n    descending_probs = np.argsort(probs)[::-1]\n    labels_by_probs = region_labels[descending_probs]\n    # Iterate over faces in video\n    rois = [] ; face_maps = []\n    for i_face in labels_by_probs:#labels_by_size:\n        # Find the first and last frame containing the face\n        frames = np.where(np.any(id_face_3d == i_face, axis=(1, 2)) == True)\n        starting_frame, ending_frame = frames[0].min(), frames[0].max()\n\n        # Iterate over the frames with faces in and find the min/max cols/rows (bounding box)\n        cols, rows = [], []\n        for i_frame in range(starting_frame, ending_frame + 1):\n            rs = np.where(np.any(id_face_3d[i_frame] == i_face, axis=1) == True)\n            rows.append((rs[0].min(), rs[0].max()))\n            cs = np.where(np.any(id_face_3d[i_frame] == i_face, axis=0) == True)\n            cols.append((cs[0].min(), cs[0].max()))\n        frame_from, frame_to = starting_frame*temporal_upsample, ((ending_frame+1)*temporal_upsample)-1\n        rows_from, rows_to = np.array(rows)[:, 0].min(), np.array(rows)[:, 1].max()\n        cols_from, cols_to = np.array(cols)[:, 0].min(), np.array(cols)[:, 1].max()\n        \n        frame_to = min(frame_to, frame_from + MAX_FRAMES_FOR_FACE)\n        \n        if frame_to - frame_from >= MIN_FRAMES_FOR_FACE:\n            tmp_face_map = id_face_3d.copy()\n            tmp_face_map[tmp_face_map != i_face] = 0\n            tmp_face_map[tmp_face_map == i_face] = 1\n            face_maps.append(tmp_face_map[frame_from//temporal_upsample:frame_to//temporal_upsample+1])\n            rois.append(((frame_from, frame_to),\n                         (int(rows_from*upsample), int(rows_to*upsample)),\n                         (int(cols_from*upsample), int(cols_to*upsample))))\n            \n    return np.array(rois), face_maps","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:18:43.030949Z","iopub.execute_input":"2022-05-31T08:18:43.031568Z","iopub.status.idle":"2022-05-31T08:18:43.073698Z","shell.execute_reply.started":"2022-05-31T08:18:43.031247Z","shell.execute_reply":"2022-05-31T08:18:43.073088Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Find faces","metadata":{}},{"cell_type":"code","source":"def get_center(bbox):\n    x1, y1, x2, y2 = bbox\n    return (x1+x2)/2, (y1+y2)/2\n\n\ndef get_coords(faces_roi):\n    coords = np.argwhere(faces_roi == 1)\n    #print(coords)\n    if coords.shape[0] == 0:\n        return None\n    y1, x1 = coords[0]\n    y2, x2 = coords[-1]\n    return x1, y1, x2, y2\n\n\ndef interpolate_center(c1, c2, length):\n    x1, y1 = c1\n    x2, y2 = c2\n    xi, yi = np.linspace(x1, x2, length), np.linspace(y1, y2, length)\n    return np.vstack([xi, yi]).transpose(1,0)\n\n\ndef get_faces(faces_roi, upsample): \n    all_faces = []\n    rows = faces_roi[0].shape[1]\n    cols = faces_roi[0].shape[2]\n    for i in range(len(faces_roi)):\n        faces = np.asarray([get_coords(faces_roi[i][j]) for j in range(len(faces_roi[i]))])\n        if faces[0] is None:  faces[0] = faces[1]\n        if faces[-1] is None: faces[-1] = faces[-2]\n        if None in faces:\n            #print(faces)\n            raise Exception('This should not have happened ...')\n        all_faces.append(faces)\n\n    extracted_faces = []\n    for face in all_faces:\n        # Get max dim size\n        max_dim = np.concatenate([face[:,2]-face[:,0],face[:,3]-face[:,1]])\n        max_dim = np.percentile(max_dim, 90)\n        # Enlarge by 1.2\n        max_dim = int(max_dim * 1.2)\n        # Get center coords\n        centers = np.asarray([get_center(_) for _ in face])\n        # Interpolate\n        centers = np.vstack([interpolate_center(centers[i], centers[i+1], length=10) for i in range(len(centers)-1)]).astype('int')\n        x1y1 = centers - max_dim // 2\n        x2y2 = centers + max_dim // 2 \n        x1, y1 = x1y1[:,0], x1y1[:,1]\n        x2, y2 = x2y2[:,0], x2y2[:,1]\n        # If x1 or y1 is negative, turn it to 0\n        # Then add to x2 y2 or y2\n        x2[x1 < 0] -= x1[x1 < 0]\n        y2[y1 < 0] -= y1[y1 < 0]\n        x1[x1 < 0] = 0\n        y1[y1 < 0] = 0\n        # If x2 or y2 is too big, turn it to max image shape\n        # Then subtract from y1\n        y1[y2 > rows] += rows - y2[y2 > rows]\n        x1[x2 > cols] += cols - x2[x2 > cols]\n        y2[y2 > rows] = rows\n        x2[x2 > cols] = cols\n        vidface = np.asarray([[x1[_],y1[_],x2[_],y2[_]] for _,c in enumerate(centers)])\n        vidface = (vidface*upsample).astype('int')\n        extracted_faces.append(vidface)\n\n    return extracted_faces\n\ndef detect_face_with_mtcnn(mtcnn_model, pil_frames, facedetection_upsample, video_shape, face_frames):\n    boxes, _probs = mtcnn_model.detect(pil_frames, landmarks=False)\n    faces, faces_roi = get_roi_for_each_face(faces_by_frame=boxes, probs=_probs, video_shape=video_shape, temporal_upsample=face_frames, upsample=facedetection_upsample)\n    coords = [] if len(faces_roi) == 0 else get_faces(faces_roi, upsample=facedetection_upsample)\n    return faces, coords\n\ndef face_detection_wrapper(mtcnn_model, videopath, every_n_frames, facedetection_downsample, max_frames_to_load):\n    video, pil_frames, rescale = load_video(videopath, every_n_frames=every_n_frames, to_rgb=True, rescale=facedetection_downsample, inc_pil=True, max_frames=max_frames_to_load)\n    if len(pil_frames):\n        try:\n            faces, coords = detect_face_with_mtcnn(mtcnn_model=mtcnn, \n                                                   pil_frames=pil_frames, \n                                                   facedetection_upsample=1/rescale, \n                                                   video_shape=video.shape, \n                                                   face_frames=every_n_frames)\n        except RuntimeError:  # Out of CUDA RAM\n            print(f\"Failed to process {videopath} ! Downsampling x2 ...\")\n            video, pil_frames, rescale = load_video(videopath, every_n_frames=every_n_frames, to_rgb=True, rescale=facedetection_downsample/2, inc_pil=True, max_frames=max_frames_to_load)\n\n            try:\n                faces, coords = detect_face_with_mtcnn(mtcnn_model=mtcnn, \n                                       pil_frames=pil_frames, \n                                       facedetection_upsample=1/rescale, \n                                       video_shape=video.shape, \n                                       face_frames=every_n_frames)\n            except RuntimeError:\n                print(f\"Failed on downsample ! Skipping...\")\n                return [], []\n                \n    else:\n        print('Failed to fetch frames ! Skipping ...')\n        return [], []\n        \n    if len(faces) == 0:\n        print('Failed to find faces ! Upsampling x2 ...')\n        try:\n            video, pil_frames, rescale = load_video(videopath, every_n_frames=every_n_frames, to_rgb=True, rescale=facedetection_downsample*2, inc_pil=True, max_frames=max_frames_to_load)\n            faces, coords = detect_face_with_mtcnn(mtcnn_model=mtcnn, \n                                                   pil_frames=pil_frames, \n                                                   facedetection_upsample=1/rescale, \n                                                   video_shape=video.shape, \n                                                   face_frames=every_n_frames)\n        except Exception as e:\n            print(e)\n            return [], []\n    \n    return faces, coords","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:18:43.077051Z","iopub.execute_input":"2022-05-31T08:18:43.077272Z","iopub.status.idle":"2022-05-31T08:18:43.111354Z","shell.execute_reply.started":"2022-05-31T08:18:43.07723Z","shell.execute_reply":"2022-05-31T08:18:43.109475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"videopaths = sorted(glob(os.path.join(INPUT_DIR, \"*.mp4\")))\nprint(f'Found {len(videopaths)} videos !')\n\nmtcnn = MTCNN(margin=0, keep_all=True, post_process=False, select_largest=False, device='cuda:0', thresholds=MTCNN_THRESHOLDS, factor=MMTNN_FACTOR)\n\nfaces_by_videopath = {}\ncoords_by_videopath = {}\n\nfor i_video, videopath in enumerate(tqdm(videopaths)):\n    faces, coords = face_detection_wrapper(mtcnn, videopath, every_n_frames=FACE_FRAMES, facedetection_downsample=FACEDETECTION_DOWNSAMPLE, max_frames_to_load=MAX_FRAMES_TO_LOAD)\n            \n    if len(faces):\n        faces_by_videopath[videopath]  = faces[:MAX_FACES_HIGHTHRESH]\n        coords_by_videopath[videopath] = coords[:MAX_FACES_HIGHTHRESH]\n    else:\n        print(f\"Found no faces for {videopath} !\")\n\n        \ndel mtcnn\nimport gc\ngc.collect()\n\ndef load_state(net,ckpt):\n    sd=net.state_dict()\n    nd={}\n    goodmatch=True\n    for i in ckpt:\n        if i in sd and sd[i].shape==ckpt[i].shape:\n            nd[i]=ckpt[i]\n            #print(i)\n        else:\n            print('fail to load %s'%i)\n            goodmatch=False\n    net.load_state_dict(nd,strict=False)\n    return goodmatch\ndef main_worker(local_rank,world_size,rank_offset,config):\n    rank=local_rank+rank_offset\n    if rank==0:\n        logging.basicConfig(\n        filename=os.path.join('runs', config.name,'train.log'),\n        filemode='w',\n        format='%(asctime)s: %(levelname)s: [%(filename)s:%(lineno)d]: %(message)s',\n        level=logging.INFO)\n    warnings.filterwarnings(\"ignore\")\n    dist.init_process_group(backend='nccl', init_method='file://%s/.pytorch_distribute'%os.getcwd(),world_size=world_size, rank=rank)\n    if rank==0:\n        try:\n            os.remove('.pytorch_distribute')\n        except:\n            pass\n    numpy.random.seed(1234567)\n    torch.manual_seed(1234567)\n    torch.cuda.manual_seed(1234567)\n    torch.cuda.set_device(local_rank)\n    train_dataset = DeepfakeDataset(phase='train',**config.train_dataset)\n    validate_dataset=DeepfakeDataset(phase='val',**config.val_dataset)\n    train_sampler=torch.utils.data.distributed.DistributedSampler(train_dataset)\n    validate_sampler=torch.utils.data.distributed.DistributedSampler(validate_dataset)\n    train_loader=torch.utils.data.DataLoader(train_dataset, batch_size=config.batch_size,sampler=train_sampler,pin_memory=True,num_workers=config.workers)\n    validate_loader=torch.utils.data.DataLoader(validate_dataset, batch_size=config.batch_size,sampler=validate_sampler,pin_memory=True,num_workers=config.workers)\n    logs = {}\n    start_epoch = 0\n    net = netrunc(config.net,config.feature_layer,config.num_classes,config.dropout_rate,config.pretrained)   \n    net=nn.SyncBatchNorm.convert_sync_batchnorm(net).to(local_rank)\n    net = nn.parallel.DistributedDataParallel(net,device_ids=[local_rank],output_device=local_rank,find_unused_parameters=True)\n    optimizer = torch.optim.AdamW(net.parameters(), lr=config.learning_rate, betas=config.adam_betas, weight_decay=config.weight_decay)\n    scheduler=torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.scheduler_step, gamma=config.scheduler_gamma)\n    if config.ckpt:\n        loc = 'cuda:{}'.format(local_rank)\n        checkpoint = torch.load(config.ckpt, map_location=loc)\n        logs = checkpoint['logs']\n        start_epoch = int(logs['epoch'])+1\n        if load_state(net.module,checkpoint['state_dict']) and config.resume_optim:\n            optimizer.load_state_dict(checkpoint['optimizer_state'])\n            try:\n                scheduler.load_state_dict(checkpoint['scheduler_state'])\n            except:\n                pass\n\n        \n    for epoch in range(start_epoch, config.epochs):\n        logs['epoch'] = epoch\n        train_sampler.set_epoch(epoch)\n        train_sampler.dataset.next_epoch()\n        run(logs=logs,data_loader=train_loader,net=net,optimizer=optimizer,local_rank=local_rank,config=config,phase='train')\n        run(logs=logs,data_loader=validate_loader,net=net,optimizer=optimizer,local_rank=local_rank,config=config,phase='valid')\n        scheduler.step()\n        if local_rank==0:\n            torch.save({\n                    'logs': logs,\n                    'state_dict': net.module.state_dict(),\n                    'optimizer_state': optimizer.state_dict(),\n                    'scheduler_state':scheduler.state_dict()}, 'checkpoints/'+config.name+'/ckpt_%s.pth'%epoch)\n        dist.barrier()\n\n    \ndef run(logs,data_loader,net,optimizer,local_rank,config,phase='train'):\n    if local_rank==0:\n        print('start ',phase)\n    recorder={}\n    record_list=['loss','acc']\n    for i in record_list:\n        recorder[i]=dist_average(local_rank)\n    # begin training\n    start_time = time.time()\n    if phase=='train':\n        net.train()\n    else: net.eval()\n    for i, (X, y) in enumerate(data_loader):\n        loss_pack={}\n        X = X.to(local_rank,non_blocking=True)\n        y = y.to(local_rank,non_blocking=True)\n        with torch.set_grad_enabled(phase=='train'):\n            logit=net(X)\n            batch_loss = F.cross_entropy(logit,y)\n        if phase=='train':\n            batch_loss.backward()\n            optimizer.step()\n            optimizer.zero_grad()\n        with torch.no_grad():\n            loss_pack['loss']=batch_loss\n            loss_pack['acc']=ACC(logit,y)\n        for i in record_list:\n            recorder[i].step(loss_pack[i])\n\n    # end of this epoch\n    batch_info=[]\n    for i in record_list:\n        mesg=recorder[i].get()\n        logs[i]=mesg\n        batch_info.append('{}:{:.4f}'.format(i,mesg))\n    end_time = time.time()\n\n    # write log for this epoch\n    if local_rank==0:\n        logging.info('{}: {}, Time {:3.2f}'.format(phase,'  '.join(batch_info), end_time - start_time))\n\n\ndef distributed_train(config,world_size=0,num_gpus=0,rank_offset=0):\n    if not num_gpus:\n        num_gpus = torch.cuda.device_count()\n    if not world_size:\n        world_size=num_gpus\n    mp.spawn(main_worker, nprocs=num_gpus, args=(world_size,rank_offset,config))\n","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:26:43.132444Z","iopub.execute_input":"2022-05-31T08:26:43.132778Z","iopub.status.idle":"2022-05-31T08:33:11.252664Z","shell.execute_reply.started":"2022-05-31T08:26:43.132724Z","shell.execute_reply":"2022-05-31T08:33:11.251855Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mtcnn = MTCNN(margin=0, keep_all=True, post_process=False, select_largest=False ,device='cuda:0', thresholds=MTCNN_THRESHOLDS_RETRY, factor=MMTNN_FACTOR)\n\nfor i_video, videopath in enumerate(tqdm(videopaths_missing_faces)):\n    faces, coords = face_detection_wrapper(mtcnn, \n                                           videopath, \n                                           every_n_frames=FACE_FRAMES, \n                                           facedetection_downsample=FACEDETECTION_DOWNSAMPLE, \n                                           max_frames_to_load=MAX_FRAMES_TO_LOAD)\n            \n    if len(faces):\n        faces_by_videopath[videopath]  = faces[:MAX_FACES_LOWTHRESH]\n        coords_by_videopath[videopath] = coords[:MAX_FACES_LOWTHRESH]\n    else:\n        print(f\"Found no faces for {videopath} !\")\n\ndel mtcnn\nimport gc\ngc.collect()\n\nfaces_by_videopath = dict(sorted(faces_by_videopath.items(), key=lambda x: x[0]))\nvideopaths_missing_faces = {p for p in videopaths if p not in faces_by_videopath}\nprint(f\"Found faces for {len(faces_by_videopath)} videos; {len(videopaths_missing_faces)} missing\")","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:11.254564Z","iopub.execute_input":"2022-05-31T08:33:11.255091Z","iopub.status.idle":"2022-05-31T08:33:18.036299Z","shell.execute_reply.started":"2022-05-31T08:33:11.255038Z","shell.execute_reply":"2022-05-31T08:33:18.035541Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"faces_by_videopath = dict(sorted(faces_by_videopath.items(), key=lambda x: x[0]))\nvideopaths_missing_faces = {p for p in videopaths if p not in faces_by_videopath}\nprint(f\"Found faces for {len(faces_by_videopath)} videos; {len(videopaths_missing_faces)} missing\")","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:18.037728Z","iopub.execute_input":"2022-05-31T08:33:18.038241Z","iopub.status.idle":"2022-05-31T08:33:18.045269Z","shell.execute_reply.started":"2022-05-31T08:33:18.038188Z","shell.execute_reply":"2022-05-31T08:33:18.044487Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 3D CNNs","metadata":{}},{"cell_type":"code","source":"class MaxPool3dSamePadding(nn.MaxPool3d):\n\n    def compute_pad(self, dim, s):\n        if s % self.stride[dim] == 0:\n            return max(self.kernel_size[dim] - self.stride[dim], 0)\n        else:\n            return max(self.kernel_size[dim] - (s % self.stride[dim]), 0)\n\n    def forward(self, x):\n        # compute 'same' padding\n        (batch, channel, t, h, w) = x.size()\n        # print t,h,w\n        out_t = np.ceil(float(t) / float(self.stride[0]))\n        out_h = np.ceil(float(h) / float(self.stride[1]))\n        out_w = np.ceil(float(w) / float(self.stride[2]))\n        # print out_t, out_h, out_w\n        pad_t = self.compute_pad(0, t)\n        pad_h = self.compute_pad(1, h)\n        pad_w = self.compute_pad(2, w)\n        # print pad_t, pad_h, pad_w\n\n        pad_t_f = pad_t // 2\n        pad_t_b = pad_t - pad_t_f\n        pad_h_f = pad_h // 2\n        pad_h_b = pad_h - pad_h_f\n        pad_w_f = pad_w // 2\n        pad_w_b = pad_w - pad_w_f\n\n        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)\n        # print x.size()\n        # print pad\n        x = F.pad(x, pad)\n        return super(MaxPool3dSamePadding, self).forward(x)\n\n\nclass Unit3D(nn.Module):\n\n    def __init__(self, in_channels,\n                 output_channels,\n                 kernel_shape=(1, 1, 1),\n                 stride=(1, 1, 1),\n                 padding=0,\n                 activation_fn=F.relu,\n                 use_batch_norm=True,\n                 use_bias=False,\n                 name='unit_3d'):\n\n        \"\"\"Initializes Unit3D module.\"\"\"\n        super(Unit3D, self).__init__()\n\n        self._output_channels = output_channels\n        self._kernel_shape = kernel_shape\n        self._stride = stride\n        self._use_batch_norm = use_batch_norm\n        self._activation_fn = activation_fn\n        self._use_bias = use_bias\n        self.name = name\n        self.padding = padding\n\n        self.conv3d = nn.Conv3d(in_channels=in_channels,\n                                out_channels=self._output_channels,\n                                kernel_size=self._kernel_shape,\n                                stride=self._stride,\n                                padding=0,\n                                # we always want padding to be 0 here. We will dynamically pad based on input size in forward function\n                                bias=self._use_bias)\n\n        if self._use_batch_norm:\n            self.bn = nn.BatchNorm3d(self._output_channels, eps=0.001, momentum=0.01)\n\n    def compute_pad(self, dim, s):\n        if s % self._stride[dim] == 0:\n            return max(self._kernel_shape[dim] - self._stride[dim], 0)\n        else:\n            return max(self._kernel_shape[dim] - (s % self._stride[dim]), 0)\n\n    def forward(self, x):\n        # compute 'same' padding\n        (batch, channel, t, h, w) = x.size()\n        # print t,h,w\n        out_t = np.ceil(float(t) / float(self._stride[0]))\n        out_h = np.ceil(float(h) / float(self._stride[1]))\n        out_w = np.ceil(float(w) / float(self._stride[2]))\n        # print out_t, out_h, out_w\n        pad_t = self.compute_pad(0, t)\n        pad_h = self.compute_pad(1, h)\n        pad_w = self.compute_pad(2, w)\n        # print pad_t, pad_h, pad_w\n\n        pad_t_f = pad_t // 2\n        pad_t_b = pad_t - pad_t_f\n        pad_h_f = pad_h // 2\n        pad_h_b = pad_h - pad_h_f\n        pad_w_f = pad_w // 2\n        pad_w_b = pad_w - pad_w_f\n\n        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)\n        # print x.size()\n        # print pad\n        x = F.pad(x, pad)\n        # print x.size()\n\n        x = self.conv3d(x)\n        if self._use_batch_norm:\n            x = self.bn(x)\n        if self._activation_fn is not None:\n            x = self._activation_fn(x)\n        return x\n\n\nclass InceptionModule(nn.Module):\n    def __init__(self, in_channels, out_channels, name):\n        super(InceptionModule, self).__init__()\n\n        self.b0 = Unit3D(in_channels=in_channels, output_channels=out_channels[0], kernel_shape=[1, 1, 1], padding=0,\n                         name=name + '/Branch_0/Conv3d_0a_1x1')\n        self.b1a = Unit3D(in_channels=in_channels, output_channels=out_channels[1], kernel_shape=[1, 1, 1], padding=0,\n                          name=name + '/Branch_1/Conv3d_0a_1x1')\n        self.b1b = Unit3D(in_channels=out_channels[1], output_channels=out_channels[2], kernel_shape=[3, 3, 3],\n                          name=name + '/Branch_1/Conv3d_0b_3x3')\n        self.b2a = Unit3D(in_channels=in_channels, output_channels=out_channels[3], kernel_shape=[1, 1, 1], padding=0,\n                          name=name + '/Branch_2/Conv3d_0a_1x1')\n        self.b2b = Unit3D(in_channels=out_channels[3], output_channels=out_channels[4], kernel_shape=[3, 3, 3],\n                          name=name + '/Branch_2/Conv3d_0b_3x3')\n        self.b3a = MaxPool3dSamePadding(kernel_size=[3, 3, 3],\n                                        stride=(1, 1, 1), padding=0)\n        self.b3b = Unit3D(in_channels=in_channels, output_channels=out_channels[5], kernel_shape=[1, 1, 1], padding=0,\n                          name=name + '/Branch_3/Conv3d_0b_1x1')\n        self.name = name\n\n    def forward(self, x):\n        b0 = self.b0(x)\n        b1 = self.b1b(self.b1a(x))\n        b2 = self.b2b(self.b2a(x))\n        b3 = self.b3b(self.b3a(x))\n        return torch.cat([b0, b1, b2, b3], dim=1)\n\n\nclass InceptionI3d(nn.Module):\n    \"\"\"Inception-v1 I3D architecture.\n    The model is introduced in:\n        Quo Vadis, Action Recognition? A New Model and the Kinetics Dataset\n        Joao Carreira, Andrew Zisserman\n        https://arxiv.org/pdf/1705.07750v1.pdf.\n    See also the Inception architecture, introduced in:\n        Going deeper with convolutions\n        Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed,\n        Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, Andrew Rabinovich.\n        http://arxiv.org/pdf/1409.4842v1.pdf.\n    \"\"\"\n\n    # Endpoints of the model in order. During construction, all the endpoints up\n    # to a designated `final_endpoint` are returned in a dictionary as the\n    # second return value.\n    VALID_ENDPOINTS = (\n        'Conv3d_1a_7x7',\n        'MaxPool3d_2a_3x3',\n        'Conv3d_2b_1x1',\n        'Conv3d_2c_3x3',\n        'MaxPool3d_3a_3x3',\n        'Mixed_3b',\n        'Mixed_3c',\n        'MaxPool3d_4a_3x3',\n        'Mixed_4b',\n        'Mixed_4c',\n        'Mixed_4d',\n        'Mixed_4e',\n        'Mixed_4f',\n        'MaxPool3d_5a_2x2',\n        'Mixed_5b',\n        'Mixed_5c',\n        'Logits',\n        'Predictions',\n    )\n\n    def __init__(self, num_classes=400, spatial_squeeze=True, output_method='per_frame',\n                 final_endpoint='Logits', name='inception_i3d', in_channels=3, dropout_keep_prob=0.5):\n        \"\"\"Initializes I3D model instance.\n        Args:\n          num_classes: The number of outputs in the logit layer (default 400, which\n              matches the Kinetics dataset).\n          spatial_squeeze: Whether to squeeze the spatial dimensions for the logits\n              before returning (default True).\n          final_endpoint: The model contains many possible endpoints.\n              `final_endpoint` specifies the last endpoint for the model to be built\n              up to. In addition to the output at `final_endpoint`, all the outputs\n              at endpoints up to `final_endpoint` will also be returned, in a\n              dictionary. `final_endpoint` must be one of\n              InceptionI3d.VALID_ENDPOINTS (default 'Logits').\n          name: A string (optional). The name of this module.\n        Raises:\n          ValueError: if `final_endpoint` is not recognized.\n        \"\"\"\n\n        if final_endpoint not in self.VALID_ENDPOINTS:\n            raise ValueError('Unknown final endpoint %s' % final_endpoint)\n\n        super(InceptionI3d, self).__init__()\n        self.output_method = output_method\n        assert output_method in ('per_frame', 'avg_pool', 'max_pool', 'dual_pool')\n        self._num_classes = num_classes\n        self._spatial_squeeze = spatial_squeeze\n        self._final_endpoint = final_endpoint\n        self.logits = None\n\n        if self._final_endpoint not in self.VALID_ENDPOINTS:\n            raise ValueError('Unknown final endpoint %s' % self._final_endpoint)\n\n        self.end_points = {}\n        end_point = 'Conv3d_1a_7x7'\n        self.end_points[end_point] = Unit3D(in_channels=in_channels, output_channels=64, kernel_shape=[7, 7, 7],\n                                            stride=(2, 2, 2), padding=(3, 3, 3), name=name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'MaxPool3d_2a_3x3'\n        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),\n                                                          padding=0)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Conv3d_2b_1x1'\n        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=64, kernel_shape=[1, 1, 1], padding=0,\n                                            name=name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Conv3d_2c_3x3'\n        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=192, kernel_shape=[3, 3, 3], padding=1,\n                                            name=name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'MaxPool3d_3a_3x3'\n        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),\n                                                          padding=0)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Mixed_3b'\n        self.end_points[end_point] = InceptionModule(192, [64, 96, 128, 16, 32, 32], name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Mixed_3c'\n        self.end_points[end_point] = InceptionModule(256, [128, 128, 192, 32, 96, 64], name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'MaxPool3d_4a_3x3'\n        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[3, 3, 3], stride=(2, 2, 2),\n                                                          padding=0)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Mixed_4b'\n        self.end_points[end_point] = InceptionModule(128 + 192 + 96 + 64, [192, 96, 208, 16, 48, 64], name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Mixed_4c'\n        self.end_points[end_point] = InceptionModule(192 + 208 + 48 + 64, [160, 112, 224, 24, 64, 64], name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Mixed_4d'\n        self.end_points[end_point] = InceptionModule(160 + 224 + 64 + 64, [128, 128, 256, 24, 64, 64], name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Mixed_4e'\n        self.end_points[end_point] = InceptionModule(128 + 256 + 64 + 64, [112, 144, 288, 32, 64, 64], name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Mixed_4f'\n        self.end_points[end_point] = InceptionModule(112 + 288 + 64 + 64, [256, 160, 320, 32, 128, 128],\n                                                     name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'MaxPool3d_5a_2x2'\n        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[2, 2, 2], stride=(2, 2, 2),\n                                                          padding=0)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Mixed_5b'\n        self.end_points[end_point] = InceptionModule(256 + 320 + 128 + 128, [256, 160, 320, 32, 128, 128],\n                                                     name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Mixed_5c'\n        self.end_points[end_point] = InceptionModule(256 + 320 + 128 + 128, [384, 192, 384, 48, 128, 128],\n                                                     name + end_point)\n        if self._final_endpoint == end_point: return\n\n        end_point = 'Logits'\n        self.avg_pool = nn.AvgPool3d(kernel_size=[2, 7, 7],\n                                     stride=(1, 1, 1))\n        self.dropout = nn.Dropout(dropout_keep_prob)\n        self.logits = Unit3D(in_channels=384 + 384 + 128 + 128, output_channels=self._num_classes,\n                             kernel_shape=[1, 1, 1],\n                             padding=0,\n                             activation_fn=None,\n                             use_batch_norm=False,\n                             use_bias=True,\n                             name='logits')\n\n        self.build()\n\n    def replace_logits(self, num_classes):\n        self._num_classes = num_classes\n        self.logits = Unit3D(in_channels=384 + 384 + 128 + 128, output_channels=self._num_classes,\n                             kernel_shape=[1, 1, 1],\n                             padding=0,\n                             activation_fn=None,\n                             use_batch_norm=False,\n                             use_bias=True,\n                             name='logits')\n\n    def build(self):\n        for k in self.end_points.keys():\n            self.add_module(k, self.end_points[k])\n\n    def forward(self, x):\n        for end_point in self.VALID_ENDPOINTS:\n            if end_point in self.end_points:\n                x = self._modules[end_point](x)  # use _modules to work with dataparallel\n\n        x = self.logits(self.dropout(self.avg_pool(x)))\n        if self._spatial_squeeze:\n            logits = x.squeeze(3).squeeze(3)\n        # logits is batch X time X classes, which is what we want to work with\n        if self.output_method == 'per_frame':\n            return F.interpolate(logits, 64, mode='linear')  # -> batch_size * 2 * 64\n        elif self.output_method == 'avg_pool':\n            avg_all_frames = F.adaptive_avg_pool1d(logits, output_size=1)\n            return avg_all_frames.squeeze(-1)  # -> batch_size * 2\n\n    def extract_features(self, x):\n        for end_point in self.VALID_ENDPOINTS:\n            if end_point in self.end_points:\n                x = self._modules[end_point](x)\n        return self.avg_pool(x)","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:18.046979Z","iopub.execute_input":"2022-05-31T08:33:18.047503Z","iopub.status.idle":"2022-05-31T08:33:18.122922Z","shell.execute_reply.started":"2022-05-31T08:33:18.047454Z","shell.execute_reply":"2022-05-31T08:33:18.122222Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def conv3x3x3(in_planes, out_planes, stride=1):\n    # 3x3x3 convolution with padding\n    return nn.Conv3d(\n        in_planes,\n        out_planes,\n        kernel_size=3,\n        stride=stride,\n        padding=1,\n        bias=False)\n\ndef downsample_basic_block(x, planes, stride):\n    out = F.avg_pool3d(x, kernel_size=1, stride=stride)\n    zero_pads = torch.Tensor(\n        out.size(0), planes - out.size(1), out.size(2), out.size(3),\n        out.size(4)).zero_()\n    if isinstance(out.data, torch.cuda.FloatTensor):\n        zero_pads = zero_pads.cuda()\n \n    out = Variable(torch.cat([out.data, zero_pads], dim=1))\n \n    return out\n\nclass BasicBlock(nn.Module):\n    expansion = 1\n \n    def __init__(self, inplanes, planes, stride=1, downsample=None):\n        super(BasicBlock, self).__init__()\n        self.conv1 = conv3x3x3(inplanes, planes, stride)\n        self.bn1 = nn.BatchNorm3d(planes)\n        self.relu = nn.ReLU(inplace=True)\n        self.conv2 = conv3x3x3(planes, planes)\n        self.bn2 = nn.BatchNorm3d(planes)\n        self.downsample = downsample\n        self.stride = stride\n \n    def forward(self, x):\n        residual = x\n \n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n \n        out = self.conv2(out)\n        out = self.bn2(out)\n \n        if self.downsample is not None:\n            residual = self.downsample(x)\n \n        out += residual\n        out = self.relu(out)\n \n        return out\n\n\nclass Bottleneck(nn.Module):\n    expansion = 4\n \n    def __init__(self, inplanes, planes, stride=1, downsample=None):\n        super(Bottleneck, self).__init__()\n        self.conv1 = nn.Conv3d(inplanes, planes, kernel_size=1, bias=False)\n        self.bn1 = nn.BatchNorm3d(planes)\n        self.conv2 = nn.Conv3d(\n            planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n        self.bn2 = nn.BatchNorm3d(planes)\n        self.conv3 = nn.Conv3d(planes, planes * 4, kernel_size=1, bias=False)\n        self.bn3 = nn.BatchNorm3d(planes * 4)\n        self.relu = nn.ReLU(inplace=True)\n        self.downsample = downsample\n        self.stride = stride\n \n    def forward(self, x):\n        residual = x\n \n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n \n        out = self.conv2(out)\n        out = self.bn2(out)\n        out = self.relu(out)\n \n        out = self.conv3(out)\n        out = self.bn3(out)\n \n        if self.downsample is not None:\n            residual = self.downsample(x)\n \n        out += residual\n        out = self.relu(out)\n \n        return out\n\n\nclass ResNet(nn.Module):\n \n    def __init__(self,\n                 block,\n                 layers,\n                 sample_size,\n                 sample_duration,\n                 shortcut_type='B',\n                 num_classes=400):\n        self.inplanes = 64\n        super(ResNet, self).__init__()\n        self.conv1 = nn.Conv3d(\n            3,\n            64,\n            kernel_size=7,\n            stride=(1, 2, 2),\n            padding=(3, 3, 3),\n            bias=False)\n        self.bn1 = nn.BatchNorm3d(64)\n        self.relu = nn.ReLU(inplace=True)\n        self.maxpool = nn.MaxPool3d(kernel_size=(3, 3, 3), stride=2, padding=1)\n        self.layer1 = self._make_layer(block, 64, layers[0], shortcut_type)\n        self.layer2 = self._make_layer(\n            block, 128, layers[1], shortcut_type, stride=2)\n        self.layer3 = self._make_layer(\n            block, 256, layers[2], shortcut_type, stride=2)\n        self.layer4 = self._make_layer(\n            block, 512, layers[3], shortcut_type, stride=2)\n        # last_duration = int(math.ceil(sample_duration / 16))\n        # last_size = int(math.ceil(sample_size / 32))\n        # self.avgpool = nn.AvgPool3d(\n        #     (last_duration, last_size, last_size), stride=1)\n        self.avgpool = nn.AdaptiveAvgPool3d(1)\n        self.fc = nn.Linear(512 * block.expansion, num_classes)\n \n        for m in self.modules():\n            if isinstance(m, nn.Conv3d):\n                m.weight = nn.init.kaiming_normal(m.weight, mode='fan_out')\n            elif isinstance(m, nn.BatchNorm3d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n \n    def _make_layer(self, block, planes, blocks, shortcut_type, stride=1):\n        downsample = None\n        if stride != 1 or self.inplanes != planes * block.expansion:\n            if shortcut_type == 'A':\n                downsample = partial(\n                    downsample_basic_block,\n                    planes=planes * block.expansion,\n                    stride=stride)\n            else:\n                downsample = nn.Sequential(\n                    nn.Conv3d(\n                        self.inplanes,\n                        planes * block.expansion,\n                        kernel_size=1,\n                        stride=stride,\n                        bias=False), nn.BatchNorm3d(planes * block.expansion))\n \n        layers = []\n        layers.append(block(self.inplanes, planes, stride, downsample))\n        self.inplanes = planes * block.expansion\n        for i in range(1, blocks):\n            layers.append(block(self.inplanes, planes))\n \n        return nn.Sequential(*layers)\n \n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n \n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n \n        x = self.avgpool(x)\n \n        x = x.view(x.size(0), -1)\n        x = self.fc(x)\n \n        return x\n\ndef get_fine_tuning_parameters(model, ft_begin_index):\n    if ft_begin_index == 0:\n        return model.parameters()\n \n    ft_module_names = []\n    for i in range(ft_begin_index, 5):\n        ft_module_names.append('layer{}'.format(i))\n    ft_module_names.append('fc')\n \n    parameters = []\n    for k, v in model.named_parameters():\n        for ft_module in ft_module_names:\n            if ft_module in k:\n                parameters.append({'params': v})\n                break\n        else:\n            parameters.append({'params': v, 'lr': 0.0})\n \n    return parameters\n\n\ndef resnet18(**kwargs):\n    \"\"\"Constructs a ResNet-18 model.\n    \"\"\"\n    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)\n    return model\n\ndef resnet34(**kwargs):\n    \"\"\"Constructs a ResNet-34 model.\n    \"\"\"\n    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)\n    return model","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:18.126778Z","iopub.execute_input":"2022-05-31T08:33:18.127083Z","iopub.status.idle":"2022-05-31T08:33:18.173976Z","shell.execute_reply.started":"2022-05-31T08:33:18.127026Z","shell.execute_reply":"2022-05-31T08:33:18.173197Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HFlipWrapper(nn.Module):\n    def __init__(self, model, flip_dim=(-1,)):\n        super(HFlipWrapper, self).__init__()\n        self.model = model\n        self.flip_dim = flip_dim\n        \n    def forward(self, x):\n        with torch.no_grad():\n            xf = torch.flip(x, self.flip_dim)\n        pred = self.model(torch.stack([x, xf]))\n        return pred","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:18.177502Z","iopub.execute_input":"2022-05-31T08:33:18.177765Z","iopub.status.idle":"2022-05-31T08:33:18.186301Z","shell.execute_reply.started":"2022-05-31T08:33:18.177719Z","shell.execute_reply":"2022-05-31T08:33:18.185686Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"models_3d = []\n\nfor modeldict in PRETRAINED_MODELS_3D:\n    if modeldict['type'] == 'i3d':\n        model = InceptionI3d(157, in_channels=3, output_method='avg_pool')\n        model.replace_logits(2)\n        model = model.cuda()\n        model.eval()\n        model.load_state_dict(torch.load(modeldict['path']))\n        models_3d.append({'norm':'i3d', 'model':model})\n        \n    elif modeldict['type'] == 'res18':\n        model = resnet18(num_classes=2, shortcut_type='A', sample_size=224, sample_duration=32)  # , last_fc=True)\n        model.load_state_dict(torch.load(modeldict['path']))\n        model = model.cuda()\n        model.eval()\n        models_3d.append({'norm':'nil', 'model':model})\n        \n    elif modeldict['type'] == 'res34':\n        model = resnet34(num_classes=2, shortcut_type='A', sample_size=224, sample_duration=32)  # , last_fc=True)\n        model.load_state_dict(torch.load(modeldict['path']))\n        model = model.cuda()\n        model.eval()\n        models_3d.append({'norm':'nil', 'model':model})\n        \n    elif modeldict['type'] == 'mc3_112':\n        model = mc3_18(num_classes=2)\n        model.load_state_dict(torch.load(modeldict['path']))\n        model = model.cuda()\n        model.eval()\n        models_3d.append({'norm':'112_imagenet', 'model':model})\n        \n    elif modeldict['type'] == 'mc3_224':\n        model = mc3_18(num_classes=2)\n        model.load_state_dict(torch.load(modeldict['path']))\n        model = model.cuda()\n        model.eval()\n        models_3d.append({'norm':'224_imagenet', 'model':model})\n        \n    elif modeldict['type'] == 'r2p1_112':\n        model = r2plus1d_18(num_classes=2)\n        model.load_state_dict(torch.load(modeldict['path']))\n        model = model.cuda()\n        model.eval()\n        models_3d.append({'norm':'112_imagenet', 'model':model})\n        \n        \n    else:\n        raise ValueError(f\"Unknown model type {modeldict['type']}\")","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:18.18902Z","iopub.execute_input":"2022-05-31T08:33:18.189247Z","iopub.status.idle":"2022-05-31T08:33:26.62787Z","shell.execute_reply.started":"2022-05-31T08:33:18.189203Z","shell.execute_reply":"2022-05-31T08:33:26.627136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys ; sys.path.insert(0, '/kaggle/input/deepfake/deepfake/deepfake/skp/')\nimport yaml\n\nwith open('/kaggle/input/deepfake/deepfake/deepfake/skp/configs/experiments/experiment001.yaml') as f:\n    CFG = yaml.load(f, Loader=yaml.FullLoader)\n\nfrom factory.builder import build_model, build_dataloader\n\nCFG['model']['params']['pretrained'] = None\nmodel2d = build_model(CFG['model']['name'], CFG['model']['params'])\nmodel2d.load_state_dict(torch.load('/kaggle/input/deepfake/SRXT50_094_VM-0.2504.PTH'))\nmodel2d = model2d.eval().cuda()\nloader = build_dataloader(CFG, data_info={'vidfiles': [], 'labels': []}, mode='predict')","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:26.629635Z","iopub.execute_input":"2022-05-31T08:33:26.630052Z","iopub.status.idle":"2022-05-31T08:33:30.103054Z","shell.execute_reply.started":"2022-05-31T08:33:26.630003Z","shell.execute_reply":"2022-05-31T08:33:30.10229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 3D Main loop","metadata":{}},{"cell_type":"code","source":"def resize_and_square_face(video, output_size):\n    input_size = max(video.shape[1], video.shape[2])  # We will square it, so this is the effective input size\n    out_video = np.empty((len(video), output_size[0], output_size[1], 3), np.dtype('uint8'))\n    \n    for i_frame, frame in enumerate(video):\n        padded_image = np.zeros((input_size, input_size, 3))\n        padded_image[0:frame.shape[0], 0:frame.shape[1]] = frame\n        if (input_size, input_size) != output_size:\n            frame = cv2.resize(padded_image, (output_size[0], output_size[1])).astype(np.uint8)\n        else:\n            frame = padded_image.astype(np.uint8)\n        out_video[i_frame] = frame\n    return out_video","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:30.104469Z","iopub.execute_input":"2022-05-31T08:33:30.104792Z","iopub.status.idle":"2022-05-31T08:33:30.113897Z","shell.execute_reply.started":"2022-05-31T08:33:30.104744Z","shell.execute_reply":"2022-05-31T08:33:30.113032Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def center_crop_video(video, crop_dimensions):\n    height, width = video.shape[1], video.shape[2]\n    crop_height, crop_width = crop_dimensions\n    \n    y1 = (height - crop_height) // 2\n    y2 = y1 + crop_height\n    x1 = (width - crop_width) // 2\n    x2 = x1 + crop_width\n        \n    video_out = np.zeros((len(video), crop_height, crop_width, 3))\n    for i_frame,frame in enumerate(video):\n        video_out[i_frame] = frame[y1:y2, x1:x2]\n        \n    return video_out","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:30.115167Z","iopub.execute_input":"2022-05-31T08:33:30.115642Z","iopub.status.idle":"2022-05-31T08:33:30.125339Z","shell.execute_reply.started":"2022-05-31T08:33:30.115557Z","shell.execute_reply":"2022-05-31T08:33:30.124438Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_last_frame_needed_across_faces(faces):\n    last_frame = 0\n    \n    for face in faces:\n        (frame_from, frame_to), (row_from, row_to), (col_from, col_to) = face\n        last_frame = max(frame_to, last_frame)\n        \n    return last_frame","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:30.126763Z","iopub.execute_input":"2022-05-31T08:33:30.127227Z","iopub.status.idle":"2022-05-31T08:33:30.134327Z","shell.execute_reply.started":"2022-05-31T08:33:30.127044Z","shell.execute_reply":"2022-05-31T08:33:30.133041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"predictions = defaultdict(list)","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:30.135872Z","iopub.execute_input":"2022-05-31T08:33:30.136634Z","iopub.status.idle":"2022-05-31T08:33:30.144904Z","shell.execute_reply.started":"2022-05-31T08:33:30.136318Z","shell.execute_reply":"2022-05-31T08:33:30.144123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_transforms_114_imagenet = A.Compose([A.Resize(height=112, width=112),\n                                 A.Normalize()])\n\ntest_transforms_224_imagenet = A.Compose([A.Normalize()])","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:30.146455Z","iopub.execute_input":"2022-05-31T08:33:30.146718Z","iopub.status.idle":"2022-05-31T08:33:30.153499Z","shell.execute_reply.started":"2022-05-31T08:33:30.146675Z","shell.execute_reply":"2022-05-31T08:33:30.152725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for videopath, faces in tqdm(faces_by_videopath.items(), total=len(faces_by_videopath)):\n    try:\n\n        if len(faces):\n            last_frame_needed = get_last_frame_needed_across_faces(faces)\n            video, rescale = load_video(videopath, every_n_frames=1, to_rgb=True, rescale=None, inc_pil=False, max_frames=last_frame_needed)\n\n        else:\n            print(f\"Skipping {videopath} as no faces found\")\n            continue\n            \n        for modeldict in models_3d:\n            preds_video = []\n            model = modeldict['model']\n            \n            model = HFlipWrapper(model=model)\n\n            for i_face, face in enumerate(faces):\n                (frame_from, frame_to), (row_from, row_to), (col_from, col_to) = face\n\n                x = video[frame_from:frame_to, row_from:row_to + 1, col_from:col_to + 1]\n                x = resize_and_square_face(x, output_size=OUTPUT_FACE_SIZE)\n\n                if PRE_INFERENCE_CROP and PRE_INFERENCE_CROP != OUTPUT_FACE_SIZE:\n                    x = center_crop_video(x, PRE_INFERENCE_CROP)\n\n                with torch.no_grad():\n                    \n                    if modeldict['norm'] == '112_imagenet':\n                        x = np.array([test_transforms_114_imagenet(image=frame)['image'] for frame in x])\n                    elif modeldict['norm'] == '224_imagenet':\n                        x = np.array([test_transforms_224_imagenet(image=frame)['image'] for frame in x])\n                        \n                    x = torch.from_numpy(x.transpose([3, 0, 1, 2])).float()\n                    \n                    if modeldict['norm'] == 'i3d':\n                        x = (x / 255.) * 2 - 1\n                    elif modeldict['norm'] == 'nil':\n                        pass\n                    elif modeldict['norm'] == '112_imagenet':\n                        pass\n                    elif modeldict['norm'] == '224_imagenet':\n                        pass\n                    else:\n                        raise ValueError(f\"Unknown normalisation mode {modeldict['norm']}\")\n\n                    y_pred = model(x.cuda())\n                    prob0, prob1 = torch.mean(torch.exp(F.log_softmax(y_pred, dim=1)),dim=0)\n                    if REVERSE_PROBS:\n                        prob1 = 1-prob1\n                    preds_video.append(float(prob1))\n                    \n            videoname = os.path.basename(videopath)\n            if preds_video:\n                predictions[videoname].extend([USE_FACE_FUNCTION(preds_video)] * RATIO_3D)\n            \n        try:\n            FRAMES2D = 32\n            # Ian's 2D model\n            coords = coords_by_videopath[videopath]\n            videoname = os.path.basename(videopath)\n            preds_video = []\n            for i_coord, coordinate in enumerate(coords):\n                (frame_from, frame_to), (row_from, row_to), (col_from, col_to) = faces[i_coord]\n                x = []\n                for coord_ind, frame_number in enumerate(range(frame_from, min(frame_from+FRAMES2D, frame_to-1))):\n                    if coord_ind >= len(coordinate):\n                        break\n                    x1, y1, x2, y2 = coordinate[coord_ind]\n                    x.append(video[frame_number, y1:y2, x1:x2])\n                x = np.asarray(x)\n                # Reverse back to BGR because it will get reversed to RGB when preprocessed\n                #x = x[...,::-1]\n                # Preprocess\n                x = loader.dataset.process_video(x)\n                #x = np.asarray([loader.dataset.process_image(_) for _ in x])\n                # Flip every other frame\n                x[:,::2] = x[:,::2,:,::-1]\n                # RGB reverse every 3rd frame\n                #x[:,::3] = x[::-1,::3]\n                with torch.no_grad():\n                    out = model2d(torch.from_numpy(np.ascontiguousarray(x)).unsqueeze(0).cuda())\n                #out = np.median(out.cpu().numpy())\n                preds_video.append(out.cpu().numpy())\n            if len(preds_video) > 0:\n                predictions[videoname].extend([USE_FACE_FUNCTION(preds_video)] * RATIO_2D)\n        except:\n            pass\n            \n    except Exception as e:\n        print(f\"ERROR: Video {videopath}: {e}\")","metadata":{"execution":{"iopub.status.busy":"2022-05-31T08:33:30.154915Z","iopub.execute_input":"2022-05-31T08:33:30.155367Z","iopub.status.idle":"2022-05-31T09:12:59.232331Z","shell.execute_reply.started":"2022-05-31T08:33:30.155316Z","shell.execute_reply":"2022-05-31T09:12:59.231546Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nx_show = x.transpose(1,2,3,0)\nx_show -= np.min(x_show)\nx_show /= np.max(x_show)\ndef view(x, nrows=4, ncols=4):\n    indices = np.linspace(0, len(x)-1, nrows*ncols) \n    for ind, i in enumerate(indices):\n        plt.subplot(nrows,ncols,ind+1)\n        plt.imshow(x[int(i),...,::-1])\n    plt.show()\n\nview(x_show, 5, 5)","metadata":{"execution":{"iopub.status.busy":"2022-05-31T09:12:59.233797Z","iopub.execute_input":"2022-05-31T09:12:59.234322Z","iopub.status.idle":"2022-05-31T09:13:01.565714Z","shell.execute_reply.started":"2022-05-31T09:12:59.234267Z","shell.execute_reply":"2022-05-31T09:13:01.564654Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for k,v in predictions.items():\n    string = '{} : {:.4f} //'.format(k, np.mean(v))\n    for proba in v:\n        string += ' {:.4f}'.format(proba)\n    print(string)","metadata":{"execution":{"iopub.status.busy":"2022-05-31T09:13:01.567526Z","iopub.execute_input":"2022-05-31T09:13:01.57171Z","iopub.status.idle":"2022-05-31T09:13:01.630198Z","shell.execute_reply.started":"2022-05-31T09:13:01.57164Z","shell.execute_reply":"2022-05-31T09:13:01.629292Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del model\nfor modeldict in models_3d:\n    del modeldict['model']\n    del modeldict\ndel models_3d\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2022-05-31T09:13:01.63156Z","iopub.execute_input":"2022-05-31T09:13:01.631861Z","iopub.status.idle":"2022-05-31T09:13:01.732393Z","shell.execute_reply.started":"2022-05-31T09:13:01.631812Z","shell.execute_reply":"2022-05-31T09:13:01.731602Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Fill in missing predictions with 0.5","metadata":{}},{"cell_type":"code","source":"for videopath in videopaths:\n    videoname = os.path.basename(videopath)\n    if videoname not in predictions:\n        predictions[videoname] = [DEFAULT_MISSING_PRED]","metadata":{"execution":{"iopub.status.busy":"2022-05-31T09:13:01.733751Z","iopub.execute_input":"2022-05-31T09:13:01.733997Z","iopub.status.idle":"2022-05-31T09:13:01.743956Z","shell.execute_reply.started":"2022-05-31T09:13:01.73395Z","shell.execute_reply":"2022-05-31T09:13:01.739691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Calculate ensembled prediction & clamp","metadata":{}},{"cell_type":"code","source":"preds_ensembled = {}\n\nfor videopath, pred in predictions.items():\n    #print(f\"{videopath} Got {len(pred)} predictions\")\n    preds_ensembled[videopath] = np.clip(np.mean(pred), PROB_MIN, PROB_MAX)","metadata":{"execution":{"iopub.status.busy":"2022-05-31T09:13:01.745456Z","iopub.execute_input":"2022-05-31T09:13:01.745928Z","iopub.status.idle":"2022-05-31T09:13:01.772Z","shell.execute_reply.started":"2022-05-31T09:13:01.74587Z","shell.execute_reply":"2022-05-31T09:13:01.768441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Save predictions","metadata":{}},{"cell_type":"code","source":"!rm -rf /kaggle/working/*","metadata":{"execution":{"iopub.status.busy":"2022-05-31T09:13:01.773345Z","iopub.execute_input":"2022-05-31T09:13:01.773711Z","iopub.status.idle":"2022-05-31T09:13:02.502295Z","shell.execute_reply.started":"2022-05-31T09:13:01.773659Z","shell.execute_reply":"2022-05-31T09:13:02.501419Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"{len(preds_ensembled)} predictions\")\nwith open(\"submission.csv\", 'w') as f:\n    f.write(\"filename,label\\n\")\n    for filename, label in preds_ensembled.items():\n        f.write(f\"{filename},{label}\\n\")","metadata":{"execution":{"iopub.status.busy":"2022-05-31T09:13:02.504067Z","iopub.execute_input":"2022-05-31T09:13:02.504396Z","iopub.status.idle":"2022-05-31T09:13:02.511958Z","shell.execute_reply.started":"2022-05-31T09:13:02.504347Z","shell.execute_reply":"2022-05-31T09:13:02.511195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!head -10 submission.csv","metadata":{"execution":{"iopub.status.busy":"2022-05-31T09:13:02.513338Z","iopub.execute_input":"2022-05-31T09:13:02.513874Z","iopub.status.idle":"2022-05-31T09:13:03.284797Z","shell.execute_reply.started":"2022-05-31T09:13:02.513786Z","shell.execute_reply":"2022-05-31T09:13:03.283605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd \nsubmission = pd.read_csv('./submission.csv')\nplt.hist(submission.label)\nplt.show()\nprint(\"ACC：\",sum(submission['label']<0.95)/400)\n","metadata":{"execution":{"iopub.status.busy":"2022-05-31T09:13:03.289813Z","iopub.execute_input":"2022-05-31T09:13:03.290451Z","iopub.status.idle":"2022-05-31T09:13:03.589464Z","shell.execute_reply.started":"2022-05-31T09:13:03.29039Z","shell.execute_reply":"2022-05-31T09:13:03.588672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}