CouchPotato101/space_network_pred

1from fastapi import FastAPI, Query2from models import LightweightTelecomNet3from schemas import PredictionResponse4import torch5import os6 7# 1. Initialize FastAPI8app = FastAPI()9 10# 2. Determine computation device11device = "cuda" if torch.cuda.is_available() else "cpu"12 13# 3. Load the pre-trained model14MODEL_PATH = "model/best_lightweight_model.pt"  # Update this path for HuggingFace Spaces15print(os.path.exists(MODEL_PATH))16try:17    model = LightweightTelecomNet().to(device)18    model.load_state_dict(torch.load(MODEL_PATH, map_location=device))19    model.eval()20    print(f"Model loaded successfully on {device}.")21except FileNotFoundError:22    print(f"Error: Model file not found. Please ensure '{MODEL_PATH}' is present.")23    model = None24except Exception as e:25    print(f"An error occurred while loading the model: {e}")26    model = None27 28 29# 4. Define class_map30class_map = {0: "Normal", 1: "Scintillation", 2: "Congestion", 3: "Rain Fade"}31 32# 5. Load synthetic data33SYNTHETIC_DATA_PATH = os.path.join("data", "synthetic_data.pt")34# try:35#     synthetic_data = torch.load(SYNTHETIC_DATA_PATH, map_location=device)36#     # Expecting dict with keys: 'temporal', 'spatial', 'operator_id' or similar37#     print(f"Synthetic data loaded: {list(synthetic_data.keys())}")38# except Exception as e:39#     print(f"Error loading synthetic data: {e}")40#     synthetic_data = None41 42try:43    synthetic_data = torch.load(SYNTHETIC_DATA_PATH, map_location=device)44 45    if isinstance(synthetic_data, dict):46        print("Synthetic data loaded (dict keys):", list(synthetic_data.keys()))47    elif isinstance(synthetic_data, list):48        print("Synthetic data loaded (list length):", len(synthetic_data))49    else:50        print("Synthetic data loaded:", type(synthetic_data))51 52except Exception as e:53    print("Error loading synthetic data:", e)54    synthetic_data = None55 56 57# Use the first available sample from synthetic_data for prediction58def get_synthetic_sample():59    if synthetic_data is None:60        return None, None, None61    # Try to support both dict of lists and list of dicts62    if isinstance(synthetic_data, dict):63        temporal = synthetic_data.get('temporal')64        spatial = synthetic_data.get('spatial')65        operator_id = synthetic_data.get('operator_id', 0)66        # If these are lists, use the first sample67        if isinstance(temporal, list) and len(temporal) > 0 and isinstance(temporal[0], list):68            temporal = temporal[0]69        if isinstance(spatial, list) and len(spatial) > 0 and isinstance(spatial[0], list):70            spatial = spatial[0]71        if isinstance(operator_id, list):72            operator_id = operator_id[0]73        return temporal, spatial, operator_id74    elif isinstance(synthetic_data, list) and len(synthetic_data) > 0:75        sample = synthetic_data[0]76        return sample.get('temporal'), sample.get('spatial'), sample.get('operator_id', 0)77    return None, None, None78 79@app.get("/ping")80def ping():81    return {"status": "alive"}82    83@app.get("/predict_segment", response_model=PredictionResponse)84async def predict_segment(85    sample_idx: int = Query(200, description="Index of the sample in the validation dataset"),86    start_time_idx: int = Query(0, description="Starting 10-min interval (0-143)"),87    duration_in_10min_intervals: int = Query(1, description="Number of consecutive 10-min intervals (1-144)")88):89    if model is None or synthetic_data is None:90        return {"predicted_event": "Model not loaded", "predicted_vulnerability": 0.0}91 92    # Try to treat synthetic_data as a list of tuples or dicts like a dataset93    val_dataset = synthetic_data94    try:95        # Try to index as a list/tuple96        sample = val_dataset[sample_idx]97        # Support tuple or dict98        if isinstance(sample, dict):99            temporal_full = torch.tensor(sample['temporal'])100            spatial_full = torch.tensor(sample['spatial'])101            operator_id = torch.tensor(sample['operator_id'])102            true_label = torch.tensor(sample.get('true_label', 0))103            true_vuln = torch.tensor(sample.get('true_vuln', 0.0))104        else:105            # Assume tuple: (temporal, spatial, operator_id, true_label, true_vuln)106            temporal_full, spatial_full, operator_id, true_label, true_vuln = sample107        # Extract the temporal segment108        temporal_segment = temporal_full[start_time_idx : start_time_idx + duration_in_10min_intervals]109        temporal_input_configured = temporal_segment.unsqueeze(0).to(device) # [1, duration, 5]110        spatial_input_configured = spatial_full.unsqueeze(0).to(device) # [1, 50, 50]111        operator_input_configured = torch.tensor([operator_id.item() if hasattr(operator_id, 'item') else int(operator_id)], device=device)112 113        with torch.no_grad():114            class_out_configured, vuln_out_configured = model(temporal_input_configured, spatial_input_configured, operator_input_configured)115        pred_label_configured = torch.argmax(class_out_configured, dim=1).item()116        pred_vuln_configured = vuln_out_configured.item()117        predicted_event_label = class_map[pred_label_configured]118 119        # Optionally, return true label/vuln for comparison120        return {121            "predicted_event": predicted_event_label,122            "predicted_vulnerability": pred_vuln_configured,123            # "true_event": class_map[true_label.item() if hasattr(true_label, 'item') else int(true_label)],124            # "true_vulnerability": float(true_vuln.item() if hasattr(true_vuln, 'item') else float(true_vuln))125        }126    except Exception as e:127        return {"predicted_event": f"Error: {e}", "predicted_vulnerability": 0.0}128 129@app.get("/predict_moment", response_model=PredictionResponse)130async def predict_moment(131    sample_idx: int = Query(200, description="Index of the sample in the validation dataset"),132    moment_idx: int = Query(0, description="Index of the moment in the temporal sequence (0-143)")133):134    if model is None or synthetic_data is None:135        return {"predicted_event": "Model not loaded", "predicted_vulnerability": 0.0}136 137    val_dataset = synthetic_data138    try:139        sample = val_dataset[sample_idx]140        if isinstance(sample, dict):141            temporal = torch.tensor(sample['temporal'])142            spatial = torch.tensor(sample['spatial'])143            operator_id = torch.tensor(sample['operator_id'])144        else:145            temporal, spatial, operator_id = sample[:3]146 147        # Select the single moment148        single_moment_temporal = temporal[moment_idx] # [5]149        single_moment_temporal_input = single_moment_temporal.unsqueeze(0).unsqueeze(0).to(device) # [1, 1, 5]150        spatial_input = spatial.unsqueeze(0).to(device)151        operator_input = torch.tensor([operator_id.item() if hasattr(operator_id, 'item') else int(operator_id)], device=device)152 153        with torch.no_grad():154            class_out_single_moment, vuln_out_single_moment = model(single_moment_temporal_input, spatial_input, operator_input)155        pred_label_single_moment = torch.argmax(class_out_single_moment, dim=1).item()156        pred_vuln_single_moment = vuln_out_single_moment.item()157        predicted_event_label = class_map[pred_label_single_moment]158 159        return {160            "predicted_event": predicted_event_label,161            "predicted_vulnerability": pred_vuln_single_moment162        }163    except Exception as e:164        return {"predicted_event": f"Error: {e}", "predicted_vulnerability": 0.0}165