Update app.py

app.py

@@ -4,6 +4,27 @@ from tensorflow.keras.models import load_model
 from tensorflow.keras import backend as K
 from PIL import Image
 from huggingface_hub import hf_hub_download
+import cv2
+
+
+# -----------------------
+# Preprocessing
+# -----------------------
+def preprocess(img):
+    img_gray = np.array(img.convert("L"))  # grayscale
+    img_resized = cv2.resize(img_gray, (256, 256))  # resize to model input
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
+    img_clahe = clahe.apply(img_resized)  # enhance contrast
+    img_norm = img_clahe / 255.0  # normalize to [0,1]
+    return img_norm
+from skimage import measure
+
+def remove_small_blobs(mask, min_size=50):
+    labels = measure.label(mask)
+    for region in measure.regionprops(labels):
+        if region.area < min_size:
+            mask[labels == region.label] = 0
+    return mask
 
 # -----------------------
 # Custom metric
@@ -46,15 +67,15 @@ def predict(img):
     orig_size = img.size  # (width, height)
     
     # Resize image for model
-    img_gray = img.convert("L").resize((224, 224))
-    x = np.expand_dims(np.expand_dims(np.array(img_gray)/255.0, -1), 0)
-    
+    img_processed = preprocess(img)          # preprocess function
+    x = np.expand_dims(np.expand_dims(img_processed, -1), 0)
+
     # Predict mask
     pred = model.predict(x)[0]
     tumor_present = pred.max() > 0.7
     
     if tumor_present:
-        mask = (pred > 0.7).astype(np.uint8) * 255
+        mask = remove_small_blobs(mask, min_size=50)
         mask_img = Image.fromarray(mask.squeeze()).convert("L")
         
         # Resize mask back to original image size