doodle / app.py2
Apex-X's picture
Rename app.py to app.py2
3d81129 verified
import pandas as pd
import numpy as np
import tensorflow as tf
from keras.models import load_model
import gradio as gr
# Extended classes and labels
classes = [
'car', 'house', 'wine bottle', 'chair', 'table',
'tree', 'camera', 'fish', 'rain', 'clock', 'hat',
'dog', 'cat', 'bicycle', 'plane', 'book', 'computer'
]
labels = {name: index for index, name in enumerate(classes)}
num_classes = len(classes)
# Load the model - update this path to the actual location of your model file
from keras.models import load_model
model = load_model('sketch_recogination_model_cnn.h5')
# Predict function for interface
def predict_fn(image):
"""
Predict the class of a drawn image.
Args:
image: The input image drawn by the user.
Returns:
The predicted class name.
"""
try:
# Preprocessing the image
resized_image = tf.image.resize(image, (28, 28)) # Resize image to (28, 28)
grayscale_image = tf.image.rgb_to_grayscale(resized_image) # Convert image to grayscale
image_array = np.array(grayscale_image) / 255.0 # Normalize the image
# Prepare image for model input
image_array = image_array.reshape(1, 28, 28, 1) # Add batch dimension
predictions = model.predict(image_array).reshape(num_classes) # 2D output to 1D
# Predict the class index
predicted_index = tf.argmax(predictions).numpy() # Get the index of the highest score
class_name = classes[predicted_index] # Retrieve the class name
return class_name
except Exception as e:
return f"Error in prediction: {str(e)}"
# Gradio application interface
gr.Interface(
fn=predict_fn,
inputs="paint",
outputs="label",
title="DoodleDecoder",
description="Draw something from: Car, House, Wine bottle, Chair, Table, Tree, Camera, Fish, Rain, Clock, Hat, Dog, Cat, Bicycle, Plane, Book, Computer",
article="Draw large with thick stroke."
).launch()