Prayesh007/perceptionX-python

1import sys2import asyncio3import os4import requests5import tempfile6import subprocess7from io import BytesIO8from fastapi import FastAPI, HTTPException9from pydantic import BaseModel10import uvicorn11 12from motor.motor_asyncio import AsyncIOMotorClient13from ultralytics import YOLO14from PIL import Image15import numpy as np16import cv217from bson import ObjectId18import imageio_ffmpeg as ffmpeg19from bson.binary import Binary20from dotenv import load_dotenv21 22# --------------------- CONFIG ------------------------23load_dotenv()24 25PORT = int(os.environ.get("PORT", 7860))26MONGO_URI = os.environ.get("MONGO_URI")27DATABASE_NAME = os.environ.get("DATABASE_NAME", "test")28COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "files")29WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "./yolov11/best.pt")30MODEL_WEIGHTS_URL = os.environ.get("MODEL_WEIGHTS_URL", "")31 32os.environ['MPLCONFIGDIR'] = '/tmp'33os.environ["ULTRALYTICS_CONFIG_DIR"] = "/tmp/ultralytics"34 35# --------------------- MONGODB ------------------------36client = None37collection = None38 39if MONGO_URI:40    try:41        client = AsyncIOMotorClient(MONGO_URI)42        db = client[DATABASE_NAME]43        collection = db[COLLECTION_NAME]44        print("✅ Connected to MongoDB")45    except Exception as e:46        print(f"❌ Failed to connect MongoDB: {e}")47else:48    print("⚠️ No MONGO_URI provided. MongoDB features disabled.")49 50# --------------------- YOLO MODEL ------------------------51_MODEL = None52 53def download_weights_if_missing():54    if os.path.exists(WEIGHTS_PATH):55        print(f"✅ Model exists at: {WEIGHTS_PATH}")56        return True57    if not MODEL_WEIGHTS_URL:58        print(f"❌ Weights not found at {WEIGHTS_PATH} and no MODEL_WEIGHTS_URL provided.")59        return False60    try:61        os.makedirs(os.path.dirname(WEIGHTS_PATH), exist_ok=True)62        print(f"⬇️ Downloading model from {MODEL_WEIGHTS_URL} ...")63        r = requests.get(MODEL_WEIGHTS_URL, stream=True, timeout=60)64        r.raise_for_status()65        with open(WEIGHTS_PATH, "wb") as f:66            for chunk in r.iter_content(chunk_size=8192):67                f.write(chunk)68        print("✅ Model downloaded successfully.")69        return True70    except Exception as e:71        print(f"❌ Error downloading model: {e}")72        return False73 74def load_model():75    global _MODEL76    if _MODEL is not None:77        return _MODEL78    try:79        if not os.path.exists(WEIGHTS_PATH):80            ok = download_weights_if_missing()81            if not ok:82                raise FileNotFoundError("YOLO model weights not found and could not be downloaded.")83        _MODEL = YOLO(WEIGHTS_PATH)84        print("🧠 YOLO model loaded.")85        return _MODEL86    except Exception as e:87        print(f"❌ Failed to load model: {e}")88        _MODEL = None89        raise90 91# --------------------- MONGO HELPERS ------------------------92async def fetch_file_from_mongo(file_id):93    if collection is None:94        return None, None95    try:96        obj_id = ObjectId(file_id)97    except Exception:98        return None, None99    doc = await collection.find_one({"_id": obj_id})100    if not doc or "data" not in doc:101        return None, None102    return bytes(doc["data"]), doc.get("mimetype", "")103 104async def save_processed_file(file_id, processed_bytes):105    if collection is None:106        return False107    try:108        obj_id = ObjectId(file_id)109    except Exception:110        return False111    result = await collection.update_one(112        {"_id": obj_id},113        {"$set": {"processedData": Binary(processed_bytes)}}114    )115    return result.modified_count > 0 or result.matched_count > 0116 117# --------------------- IMAGE PROCESSING ------------------------118async def process_image(file_id, model):119    data, _ = await fetch_file_from_mongo(file_id)120    if data is None:121        return False122    try:123        img = Image.open(BytesIO(data)).convert("RGB")124        results = await asyncio.to_thread(model, np.array(img))125        annotated = results[0].plot()126        bgr = cv2.cvtColor(annotated, cv2.COLOR_RGB2BGR)127        success, buffer = cv2.imencode(".jpg", bgr)128        if not success:129            return False130        return await save_processed_file(file_id, buffer.tobytes())131    except Exception as e:132        print(f"❌ Image processing error: {e}")133        return False134 135# --------------------- VIDEO PROCESSING ------------------------136async def process_video(file_id, model):137    data, _ = await fetch_file_from_mongo(file_id)138    if data is None:139        return False140    tmp_in = os.path.join(tempfile.gettempdir(), f"input_{file_id}.mp4")141    tmp_out = os.path.join(tempfile.gettempdir(), f"out_{file_id}.avi")142    mp4_path = tmp_out.replace(".avi", ".mp4")143    try:144        with open(tmp_in, "wb") as f:145            f.write(data)146        cap = cv2.VideoCapture(tmp_in)147        if not cap.isOpened():148            return False149        fps = cap.get(cv2.CAP_PROP_FPS) or 25150        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))151        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))152        out = cv2.VideoWriter(tmp_out, cv2.VideoWriter_fourcc(*"XVID"), fps, (width, height))153        while True:154            ret, frame = cap.read()155            if not ret:156                break157            results = await asyncio.to_thread(model, frame)158            annotated = results[0].plot()159            try:160                annotated_bgr = cv2.cvtColor(annotated, cv2.COLOR_RGB2BGR)161            except:162                annotated_bgr = annotated163            out.write(annotated_bgr)164        cap.release()165        out.release()166        ffmpeg_exe = ffmpeg.get_ffmpeg_exe()167        subprocess.run(168            [ffmpeg_exe, "-y", "-i", tmp_out, "-vcodec", "libx264", "-crf", "23", "-preset", "fast", mp4_path],169            stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True, timeout=300170        )171        with open(mp4_path, "rb") as f:172            processed_video = f.read()173        return await save_processed_file(file_id, processed_video)174    except subprocess.CalledProcessError as e:175        print(f"❌ FFmpeg conversion failed. Stderr: {e.stderr.decode()}")176        return False177    except subprocess.TimeoutExpired:178        print("❌ FFmpeg conversion timed out.")179        return False180    except Exception as e:181        print(f"❌ Video processing error: {e}")182        return False183    finally:184        for f in [tmp_in, tmp_out, mp4_path]:185            try:186                os.remove(f)187            except:188                pass189 190# --------------------- PROCESS ROUTINE ------------------------191async def run_process(file_id, file_type):192    if _MODEL is None:193        raise RuntimeError("YOLO model is not loaded")194    if file_type.startswith("image"):195        return await process_image(file_id, _MODEL)196    elif file_type.startswith("video"):197        return await process_video(file_id, _MODEL)198    else:199        return False200 201# --------------------- FASTAPI ------------------------202app = FastAPI(title="YOLO Processing Service")203 204class ProcessRequest(BaseModel):205    fileId: str206    fileType: str207 208@app.on_event("startup")209async def startup_event():210    try:211        load_model()212        print("Service startup complete.")213    except Exception as e:214        print(f"❌ CRITICAL ERROR: {e}")215 216@app.get("/")217async def root():218    return {"message": "YOLO backend is running!"}219 220@app.get("/health")221async def health():222    if _MODEL is None:223        raise HTTPException(status_code=503, detail="Model not ready")224    return {"status": "ok"}225 226@app.get("/warmup")227async def warmup():228    if _MODEL is None:229        load_model()230    return {"status": "ready"}231 232@app.post("/process")233async def process_endpoint(payload: ProcessRequest):234    if _MODEL is None:235        raise HTTPException(status_code=503, detail="YOLO model not ready")236    try:237        ok = await run_process(payload.fileId, payload.fileType)238        if not ok:239            raise HTTPException(status_code=500, detail="Processing failed")240        return {"status": "ok", "fileId": payload.fileId}241    except Exception as e:242        print(f"❌ Processing failed for file {payload.fileId}: {e}")243        raise HTTPException(status_code=500, detail=f"Processing failed. Reason: {e}")244 245if __name__ == "__main__":246    uvicorn.run("app:app", host="0.0.0.0", port=PORT)247