Initial AI service commit

.gitignore

@@ -0,0 +1,44 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual Environments
+venv/
+env/
+ENV/
+.env
+.venv
+
+# Environment Variables
+.env
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Large Files / Models
+*.pth
+*.pt
+*.onnx
+models/
+data/
+hf_cache/

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+COPY requirements.txt /app/requirements.txt
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential git ffmpeg libgl1 \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . /app
+
+ENV PYTHONUNBUFFERED=1
+ENV PYTHONDONTWRITEBYTECODE=1
+
+# For GPU: use an nvidia/cuda base image and install the correct torch wheel with CUDA.
+# When running with nvidia runtime, pass --gpus=all to docker run or set deploy settings in compose.

app.py

@@ -0,0 +1,102 @@
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import FileResponse, JSONResponse
+from uuid import uuid4
+from pathlib import Path
+import shutil
+import os
+import json
+import redis
+from celery import Celery
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Directories (mounted by docker-compose)
+UPLOAD_DIR = Path(os. environ. get("UPLOAD_DIR", "/data/uploads"))
+RESULT_DIR = Path(os.environ. get("RESULT_DIR", "/data/results"))
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+RESULT_DIR.mkdir(parents=True, exist_ok=True)
+
+# Redis for job status
+REDIS_HOST = os. environ.get("REDIS_HOST", "redis")
+REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))
+CELERY_BROKER_URL = os.environ.get("CELERY_BROKER_URL", "redis://redis:6379/0")
+
+# Lazy connections
+_rdb = None
+_celery_client = None
+
+def get_redis():
+    global _rdb
+    if _rdb is None:
+        _rdb = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0, decode_responses=True)
+    return _rdb
+
+def get_celery():
+    global _celery_client
+    if _celery_client is None:
+        _celery_client = Celery(broker=CELERY_BROKER_URL)
+    return _celery_client
+
+app = FastAPI(title="Depth->STL processing service (API)")
+
+
+def set_status(job_id: str, state: str, detail: str = "", result: str = ""):
+    payload = {"state": state, "detail": detail, "result": result}
+    get_redis().set(job_id, json.dumps(payload))
+
+
+@app.post("/upload/")
+async def upload_image(file: UploadFile = File(... )):
+    # Basic validation
+    if not file.content_type. startswith("image/"):
+        raise HTTPException(status_code=400, detail="File must be an image")
+
+    job_id = str(uuid4())
+    safe_name = Path(file.filename).name
+    fname = f"{job_id}_{safe_name}"
+    save_path = UPLOAD_DIR / fname
+
+    # Save uploaded file to mounted volume
+    try:
+        with save_path.open("wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to save upload: {e}")
+
+    # Mark queued and enqueue Celery task
+    set_status(job_id, "QUEUED", "Job received and queued")
+    try:
+        async_result = get_celery().send_task(
+            "tasks.process_image_task",
+            args=[str(save_path), str(RESULT_DIR), job_id],
+            kwargs={},
+            queue=os.environ.get("CELERY_QUEUE", None),
+        )
+    except Exception as e:
+        set_status(job_id, "FAILURE", f"Failed to enqueue task: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to enqueue task: {e}")
+
+    return {"job_id": job_id, "celery_id": str(async_result.id)}
+
+
+@app.get("/status/{job_id}")
+def status(job_id: str):
+    raw = get_redis().get(job_id)
+    if not raw:
+        return JSONResponse({"state": "UNKNOWN", "detail": "No such job_id"}, status_code=404)
+    return JSONResponse(json.loads(raw))
+
+
+@app.get("/download/{job_id}")
+def download(job_id: str):
+    raw = get_redis().get(job_id)
+    if not raw:
+        raise HTTPException(status_code=404, detail="No such job")
+    info = json.loads(raw)
+    if info.get("state") != "SUCCESS":
+        raise HTTPException(status_code=404, detail="Result not ready")
+    stl_path = info.get("result")
+    if not stl_path or not Path(stl_path).exists():
+        raise HTTPException(status_code=404, detail="Result file missing")
+    return FileResponse(path=stl_path, filename=Path(stl_path).name, media_type="application/sla")

docker-compose.dev.yml

