Update app.py

app.py

@@ -14,15 +14,15 @@ from huggingface_hub import hf_hub_download
 HF_MODEL_REPO_ID = "LeafNet75/Leaf-Annotate-v2"
 DEVICE = "cpu"
 IMG_SIZE = 256
-CONFIDENCE_THRESHOLD = 0.5
+CONFIDENCE_THRESHOLD = 0.298
 
 # --- DATA MODELS FOR API ---
 class InferenceRequest(BaseModel):
-    image: str          # base64 encoded image string
-    scribble_mask: str  # base64 encoded scribble mask string
+    image: str
+    scribble_mask: str
 
 class InferenceResponse(BaseModel):
-    predicted_mask: str # base64 encoded raw binary mask string
+    predicted_mask: str
 
 # --- INITIALIZE FASTAPI APP ---
 app = FastAPI()
@@ -30,8 +30,16 @@ app = FastAPI()
 # --- LOAD MODEL ON STARTUP ---
 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")
-    
+    try:
+        model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename="best_model.pth")
+    except Exception as e:
+        # Fallback for local testing if the model file is in the same directory
+        if os.path.exists("best_model.pth"):
+            print("Could not download from Hub, using local 'best_model.pth'.")
+            model_path = "best_model.pth"
+        else:
+            raise e
+
     model = smp.Unet(
         encoder_name="mobilenet_v2",
         encoder_weights=None,
@@ -48,35 +56,15 @@ model = load_model()
 
 # --- HELPER FUNCTIONS ---
 def base64_to_cv2_rgba(base64_string: str):
-    try:
-        # Handle data URL format
-        if "," in base64_string:
-            header, encoded = base64_string.split(",", 1)
-        else:
-            encoded = base64_string
-            
-        img_data = base64.b64decode(encoded)
-        pil_image = Image.open(io.BytesIO(img_data))
-        
-        # Convert to RGBA if not already
-        if pil_image.mode != 'RGBA':
-            pil_image = pil_image.convert('RGBA')
-            
-        # Convert PIL to numpy array and then to OpenCV format
-        np_array = np.array(pil_image)
-        return cv2.cvtColor(np_array, cv2.COLOR_RGBA2BGRA)
-    except Exception as e:
-        print(f"Error in base64_to_cv2_rgba: {e}")
-        raise
+    header, encoded = base64_string.split(",", 1)
+    img_data = base64.b64decode(encoded)
+    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):
-    try:
-        _, buffer = cv2.imencode('.png', image)
-        png_as_text = base64.b64encode(buffer).decode('utf-8')
-        return f"data:image/png;base64,{png_as_text}"
-    except Exception as e:
-        print(f"Error in cv2_to_base64: {e}")
-        raise
+    _, 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("/")
@@ -85,62 +73,30 @@ def read_root():
 
 @app.post("/predict", response_model=InferenceResponse)
 async def predict(request: InferenceRequest):
-    try:
-        # Convert base64 images to OpenCV format
-        image_cv = base64_to_cv2_rgba(request.image)
-        scribble_cv = base64_to_cv2_rgba(request.scribble_mask)
-        
-        # Convert scribble mask to grayscale if it has multiple channels
-        if len(scribble_cv.shape) > 2 and scribble_cv.shape[2] > 1:
-            scribble_cv = cv2.cvtColor(scribble_cv, cv2.COLOR_BGRA2GRAY)
-        
-        # Get original dimensions - FIXED SYNTAX ERROR
-        h, w, *_ = image_cv.shape
-        
-        # Resize images to model input size
-        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
-        )
-        
-        # Convert to tensors and normalize
-        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
-        
-        # Concatenate image and scribble mask as 4-channel input
-        input_tensor = torch.cat([image_tensor, scribble_tensor], dim=0).unsqueeze(0).to(DEVICE)
-        
-        # Run inference
-        with torch.no_grad():
-            output = model(input_tensor)
-        
-        # Post-process output
-        probs = torch.sigmoid(output)
-        binary_mask = (probs > CONFIDENCE_THRESHOLD).float().squeeze().cpu().numpy()
-        
-        # Resize output mask back to original image dimensions
-        output_mask_resized = cv2.resize(
-            binary_mask, 
-            (w, h), 
-            interpolation=cv2.INTER_NEAREST
-        )
-        
-        # Convert to uint8 format
-        output_mask_uint8 = (output_mask_resized * 255).astype(np.uint8)
-        
-        # Convert to base64 for response
-        result_base64 = cv2_to_base64(output_mask_uint8)
-        
-        return InferenceResponse(predicted_mask=result_base64)
-        
-    except Exception as e:
-        print(f"Error in predict endpoint: {e}")
-        import traceback
-        traceback.print_exc()
-        raise
+    image_cv = base64_to_cv2_rgba(request.image)
+    scribble_cv = base64_to_cv2_rgba(request.scribble_mask)
+    
+    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
+
+    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)
+
+    with torch.no_grad():
+        output = model(input_tensor)
+
+    probs = torch.sigmoid(output)
+    binary_mask = (probs > CONFIDENCE_THRESHOLD).float().squeeze().cpu().numpy()
+    
+    output_mask_resized = cv2.resize(binary_mask, (w, h), interpolation=cv2.INTER_NEAREST)
+    output_mask_uint8 = (output_mask_resized * 255).astype(np.uint8)
+
+    result_base64 = cv2_to_base64(output_mask_uint8)
+    
+    return InferenceResponse(predicted_mask=result_base64)