Spaces:
Sleeping
Sleeping
| 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) | |