muhyousf's picture
4th update app,py
20e58ac verified
Raw
History Blame Contribute Delete
2.46 kB
import gradio as gr
import tensorflow as tf
import numpy as np
import pandas as pd
from tensorflow.keras.preprocessing import image
import os
# Check if the model file exists before loading
MODEL_PATH = "civil_tool_classifier.h5"
if not os.path.exists(MODEL_PATH):
print(f"Error: Model file '{MODEL_PATH}' not found.")
model = None
else:
try:
model = tf.keras.models.load_model(MODEL_PATH)
except Exception as e:
print(f"Error loading model: {e}")
model = None
# Check if the CSV file exists before loading
CSV_PATH = "tools.csv"
if not os.path.exists(CSV_PATH):
print(f"Error: CSV file '{CSV_PATH}' not found.")
tools_df = None
CLASS_NAMES = []
else:
try:
tools_df = pd.read_csv(CSV_PATH)
CLASS_NAMES = list(tools_df["tool_name"].values)
except Exception as e:
print(f"Error loading CSV file: {e}")
tools_df = None
CLASS_NAMES = []
def predict_tool(img):
if model is None or tools_df is None:
return {"tool": "Error", "usage": "Model/CSV not loaded", "safety": "N/A"}
# Pre-process the image
img = image.array_to_img(img).resize((224, 224))
img_array = image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0) / 255.0
# Predict
predictions = model.predict(img_array)
class_idx = np.argmax(predictions[0])
# Get the tool info from the DataFrame
if class_idx < len(CLASS_NAMES):
tool_name = CLASS_NAMES[class_idx]
try:
row = tools_df[tools_df["tool_name"] == tool_name].iloc[0]
usage = row.get("tool_usage", "Usage not available.")
safety = row.get("tool_safety", "Safety not available.")
except Exception:
usage = "Usage information not available."
safety = "Safety guidance not available."
else:
tool_name = "Unknown Tool"
usage = "Not found in the list."
safety = "Not found in the list."
# πŸ‘‡ Return JSON instead of multiple textboxes
return {"tool": tool_name, "usage": usage, "safety": safety}
# Gradio interface with JSON output
demo = gr.Interface(
fn=predict_tool,
inputs=gr.Image(type="numpy"),
outputs=gr.JSON(), # πŸ‘ˆ JSON output for API
title="Civil Tool Classifier",
description="Upload an image of a civil tool to get its details."
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860)