Upload landmarkdiff/api_client.py with huggingface_hub

landmarkdiff/api_client.py

@@ -0,0 +1,242 @@
+"""Python client for the LandmarkDiff REST API.
+
+Provides a clean interface for interacting with the FastAPI server,
+handling image encoding/decoding, error handling, and session management.
+
+Usage:
+    from landmarkdiff.api_client import LandmarkDiffClient
+
+    client = LandmarkDiffClient("http://localhost:8000")
+
+    # Single prediction
+    result = client.predict("patient.png", procedure="rhinoplasty", intensity=65)
+    result.save("output.png")
+
+    # Face analysis
+    analysis = client.analyze("patient.png")
+    print(f"Fitzpatrick type: {analysis['fitzpatrick_type']}")
+
+    # Batch processing
+    results = client.batch_predict(
+        ["patient1.png", "patient2.png"],
+        procedure="blepharoplasty",
+    )
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import cv2
+import numpy as np
+
+
+@dataclass
+class PredictionResult:
+    """Result from a single prediction."""
+
+    output_image: np.ndarray
+    procedure: str
+    intensity: float
+    confidence: float = 0.0
+    landmarks_before: list | None = None
+    landmarks_after: list | None = None
+    metrics: dict[str, float] = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def save(self, path: str | Path, fmt: str = ".png") -> None:
+        """Save the output image to a file."""
+        cv2.imwrite(str(path), self.output_image)
+
+    def show(self) -> None:
+        """Display the output image (requires GUI)."""
+        cv2.imshow("LandmarkDiff Prediction", self.output_image)
+        cv2.waitKey(0)
+        cv2.destroyAllWindows()
+
+
+class LandmarkDiffClient:
+    """Client for the LandmarkDiff REST API.
+
+    Args:
+        base_url: Server URL (e.g. "http://localhost:8000").
+        timeout: Request timeout in seconds.
+    """
+
+    def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 60.0) -> None:
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self._session = None
+
+    def _get_session(self):
+        """Lazy-initialize requests session."""
+        if self._session is None:
+            try:
+                import requests
+            except ImportError:
+                raise ImportError("requests required. Install with: pip install requests")
+            self._session = requests.Session()
+            self._session.timeout = self.timeout
+        return self._session
+
+    def _read_image(self, image_path: str | Path) -> bytes:
+        """Read image file as bytes."""
+        path = Path(image_path)
+        if not path.exists():
+            raise FileNotFoundError(f"Image not found: {path}")
+        return path.read_bytes()
+
+    def _decode_base64_image(self, b64_string: str) -> np.ndarray:
+        """Decode a base64-encoded image to numpy array."""
+        img_bytes = base64.b64decode(b64_string)
+        arr = np.frombuffer(img_bytes, np.uint8)
+        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+        if img is None:
+            raise ValueError("Failed to decode base64 image")
+        return img
+
+    # ------------------------------------------------------------------
+    # API methods
+    # ------------------------------------------------------------------
+
+    def health(self) -> dict[str, Any]:
+        """Check server health.
+
+        Returns:
+            Dict with status and version info.
+        """
+        session = self._get_session()
+        resp = session.get(f"{self.base_url}/health")
+        resp.raise_for_status()
+        return resp.json()
+
+    def procedures(self) -> list[str]:
+        """List available surgical procedures.
+
+        Returns:
+            List of procedure names.
+        """
+        session = self._get_session()
+        resp = session.get(f"{self.base_url}/procedures")
+        resp.raise_for_status()
+        return resp.json().get("procedures", [])
+
+    def predict(
+        self,
+        image_path: str | Path,
+        procedure: str = "rhinoplasty",
+        intensity: float = 65.0,
+        seed: int = 42,
+    ) -> PredictionResult:
+        """Run surgical outcome prediction.
+
+        Args:
+            image_path: Path to input face image.
+            procedure: Surgical procedure type.
+            intensity: Intensity of the modification (0-100).
+            seed: Random seed for reproducibility.
+
+        Returns:
+            PredictionResult with output image and metadata.
+        """
+        session = self._get_session()
+        image_bytes = self._read_image(image_path)
+
+        files = {"image": ("image.png", image_bytes, "image/png")}
+        data = {
+            "procedure": procedure,
+            "intensity": str(intensity),
+            "seed": str(seed),
+        }
+
+        resp = session.post(f"{self.base_url}/predict", files=files, data=data)
+        resp.raise_for_status()
+        result = resp.json()
+
+        # Decode output image
+        output_img = self._decode_base64_image(result["output_image"])
+
+        return PredictionResult(
+            output_image=output_img,
+            procedure=procedure,
+            intensity=intensity,
+            confidence=result.get("confidence", 0.0),
+            metrics=result.get("metrics", {}),
+            metadata=result.get("metadata", {}),
+        )
+
+    def analyze(self, image_path: str | Path) -> dict[str, Any]:
+        """Analyze a face image without generating a prediction.
+
+        Returns face landmarks, Fitzpatrick type, pose estimation, etc.
+
+        Args:
+            image_path: Path to input face image.
+
+        Returns:
+            Dict with analysis results.
+        """
+        session = self._get_session()
+        image_bytes = self._read_image(image_path)
+
+        files = {"image": ("image.png", image_bytes, "image/png")}
+        resp = session.post(f"{self.base_url}/analyze", files=files)
+        resp.raise_for_status()
+        return resp.json()
+
+    def batch_predict(
+        self,
+        image_paths: list[str | Path],
+        procedure: str = "rhinoplasty",
+        intensity: float = 65.0,
+        seed: int = 42,
+    ) -> list[PredictionResult]:
+        """Run batch prediction on multiple images.
+
+        Args:
+            image_paths: List of image file paths.
+            procedure: Procedure to apply to all images.
+            intensity: Intensity for all images.
+            seed: Base random seed.
+
+        Returns:
+            List of PredictionResult objects.
+        """
+        results = []
+        for i, path in enumerate(image_paths):
+            try:
+                result = self.predict(
+                    path,
+                    procedure=procedure,
+                    intensity=intensity,
+                    seed=seed + i,
+                )
+                results.append(result)
+            except Exception as e:
+                # Create a failed result
+                results.append(PredictionResult(
+                    output_image=np.zeros((512, 512, 3), dtype=np.uint8),
+                    procedure=procedure,
+                    intensity=intensity,
+                    metadata={"error": str(e), "path": str(path)},
+                ))
+        return results
+
+    def close(self) -> None:
+        """Close the HTTP session."""
+        if self._session is not None:
+            self._session.close()
+            self._session = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+    def __repr__(self) -> str:
+        return f"LandmarkDiffClient(base_url='{self.base_url}')"