predict.py

import argparse
import cv2
import numpy as np
import yaml
import random

from typing import Type
import preprocessing.face_detector as face_detector
from preprocessing.face_detector import VideoDataset, VideoFaceDetector
from torch.utils.data.dataloader import DataLoader

from PIL import Image

import torch
from preprocessing.utils import preprocess_images, _generate_connected_components
from facenet_pytorch import InceptionResnetV1, fixed_image_standardization

from statistics import mean

from albumentations import Compose, RandomBrightnessContrast, HorizontalFlip, FancyPCA, HueSaturationValue, OneOf, ToGray, ShiftScaleRotate, ImageCompression, PadIfNeeded, GaussNoise, GaussianBlur, Rotate, Normalize, Resize
from transforms.albu import IsotropicResize

from models.size_invariant_timesformer import SizeInvariantTimeSformer
from models.efficientnet.efficientnet_pytorch import EfficientNet
from models.baseline import Baseline
import os
from einops import rearrange
from utils import aggregate_attentions, draw_border, save_attention_plots
from models.xception import xception



RANGE_SIZE = 5
SIZE_EMB_DICT = [(1+i*RANGE_SIZE, (i+1)*RANGE_SIZE) if i != 0 else (0, RANGE_SIZE) for i in range(20)]

def detect_faces(video_path, detector_cls: Type[VideoFaceDetector], opt):
    # Init the face detector
    detector = face_detector.__dict__[detector_cls](device=opt.gpu_id)

    # Read the video and its information
    dataset = VideoDataset([video_path])
    loader = DataLoader(dataset, shuffle=False, num_workers=opt.workers, batch_size=1, collate_fn=lambda x: x)
    
    # Detect the faces
    for item in loader: 
        bboxes = {}
        video, indices, fps, frames = item[0]
        bboxes.update({i : b for i, b in zip(indices, detector._detect_faces(frames))})
        found_faces = False
        for key in bboxes:
            if type(bboxes[key]) == list:
                found_faces = True
                break

        if not found_faces:
            raise Exception("No faces found.")

    return bboxes

def extract_crops(video_path, bboxes_dict):

    # Read video frames
    frames = []
    
    capture = cv2.VideoCapture(video_path)
    frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = int(capture.get(5))

    for i in range(frames_num):
        capture.grab()
        success, frame = capture.retrieve()
        if not success:
            continue
        frames.append(frame)

    # Extract the faces crops
    explored_indexes = []
    crops = []

    for i in range(0, len(frames), fps):
        while str(i) not in bboxes_dict:
            if i == frames_num - 1:
                i -= 1
            if i in explored_indexes:
                break
            else:
                explored_indexes.append(i)

        frame = frames[i]
        index = i
        limit = i + fps - 1
        keys = [int(x) for x in list(bboxes_dict.keys())]

        while index < limit:
            index += 1
            if index in keys and bboxes_dict[index] is not None:
                break
        if index == limit:
            continue

        bboxes = bboxes_dict[index]
        for bbox in bboxes:
            xmin, ymin, xmax, ymax = [int(b * 2) for b in bbox]
            w = xmax - xmin
            h = ymax - ymin

            # Add some padding to catch background too
            p_h = h // 3
            p_w = w // 3
            
            crop_h = (ymax + p_h) - max(ymin - p_h, 0)
            crop_w = (xmax + p_w) - max(xmin - p_w, 0)

            # Make the image square
            if crop_h > crop_w:
                p_h -= int(((crop_h - crop_w)/2))
            else:
                p_w -= int(((crop_w - crop_h)/2))

            # Extract the face from the frame
            crop = frame[max(ymin - p_h, 0):ymax + p_h, max(xmin - p_w, 0):xmax + p_w]
            
            # Check if out of bound and correct
            h, w = crop.shape[:2]
            if h > w:
                diff = int((h - w)/2)
                if diff > 0:         
                    crop = crop[diff:-diff,:]
                else:
                    crop = crop[1:,:]
            elif h < w:
                diff = int((w - h)/2)
                if diff > 0:
                    crop = crop[:,diff:-diff]
                else:
                    crop = crop[:,:-1]

            # Add the extracted face to the list
            crops.append((i, Image.fromarray(crop), bbox))

    return crops

