Update app.py

app.py

@@ -1,62 +1,58 @@
 import pandas as pd
 import numpy as np
 import tensorflow as tf
+from keras.models import load_model
+import gradio as gr
 
-# classes:
+# Extended classes and labels
 classes = [
-'car',
-'house',
-'wine bottle',
-'chair',
-'table',
-'tree',
-'camera',
-'fish',
-'rain',
-'clock',
-'hat'
+    'car', 'house', 'wine bottle', 'chair', 'table',
+    'tree', 'camera', 'fish', 'rain', 'clock', 'hat',
+    'dog', 'cat', 'bicycle', 'plane', 'book', 'computer'
 ]
 
-# labels :
-labels = {
-'car': 0,
-'house': 1,
-'wine bottle': 2,
-'chair': 3,
-'table': 4,
-'tree': 5,
-'camera': 6,
-'fish': 7,
-'rain': 8,
-'clock': 9,
-'hat': 10
-}
-
+labels = {name: index for index, name in enumerate(classes)}
 num_classes = len(classes)
 
-# load the model:
-from keras.models import load_model
-model = load_model('sketch_recogination_model_cnn.h5')
+# Load the model
+model = load_model('sketch_recognition_model_cnn.h5')
 
-# Predict function for interface:
+# Predict function for interface
 def predict_fn(image):
-
-  # preprocessing the size:
-  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 = np.array(grayscale_image)
-
-  # model requirements:
-  image = image.reshape(1,28,28,1)
-  label = tf.constant(model.predict(image).reshape(num_classes))   # giving 2D output so 1D
-
-  # predict:
-  predicted_index = tf.argmax(label)
-  class_name = [name for name, index in labels.items() if predicted_index == index][0]
-  return class_name
-
-
-# application interface:
-import gradio as gr
-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", interpretation='default', article="Draw large with thick stroke.").launch()
+    """
+    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",
+    interpretation='default',
+    article="Draw large with thick stroke."
+).launch()