Update app.py

app.py

@@ -14,20 +14,20 @@ from huggingface_hub import hf_hub_download
 HF_MODEL_REPO_ID = "LeafNet75/Leaf-Annotate-v2"
 DEVICE = "cpu"
 IMG_SIZE = 256
+CONFIDENCE_THRESHOLD = 0.5
 
-# --- DATA MODELS FOR API (using Pydantic) ---
+# --- DATA MODELS FOR API ---
 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
+    predicted_mask: str # base64 encoded raw binary 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")
@@ -47,62 +47,48 @@ def load_model():
 model = load_model()
 
 # --- HELPER FUNCTIONS ---
-def base64_to_cv2(base64_string: str):
-    # Remove the "data:image/..." header
+def base64_to_cv2_rgba(base64_string: str):
     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)
+    image_cv = base64_to_cv2_rgba(request.image)
+    scribble_cv = base64_to_cv2_rgba(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()
+    binary_mask = (probs > CONFIDENCE_THRESHOLD).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)