Spaces:
Running
Running
| # Final, API-only app.py for Hugging Face Space | |
| import os | |
| import cv2 | |
| import tempfile | |
| import numpy as np | |
| import uvicorn | |
| from PIL import Image | |
| from inference_sdk import InferenceHTTPClient | |
| from fastapi import FastAPI, File, UploadFile | |
| from fastapi.responses import JSONResponse | |
| import tensorflow as tf | |
| from huggingface_hub import hf_hub_download | |
| # --- 1. Configuration and Model Loading --- | |
| ROBOFLOW_API_KEY = os.environ.get("ROBOFLOW_API_KEY") | |
| CLIENT_FACE = InferenceHTTPClient(api_url="https://detect.roboflow.com", api_key=ROBOFLOW_API_KEY) | |
| CLIENT_EYES = InferenceHTTPClient(api_url="https://detect.roboflow.com", api_key=ROBOFLOW_API_KEY) | |
| CLIENT_IRIS = InferenceHTTPClient(api_url="https://detect.roboflow.com", api_key=ROBOFLOW_API_KEY) | |
| leuko_model = None | |
| try: | |
| model_path = hf_hub_download("skibi11/leukolook-eye-detector", "MobileNetV1_best.keras") | |
| leuko_model = tf.keras.models.load_model(model_path) | |
| print("--- LEUKOCORIA MODEL LOADED SUCCESSFULLY! ---") | |
| except Exception as e: | |
| print(f"--- FATAL ERROR: COULD NOT LOAD LEUKOCORIA MODEL: {e} ---") | |
| raise RuntimeError(f"Could not load leukocoria model: {e}") | |
| # --- 2. All Helper Functions --- | |
| def enhance_image_unsharp_mask(image, strength=0.5, radius=5): | |
| blur = cv2.GaussianBlur(image, (radius, radius), 0) | |
| return cv2.addWeighted(image, 1.0 + strength, blur, -strength, 0) | |
| def detect_faces_roboflow(image_path): | |
| return CLIENT_FACE.infer(image_path, model_id="face-detector-v4liw/2").get("predictions", []) | |
| def detect_eyes_roboflow(image_path, raw_image): | |
| resp = CLIENT_EYES.infer(image_path, model_id="eye-detection-kso3d/3") | |
| crops = [] | |
| for p in resp.get("predictions", []): | |
| x1 = int(p['x'] - p['width'] / 2) | |
| y1 = int(p['y'] - p['height'] / 2) | |
| x2 = int(p['x'] + p['width'] / 2) | |
| y2 = int(p['y'] + p['height'] / 2) | |
| crop = raw_image[y1:y2, x1:x2] | |
| if crop.size > 0: | |
| crops.append(crop) | |
| return crops | |
| def get_largest_iris_prediction(eye_crop): | |
| is_success, buffer = cv2.imencode(".jpg", eye_crop) | |
| if not is_success: return None | |
| resp = CLIENT_IRIS.infer(buffer, model_id="iris_120_set/7") | |
| preds = resp.get("predictions", []) | |
| return max(preds, key=lambda p: p["width"] * p["height"]) if preds else None | |
| def run_leukocoria_prediction(iris_crop): | |
| if leuko_model is None: return {"error": "Leukocoria model not loaded"}, 0.0 | |
| img_pil = Image.fromarray(cv2.cvtColor(iris_crop, cv2.COLOR_BGR2RGB)) | |
| enh = enhance_image_unsharp_mask(np.array(img_pil)) | |
| enh_rs = cv2.resize(enh, (224, 224)) | |
| img_array = np.array(enh_rs) / 255.0 | |
| img_array = np.expand_dims(img_array, axis=0) | |
| prediction = leuko_model.predict(img_array) | |
| confidence = float(prediction[0][0]) | |
| has_leuko = confidence > 0.5 | |
| return has_leuko, confidence | |
| # --- 3. FastAPI Application --- | |
| app = FastAPI() | |
| async def full_detection_pipeline(image: UploadFile = File(...)): | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp: | |
| contents = await image.read() | |
| tmp.write(contents) | |
| temp_image_path = tmp.name | |
| try: | |
| if not detect_faces_roboflow(temp_image_path): | |
| return JSONResponse(status_code=400, content={"error": "No face detected."}) | |
| raw_image = cv2.imread(temp_image_path) | |
| eye_crops = detect_eyes_roboflow(temp_image_path, raw_image) | |
| if len(eye_crops) != 2: | |
| return JSONResponse(status_code=400, content={"error": "Exactly two eyes not detected."}) | |
| eye_crops.sort(key=lambda c: cv2.boundingRect(cv2.cvtColor(c, cv2.COLOR_BGR2GRAY))[0]) | |
| flags = {} | |
| for i, eye_crop in enumerate(eye_crops): | |
| side = "left" if i == 0 else "right" | |
| pred = get_largest_iris_prediction(eye_crop) | |
| if pred: | |
| x1, y1 = int(pred['x'] - pred['width'] / 2), int(pred['y'] - pred['height'] / 2) | |
| x2, y2 = int(pred['x'] + pred['width'] / 2), int(pred['y'] + pred['height'] / 2) | |
| iris_crop = eye_crop[y1:y2, x1:x2] | |
| has_leuko, confidence = run_leukocoria_prediction(iris_crop) | |
| flags[side] = has_leuko | |
| else: | |
| flags[side] = None | |
| return JSONResponse(content={"leukocoria": flags, "warnings": []}) | |
| finally: | |
| os.remove(temp_image_path) | |
| # --- 4. Run the Server --- | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) |