Update stenosis.py

stenosis.py

@@ -1,108 +1,44 @@
-import cv2
+# stenosis.py
 import numpy as np
-import base64
+import cv2
 from skimage.morphology import skeletonize
 from scipy.ndimage import distance_transform_edt
 
+def detect_stenosis(mask: np.ndarray):
 
-# -------------------------------
-# Severity classification
-# -------------------------------
-def classify_severity(percent):
-    if percent < 30:
-        return "Mild"
-    elif percent < 70:
-        return "Moderate"
-    else:
-        return "Severe"
-
-
-# -------------------------------
-# Core stenosis detection
-# -------------------------------
-def detect_stenosis(binary_mask):
-
-    mask = (binary_mask > 0).astype(np.uint8)
+    mask = (mask > 0).astype(np.uint8)
     skeleton = skeletonize(mask).astype(np.uint8)
     dist = distance_transform_edt(mask)
 
     coords = np.column_stack(np.where(skeleton > 0))
-    if len(coords) < 20:
+    if len(coords) < 30:
         return 0.0, None
 
     radii = np.array([dist[y, x] for y, x in coords])
 
-    ref_radius = np.percentile(radii, 90)
+    # smooth radius (VERY IMPORTANT)
+    radii = np.convolve(radii, np.ones(7)/7, mode="same")
+
+    ref_radius = np.percentile(radii, 85)
     min_radius = radii.min()
 
-    narrowing_percent = (1 - min_radius / ref_radius) * 100
+    severity = float((1 - min_radius / ref_radius) * 100)
+
+    # pick worst region center
+    worst = np.where(radii < ref_radius * 0.6)[0]
+    if len(worst) == 0:
+        return severity, None
 
-    idx = np.argmin(radii)
+    idx = worst[len(worst)//2]
     y, x = coords[idx]
 
-    box = int(ref_radius * 4)
-    h, w = mask.shape
+    # force visible bounding box
+    box = max(int(ref_radius * 8), 30)
 
+    h, w = mask.shape
     x1 = max(0, x - box)
     y1 = max(0, y - box)
     x2 = min(w, x + box)
     y2 = min(h, y + box)
 
-    return float(narrowing_percent), (x1, y1, x2, y2)
-
-
-# -------------------------------
-# Draw bounding box
-# -------------------------------
-def draw_stenosis_box(image, bbox, percent):
-
-    if len(image.shape) == 2:
-        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
-
-    severity = classify_severity(percent)
-    x1, y1, x2, y2 = bbox
-
-    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 3)
-
-    label = f"{severity} ({percent:.1f}%)"
-    (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
-
-    cv2.rectangle(image, (x1, y1 - h - 10), (x1 + w + 6, y1), (0, 0, 255), -1)
-    cv2.putText(image, label, (x1 + 3, y1 - 4),
-                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
-
-    return image
-
-
-# -------------------------------
-# Convert image → base64
-# -------------------------------
-def image_to_base64(image):
-    _, buffer = cv2.imencode(".png", image)
-    return base64.b64encode(buffer).decode("utf-8")
-
-
-# -------------------------------
-# MAIN API FUNCTION
-# -------------------------------
-def analyze_stenosis(original_image, vessel_mask):
-
-    percent, bbox = detect_stenosis(vessel_mask)
-
-    if bbox is None:
-        return {
-            "stenosis_detected": False,
-            "max_narrowing_percent": 0.0,
-            "severity": "None",
-            "annotated_image_base64": None
-        }
-
-    severity = classify_severity(percent)
-    annotated = draw_stenosis_box(original_image, bbox, percent)
-
-    return {
-        "stenosis_detected": True,
-        "max_narrowing_percent": round(percent, 2),
-        "severity": severity,
-        "annotated_image_base64": image_to_base64(annotated)
-    }
+    return severity, (x1, y1, x2, y2)