Spaces:
Running on Zero
Running on Zero
| import json | |
| import os | |
| import gradio as gr | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| import torch.nn as nn | |
| from joblib import load | |
| from scipy.signal import find_peaks | |
| from scipy.stats import entropy | |
| try: | |
| import spaces | |
| except ImportError: | |
| class _LocalSpaces: | |
| def GPU(function): | |
| return function | |
| spaces = _LocalSpaces() | |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| MODEL_DIR = os.path.join(BASE_DIR, "best_models") | |
| ARTIFACT_DIR = os.path.join(BASE_DIR, "artifacts") | |
| DEVICE = torch.device("cpu") | |
| WINDOW_SIZE = 60 | |
| PHYSICAL_LABELS = {0: "STRESS", 1: "PHYSICAL ACTIVITY", 2: "NORMAL"} | |
| DISTRACTION_LABELS = {0: "FOCUSED", 1: "DISTRACTED"} | |
| class WearableMLP(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Linear(13, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.3), | |
| nn.Linear(256, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(0.3), | |
| nn.Linear(128, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(0.3), | |
| nn.Linear(64, 3), | |
| ) | |
| def forward(self, inputs): | |
| return self.net(inputs) | |
| class DistractionNet(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Linear(7, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(0.3), | |
| nn.Linear(64, 32), nn.BatchNorm1d(32), nn.ReLU(), nn.Dropout(0.2), | |
| nn.Linear(32, 1), nn.Sigmoid(), | |
| ) | |
| def forward(self, inputs): | |
| return self.net(inputs).squeeze(1) | |
| def load_model(model, filename): | |
| path = os.path.join(MODEL_DIR, filename) | |
| if not os.path.exists(path): | |
| raise FileNotFoundError(f"Missing model file: {path}") | |
| model.load_state_dict(torch.load(path, map_location=DEVICE)) | |
| model.eval() | |
| return model | |
| physical_model = load_model(WearableMLP(), "model_1c_state_scratch.pt") | |
| distraction_model = load_model(DistractionNet(), "model_2c_distraction_scratch.pt") | |
| physical_scaler = load(os.path.join(ARTIFACT_DIR, "physical_scaler.joblib")) | |
| distraction_scaler = load(os.path.join(ARTIFACT_DIR, "distraction_scaler.joblib")) | |
| def check_values(values): | |
| values = np.asarray(values, dtype=np.float32) | |
| if not np.isfinite(values).all(): | |
| raise gr.Error("All input values must be finite numbers.") | |
| return values | |
| def predict_physical_state( | |
| heart_rate, spo2, skin_temperature, accel_x, accel_y, accel_z, | |
| gyro_x, gyro_y, gyro_z, lat_delta, lon_delta, altitude, time_delta, | |
| ): | |
| values = check_values([ | |
| heart_rate, spo2, skin_temperature, accel_x, accel_y, accel_z, | |
| gyro_x, gyro_y, gyro_z, lat_delta, lon_delta, altitude, time_delta, | |
| ]).reshape(1, -1) | |
| scaled = physical_scaler.transform(values).astype(np.float32) | |
| with torch.inference_mode(): | |
| probabilities = torch.softmax( | |
| physical_model(torch.from_numpy(scaled)), dim=1 | |
| )[0].numpy() | |
| class_id = int(np.argmax(probabilities)) | |
| return { | |
| "model": "model_1c_state_scratch", | |
| "task": "physical_state", | |
| "class_id": class_id, | |
| "state": PHYSICAL_LABELS[class_id], | |
| "probability": round(float(probabilities[class_id]), 6), | |
| "probabilities": { | |
| PHYSICAL_LABELS[index]: round(float(value), 6) | |
| for index, value in enumerate(probabilities) | |
| }, | |
| "alert": class_id == 0, | |
| } | |
| def calculate_window_features(samples): | |
| if isinstance(samples, dict): | |
| samples = samples.get("samples") | |
| if not isinstance(samples, list) or len(samples) != WINDOW_SIZE: | |
| raise gr.Error("The distraction input must contain exactly 60 samples.") | |
| frame = pd.DataFrame(samples) | |
| required = [ | |
| "heart_rate", "accel_x", "accel_y", "accel_z", | |
| "gyro_x", "gyro_y", "gyro_z", "lat_delta", "lon_delta", "time_delta", | |
| ] | |
| missing = [name for name in required if name not in frame.columns] | |
| if missing: | |
| raise gr.Error(f"Missing sample fields: {', '.join(missing)}") | |
| frame = frame[required].astype(np.float32) | |
| if not np.isfinite(frame.to_numpy()).all(): | |
| raise gr.Error("All sample values must be finite numbers.") | |
| acceleration = np.sqrt( | |
| frame["accel_x"] ** 2 + frame["accel_y"] ** 2 + frame["accel_z"] ** 2 | |
| ).to_numpy() | |
| gyroscope = np.sqrt( | |
| frame["gyro_x"] ** 2 + frame["gyro_y"] ** 2 + frame["gyro_z"] ** 2 | |
| ).to_numpy() | |
| displacement = np.sqrt(frame["lat_delta"] ** 2 + frame["lon_delta"] ** 2) | |
| speed = np.mean(displacement / (frame["time_delta"] + 1e-6)) | |
| histogram, _ = np.histogram(acceleration, bins=10, density=True) | |
| histogram += 1e-8 | |
| return np.asarray([ | |
| np.var(acceleration), | |
| np.var(gyroscope), | |
| entropy(histogram), | |
| len(find_peaks(acceleration)[0]), | |
| np.var(frame["heart_rate"].to_numpy()), | |
| (frame["heart_rate"].iloc[-1] - frame["heart_rate"].iloc[0]) / len(frame), | |
| speed, | |
| ], dtype=np.float32) | |
| def predict_distraction(samples_json): | |
| try: | |
| samples = json.loads(samples_json) | |
| except (TypeError, json.JSONDecodeError) as error: | |
| raise gr.Error("Enter valid JSON containing a samples array.") from error | |
| features = calculate_window_features(samples).reshape(1, -1) | |
| scaled = distraction_scaler.transform(features).astype(np.float32) | |
| with torch.inference_mode(): | |
| probability = float(distraction_model(torch.from_numpy(scaled))[0].item()) | |
| class_id = int(probability >= 0.5) | |
| return { | |
| "model": "model_2c_distraction_scratch", | |
| "task": "focus_distraction", | |
| "class_id": class_id, | |
| "state": DISTRACTION_LABELS[class_id], | |
| "probability": round(probability, 6), | |
| "alert": class_id == 1, | |
| "window_size": WINDOW_SIZE, | |
| } | |
| physical_inputs = [ | |
| gr.Number(label="Heart rate", value=80), | |
| gr.Number(label="SpO2", value=97), | |
| gr.Number(label="Skin temperature", value=36.5), | |
| gr.Number(label="Accel X", value=0), | |
| gr.Number(label="Accel Y", value=0), | |
| gr.Number(label="Accel Z", value=9.81), | |
| gr.Number(label="Gyro X", value=0), | |
| gr.Number(label="Gyro Y", value=0), | |
| gr.Number(label="Gyro Z", value=0), | |
| gr.Number(label="Latitude delta", value=0), | |
| gr.Number(label="Longitude delta", value=0), | |
| gr.Number(label="Altitude", value=20), | |
| gr.Number(label="Time delta", value=30), | |
| ] | |
| with gr.Blocks(title="Spandan Smartwatch AI") as demo: | |
| gr.Markdown("# Spandan Smartwatch AI\nTwo-model inference for physical state and focus monitoring.") | |
| with gr.Tab("Physical state"): | |
| gr.Markdown("Enter one cleaned telemetry row. The endpoint returns STRESS, PHYSICAL ACTIVITY, or NORMAL.") | |
| physical_button = gr.Button("Run physical-state model", variant="primary") | |
| physical_output = gr.JSON(label="Physical-state output") | |
| physical_button.click( | |
| predict_physical_state, | |
| inputs=physical_inputs, | |
| outputs=physical_output, | |
| api_name="predict_physical_state", | |
| ) | |
| with gr.Tab("Focus and distraction"): | |
| gr.Markdown("Paste JSON with exactly 60 ordered samples from one officer.") | |
| distraction_input = gr.Textbox( | |
| label="60-sample JSON", | |
| lines=12, | |
| placeholder='{"samples": [{"heart_rate": 80, "accel_x": 0, ...}]}' | |
| ) | |
| distraction_button = gr.Button("Run distraction model", variant="primary") | |
| distraction_output = gr.JSON(label="Distraction output") | |
| distraction_button.click( | |
| predict_distraction, | |
| inputs=distraction_input, | |
| outputs=distraction_output, | |
| api_name="predict_distraction", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |