scanner/app.py

from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import FileResponse, JSONResponse
import shutil
import os
import sys
from io import StringIO
import contextlib

# ඔයාගේ original ෆයිල් දෙක මෙතනට import කරනවා
import scan

app = FastAPI()

@app.get("/")
def home():
    return {"message": "Cipher-MD Scanner Engine is Running! 🚀"}


# =================================================================
# 2️⃣ SCAN / REVEAL ENDPOINT
# =================================================================
@app.post("/scan")
async def scan_watermark(image: UploadFile = File(...)):
    in_file = f"in_{image.filename}"
    out_file = f"out_{image.filename}"

    # 1. Image එක Save කරගන්නවා
    with open(in_file, "wb") as buffer:
        shutil.copyfileobj(image.file, buffer)

    try:
        # 2. ඔයාගේ scan script (reveal_watermark) එක run කරනවා
        scan.reveal_watermark(in_file, out_file)

        # 3. හැදුනු ෆොටෝ එක ආපහු යවනවා
        if os.path.exists(out_file):
            return FileResponse(out_file)
        else:
            return {"error": "Processing failed"}
            
    except Exception as e:
        return {"error": str(e)}
        
    finally:
        # Input ෆයිල් එක විතරක් මකනවා (Output එක යැව්වට පස්සේ මැකෙන විදිහට HuggingFace එකේ auto clean වෙනවා)
        if os.path.exists(in_file):
            os.remove(in_file)

# Docker එකේ Run වෙන්න ඕන නිසා
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7864)