AI_Image_Detector / pipeline.py
truens66's picture
Update pipeline.py
4b2a093 verified
raw
history blame
956 Bytes
import cv2
import numpy as np
import tensorflow as tf
# Load model
model = tf.keras.models.load_model("efficientnet-b0/")
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