def cluster_faces(crops, valid_cluster_size_ratio = 0.20, similarity_threshold = 0.45):

    # Convert crops to PIL images
    crops_images = [row[1] for row in crops]
    
    # Extract the embeddings
    embeddings_extractor = InceptionResnetV1(pretrained='vggface2').eval().to(device)
    faces = [preprocess_images(face) for face in crops_images]
    faces = np.stack([np.uint8(face) for face in faces])
    faces = torch.as_tensor(faces)
    faces = faces.permute(0, 3, 1, 2).float()
    faces = fixed_image_standardization(faces)
    face_recognition_input = faces.cuda()
    embeddings = []
    embeddings = embeddings_extractor(face_recognition_input).detach().cpu().numpy()

    # Clustering
    valid_cluster_size = int(len(faces) * valid_cluster_size_ratio)
    similarities = np.dot(np.array(embeddings), np.array(embeddings).T)
    
    components = _generate_connected_components(
        similarities, similarity_threshold=similarity_threshold
    )
    components = [sorted(component) for component in components]

    clustered_faces = {}
    for identity_index, component in enumerate(components):
        for index, face_index in enumerate(component):
            component[index] = crops[face_index]
        
        clustered_faces[identity_index] = component

    return clustered_faces

def get_identity_information(identity, faces):
    mean_side = mean([row[1].size[0] for row in faces])   
    number_of_faces = len(faces)
    return [identity, mean_side, number_of_faces, faces]

def get_sorted_identities(identities, discarded_faces, max_identities = 2, num_frames = 16):
    sorted_identities = []
    discarded_faces = []
    for identity in identities:
        sorted_identities.append(get_identity_information(identity, identities[identity]))

    '''
    # If no faces have been found, use the discarded faces
    if len(sorted_identities) == 0:
        sorted_identities.append(self.get_identity_information(identities))
        discarded_faces = []
    '''

    # Sort identities based on faces size
    sorted_identities = sorted(sorted_identities, key=lambda x:x[1], reverse=True)

    if len(sorted_identities) > max_identities:
        sorted_identities = sorted_identities[:max_identities]

    # Adjust the identities list faces number
    identities_number = len(sorted_identities)
    available_additional_faces = []
    if identities_number > 1:
        max_faces_per_identity = {1: [num_frames], 
                  2:  [int(num_frames/2), int(num_frames/2)],
                  3:  [int(num_frames/3), int(num_frames/3), int(num_frames/4)],
                  4:  [int(num_frames/3), int(num_frames/3), int(num_frames/8), int(num_frames/8)]}

        max_faces_per_identity = max_faces_per_identity[identities_number]
        for i in range(identities_number):
            if sorted_identities[i][2] < max_faces_per_identity[i] and i < identities_number - 1:
                sorted_identities[i+1][2] += max_faces_per_identity[i] - sorted_identities[i][2] 
                available_additional_faces.append(0)
            elif sorted_identities[i][2] > max_faces_per_identity[i]:
                available_additional_faces.append(sorted_identities[i][2] - max_faces_per_identity[i])
                sorted_identities[i][2] = max_faces_per_identity[i]
            else:
                available_additional_faces.append(0)

    else: # If only one identity is in the video, all the frames are assigned to this identity
        sorted_identities[0][2] = num_frames
        available_additional_faces.append(0)


    # Check if we found enough faces to fullfill the input sequence, otherwise go back and add some faces from previous identities
    input_sequence_length = sum(faces_number for _, _, faces_number, _ in sorted_identities)
    if input_sequence_length < num_frames:
        for i in range(identities_number):
            needed_faces = num_frames - input_sequence_length
            if available_additional_faces[i] > 0:
                added_faces = min(available_additional_faces[i], needed_faces)
                sorted_identities[i][2] += added_faces
                input_sequence_length += added_faces
                if input_sequence_length == num_frames:
                    break
        # If not enough faces have been found, add some "dummy" images in the last identity
        if input_sequence_length < num_frames:
            needed_faces = num_frames - input_sequence_length
            sorted_identities[-1][2] += needed_faces
            input_sequence_length += needed_faces
    
    return sorted_identities, discarded_faces

def create_val_transform(size, additional_targets):
    return Compose([
        IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC),
        PadIfNeeded(min_height=size, min_width=size, border_mode=cv2.BORDER_CONSTANT),
        Resize(height=size, width=size)
    ],  additional_targets = additional_targets
    )