@@ -0,0 +1,85 @@
+services:
+  redis:
+    image: redis:7
+    ports:
+      - "6379:6379"
+    healthcheck:
+      test: ["CMD", "redis-cli", "ping"]
+      interval: 5s
+      timeout: 3s
+      retries: 5
+
+  python-api:
+    build:
+      context: ./python
+    command: ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
+    ports:
+      - "8000:8000"
+    environment:
+      - CELERY_BROKER_URL=redis://redis:6379/0
+      - CELERY_RESULT_BACKEND=redis://redis:6379/0
+      - REDIS_HOST=redis
+      - REDIS_PORT=6379
+      - UPLOAD_DIR=/data/uploads
+      - RESULT_DIR=/data/results
+      - DEPTH_CHECKPOINT=/models/depth-anything-Large-hf
+      - HF_HUB_OFFLINE=1
+      - TRANSFORMERS_OFFLINE=1
+      - HF_HUB_DISABLE_TELEMETRY=1
+      - OMP_NUM_THREADS=4
+      - MKL_NUM_THREADS=4
+    volumes:
+      - ./python:/app
+      - ./data:/data
+      - ./models:/models:ro
+    depends_on:
+      redis:
+        condition: service_healthy
+
+  python-worker:
+    build:
+      context: ./python
+    # GPU will be added by running docker run --gpus all manually (Option 1)
+    command: ["celery", "-A", "tasks.celery", "worker", "--loglevel=info", "--pool=solo"]
+    environment:
+      - CELERY_BROKER_URL=redis://redis:6379/0
+      - CELERY_RESULT_BACKEND=redis://redis:6379/0
+      - REDIS_HOST=redis
+      - REDIS_PORT=6379
+      - UPLOAD_DIR=/data/uploads
+      - RESULT_DIR=/data/results
+      - DEPTH_CHECKPOINT=/models/depth-anything-Large-hf
+      - USE_GPU=0
+      - POISSON_DEPTH=9
+      - OUTLIER_NEIGHBORS=15
+      - OUTLIER_STD_RATIO=1.0
+      - ORTHO_SCALE_FACTOR=255
+      - INFERENCE_RESIZE=0
+      - RESULT_PREFIX=
+      - HF_HUB_OFFLINE=1
+      - TRANSFORMERS_OFFLINE=1
+      - HF_HUB_DISABLE_TELEMETRY=1
+      - OMP_NUM_THREADS=4
+      - MKL_NUM_THREADS=4
+    volumes:
+      - ./python:/app
+      - ./data:/data:rw
+      - ./models:/models:ro
+      - ./hf_cache:/root/.cache/huggingface:rw
+    depends_on:
+      redis:
+        condition: service_healthy
+
+  node:
+    build:
+      context: ./node
+    command: ["node", "server.js"]
+    ports:
+      - "3000:3000"
+    environment:
+      PYTHON_URL: http://python-api:8000
+      MONGODB_URI: mongodb+srv://ironman88103102_db_user:[email protected]/teethnet?retryWrites=true&w=majority&appName=Cluster0
+      JWT_SECRET: 3fe4191be8414cac9a2185511b0045400be14cfb2a181cad3969a61594a2246d
+    depends_on:
+      python-api:
+        condition: service_started

model_loader.py

