app.py

import base64
import cv2
import numpy as np
import urllib.request
import tempfile
from pathlib import Path
from fastapi import FastAPI, UploadFile, File, Request, Form
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.templating import Jinja2Templates
from rembg import remove

CASCADE_URL = "https://raw.githubusercontent.com/opencv/opencv/master/data/haarcascades/haarcascade_frontalface_default.xml"
CASCADE_PATH = Path("haarcascade_frontalface_default.xml")
if not CASCADE_PATH.exists():
    urllib.request.urlretrieve(CASCADE_URL, CASCADE_PATH.name)

face_cascade = cv2.CascadeClassifier(str(CASCADE_PATH))

MODEL_PROTO = "deploy.prototxt"
MODEL_WEIGHTS = "res10_300x300_ssd_iter_140000.caffemodel"
face_net = cv2.dnn.readNetFromCaffe(MODEL_PROTO, MODEL_WEIGHTS)

app = FastAPI(title="Pipeline Tête + Fond")
templates = Jinja2Templates(directory="templates")


def detect_faces_dnn(img):
    (h, w) = img.shape[:2]
    blob = cv2.dnn.blobFromImage(cv2.resize(img, (300, 300)), 1.0,
                                 (300, 300), (104.0, 177.0, 123.0))
    face_net.setInput(blob)
    detections = face_net.forward()

    faces = []
    for i in range(0, detections.shape[2]):
        confidence = detections[0, 0, i, 2]
        if confidence > 0.6:  
            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            (x1, y1, x2, y2) = box.astype("int")
            faces.append((x1, y1, x2 - x1, y2 - y1))
    return faces


def remove_head_bytes(image_bytes: bytes):
    img_arr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(img_arr, cv2.IMREAD_UNCHANGED)
    if img is None:
        return None

    if img.shape[2] == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)


    faces = detect_faces_dnn(img[:, :, :3])


    if len(faces) == 0:
        gray = cv2.cvtColor(img[:, :, :3], cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(
            gray, scaleFactor=1.1, minNeighbors=6, minSize=(40, 40)
        )

    if len(faces) == 0:
        return None

   
    mask = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
    for (x, y, fw, fh) in faces:
        pad_top = int(fh * 2)
        pad_side = int(fw * 2)
        pad_bottom = int(fh * 0.1)

        x1 = max(0, x - pad_side)
        x2 = min(img.shape[1], x + fw + pad_side)
        y1 = max(0, y - pad_top)
        y2 = min(img.shape[0], y + fh + pad_bottom)

        cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
        ax, ay = (x2 - x1) // 2, (y2 - y1) // 2

        ellipse_mask = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
        cv2.ellipse(ellipse_mask, (cx, cy), (ax, ay), 0, 0, 360, 255, -1)
        mask = cv2.bitwise_or(mask, ellipse_mask)

    img[mask == 255] = [0, 0, 0, 0]

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    cv2.imwrite(tmp.name, img)
    return tmp.name



def remove_bg_bytes(image_bytes: bytes):
    try:
        output_bytes = remove(image_bytes)
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        tmp.write(output_bytes)
        return tmp.name
    except Exception as e:
        print("Erreur rembg lib:", e)
        return None



@app.get("/", response_class=HTMLResponse)
def home(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})



@app.post("/process")
async def process_image(request: Request, image: UploadFile = File(...), action: str = Form(...)):
    image_bytes = await image.read()

    if action == "pipeline":
        bgless_path = remove_bg_bytes(image_bytes)
        if not bgless_path:
            return {"error": "Erreur lors de la suppression du fond"}

        with open(bgless_path, "rb") as f:
            output_path = remove_head_bytes(f.read())

        if not output_path:
            output_path = bgless_path

    else:
        return {"error": "Action inconnue"}

    with open(output_path, "rb") as f:
        img_b64 = base64.b64encode(f.read()).decode("utf-8")

    return templates.TemplateResponse("result.html", {
        "request": request,
        "img_b64": img_b64,
        "filename": Path(output_path).name
    })



@app.post("/pipeline-file")
async def pipeline_file(image: UploadFile = File(...)):
    image_bytes = await image.read()

    bgless_path = remove_bg_bytes(image_bytes)
    if not bgless_path:
        return {"error": "Erreur lors de la suppression du fond"}

    with open(bgless_path, "rb") as f:
        output_path = remove_head_bytes(f.read())

    if not output_path:
        output_path = bgless_path

    return FileResponse(output_path, media_type="image/png", filename="result.png")