pipeline.py

import cv2
import numpy as np
import tensorflow as tf
import zipfile
import tensorflow_addons as tfa

# Set random seed for reproducibility.
tf.random.set_seed(42)

# Extract model if zipped.
local_zip = "./efficientnet-b0.zip"
zip_ref = zipfile.ZipFile(local_zip, 'r')
zip_ref.extractall()
zip_ref.close()

# Load model with custom objects
custom_objects = {"RectifiedAdam": tfa.optimizers.RectifiedAdam}
model = tf.keras.models.load_model("efficientnet-b0/", custom_objects=custom_objects)

class DetectionPipeline:
    def __init__(self, batch_size=1):
        self.batch_size = batch_size
    
    def __call__(self, filename):
        print('Processing image...')
        image = cv2.cvtColor(filename, cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, (224, 224))
        return image

# Initialize image detection pipeline
detection_image_pipeline = DetectionPipeline()

def deepfakes_image_predict(input_image):
    face = detection_image_pipeline(input_image)
    face2 = face / 255.0
    pred = model.predict(np.expand_dims(face2, axis=0))[0]
    real, fake = pred[0], pred[1]
    
    if real > 0.5:
        text = f"The image is REAL. \n Deepfakes Confidence: {round(100 - (real * 100), 3)}%"
    else:
        text = f"The image is FAKE. \n Deepfakes Confidence: {round(fake * 100, 3)}%"
    
    return text