def generate_masks(video_path, identities, discarded_faces, num_frames, image_size, num_patches):
    mask = []
    last_range_end = 0
    sequence = []
    size_embeddings = []
    
    images_frames = []
    for identity_index, identity in enumerate(identities):
        max_faces = identity[2]
        identity_images = identity[3]
        '''
        # If no faces were considered for a frame during clustering, probably it is inside the discarded faces
        if identity_index == 0 and len(discarded_faces) > 0:
            frames = [int(os.path.basename(image_path).split("_")[0]) for image_path in identity_faces]
            discarded_frames = [int(os.path.basename(image_path).split("_")[0]) for image_path in discarded_faces]
            missing_frames = list(set(discarded_frames) - set(frames))
            missing_faces = [discarded_faces[discarded_frames.index(missing_frame)] for missing_frame in missing_frames]
            
            if len(missing_faces) > 0:
                identity_faces = identity_faces + missing_faces # Add the missing faces to the identity
        '''


        # Select uniformly the frames in an alternate way
        if len(identity_images) > max_faces:
            idx = np.round(np.linspace(0, len(identity_images) - 2, max_faces)).astype(int)
            identity_images = np.asarray(identity_images)[idx]
            
        images_frames.extend(identity_image[0] for identity_image in identity_images)
        identity_images = [identity_image[1] for identity_image in identity_images]

        # Generate size embeddings
        capture = cv2.VideoCapture(video_path)
        width  = capture.get(3)  
        height = capture.get(4) 
        video_area = width*height/2
        identity_size_embeddings = []
        
        for image_index, image in enumerate(identity_images):
            # Get face-frame area ratio for size embedding
            face_area = image.size[0] * image.size[1]
            ratio = int(face_area * 100 / video_area)
            side_ranges = list(map(lambda a_: ratio in range(a_[0], a_[1] + 1), SIZE_EMB_DICT))
            identity_size_embeddings.append(np.where(side_ranges)[0][0]+1)
      

        # If the readed faces are less than max_faces we need to add empty images and generate the mask
        if len(identity_images) < max_faces: 
            diff = max_faces - len(identity_size_embeddings)
            identity_size_embeddings = np.concatenate((identity_size_embeddings, np.zeros(diff)))
            identity_images.extend([np.zeros((image_size, image_size, 3), dtype=np.uint8) for i in range(diff)])
            mask.extend([1 if i < max_faces - diff else 0 for i in range(max_faces)])
            images_frames.extend([max(images_frames) for i in range(diff)])
        else: # Otherwise all the faces are valid
            mask.extend([1 for i in range(max_faces)])

        # Compose the size_embedding and sequence list
        size_embeddings.extend(identity_size_embeddings)
        sequence.extend(identity_images)

    # Transform the images, the same transformation is applied to all the faces in the same video
    sequence = [np.asarray(image) for image in sequence]
    additional_targets_keys = ["image" + str(i) for i in range(num_frames)]
    additional_targets_values = ["image" for i in range(num_frames)]
    additional_targets = dict(zip(additional_targets_keys, additional_targets_values))

 
    transform = create_val_transform(image_size, additional_targets)  
    if len(sequence) == 8:
        transformed_images = transform(image=sequence[0], image1=sequence[1], image2=sequence[2], image3=sequence[3], image4=sequence[4], image5=sequence[5], image6=sequence[6], image7=sequence[7])
    elif len(sequence) == 16:
        transformed_images = transform(image=sequence[0], image1=sequence[1], image2=sequence[2], image3=sequence[3], image4=sequence[4], image5=sequence[5], image6=sequence[6], image7=sequence[7], image8=sequence[8], image9=sequence[9], image10=sequence[10], image11=sequence[11], image12=sequence[12], image13=sequence[13], image14=sequence[14], image15=sequence[15])
    else:
        raise Exception("Invalid number of frames.")

    sequence = [transformed_images[key] for key in transformed_images]
        
    # Generate the identities_mask telling to the model which faces attend to an identity and which to another one
    identities_mask = []
    last_range_end = 0
    for identity_index in range(len(identities)):
        identity_mask = [True if i >= last_range_end and i < last_range_end + identities[identity_index][2] else False for i in range(0, num_frames)]
        for k in range(identities[identity_index][2]):
            identities_mask.append(identity_mask)
        last_range_end += identities[identity_index][2]

    # Generate coherent temporal-positional embedding
    images_frames_positions = {k: v+1 for v, k in enumerate(sorted(set(images_frames)))}
    frame_positions = [images_frames_positions[frame] for frame in images_frames]   
    if num_patches != None: 
        positions = [[i+1 for i in range(((frame_position-1)*num_patches), num_patches*(frame_position))] for frame_position in frame_positions]
        positions = sum(positions, []) # Merge the lists
        positions.insert(0,0) # Add CLS
    else:
        positions = []

    tokens_per_identity = [(identities[i][0], identities[i][2]*num_patches + identities[i-1][2]*num_patches) if i > 0 else (identities[i][0], identities[i][2]*num_patches) for i in range(len(identities))]     

    return torch.tensor([sequence]).float(), torch.tensor([size_embeddings]).int(), torch.tensor([mask]).bool(), torch.tensor([identities_mask]).bool(), torch.tensor([positions]), tokens_per_identity


