storesource's picture
Update app.py
c0b6df5 verified
import gradio as gr
import tensorflow as tf
from PIL import Image, ImageOps
import numpy as np
import os
import cv2
model_path = "IM_417_128.keras"
print("Loading Model...")
model = tf.keras.models.load_model(model_path)
model.summary()
labelled_images = [f"labelled/{i}.png" for i in range(1, 418)]
# Define regression function
def predict_regression(image):
# Preprocess image
print(f" {type(image)} image layers ===> { len(image['layers']) }")
image = np.squeeze(image["layers"][0][:, :, -1:], axis=-1)
image = np.clip(image, 0, 255).astype('uint8')
resized_image = cv2.resize(image, (128, 128), interpolation=cv2.INTER_CUBIC)
# Convert to grayscale
# gray_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2GRAY)
# Create a mask for values greater than 10
mask = resized_image > 10
# Use the mask to assign values
gray_image = np.where(mask, 255, 0)
print(f"gray_image image TEST :: {gray_image.shape}")
# Convert the image to a NumPy array
image_array = np.array(gray_image)
# Invert the colors
inverted_image_array = 255 - image_array
# Predict
predictions = model.predict(inverted_image_array[None, ...]) # Assuming single regression value
predicted_output = np.argmax(predictions, axis=1)
predicted_label = f"labelled/{str(predicted_output[0])}.png"
print(f"prediction IM-417 character :: {predicted_label} :: source {predicted_output}")
output_data = cv2.imread(predicted_label, cv2.COLOR_BGR2GRAY)
return cv2.resize(output_data, (128, 128)) , image_array
# Create Gradio interface
output_image = gr.Image(height=150, width=150)
i_0 = gr.Image(height=150, width=150)
# sketchpad = gr.Sketchpad(flatten=True)
interface = gr.Interface(fn=predict_regression,
inputs="sketchpad",
outputs=[output_image, i_0],
live=True,
description="A simple mlp classification model for image classification using the sign list from 'The Indus Scripts: Texts, Concordance and Tables'"
)
interface.launch()