codewithharsha's picture
Update app.py
55b44d3 verified
import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
import os
import io
# --- Configuration ---
PHOTO_SIZE = 224
MODEL_FILE_NAME = 'vgg_model50.h5' # Make sure this matches your uploaded model file
CLASS_NAMES = ["Non-Autstic", "Autstic"] # Match indices 0 and 1 from your notebook
# --- Load the Keras model ---
# Ensure the model file is in the same directory in the Space repository
model_path = os.path.join(os.path.dirname(__file__), MODEL_FILE_NAME)
if not os.path.exists(model_path):
# Basic error handling if model is missing
raise FileNotFoundError(f"Model file '{MODEL_FILE_NAME}' not found. Please upload it to the Space.")
try:
model = tf.keras.models.load_model(model_path)
print("Model loaded successfully.")
except Exception as e:
print(f"Error loading model: {e}")
# You might want more robust error handling or display in the Gradio interface
model = None
def preprocess_image(pil_image):
"""
Preprocesses the PIL image object for the VGG16 model.
Resizes to (PHOTO_SIZE, PHOTO_SIZE), normalizes to [0, 1].
"""
try:
# Gradio provides a PIL image object directly
img = pil_image.convert('RGB') # Ensure image is RGB
img = img.resize((PHOTO_SIZE, PHOTO_SIZE))
np_image = np.array(img).astype('float32') / 255.0
# Expand dimensions to create batch size of 1
np_image = np.expand_dims(np_image, axis=0)
print(f"Image preprocessed successfully. Shape: {np_image.shape}")
return np_image
except Exception as e:
print(f"Error preprocessing image: {e}")
return None
def predict_autism(image_input):
"""
Takes a PIL Image input from Gradio, preprocesses, predicts, and returns the class name.
"""
if model is None:
return "Error: Model not loaded." # Or raise an error
print(f"Received image of type: {type(image_input)}") # Should be PIL Image
# Preprocess the image (Gradio image input provides PIL image)
processed_image = preprocess_image(image_input)
if processed_image is None:
return "Error: Image preprocessing failed."
# Make prediction
print("Making prediction...")
prediction = model.predict(processed_image)
predicted_class_index = np.argmax(prediction, axis=1)[0]
predicted_class_name = CLASS_NAMES[predicted_class_index]
confidence = float(np.max(prediction)) # Get the confidence score
print(f"Prediction result index: {predicted_class_index}, Class: {predicted_class_name}, Confidence: {confidence:.4f}")
# Return prediction as a dictionary (Gradio handles JSON conversion for API)
# The key 'label' often works well with Gradio output components
# Or return just the string if using a simple Textbox output
# return {CLASS_NAMES[0]: float(1-confidence), CLASS_NAMES[1]: confidence} # Example for Label output
return predicted_class_name # Simpler for Textbox output
# --- Create Gradio Interface ---
# Input: Image Upload
# Output: Textbox to display the predicted class
# Allow flagging for feedback (optional but good practice)
# Add title and description
# Provide example images if available in your Space repo (e.g., in an 'examples' folder)
# examples_folder = os.path.join(os.path.dirname(__file__), "examples")
# example_images = [os.path.join(examples_folder, img) for img in os.listdir(examples_folder)] if os.path.exists(examples_folder) else None
iface = gr.Interface(
fn=predict_autism,
inputs=gr.Image(type="pil", label="Upload Image"), # Input is PIL format
outputs=gr.Textbox(label="Prediction Result"), # Output is simple text
# outputs=gr.Label(num_top_classes=2), # Alternative: Label output shows confidences
title="Autism Classification from Facial Images (VGG16)",
description="Upload a facial image to classify as Autistic or Non-Autistic using a VGG16 model.",
allow_flagging="never",
# examples=example_images # Uncomment if you add example images
)
# --- Launch the Gradio app ---
# When run on Hugging Face Spaces, it automatically uses the Space's URL
if __name__ == "__main__":
iface.launch() # share=True is not needed on Spaces