Update app.py

app.py

@@ -1,103 +1,93 @@
-import io
-import base64
 import numpy as np
 from PIL import Image
-import torch
-import uvicorn
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import JSONResponse
-from contextlib import asynccontextmanager
-from depth_pro import create_model_and_transforms
-from depth_pro.depth_pro import DepthProConfig
-
-CHECKPOINT_PATH = "/app/ml-depth-pro/checkpoints/depth_pro.pt"
-
-model = None
-transform = None
-
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    global model, transform
-    print("Chargement du modèle DepthPro...")
-    config = DepthProConfig(
-        patch_encoder_preset="dinov2l16_384",
-        image_encoder_preset="dinov2l16_384",
-        checkpoint_uri=CHECKPOINT_PATH,
-        decoder_features=256,
-        use_fov_head=True,
-        fov_encoder_preset="dinov2l16_384",
-    )
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    precision = torch.half if torch.cuda.is_available() else torch.float32
-    model, transform = create_model_and_transforms(
-        config=config,
-        device=device,
-        precision=precision,
-    )
-    model.eval()
-    print(f"Modèle chargé sur {device}.")
-    yield
-
-
-app = FastAPI(
-    title="DepthPro API",
-    description="API de prédiction de profondeur monoculaire avec Apple DepthPro",
-    version="1.0.0",
-    lifespan=lifespan,
-)
+import io
+import traceback
+import logging
+import torch
+from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-@app.get("/")
-def root():
-    return {"message": "DepthPro API is running", "docs": "/docs"}
+# ──────────────────────────────────────────────
+app = FastAPI(title="Depth Pro (Apple) — Metric Depth API")
 
+DEVICE = "cpu"
+MODEL_ID = "apple/DepthPro-hf"
 
-@app.get("/health")
-def health():
-    return {
-        "status": "ok",
-        "device": "cuda" if torch.cuda.is_available() else "cpu",
-        "model_loaded": model is not None,
-    }
+logger.info(f"Chargement {MODEL_ID} sur {DEVICE} ...")
+processor = DepthProImageProcessorFast.from_pretrained(MODEL_ID)
+model = DepthProForDepthEstimation.from_pretrained(MODEL_ID)
+model = model.to(DEVICE).eval()
+logger.info("Modèle Depth Pro prêt.")
 
 
+# ──────────────────────────────────────────────
 @app.post("/predict")
-async def predict_depth(file: UploadFile = File(...)):
+async def predict(file: UploadFile = File(...)):
+
+    # 1) Lecture image
+    logger.info(f"Fichier reçu : {file.filename} | type : {file.content_type}")
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail=f"Content-type invalide : {file.content_type}")
+
+    contents = await file.read()
+    logger.info(f"Taille : {len(contents)} octets")
+
     try:
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        pil_img = Image.open(io.BytesIO(contents)).convert("RGB")
+        logger.info(f"Image : {pil_img.size[0]}x{pil_img.size[1]}")
     except Exception as e:
+        logger.error(traceback.format_exc())
         raise HTTPException(status_code=400, detail=f"Impossible de lire l'image : {e}")
 
+    # 2) Prétraitement
     try:
-        image_tensor = transform(image)
-        with torch.no_grad():
-            prediction = model.infer(image_tensor)
+        inputs = processor(images=pil_img, return_tensors="pt").to(DEVICE)
+        logger.info(f"Inputs prêts : {list(inputs.keys())}")
+    except Exception as e:
+        logger.error(traceback.format_exc())
+        raise HTTPException(status_code=500, detail=f"Erreur prétraitement : {e}")
 
-        depth = prediction["depth"].squeeze().cpu().numpy()
-        focal_length_px = float(prediction["focallength_px"].cpu())
+    # 3) Inférence
+    try:
+        logger.info("Inférence en cours ...")
+        with torch.no_grad():
+            outputs = model(**inputs)
+        logger.info("Inférence terminée.")
+    except Exception as e:
+        logger.error(traceback.format_exc())
+        raise HTTPException(status_code=500, detail=f"Erreur inférence : {e}")
 
-        depth_min = float(depth.min())
-        depth_max = float(depth.max())
-        depth_normalized = ((depth - depth_min) / (depth_max - depth_min + 1e-8) * 65535).astype(np.uint16)
+    # 4) Post-traitement → profondeur en mètres à la résolution originale
+    try:
+        post = processor.post_process_depth_estimation(
+            outputs,
+            target_sizes=[(pil_img.height, pil_img.width)],
+        )
+        depth_map = post[0]["predicted_depth"].squeeze().cpu().numpy()  # [H, W] float32, mètres
+        logger.info(f"depth_map shape={depth_map.shape} min={depth_map.min():.3f} max={depth_map.max():.3f}")
+    except Exception as e:
+        logger.error(traceback.format_exc())
+        raise HTTPException(status_code=500, detail=f"Erreur post-traitement : {e}")
 
-        depth_img = Image.fromarray(depth_normalized, mode="I;16")
-        buf = io.BytesIO()
-        depth_img.save(buf, format="PNG")
-        depth_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    # 5) Résultats
+    H, W = depth_map.shape
+    closest_distance = float(np.min(depth_map))
+    cy, cx = H // 2, W // 2
+    center_distance = float(depth_map[cy, cx])
 
-        return JSONResponse({
-            "depth_map_base64": depth_b64,
-            "focal_length_px": focal_length_px,
-            "depth_min_meters": depth_min,
-            "depth_max_meters": depth_max,
-            "image_size": {"width": image.width, "height": image.height},
-        })
+    logger.info(f"closest={closest_distance:.3f}m | center={center_distance:.3f}m")
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Erreur lors de l'inférence : {e}")
+    return JSONResponse(content={
+        "closest_distance": closest_distance,
+        "center_distance":  center_distance,
+    })
 
 
-if __name__ == "__main__":
-    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)
+# ──────────────────────────────────────────────
+@app.get("/")
+def root():
+    return {"status": "ok", "model": MODEL_ID, "device": DEVICE}