@@ -0,0 +1,32 @@
+import os
+import torch
+from transformers import AutoImageProcessor, AutoModelForDepthEstimation
+
+# Default to the small model; overridden by DEPTH_CHECKPOINT env if set
+_CHECKPOINT = os.environ.get("DEPTH_CHECKPOINT", "/models/depth-anything-small-hf")
+
+_MODEL = None
+_PROCESSOR = None
+_DEVICE = None
+
+def get_model_and_processor():
+    global _MODEL, _PROCESSOR, _DEVICE
+    if _MODEL is None or _PROCESSOR is None:
+        print("Loading model:", _CHECKPOINT, flush=True)
+        # Strongly limit CPU threads to avoid WSL2/Docker oversubscription
+        try:
+            torch.set_num_threads(max(1, (os.cpu_count() or 2) // 2))
+        except Exception:
+            pass
+
+        _PROCESSOR = AutoImageProcessor.from_pretrained(_CHECKPOINT)
+        _MODEL = AutoModelForDepthEstimation.from_pretrained(_CHECKPOINT)
+
+        if torch.cuda.is_available():
+            _DEVICE = torch.device("cuda")
+        else:
+            _DEVICE = torch.device("cpu")
+
+        _MODEL = _MODEL.to(_DEVICE)
+        _MODEL.eval()
+    return _MODEL, _PROCESSOR, _DEVICE

requirements.txt

@@ -0,0 +1,12 @@
+python-multipart
+fastapi
+uvicorn[standard]
+celery[redis]
+redis
+transformers
+torch
+opencv-python-headless
+open3d==0.19.0
+trimesh
+numpy
+python-dotenv

tasks.py

@@ -0,0 +1,186 @@
+import os
+from celery import Celery
+from pathlib import Path
+import traceback
+import json
+import redis
+import time
+import sys
+
+import numpy as np
+import cv2
+import open3d as o3d
+import torch
+from PIL import Image
+import trimesh
+from transformers import AutoImageProcessor, AutoModelForDepthEstimation
+
+# Celery / Redis config
+CELERY_BROKER = os.environ.get("CELERY_BROKER_URL", "redis://redis:6379/0")
+CELERY_BACKEND = os.environ.get("CELERY_RESULT_BACKEND", "redis://redis:6379/0")
+REDIS_HOST = os.environ.get("REDIS_HOST", "redis")
+REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))
+
+# Pipeline settings (fixed to orthographic + Poisson depth=9 to match your notebook)
+DEPTH_CHECKPOINT = os.environ.get("DEPTH_CHECKPOINT", "/models/depth-anything-Large-hf")
+USE_GPU = int(os.environ.get("USE_GPU", "1"))
+POISSON_DEPTH = int(os.environ.get("POISSON_DEPTH", "9"))
+OUTLIER_NEIGHBORS = int(os.environ.get("OUTLIER_NEIGHBORS", "15"))
+OUTLIER_STD_RATIO = float(os.environ.get("OUTLIER_STD_RATIO", "1.0"))
+ORTHO_SCALE_FACTOR = float(os.environ.get("ORTHO_SCALE_FACTOR", "255"))  # same as your function
+INFERENCE_RESIZE = int(os.environ.get("INFERENCE_RESIZE", "0"))  # 0 keeps original
+RESULT_PREFIX = os.environ.get("RESULT_PREFIX", "")
+
+try:
+    torch.set_num_threads(max(1, (os.cpu_count() or 2) // 2))
+except Exception:
+    pass
+
+celery = Celery("tasks", broker=CELERY_BROKER, backend=CELERY_BACKEND)
+rdb = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0, decode_responses=True)
+
+_model = None
+_processor = None
+_device = "cpu"
+
+def log(msg):
+    print(msg, flush=True)
+    sys.stdout.flush()
+
+def set_status(job_id: str, state: str, detail: str = "", result: str = ""):
+    payload = {"state": state, "detail": detail, "result": result}
+    rdb.set(job_id, json.dumps(payload))
+
+def load_model():
+    global _model, _processor, _device
+    if _model is None:
+        log(f"Loading model: {DEPTH_CHECKPOINT}")
+        _processor = AutoImageProcessor.from_pretrained(DEPTH_CHECKPOINT)
+        _model = AutoModelForDepthEstimation.from_pretrained(DEPTH_CHECKPOINT)
+        if USE_GPU and torch.cuda.is_available():
+            _device = "cuda"
+            _model = _model.to("cuda")
+        else:
+            _device = "cpu"
+        _model.eval()
+    return _model, _processor, _device
+
+def normalize_depth_uint8(depth_np: np.ndarray) -> np.ndarray:
+    m = np.max(depth_np)
+    if m <= 0:
+        return np.zeros_like(depth_np, dtype=np.uint8)
+    return (depth_np * 255.0 / m).astype("uint8")
+
+def build_orthographic_point_cloud(depth_u8: np.ndarray, color_rgb: np.ndarray) -> o3d.geometry.PointCloud:
+    depth_map = depth_u8.astype(np.float32)
+    h, w = depth_map.shape
+    y, x = np.meshgrid(np.arange(h), np.arange(w), indexing='ij')
+    z = (depth_map / ORTHO_SCALE_FACTOR) * (h / 2.0)
+    points = np.stack((x, y, z), axis=-1).reshape(-1, 3)
+    mask = points[:, 2] != 0
+    points = points[mask]
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points)
+    colors = color_rgb.reshape(-1, 3)[mask] / 255.0
+    pcd.colors = o3d.utility.Vector3dVector(colors)
+    return pcd
+
+@celery.task(bind=True)
+def process_image_task(self, image_path: str, result_dir: str, job_id: str):
+    start = time.time()
+    try:
+        set_status(job_id, "RUNNING", "Loading model")
+        model, processor, device = load_model()
+        log(f"[{job_id}] Model loaded on {device}")
+
+        img_bgr = cv2.imread(image_path)
+        if img_bgr is None:
+            raise RuntimeError("Failed to read image")
+        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+        orig_h, orig_w = img_rgb.shape[:2]
+
+        # Optional resize (not used in your notebook; keep 0 for fidelity)
+        if INFERENCE_RESIZE and INFERENCE_RESIZE > 0:
+            scale = INFERENCE_RESIZE / max(orig_h, orig_w)
+            new_w = int(orig_w * scale)
+            new_h = int(orig_h * scale)
+            img_proc = cv2.resize(img_rgb, (new_w, new_h), interpolation=cv2.INTER_AREA)
+        else:
+            img_proc = img_rgb
+
+        set_status(job_id, "RUNNING", "Running depth inference")
+        depth_inputs = processor(images=img_proc, return_tensors="pt").to(device)
+        with torch.no_grad():
+            outputs = model(**depth_inputs)
+        depth = outputs.predicted_depth.squeeze().detach().cpu().numpy()
+
+        # Match notebook: use depth resolution, resize color to depth size
+        dh, dw = depth.shape
+        color_resized = cv2.resize(img_proc, (dw, dh), interpolation=cv2.INTER_LINEAR)
+
+        depth_u8 = normalize_depth_uint8(depth)
+
+        set_status(job_id, "RUNNING", "Building orthographic point cloud")
+        pcd = build_orthographic_point_cloud(depth_u8, color_resized)
+
+        # Outlier removal (nb=15, std_ratio=1.0)
+        try:
+            cl, ind = pcd.remove_statistical_outlier(nb_neighbors=OUTLIER_NEIGHBORS,
+                                                     std_ratio=OUTLIER_STD_RATIO)
+            pcd = pcd.select_by_index(ind)
+        except Exception as e:
+            log(f"[{job_id}] Outlier removal warning: {e}")
+
+        # Normals (your notebook: estimate_normals + orient_normals_to_align_with_direction)
+        if len(pcd.points) >= 10:
+            try:
+                pcd.estimate_normals()
+                pcd.orient_normals_to_align_with_direction()
+            except Exception as e:
+                log(f"[{job_id}] Normal estimation warning: {e}")
+
+        num_pts = np.asarray(pcd.points).shape[0]
+        log(f"[{job_id}] Point cloud size after cleanup: {num_pts}")
+        if num_pts == 0:
+            raise RuntimeError("Empty point cloud after cleanup")
+
+        set_status(job_id, "RUNNING", f"Poisson reconstruction depth={POISSON_DEPTH}")
+        mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
+            pcd, depth=POISSON_DEPTH
+        )
+
+        # Compute normals
+        try:
+            mesh.compute_vertex_normals()
+        except Exception:
+            pass
+        mesh.compute_triangle_normals()
+
+        num_vertices = np.asarray(mesh.vertices).shape[0]
+        num_tris = np.asarray(mesh.triangles).shape[0]
+        log(f"[{job_id}] Mesh stats vertices={num_vertices} triangles={num_tris}")
+        if num_tris == 0:
+            raise RuntimeError("Poisson produced empty mesh")
+
+        Path(result_dir).mkdir(parents=True, exist_ok=True)
+        stl_path = Path(result_dir) / f"{RESULT_PREFIX}{job_id}.stl"
+
+        set_status(job_id, "RUNNING", "Exporting STL")
+        tm = trimesh.Trimesh(vertices=np.asarray(mesh.vertices),
+                             faces=np.asarray(mesh.triangles),
+                             process=True)
+        tm.export(str(stl_path), file_type="stl")
+
+        total = time.time() - start
+        set_status(job_id, "SUCCESS", f"Done in {total:.2f}s", str(stl_path))
+        log(f"[{job_id}] SUCCESS total={total:.2f}s STL={stl_path}")
+        return {
+            "status": "success",
+            "stl": str(stl_path),
+            "mesh_stats": {"vertices": int(num_vertices), "triangles": int(num_tris)}
+        }
+    except Exception as e:
+        traceback.print_exc()
+        set_status(job_id, "FAILURE", str(e))
+        log(f"[{job_id}] FAILURE: {e}")
+        raise