def predict(video_path, clustered_faces, config, opt, discarded_faces = None):

    # Load required weights for feature extractor
    if opt.extractor_model == 0: # EfficientNet-B0
        if opt.extractor_weights.lower() == 'imagenet':
            features_extractor = EfficientNet.from_pretrained('efficientnet-b0')
        else:
            features_extractor = EfficientNet.from_name('efficientnet-b0')
            features_extractor.load_matching_state_dict(torch.load(opt.extractor_weights, map_location=torch.device('cpu')))
            print("Custom features extractor weights loaded.")
    else: # XceptionNet
        if opt.extractor_weights.lower() == 'pretrained':
            features_extractor = xception(num_classes=1, pretrain_path="weights/ckpt_iter.pth.tar")
        else:
            features_extractor = xception(num_classes=1, pretrain_path=opt.extractor_weights)



    # Init the model
    model = SizeInvariantTimeSformer(config=config, require_attention=True)
    num_patches = config['model']['num-patches']

    
    features_extractor = torch.nn.DataParallel(features_extractor)
    model = torch.nn.DataParallel(model)

    # Move into GPU
    features_extractor = features_extractor.to(device)    
    model = model.to(device)
    features_extractor.eval()
    model.eval()

    if os.path.exists(opt.model_weights):
        model.load_state_dict(torch.load(opt.model_weights))
    else:
        raise Exception("No checkpoint loaded for the model.")    
    
    identities, discarded_faces  = get_sorted_identities(clustered_faces, discarded_faces)
    videos, size_embeddings, mask, identities_mask, positions, tokens_per_identity = generate_masks(video_path, identities, discarded_faces, config["model"]["num-frames"], config["model"]["image-size"], config["model"]["num-patches"])
    b, f, h, w, c = videos.shape
    videos = videos.to(device)    
    identities_mask = identities_mask.to(device)
    mask = mask.to(device)
    positions = positions.to(device)
    

    with torch.no_grad():
        video = rearrange(videos, "b f h w c -> (b f) c h w")
        features = features_extractor(video)  

        features = rearrange(features, '(b f) c h w -> b f c h w', b = b, f = f)   
        test_pred, attentions = model(features, mask=mask, size_embedding=size_embeddings, identities_mask=identities_mask, positions=positions)
        
        identity_names = [row[0] for row in tokens_per_identity]
        frames_per_identity = [int(row[1] / config["model"]["num-patches"]) for row in tokens_per_identity]
        
        if opt.save_attentions:
            aggregated_attentions, identity_attentions = aggregate_attentions(attentions, config['model']['heads'], config['model']['num-frames'], frames_per_identity)
            save_attention_plots(aggregated_attentions, identity_names, frames_per_identity, config['model']['num-frames'], os.path.basename(video_path))
        else:
            identity_attentions = []
            aggregated_attentions = []
        return torch.sigmoid(test_pred[0]).item(), identity_attentions, aggregated_attentions, identities, frames_per_identity

def get_identities_bboxes(identities):
    identities_bboxes = {}
    for row in identities:
        identity = row[3]
        for face in identity:
            frame = face[0]
            if frame in identities_bboxes:
                identities_bboxes[frame].append(face[2])
            else:
                identities_bboxes[frame] = [face[2]]
    return identities_bboxes


