Spaces:
Paused
Paused
| # ============================== | |
| # Imports | |
| # ============================== | |
| from fastapi import FastAPI, File, UploadFile | |
| from fastapi.responses import JSONResponse | |
| import numpy as np | |
| import cv2 | |
| import tensorflow as tf | |
| from tensorflow.keras.models import Sequential | |
| from tensorflow.keras.layers import ( | |
| InputLayer, TimeDistributed, GlobalAveragePooling2D, | |
| LSTM, Dense, Dropout | |
| ) | |
| from tensorflow.keras.applications import MobileNetV2 | |
| import os | |
| # ============================== | |
| # Parameters | |
| # ============================== | |
| IMG_H, IMG_W, CHANNELS = 224, 224, 3 | |
| NUM_FRAMES = 30 | |
| MODEL_PATH = "video_mobilenetv2_final_backup.keras" | |
| model = None # Lazy loading | |
| # ============================== | |
| # Model Loader | |
| # ============================== | |
| def load_model(): | |
| global model | |
| if model is not None: | |
| return model | |
| print("🔄 Loading model...") | |
| base_cnn = MobileNetV2( | |
| include_top=False, | |
| weights='imagenet', | |
| input_shape=(IMG_H, IMG_W, CHANNELS) | |
| ) | |
| for layer in base_cnn.layers[:-20]: | |
| layer.trainable = False | |
| model = Sequential([ | |
| InputLayer(shape=(NUM_FRAMES, IMG_H, IMG_W, CHANNELS)), | |
| TimeDistributed(base_cnn), | |
| TimeDistributed(GlobalAveragePooling2D()), | |
| LSTM(128), | |
| Dropout(0.5), | |
| Dense(128, activation='relu'), | |
| Dropout(0.5), | |
| Dense(1, activation='sigmoid') | |
| ]) | |
| model.load_weights(MODEL_PATH) | |
| print("✅ Model loaded successfully!") | |
| return model | |
| # ============================== | |
| # Frame Extraction | |
| # ============================== | |
| def extract_video_frames(video_path, n_frames=NUM_FRAMES): | |
| cap = cv2.VideoCapture(video_path) | |
| total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| if total <= 0: | |
| blank = np.zeros((IMG_H, IMG_W, 3), dtype=np.float32) | |
| return np.stack([blank]*n_frames) | |
| idxs = np.linspace(0, total - 1, n_frames).astype(int) | |
| frames = [] | |
| for i in idxs: | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, i) | |
| ok, frame = cap.read() | |
| if not ok: | |
| frames.append(np.zeros((IMG_H, IMG_W, 3), dtype=np.float32)) | |
| continue | |
| frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| frame = cv2.resize(frame, (IMG_W, IMG_H)) | |
| frame = frame.astype(np.float32) / 255.0 | |
| frames.append(frame) | |
| cap.release() | |
| return np.array(frames, dtype=np.float32) | |
| # ============================== | |
| # Predict | |
| # ============================== | |
| def predict_video(path): | |
| model = load_model() | |
| frames = extract_video_frames(path) | |
| frames = np.expand_dims(frames, 0) | |
| score = model.predict(frames, verbose=0)[0][0] | |
| label = "Violence" if score >= 0.5 else "NonViolence" | |
| conf = score*100 if score >= 0.5 else (1 - score)*100 | |
| return label, conf | |
| # ============================== | |
| # FastAPI App | |
| # ============================== | |
| app = FastAPI( | |
| title="Video Violence Detector", | |
| description="API for detecting violence in videos using deep learning", | |
| version="1.0.0" | |
| ) | |
| # ============================== | |
| # Startup Event - Preload Model | |
| # ============================== | |
| async def startup_event(): | |
| """Preload model on startup to avoid cold start delays""" | |
| print("🚀 Starting up application...") | |
| try: | |
| load_model() | |
| print("✅ Model preloaded successfully!") | |
| except Exception as e: | |
| print(f"⚠️ Warning: Could not preload model: {e}") | |
| # ============================== | |
| # Health Check Endpoint | |
| # ============================== | |
| def health_check(): | |
| """Health check endpoint for monitoring""" | |
| return { | |
| "status": "healthy", | |
| "model_loaded": model is not None | |
| } | |
| # ============================== | |
| # Root Endpoint | |
| # ============================== | |
| def home(): | |
| """Root endpoint with API information""" | |
| return { | |
| "message": "Video Violence Detector API", | |
| "endpoints": { | |
| "POST /predict": "Upload a video file for violence detection", | |
| "GET /health": "Check API health status" | |
| }, | |
| "usage": "Send a POST request to /predict with a video file" | |
| } | |
| # ============================== | |
| # Prediction Endpoint | |
| # ============================== | |
| async def upload(file: UploadFile = File(...)): | |
| """ | |
| Upload a video file and get violence prediction | |
| Returns: | |
| JSON with prediction (Violence/NonViolence) and confidence score | |
| """ | |
| # Validate file type | |
| if not file.content_type.startswith('video/'): | |
| return JSONResponse( | |
| status_code=400, | |
| content={"error": "File must be a video"} | |
| ) | |
| temp = f"temp_{file.filename}" | |
| try: | |
| # Save uploaded file | |
| with open(temp, "wb") as f: | |
| f.write(await file.read()) | |
| # Make prediction | |
| label, conf = predict_video(temp) | |
| # Clean up temp file | |
| if os.path.exists(temp): | |
| os.remove(temp) | |
| return { | |
| "prediction": label, | |
| "confidence": f"{conf:.2f}%", | |
| "filename": file.filename | |
| } | |
| except Exception as e: | |
| # Clean up on error | |
| if os.path.exists(temp): | |
| os.remove(temp) | |
| return JSONResponse( | |
| status_code=500, | |
| content={ | |
| "error": "Prediction failed", | |
| "details": str(e) | |
| } | |
| ) | |
| # ============================== | |
| # Run locally (not used in HF) | |
| # ============================== | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True) | |