ELM / app.py
LordChristoff's picture
Create app.py
c10e38a verified
import gradio as gr
import zipfile
import tensorflow as tf
import numpy as np
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import os
from tqdm import tqdm
def unzip_and_load(zip_file_path, data_dir):
with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
zip_ref.extractall(data_dir)
unzip_and_load('realfake.zip', 'unzipped_data')
train_datagen = ImageDataGenerator(
rescale=1./255,
)
batch_size = 50 # Change Batch Size (Default 32)
train_generator = train_datagen.flow_from_directory(
'unzipped_data',
target_size=(150, 150),
batch_size=batch_size,
class_mode='binary'
)
class ELM(object):
def __init__(self, input_size, output_size, hidden_size):
self.input_size = input_size
self.output_size = output_size
self.hidden_size = hidden_size
self.weight = np.random.normal(size=[self.hidden_size, self.input_size])
self.bias = np.random.normal(size=[self.hidden_size])
self.beta = np.random.normal(size=[self.hidden_size, self.output_size])
def sigmoid(self, x):
return 1.0 / (1.0 + np.exp(-x))
def relu(self, x):
return tf.nn.relu(x)
def predict(self, X):
X = tf.convert_to_tensor(X, dtype=tf.float32)
X = tf.reshape(X, [X.shape[0], -1]) # Flatten the input data
y = self.relu((X @ self.weight.T) + self.bias) @ self.beta
return y
def train(self, X, y):
X = tf.convert_to_tensor(X, dtype=tf.float32)
y = tf.convert_to_tensor(y, dtype=tf.float32)
X = tf.reshape(X, [X.shape[0], -1])
H = self.relu((X @ self.weight.T) + self.bias)
H_inv = tf.linalg.pinv(H)
# Add a new dimension to y to make it a column vector
y = tf.expand_dims(y, axis=-1) # Now y has shape (32, 1)
self.beta = H_inv @ y
loss = tf.reduce_mean(tf.square(self.predict(X) - y)) # Replace with your loss function
return loss
def calculate_loss(self, X, y): # define the missing function
X = tf.convert_to_tensor(X, dtype=tf.float32)
y = tf.convert_to_tensor(y, dtype=tf.float32)
y_pred = self.predict(X)
loss = tf.reduce_mean(tf.square(y_pred - y))
return loss
img_width = 150
img_height = 150
hidden_size = 100
elm = ELM(img_width * img_height * 3, 1, hidden_size)
num_epochs = 10 # Change the amount of Epochs (Default 10)
steps_per_epoch = len(train_generator)
for epoch in range(num_epochs):
train_generator.reset()
with tqdm(total=steps_per_epoch, desc=f"Training progress Epoch {epoch+1}/{num_epochs}", unit="batch", colour="green") as pbar:
for batch_x, batch_y in train_generator:
elm.train(batch_x, batch_y)
pbar.update(1)
pbar.set_postfix(loss=elm.calculate_loss(batch_x, batch_y))
if pbar.n == pbar.total:
break
val_datagen = ImageDataGenerator(rescale=1./255)
val_generator = val_datagen.flow_from_directory(
'unzipped_data',
target_size=(150, 150),
batch_size=batch_size,
class_mode='binary'
)
train_acc = []
val_acc = []
losses = []
import gradio as gr
from tensorflow.keras.preprocessing.image import img_to_array
from PIL import Image
import numpy as np
def predict_image(image):
"""Preprocesses and predicts on a single image."""
img_width = 150
img_height = 150
img = Image.fromarray(np.uint8(image)).convert(
"RGB"
) # Convert to PIL Image and ensure RGB format
img = img.resize((img_width, img_height)) # Resize using PIL
if img is None:
return "Invalid image: Resizing failed"
x = img_to_array(img)
x = np.expand_dims(x, axis=0) # Add batch dimension
x = x / 255.0 # Normalize
prediction = elm.predict(x)
# Ensure prediction is a NumPy array and handle potential shape issues
prediction = np.array(prediction)
if prediction.size > 0:
# Calculate percentages based on prediction value
real_percentage = (1 - prediction.item()) * 100
fake_percentage = prediction.item() * 100
return f"Real: {real_percentage:.2f}% Generated: {fake_percentage:.2f}%"
else:
return "Prediction not available"
interface = gr.Interface(
fn=predict_image,
inputs="image",
outputs="text",
allow_flagging="manual", # Allow users to flag uncertain predictions
flagging_options=[
"incorrect",
"other",
], # specify the options the user can select when flagging
css="""
.gradio-component-image {
width: 300px;
}
""", # Add your CSS here within the gr.Interface constructor
)
interface.launch(share=True, debug=True)