Spaces:
Sleeping
Sleeping
File size: 3,501 Bytes
a61e6d6 4838a1d a61e6d6 4838a1d a61e6d6 4838a1d 3b04332 4838a1d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
from PIL import Image
# --- CONFIGURATION (MUST MATCH YOUR TRAINING) ---
# The model file is automatically available in the Space's file system
MODEL_PATH = 'weapon_classifier_final_tuned.keras'
IMG_SIZE = (224, 224)
# The names of your classes, corresponding to the order they were generated (0=Not_Weapon, 1=Weapon)
CLASS_NAMES = ['Not a Weapon', 'Weapon']
# --- END CONFIGURATION ---
# 1. LOAD MODEL GLOBALLY
# Loading the model here ensures it only happens once when the app starts.
try:
classifier_model = load_model(MODEL_PATH)
print("Model loaded successfully for Gradio interface.")
except Exception as e:
print(f"Error loading model: {e}")
# Fallback in case of model load failure
classifier_model = None
# 2. DEFINE THE PREDICTION FUNCTION
# Gradio automatically passes the uploaded image as a NumPy array (or PIL Image, depending on 'type').
def classify_weapon(input_img_array):
"""
Takes a NumPy image array, preprocesses it, and returns the classification.
"""
if classifier_model is None:
return "Error: Model failed to load."
# Gradio passes the image as a NumPy array with shape (H, W, 3) and values 0-255.
# a) Resize: Must resize to the model's required input size (224x224 for VGG16).
# We use PIL/Image.fromarray and resize before conversion to prevent distortion
img = Image.fromarray(input_img_array.astype('uint8'))
img = img.resize(IMG_SIZE)
# b) Convert to array and add batch dimension (1, 224, 224, 3)
img_array = np.array(img).astype('float32')
img_array = np.expand_dims(img_array, axis=0)
# c) Normalize (Rescale): Must match the 1./255 scaling used during training
processed_image = img_array / 255.0
# d) Make Prediction
# Prediction returns a 1x1 array, e.g., [[0.95]]
prediction = classifier_model.predict(processed_image)
# e) Interpret Output for Gradio Label Component
probability = prediction[0][0]
# Gradio's Label component expects a dictionary mapping labels to probabilities.
# We calculate the confidence for both classes based on the Sigmoid output.
confidences = {
CLASS_NAMES[1]: float(probability), # Weapon confidence
CLASS_NAMES[0]: float(1 - probability) # Not a Weapon confidence
}
return confidences
# 3. CREATE THE GRADIO INTERFACE
# Gradio will handle the 'Submit' button click automatically.
gr_interface = gr.Interface(
# The function to run when the user submits an image
fn=classify_weapon,
# Input component: Image component set to resize input to 224x224 for display
# Gradio automatically converts the image to a NumPy array for the function.
inputs=gr.Image(
label="Upload Image for Classification",
type="numpy" # Pass NumPy array to the function
),
# Output component: Label displays the class names and confidence scores.
outputs=gr.Label(
num_top_classes=2,
label="Classification Result"
),
title="Weapon or Fruit Detector",
description="Upload an image to classify it as 'Weapon' or 'Not a Weapon'."
)
# 4. LAUNCH THE INTERFACE
# In a Hugging Face Space, the launch() call is automatically handled by the environment.
# If running locally, you would use: demo.launch()
gr_interface.launch() |