{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Fast and stable generator for training set\n\nThis notebook uses \"facenet-pytorch\" to detect faces and mekes face images from every frames of training movies. \nThanks to timesler for his usful benchmark. https://www.kaggle.com/timesler/comparison-of-face-detection-packages  \n\nDetected faces are checked by seeing relatively short movies for fake and real, respectivly. \n\nUsing a GPU it takes a few seconds on average for every movie to extract faces from all frames.  \nThis example is only for train_sample_videos but it is easy to change it for train_videos. \n\nProblem:  \nAlthough movies are created successfuly, these are not displayed on this notebook. \nPlease tell me how to display a local mp4 movie. "},{"metadata":{"trusted":false},"cell_type":"code","source":"import cv2\nimport numpy as np\nfrom matplotlib import pyplot as plt\nimport pandas as pd\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\nimport time\nimport glob\nimport copy\nimport torch\n!pip install facenet-pytorch\n#!git clone https://github.com/timesler/facenet-pytorch.git facenet_pytorch\nfrom facenet_pytorch import MTCNN\nimport pickle\nimport itertools\nimport os\nimport gc\n\ndevice = 'cuda:0' if torch.cuda.is_available() else 'cpu'\ndevice = 'cpu'\ndir='/kaggle/input'\ndir2='/kaggle/working'","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def detect_face(sample,interval=20,normalise=False):\n    reader = cv2.VideoCapture(sample)\n    images=[]\n    tmp0,tmp1=1000,-1000\n    for i in range(int(reader.get(cv2.CAP_PROP_FRAME_COUNT))):\n        _, image = reader.read()\n        try:\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        except:\n            continue\n        images.append(image)\n        if normalise:\n            tmp0=min(tmp0,np.min(image))\n            tmp1=max(tmp1,np.max(image))\n    rd_frames = np.linspace(0, len(images)-1, interval).astype(\"int32\")\n    if normalise:\n        tmp2=1/(tmp1-tmp0)\n        tmp=[(255.99*(images[e]-tmp0)*tmp2).astype(\"uint8\") for e in rd_frames]\n    else:\n        tmp=[images[e] for e in rd_frames]        \n    imgs_pil = [Image.fromarray(e) for e in tmp]\n    box,p=[],[]\n    for i in range(len(imgs_pil)):\n        try:\n            device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n            detector = MTCNN(device=device, keep_all=True, post_process=False)\n            box2,p2=detector.detect(imgs_pil[i])\n        except:\n            device = 'cpu'\n            detector = MTCNN(device=device, keep_all=True, post_process=False)\n            box2,p2=detector.detect(imgs_pil[i])\n        box.append(box2)\n        p.append(p2)       \n    reader.release()\n    return images,box,p,rd_frames","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def divide_face(box,p,t,w0=10,pt=0.99,nframe0=0.5):\n    p_ave, ns = 0., 0\n    for i in range(len(p)):\n        if len(p[i])<1:\n            continue\n        if p[i][0] is not None:\n                p_ave+=max(p[i])\n                ns+=1\n    p_ave=p_ave/len(p)\n    \n    p0, nface = p_ave*pt, 0\n    for i in range(len(p)):\n        if len(p[i])<1:\n            continue\n        tmp=[]\n        if p[i][0] is not None:\n            dx, dy = box[i][:,2]-box[i][:,0], box[i][:,3]-box[i][:,1]\n            for j in range(len(p[i])):\n                if (p[i][j]>=p0) & (dx[j]>=w0) & (dy[j]>=w0):\n                    tmp.append(j)\n        if p[i][0] is None or len(tmp)<1:\n            box[i], p[i] = [], []\n        else:\n            box[i], p[i] = box[i][tmp], p[i][tmp]\n        if len(p[i])>nface:\n            nface=len(p[i])\n\n#    xp, yp = (np.arange(nface)-nface)*1000, np.zeros(nface)\n    xp, yp = np.zeros(nface,dtype=\"float32\"), np.zeros(nface,dtype=\"float32\")\n    face_on=np.zeros(nface,dtype=\"int32\")\n    f_par = [[] for i in range(nface)]\n    for i in range(len(p)):\n        if len(p[i])<1:\n            continue\n        xc, yc = (box[i][:,0]+box[i][:,2])/2, (box[i][:,1]+box[i][:,3])/2\n        #perms = itertools.permutations(range(nface),len(p[i]))\n        nface2=max(np.sum(face_on),len(p[i]))\n        perms = itertools.permutations(range(nface2),len(p[i]))\n        d0 = 10.**10\n        for perm in perms:\n            tmp=list(perm)\n            d = sum(((xc-xp[tmp])**2+(yc-yp[tmp])**2)*face_on[tmp])\n            if d<d0:\n                d0, perms0 = d, tmp\n        if d0 > 700**2*np.sum(face_on):\n            continue\n        xp[perms0],yp[perms0],face_on[perms0]=xc,yc,1\n        for j, k in enumerate(perms0):\n            f_par[k].append([box[i][j,0],box[i][j,1],box[i][j,2],box[i][j,3],t[i]])\n    f_par = [e for e in f_par if len(e)>=max(len(p)*nframe0,1)]\n    for k in range(len(f_par)):\n        f_par[k] = np.array(f_par[k],dtype=np.float32).T\n    #print(p)\n    return f_par","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def mk_input(images,f_par,int_point=1):\n    t=np.arange(len(images),dtype=\"int32\")\n    xsize,ysize=images[0].shape[1],images[0].shape[0]\n    nface=len(f_par)\n    face=[]\n    for i in range(nface):\n        ts = (f_par[i][-1,:]).astype(\"int32\")\n        p0,p2=np.zeros(len(t),dtype=\"int32\"),np.zeros(len(t),dtype=\"int32\")\n        p1,p3=np.zeros(len(t),dtype=\"int32\"),np.zeros(len(t),dtype=\"int32\")\n        for n in range(len(ts)+1):\n            if n == 0:\n                if ts[n]==0:\n                    continue\n                tl=t[:ts[n]]\n            elif n == len(ts):\n                tl=t[ts[n-1]:]\n            else:\n                tl=t[ts[n-1]:ts[n]]\n            if len(tl)<=20 and n>0 and n<len(ts):\n                n0,n1=max(n-int_point,0),min(n+int_point,len(ts))\n            else:\n                n0,n1=max(n-1,0),min(n+1,len(ts))\n            coe=np.polyfit(ts[n0:n1],f_par[i][:,n0:n1].T,n1-n0-1).T                \n            q0,q2=np.poly1d(coe[0])(tl).T, np.poly1d(coe[2])(tl).T\n            q1,q3=np.poly1d(coe[1])(tl).T, np.poly1d(coe[3])(tl).T\n            p0[tl],p2[tl]=(q0+0.5).astype(\"int32\"),(q2+0.5).astype(\"int32\")\n            p1[tl],p3[tl]=(q1+0.5).astype(\"int32\"),(q3+0.5).astype(\"int32\")\n\n        p0,p1=np.where(p0<0,0,p0),np.where(p1<0,0,p1)\n        p2,p3=np.where(p2>xsize-1,xsize-1,p2),np.where(p3>ysize-1,ysize-1,p3)\n\n        tmp,tmp0,tmp1=[],np.ones(3)*1000,-np.ones(3)*1000\n        for n in range(len(t)):\n            tmp.append(images[n][p1[n]:p3[n],p0[n]:p2[n]])\n            for j in range(3):\n                tmp0[j]=min(tmp0[j],np.min(tmp[n][:,:,j]))\n                tmp1[j]=max(tmp1[j],np.max(tmp[n][:,:,j]))\n        for n in range(len(t)):\n            tmp2=tmp[n].astype(\"float32\")\n            for j in range(3):\n                tmp2[:,:,j]=((tmp2[:,:,j]-tmp0[j])/(tmp1[j]-tmp0[j]))\n            tmp[n]=cv2.resize(tmp2,(150, 200))\n        face.append((255.99*np.stack(tmp)).astype(\"uint8\"))\n    return face","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def mk_movie(face,filename):\n    fourcc = cv2.VideoWriter_fourcc('m','p','4','v')\n    video = cv2.VideoWriter(filename, fourcc, 30.0, (face.shape[2],face.shape[1]))\n    for i in range(len(face)):\n        img=face[i][:,:,[2,1,0]]\n        video.write(img)\n    video.release()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def mk_dataset(sample):\n    t0 = time.time()\n    images,box,p,rd_frames=detect_face(sample,interval=20,normalise=False)\n    t1 = time.time()\n    f_par = divide_face(box,p,rd_frames,pt=0.99)\n    gc.collect()\n    if len(f_par)<1:\n        print('no face')\n        t0 = time.time()\n        del images\n        images,box,p,rd_frames=detect_face(sample,interval=60,normalise=True)\n        t1 = time.time()\n        f_par = divide_face(box,p,rd_frames,pt=0.9,nframe0=0.)\n        #f_par = divide_face(box,p,rd_frames,pt=0.2,nframe0=0.)\n    elif len(f_par[0][0])<10:\n        print('few faces',len(f_par[0][0]))\n        t0 = time.time()\n        del images\n        images,box,p,rd_frames=detect_face(sample,interval=60,normalise=True)\n        t1 = time.time()\n        f_par = divide_face(box,p,rd_frames,pt=0.9,nframe0=0.2)\n    face = mk_input(images,f_par)\n    del images\n    t2 = time.time()\n    print(sample[-14:-4],f'{t1-t0:.3f} seconds',f'{t2-t1:.3f} seconds')\n    return face","execution_count":null,"outputs":[]},{"metadata":{"scrolled":true,"trusted":false},"cell_type":"code","source":"os.makedirs(dir2+'/train_sample',exist_ok=True)\nfiles = sorted(glob.glob(dir+'/train_sample/*.pickle'))\ndf = pd.read_json(dir+'/deepfake-detection-challenge/train_sample_videos/metadata.json').T\nsamples = (dir+'/deepfake-detection-challenge/train_sample_videos/'+df.index.values).tolist()\nvideo1 = cv2.VideoWriter(dir2+\"/train_sample/fake.mp4\", cv2.VideoWriter_fourcc('m','p','4','v'), 30.0, (150,200))\nvideo0 = cv2.VideoWriter(dir2+\"/train_sample/real.mp4\", cv2.VideoWriter_fourcc('m','p','4','v'), 30.0, (150,200))\nfor n in tqdm(range(len(df))):\n    filename=dir2+'/train_sample/'+samples[n][-14:-4]+'-'+df[\"label\"][n]\n    if df[\"label\"][n]=='FAKE':\n        filename+='-'+df[\"original\"][n][0:10]\n    face=mk_dataset(samples[n])\n    torch.cuda.empty_cache()\n    gc.collect()\n    #with open(filename+'.pickle', 'wb') as f:\n    #    pickle.dump(face , f)\n    for i in range(len(face)):\n        t_idx=(np.linspace(0,len(face[i])-1,30)+0.5).astype('int32')\n        imgs=face[i][t_idx]\n        for j in range(len(imgs)):\n            img = cv2.cvtColor(imgs[j], cv2.COLOR_BGR2RGB)\n            if df[\"label\"][n]=='FAKE':\n                video1.write(img)\n            else:\n                video0.write(img)\nvideo1.release()\nvideo0.release()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"from IPython.display import Video\nVideo(dir2+\"/train_sample/fake.mp4\", \"mp4\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"Video(dir2+\"/train_sample/real.mp4\", \"mp4\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.7.3"}},"nbformat":4,"nbformat_minor":4}