{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Embeddings: grouping similar faces automagically\nThis notebook demonstrates how embeddings can be used to group similar faces automagically. The idea is to develop this code into a validation strategy that ensures all actors (similar faces) are all in either train or validation sets."},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install facenet_pytorch\n!pip install pretrainedmodels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport pretrainedmodels\nimport pretrainedmodels.utils as utils\nfrom shutil import copyfile\nos.environ['TORCH_HOME'] = '/kaggle/working/pretrained-model-weights-pytorch'\n\ndef copy_weights(model_name):\n    found = False\n    for dirname, _, filenames in os.walk('/kaggle/input/'):\n        for filename in filenames:\n            full_path = os.path.join(dirname, filename)\n            if filename.startswith(model_name):\n                found = True\n                break\n        if found:\n            break\n            \n    base_dir = \"/kaggle/working/pretrained-model-weights-pytorch/checkpoints\"\n    os.makedirs(base_dir, exist_ok=True)\n    filename = os.path.basename(full_path)\n    copyfile(full_path, os.path.join(base_dir, filename))\n    \ncopy_weights('xception')","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nfrom pathlib import Path\nfrom facenet_pytorch import MTCNN\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport cv2\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.axes_grid1 import ImageGrid\nfrom PIL import Image\nfrom tqdm.auto import tqdm\nfrom time import time\nimport shutil","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"markdown","source":"# 1. Extract first face from all sample videos"},{"metadata":{"trusted":true},"cell_type":"code","source":"list_files = [str(x) for x in Path('/kaggle/input/deepfake-detection-challenge/test_videos').glob('*.mp4')] + \\\n             [str(x) for x in Path('/kaggle/input/deepfake-detection-challenge/train_sample_videos').glob('*.mp4')]\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\nmtcnn = MTCNN(keep_all=False, select_largest=False, device=device, min_face_size = 60)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_frame(file, folder):\n    reader = cv2.VideoCapture(file)\n    _, image = reader.read()\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    pilimg = Image.fromarray(image)\n    boxes, probs = mtcnn.detect(pilimg)\n    if boxes is None:\n        return\n    if len(boxes) > 0:\n        try:\n            best_index = probs.argmax()\n            box = [int(x) for x in boxes[best_index].tolist()]\n            face_image = image[box[1]:box[3], box[0]:box[2]]\n            pilface = Image.fromarray(face_image)\n            imgfile = f'{Path(file).stem}.jpg'\n            pilface.save(Path(folder)/imgfile)\n        except:\n            return","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"folder = '/kaggle/working/faces'\nPath(folder).mkdir(parents=True, exist_ok=True)\nfor file in tqdm(list_files):\n    save_frame(file, folder)\n\nface_files = [str(x) for x in Path(folder).glob('*')]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 2. Calculate embedding vectors from all images\nHere I'm using one of my pre-trained models, as the point of this notebook is not training/generating fake/real predictions, but group similar faces using embeddings."},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nmodel = pretrainedmodels.__dict__['xception'](num_classes=1000, pretrained='imagenet')\nmodel.eval();\nnum_ftrs = model.last_linear.in_features\nmodel.last_linear = nn.Linear(num_ftrs, 2)\nmodel = model.to(device)\n\ns = torch.load('/kaggle/input/deepfakemodels/london.pt', map_location=device)\nmodel.load_state_dict(s)\nmodel.eval();\n\ntf_img = utils.TransformImage(model)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### The magic happens here\nThis cell below calculates a forward pass through almost all NN. It stops in the 2nd last fully connected layer. The output for each image is a vector with 2048 dimensions. Later, the algorithm will find similar faces by grouping these vectors, finding which are closest to each other."},{"metadata":{"trusted":true},"cell_type":"code","source":"def embeddings(model, input):\n    f = model.features(input)\n    x = nn.ReLU(inplace=True)(f)\n    x = F.adaptive_avg_pool2d(x, (1, 1))\n    x = x.view(x.size(0), -1)\n    return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"list_embs = []\nfor face in tqdm(face_files):\n    t = tf_img(Image.open(face)).to(device)\n    e = embeddings(model, t.unsqueeze(0)).squeeze().cpu().detach().numpy().tolist()\n    list_embs.append(e)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.DataFrame({'faces': face_files, 'embeddings': list_embs})\ndf['videos'] = df['faces'].apply(lambda x: f'{Path(x).stem}.mp4')\ndf = df[['videos', 'faces', 'embeddings']]\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 4. Get similar faces using Spotify's Annoy"},{"metadata":{"trusted":true},"cell_type":"code","source":"from annoy import AnnoyIndex\n\nf = len(df['embeddings'][0])\nt = AnnoyIndex(f, metric='euclidean')\nntree = 50\n\nfor i, vector in enumerate(df['embeddings']):\n    t.add_item(i, vector)\n_  = t.build(ntree)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_similar_images_annoy(img_index):\n    t0 = time()\n    v, f  = df.iloc[img_index, [0, 1]]\n    similar_img_ids = t.get_nns_by_item(img_index, 8)\n    return v, f, df.iloc[similar_img_ids]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_idx = np.random.choice(len(df))  # 166, # 302\nv, f, s = get_similar_images_annoy(sample_idx)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = plt.figure(figsize=(15, 7))\ngs = fig.add_gridspec(2, 6)\nax1 = fig.add_subplot(gs[0:2, 0:2])\nax2 = fig.add_subplot(gs[0, 2])\nax3 = fig.add_subplot(gs[0, 3])\nax4 = fig.add_subplot(gs[0, 4])\nax5 = fig.add_subplot(gs[0, 5])\nax6 = fig.add_subplot(gs[1, 2])\nax7 = fig.add_subplot(gs[1, 3])\nax8 = fig.add_subplot(gs[1, 4])\nax9 = fig.add_subplot(gs[1, 5])\naxx = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9]\nlist_plot = [face_files[sample_idx]] + s['faces'].values.tolist()\nfor i, ax in enumerate(axx):\n    ax.imshow(plt.imread(list_plot[i]))\n    ax.xaxis.set_visible(False)\n    ax.yaxis.set_visible(False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# All these frames are from different videos!\n# From a given reference (image on the right), the algorithm finds the 8 most similar faces in all samples (8 small images on the left)!"},{"metadata":{},"cell_type":"markdown","source":"To be continued..."},{"metadata":{"trusted":true},"cell_type":"code","source":"shutil.rmtree(Path('/kaggle/working/faces'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# References:\n- https://blog.usejournal.com/fastai-image-similarity-search-pytorch-hooks-spotifys-annoy-9161bf517aaf\n- https://towardsdatascience.com/similar-images-recommendations-using-fastai-and-annoy-16d6ceb3b809\n- https://towardsdatascience.com/finding-similar-images-using-deep-learning-and-locality-sensitive-hashing-9528afee02f5"},{"metadata":{"trusted":true},"cell_type":"code","source":"","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":1}