initial commit

app.py

@@ -0,0 +1,231 @@
+# -*- coding: utf-8 -*-
+"""
+Robotic Grasp Predictor - Gradio Demo
+This script loads (or trains, if necessary) a model to predict the robustness of a robotic grasp.
+It provides a Gradio interface to generate random sensor values, compute a prediction,
+and visualize the grasp with a simple robotic arm drawing.
+
+Dependencies:
+    gradio
+    pandas
+    numpy
+    matplotlib
+    scikit-learn
+    joblib
+
+IMPORTANT:
+    Ensure that the file 'shadow_robot_dataset.csv' is included in your repository.
+    If the pickle files (final_model.pkl and feature_names.pkl) do not exist, the app will train the model,
+    save them (using compression to reduce file size), and use them for subsequent inference.
+"""
+
+import warnings
+warnings.filterwarnings("ignore")
+import os
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from joblib import dump, load
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+from sklearn.ensemble import ExtraTreesRegressor
+from sklearn.pipeline import Pipeline
+import gradio as gr
+
+# Define paths for the pre-trained model, feature names, and dataset
+MODEL_PATH = "final_model.pkl"
+FEATURE_NAMES_PATH = "feature_names.pkl"
+DATASET_PATH = "shadow_robot_dataset.csv"
+
+def load_dataset(dataset_path):
+    """
+    Attempts to load the dataset CSV file.
+    Raises an error if the file does not exist or cannot be read.
+    """
+    if not os.path.exists(dataset_path):
+        raise FileNotFoundError(
+            f"Dataset file '{dataset_path}' not found. "
+            "Please ensure it is included in the repository with the correct permissions."
+        )
+    try:
+        df = pd.read_csv(dataset_path)
+    except Exception as e:
+        raise IOError(
+            f"An error occurred while reading '{dataset_path}': {e}\n"
+            "Please check file permissions and format."
+        )
+    return df
+
+def load_or_train_model():
+    """
+    Loads a pre-trained model and corresponding feature names if available.
+    Otherwise, trains the model from the dataset, saves the model and feature names as compressed pickle files,
+    and returns them.
+    """
+    if os.path.exists(MODEL_PATH) and os.path.exists(FEATURE_NAMES_PATH):
+        final_model = load(MODEL_PATH)
+        feature_names = load(FEATURE_NAMES_PATH)
+        print("Loaded pre-trained model and feature names.")
+    else:
+        print("Pre-trained model not found. Training model; please wait...")
+        # Load and preprocess data
+        df = load_dataset(DATASET_PATH)
+        # Clean up column names and drop identifier columns if present
+        df.columns = df.columns.str.replace(r'\s+', '', regex=True)
+        for col in ['experiment_number', 'measurement_number']:
+            if col in df.columns:
+                df.drop(col, axis=1, inplace=True)
+        
+        # Define target and features
+        y = df['robustness']
+        X = df.drop('robustness', axis=1)
+        
+        # Split the data (use only the training split for model fitting)
+        X_train, _, y_train, _ = train_test_split(X, y, test_size=0.2, random_state=42)
+        
+        # Create a pipeline without caching to avoid large file size.
+        final_model = Pipeline([
+            ('scaler', StandardScaler()),
+            ('model', ExtraTreesRegressor(random_state=42, n_jobs=-1, max_depth=None, n_estimators=100))
+        ])
+        
+        # Train the model on the training data
+        final_model.fit(X_train, y_train)
+        
+        # Save the trained model and feature names with compression
+        feature_names = sorted(X_train.columns)  # Sorted to ensure consistent feature order
+        dump(final_model, MODEL_PATH, compress=3)
+        dump(feature_names, FEATURE_NAMES_PATH, compress=3)
+        print("Model training complete and saved.")
+    return final_model, feature_names
+
+# Load or train the model at startup
+final_model, feature_names = load_or_train_model()
+
+############################################
+# Utility Functions for the Gradio Demo
+############################################
+
+def generate_random_sensors(feature_names, lower_bound, upper_bound):
+    """Generates random sensor values within [lower_bound, upper_bound] for each feature."""
+    sensor_values = {}
+    for feature in feature_names:
+        sensor_values[feature] = round(np.random.uniform(lower_bound, upper_bound), 4)
+    return sensor_values
+
+def draw_robotic_arm(stability):
+    """
+    Draws a simple two-link robotic arm with three fingers and a ball.
+    The drawing changes based on the 'stability' value.
+    """
+    fig, ax = plt.subplots(figsize=(6, 6))
+    
+    # Draw Link 1 (Blue)
+    shoulder = (0, 0)
+    elbow = (0.5, 0.5)
+    ax.plot([shoulder[0], elbow[0]], [shoulder[1], elbow[1]], color='blue', linewidth=3, label='Link 1')
+    ax.plot(shoulder[0], shoulder[1], 'o', color='blue', markersize=8)
+    ax.plot(elbow[0], elbow[1], 'o', color='blue', markersize=8)
+    
+    # Draw Link 2 (Green)
+    if stability == "Stable":
+        end_effector = (0.8, 1.2)
+    else:
+        end_effector = (0.7, 0.8)
+    ax.plot([elbow[0], end_effector[0]], [elbow[1], end_effector[1]], color='green', linewidth=3, label='Link 2')
+    ax.plot(end_effector[0], end_effector[1], 'o', color='green', markersize=8)
+    
+    # Draw Ball (Red Circle)
+    if stability == "Stable":
+        ball_center = (end_effector[0], end_effector[1] + 0.4)
+    else:
+        ball_center = (end_effector[0] + 0.5, end_effector[1])
+    ball = plt.Circle(ball_center, 0.1, color='red', alpha=0.6)
+    ax.add_artist(ball)
+    ax.plot(ball_center[0], ball_center[1], 'o', color='red')
+    
+    # Draw Fingers (Black)
+    finger_base = end_effector
+    if stability == "Stable":
+        offsets = [(-0.05, 0.3), (0.0, 0.3), (0.05, 0.3)]
+    else:
+        offsets = [(-0.2, 0.2), (0.0, 0.3), (0.2, 0.2)]
+    for dx, dy in offsets:
+        tip_x = end_effector[0] + dx
+        tip_y = end_effector[1] + dy
+        ax.plot([finger_base[0], tip_x], [finger_base[1], tip_y], color='black', linewidth=3)
+        
+    ax.set_xlim(-0.5, 2.0)
+    ax.set_ylim(-0.2, 2.5)
+    ax.set_aspect('equal')
+    ax.grid(True)
+    
+    # Create a custom legend
+    from matplotlib.lines import Line2D
+    from matplotlib.patches import Patch
+    legend_elements = [
+        Line2D([0], [0], color='blue', lw=3, label='Link 1'),
+        Line2D([0], [0], color='green', lw=3, label='Link 2'),
+        Patch(facecolor='red', alpha=0.6, label='Ball'),
+        Line2D([0], [0], color='black', lw=3, label='Fingers')
+    ]
+    ax.legend(handles=legend_elements, loc='upper left')
+    ax.set_title("2-Link Robot Arm with 3-Finger Grasp")
+    return fig
+
+def predict_robustness_range(lower_bound, upper_bound):
+    """
+    1) Generates random sensor values within [lower_bound, upper_bound].
+    2) Predicts the grasp's robustness score using the pre-trained model.
+    3) Categorizes the grasp as 'Stable' or 'Unstable' (using a threshold of 100).
+    4) Returns the sensor values, a summary, and a visualization of the robotic arm.
+    """
+    if lower_bound >= upper_bound:
+        return {}, "Error: Lower bound must be less than upper bound.", None
+
+    # Generate sensor values
+    sensor_values = generate_random_sensors(feature_names, lower_bound, upper_bound)
+    
+    # Prepare model input in the same order as stored feature names
+    X_new = np.array([sensor_values[f] for f in feature_names]).reshape(1, -1)
+    pred_score = round(final_model.predict(X_new)[0], 2)
+    
+    threshold = 100
+    stability = "Stable" if pred_score >= threshold else "Unstable"
+    result_text = f"Predicted Robustness Score: {pred_score}\nStability Category: {stability}"
+    
+    # Create visualization based on prediction
+    fig = draw_robotic_arm(stability)
+    
+    return sensor_values, result_text, fig
+
+############################################
+# Gradio Interface
+############################################
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Robotic Grasp Predictor")
+    gr.Markdown(
+        "Enter lower and upper bounds for random sensor values.\n"
+        "The system will predict the grasp's robustness and categorize it as **Stable** "
+        "(if the predicted score >= 100) or **Unstable** (if the predicted score < 100)."
+    )
+    
+    with gr.Row():
+        with gr.Column():
+            lower_bound = gr.Number(label="Sensor Values Lower Bound", value=0.0)
+            upper_bound = gr.Number(label="Sensor Values Upper Bound", value=1.0)
+            button = gr.Button("Generate & Predict")
+            sensor_json = gr.JSON(label="Generated Sensor Values")
+        with gr.Column():
+            prediction_box = gr.Textbox(label="Prediction Result")
+            plot_output = gr.Plot(label="Arm Visualization")
+            
+    button.click(
+        fn=predict_robustness_range,
+        inputs=[lower_bound, upper_bound],
+        outputs=[sensor_json, prediction_box, plot_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)