correct segm

app.py

@@ -149,6 +149,8 @@ def run_detection(model, image_input, original_image, confidence_threshold):
     
     return annotated, analytics_text, {"count": len(boxes), "objects": class_counts}
 
+
+'''
 def run_segmentation(model, image_input, original_image):
     mask_tensor = model.predict(image_input)
     mask_np = mask_tensor.cpu().numpy().astype(np.uint8)
@@ -195,6 +197,38 @@ def run_segmentation(model, image_input, original_image):
 
 
     return Image.fromarray(blended), analytics_text, {"classes_found": list(found_classes)}
+'''
+def run_segmentation(model, image):
+    """
+    Handles Segmentation: Returns Tensor of shape (H, W) with class IDs.
+    """
+    # 1. Run Prediction
+    mask_tensor = model.predict(image)
+    
+    # 2. Convert to Numpy
+    mask_np = mask_tensor.cpu().numpy().astype(np.uint8)
+
+    # 3. Create a Colored Mask
+    h, w = mask_np.shape
+    colored_mask = np.zeros((h, w, 3), dtype=np.uint8)
+    
+    unique_classes = np.unique(mask_np)
+    
+    for cls_id in unique_classes:
+        if cls_id == -1: continue 
+        
+        np.random.seed(int(cls_id)) 
+        color = np.random.randint(50, 255, size=3) 
+        colored_mask[mask_np == cls_id] = color
+
+    # 4. Blend with Original Image
+    image_np = np.array(image)
+    if image_np.shape[:2] != colored_mask.shape[:2]:
+        colored_mask = cv2.resize(colored_mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST)
+
+    blended = cv2.addWeighted(image_np, 0.6, colored_mask, 0.4, 0)
+    return Image.fromarray(blended)
+
 
 # --- GRADIO UI ---