Create app.py

app.py

@@ -0,0 +1,108 @@
+import base64
+import io
+import cv2
+import numpy as np
+import torch
+from fastapi import FastAPI
+from fastapi.responses import FileResponse
+from pydantic import BaseModel
+from PIL import Image
+import segmentation_models_pytorch as smp
+from huggingface_hub import hf_hub_download
+
+# --- CONFIGURATION ---
+HF_MODEL_REPO_ID = "LeafNet75/Leaf-Annotate-v2"
+DEVICE = "cpu"
+IMG_SIZE = 256
+
+# --- DATA MODELS FOR API (using Pydantic) ---
+class InferenceRequest(BaseModel):
+    image: str          # base64 encoded image string
+    scribble_mask: str  # base64 encoded scribble mask string
+
+class InferenceResponse(BaseModel):
+    predicted_mask: str # base64 encoded predicted mask string
+
+# --- INITIALIZE FASTAPI APP ---
+app = FastAPI()
+
+# --- LOAD MODEL ON STARTUP ---
+# The model is loaded once when the application starts to ensure fast inference times.
+def load_model():
+    print(f"Loading model '{HF_MODEL_REPO_ID}'...")
+    model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename="best_model.pth")
+    
+    model = smp.Unet(
+        encoder_name="mobilenet_v2",
+        encoder_weights=None,
+        in_channels=4,
+        classes=1,
+    )
+    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+    model.to(DEVICE)
+    model.eval()
+    print("Model loaded successfully.")
+    return model
+
+model = load_model()
+
+# --- HELPER FUNCTIONS ---
+def base64_to_cv2(base64_string: str):
+    # Remove the "data:image/..." header
+    header, encoded = base64_string.split(",", 1)
+    img_data = base64.b64decode(encoded)
+    
+    # Use Pillow to open the image data and convert to OpenCV format
+    pil_image = Image.open(io.BytesIO(img_data))
+    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGBA2BGRA)
+
+def cv2_to_base64(image: np.ndarray):
+    # Convert image back to a base64 string to send to the frontend
+    _, buffer = cv2.imencode('.png', image)
+    png_as_text = base64.b64encode(buffer).decode('utf-8')
+    return f"data:image/png;base64,{png_as_text}"
+
+
+# --- API ENDPOINTS ---
+@app.get("/")
+def read_root():
+    # Serve the frontend HTML file
+    return FileResponse('index.html')
+
+@app.post("/predict", response_model=InferenceResponse)
+async def predict(request: InferenceRequest):
+    # 1. Decode input data
+    image_cv = base64_to_cv2(request.image)
+    scribble_cv = base64_to_cv2(request.scribble_mask)
+    
+    # Ensure scribble is grayscale
+    if len(scribble_cv.shape) > 2 and scribble_cv.shape[2] > 1:
+        scribble_cv = cv2.cvtColor(scribble_cv, cv2.COLOR_BGRA2GRAY)
+
+    h, w, _ = image_cv.shape
+
+    # 2. Preprocess the data for the model
+    image_resized = cv2.resize(cv2.cvtColor(image_cv, cv2.COLOR_BGRA2RGB), (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_AREA)
+    scribble_resized = cv2.resize(scribble_cv, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_NEAREST)
+
+    image_tensor = torch.from_numpy(image_resized.astype(np.float32)).permute(2, 0, 1) / 255.0
+    scribble_tensor = torch.from_numpy(scribble_resized.astype(np.float32)).unsqueeze(0) / 255.0
+
+    input_tensor = torch.cat([image_tensor, scribble_tensor], dim=0).unsqueeze(0).to(DEVICE)
+
+    # 3. Run Inference
+    with torch.no_grad():
+        output = model(input_tensor)
+
+    # 4. Post-process the output
+    probs = torch.sigmoid(output)
+    binary_mask = (probs > 0.5).float().squeeze().cpu().numpy()
+    
+    # Resize mask to the original input canvas size
+    output_mask_resized = cv2.resize(binary_mask, (w, h), interpolation=cv2.INTER_NEAREST)
+    output_mask_uint8 = (output_mask_resized * 255).astype(np.uint8)
+
+    # 5. Encode the result and return
+    result_base64 = cv2_to_base64(output_mask_uint8)
+    
+    return InferenceResponse(predicted_mask=result_base64)