def generate_output_video(video_path, pred, identity_attentions, aggregated_attentions, identities, frames_per_identity):

    identities_bboxes = get_identities_bboxes(identities)
    available_frames_keys = [frame for frame in identities_bboxes]

    cap = cv2.VideoCapture(video_path)
    width  = cap.get(3)  
    height = cap.get(4) 
    fps = int(cap.get(5))
    fourcc = hex(int(cap.get(cv2.CAP_PROP_FOURCC)))
    output = cv2.VideoWriter("examples/preds/"+str(os.path.basename(video_path).replace(".mp4", ".avi")), cv2.VideoWriter_fourcc("X", "V", "I", "D"), fps, (int(width), int(height)))
    frame_index = 0
    while True:
        ret, frame = cap.read()
        if ret:
            nearest_frame_index = min(available_frames_keys, key=lambda x:abs(x - frame_index))
            if nearest_frame_index - frame_index > fps: 
                continue

            bbox = identities_bboxes[nearest_frame_index] 
                
            for identity_index, identity_bbox in enumerate(bbox):
                
                xmin, ymin, xmax, ymax = [int(b * 2) for b in identity_bbox]
                if pred > 0.5:
                    red = 255 * identity_attentions[identity_index]
                    green = 255 - red

                    if red > green:
                        text = 'Fake ' + str(round(pred*100,2)) + "%"
                    else:
                        text = 'Pristine'
                else:
                    green = int(255 * (1 - pred))
                    red = 255 - green
                    text = 'Pristine '  + str(round((1-pred)*100,2)) + "%"

                color = (0, green, red)
                frame = draw_border(frame, (xmin,ymin), (xmax,ymax), color, 2, 10, 20)
                cv2.putText(frame, text, (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)

            output.write(frame)        
        else:
            break
        
        frame_index += 1
    output.release()
    cap.release()

        


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--video_path', type=str,
                        help='Path to the video file')
    parser.add_argument("--detector_type", help="type of the detector", default="FacenetDetector",
                        choices=["FacenetDetector"])
    parser.add_argument('--random_state', default=42, type=int,
                        help='Random state value')
    parser.add_argument('--gpu_id', default=0, type=int,
                        help='ID of GPU to be used')
    parser.add_argument('--workers', default=1, type=int,
                        help='Number of data loader workers.')
    parser.add_argument('--config', type=str, 
                        help="Which configuration to use. See into 'config' folder.")
    parser.add_argument('--model_weights', type=str,
                        help='Model weights.')
    parser.add_argument('--extractor_model', type=int, default=0, 
                        help="Which model use for features extraction (0: EfficientNet; 1: XceptionNet).")
    parser.add_argument('--extractor_weights', default='ImageNet', type=str,
                        help='Path to extractor weights or "imagenet".')
    parser.add_argument('--output_type', default=0, type=int,
                        help='Specify which type of output is requested (0: Prediction; 1: Video)".')
    parser.add_argument('--save_attentions', default=False, action="store_true",
                        help='Save attentions plots.')

    opt = parser.parse_args()
    print(opt)
    
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    with open(opt.config, 'r') as ymlfile:
        config = yaml.safe_load(ymlfile)

    # Check for integrity
    if config['model']['num-frames'] != 8 and config['model']['num-frames'] != 16:
        raise Exception("Invalid number of frames.")
        
    if not os.path.exists(opt.video_path):
        raise Exception("Invalid video path.")
 

    # Setup CUDA settings
    torch.cuda.set_device(opt.gpu_id) 
    torch.backends.cudnn.deterministic = True
    random.seed(opt.random_state)
    torch.manual_seed(opt.random_state)
    torch.cuda.manual_seed(opt.random_state)
    np.random.seed(opt.random_state)


    print("Detecting faces...")
    bboxes_dict = detect_faces(opt.video_path, opt.detector_type, opt)
    print("Face detection completed.")


    print("Cropping faces from the video...")
    crops = extract_crops(opt.video_path, bboxes_dict)
    print("Faces cropping completed.")

    '''
    for j, crop in enumerate(crops):
        cv2.imwrite("outputs/faces/face_{}.png".format(j), np.asarray(crop[1]))
    '''
    
    print("Clustering faces...")
    clustered_faces = cluster_faces(crops)
    print("Faces clustering completed.")


    print("Searching for fakes in the video...")
    pred, identity_attentions, aggregated_attentions, identities, frames_per_identity = predict(opt.video_path, clustered_faces, config, opt)
    if pred > 0.5:
        print("The video is fake ("+str(round(pred*100,2)) + "%), showing video result...")
    else:
        print("The video is pristine ("+str(round((1-pred)*100,2)) + "%), showing video result...")
    if opt.output_type == 0:
        print("Prediction", pred)
    else:
        generate_output_video(opt.video_path, pred, identity_attentions, aggregated_attentions, identities, frames_per_identity)