Update app/services/segmenter.py

app/services/segmenter.py

@@ -1,88 +1,54 @@
+# app/services/segmenter.py
 import numpy as np
-from datetime import datetime
-from typing import Dict, Any
-from PIL import Image
-import io
 import tensorflow as tf
+from typing import Dict, Any
 
-from app.configs import get_segmentation_model, SEGMENTS_DIR
-from app.storage import update_history
+from app.configs import get_segmentation_model
 from app.core.preprocessing import preprocess_image
 
 
-async def run_segmentation(request_id: str, image_filename: str, image_path: str):
-    """Run segmentation on image (background task)"""
-    print(f"🎨 SEGMENTING {request_id}.......")
-
-    seg_result: Dict[str, Any] = {
-        "request_id": request_id,
-        "image_filename": image_filename,
-        "timestamp": datetime.now().isoformat(),
-        "detections": []
-    }
-
-    # Get model
+def segment_image(image_bytes: bytes) -> Dict[str, Any]:
+    """
+    Run segmentation on image bytes.
+    
+    Returns:
+        Dict with keys: status, masks_shape, max_confidence, error (optional)
+    """
     model = get_segmentation_model()
-
+    
     if model is None:
-        seg_result["error"] = "Segmentation model not loaded"
-        seg_path = None
-    else:
-        try:
-            # Load image
-            with open(image_path, 'rb') as f:
-                img_bytes = f.read()
-
-            # Preprocess for segmentation (256x256)
-            img_array = preprocess_image(img_bytes, target_size=(256, 256))
-
-            # 🔥 Use SavedModel signature
-            result = model(tf.constant(img_array))
-            
-            if isinstance(result, dict):
-                masks = list(result.values())[0].numpy()
-            else:
-                masks = result.numpy()
-
-            # Post-process mask
-            mask = masks[0]  # Remove batch dimension
-            
-            if mask.ndim == 3:
-                mask = mask[:, :, 0]  # Take first channel
-            
-            # Normalize to 0-255
-            mask = (mask * 255).astype(np.uint8)
-            
-            # 🔥 Save mask as PNG image
-            mask_image = Image.fromarray(mask)
-            
-            # Resize to original image size if needed
-            original_image = Image.open(io.BytesIO(img_bytes))
-            original_size = original_image.size
-            mask_image = mask_image.resize(original_size)
-            
-            # Save to segments directory
-            seg_filename = f"{request_id}_mask.png"
-            seg_path = str(SEGMENTS_DIR / seg_filename)
-            mask_image.save(seg_path, format='PNG')
-
-            seg_result.update({
-                "model_used": True,
-                "masks_shape": list(masks.shape),
-                "max_confidence": float(np.max(masks)),
-                "mask_path": seg_path
-            })
-
-            print(f"✅ SEGMENTATION DONE {request_id}: {seg_path}")
-
-        except Exception as e:
-            seg_result["error"] = str(e)
-            seg_path = None
-            print(f"❌ SEGMENTATION FAILED {request_id}: {e}")
-
-    # Update history
-    update_history(
-        request_id,
-        segmentation_path=seg_path,
-        status="segmented" if seg_path else "segmentation_failed"
-    )
+        return {
+            "status": "failed",
+            "error": "Segmentation model not loaded",
+        }
+    
+    try:
+        # Preprocess for segmentation (256x256)
+        img_array = preprocess_image(image_bytes, target_size=(256, 256))
+
+        # 🔥 Call SavedModel signature
+        result = model(tf.constant(img_array))
+        
+        if isinstance(result, dict):
+            masks = list(result.values())[0].numpy()
+        else:
+            masks = result.numpy()
+
+        # Extract stats
+        max_confidence = float(np.max(masks))
+        masks_shape = list(masks.shape)
+
+        print(f"🎨 SEGMENTATION: shape={masks_shape}, max_conf={max_confidence:.3f}")
+
+        return {
+            "status": "completed",
+            "masks_shape": masks_shape,
+            "max_confidence": max_confidence,
+        }
+
+    except Exception as e:
+        print(f"❌ SEGMENTATION FAILED: {e}")
+        return {
+            "status": "failed",
+            "error": str(e),
+        }