Update app.py

app.py

@@ -7,6 +7,8 @@ from datetime import datetime
 import base64
 from io import BytesIO
 from PIL import Image
+import SimpleITK as sitk
+from skimage.transform import resize
 
 # Suppress TensorFlow warnings
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
@@ -17,7 +19,7 @@ from tensorflow.keras.models import load_model
 app = Flask(__name__)
 app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max
 app.config['UPLOAD_FOLDER'] = 'uploads'
-app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'mha'}
+app.config['ALLOWED_EXTENSIONS'] = {'mha'}  # Only MRI .mha files
 
 # Create uploads folder with proper permissions
 os.makedirs(app.config['UPLOAD_FOLDER'], mode=0o777, exist_ok=True)
@@ -45,28 +47,42 @@ def allowed_file(filename):
            filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']
 
 def preprocess_image(image_path):
-    """Preprocess image for brain segmentation"""
-    # Read image in grayscale
-    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
-    
-    if img is None:
-        raise ValueError("Could not read image")
-    
-    # Get original shape
-    original_shape = img.shape
-    
-    # Resize to model input size: 160x160 (NOT 256x256!)
-    img_resized = cv2.resize(img, (160, 160))
-    
-    # Normalize to [0, 1]
-    img_normalized = img_resized.astype(np.float32) / 255.0
-    
-    # Add batch and channel dimensions in channels_first format (NCHW)
-    # Model expects: (batch, channels, height, width) = (None, 1, 160, 160)
-    img_input = np.expand_dims(img_normalized, axis=0)  # (1, 160, 160)
-    img_input = np.expand_dims(img_input, axis=0)  # (1, 1, 160, 160)
-    
-    return img_input, original_shape
+    """Preprocess MHA image for brain segmentation (same as training)"""
+    try:
+        # Read MHA file using SimpleITK
+        img = sitk.ReadImage(image_path)
+        img = sitk.GetArrayFromImage(img)
+        
+        print(f"Original MHA shape: {img.shape}")
+        
+        # Resize to (155, 160, 160) - same as training
+        img_resized = resize(img, (155, 160, 160), preserve_range=True)
+        
+        # Select middle slice (same as training uses slice 60-130)
+        # For single prediction, use slice 95 (middle of range)
+        middle_slice = 95
+        img_slice = img_resized[middle_slice, :, :]  # (160, 160)
+        
+        # Keep original for visualization
+        original_slice = img_slice.copy()
+        
+        # Z-score normalization (same as training)
+        img_normalized = (img_slice - img_slice.mean()) / (img_slice.std() + 1e-8)
+        img_normalized = img_normalized.astype(np.float32)
+        
+        print(f"Slice shape: {img_normalized.shape}")
+        
+        # Add batch and channel dimensions in channels_first format (NCHW)
+        # Model expects: (batch, channels, height, width) = (None, 1, 160, 160)
+        img_input = np.expand_dims(img_normalized, axis=0)  # (1, 160, 160)
+        img_input = np.expand_dims(img_input, axis=0)  # (1, 1, 160, 160)
+        
+        print(f"Model input shape: {img_input.shape}")
+        
+        return img_input, original_slice
+        
+    except Exception as e:
+        raise ValueError(f"Failed to read MHA file: {str(e)}")
 
 def postprocess_mask(mask, original_shape):
     """Postprocess segmentation mask"""
@@ -144,7 +160,7 @@ def predict():
             return jsonify({'error': 'No file selected'}), 400
         
         if not allowed_file(file.filename):
-            return jsonify({'error': 'Invalid file type. Please upload PNG, JPG, or JPEG'}), 400
+            return jsonify({'error': 'Invalid file type. Please upload .mha MRI file'}), 400
         
         # Save uploaded file
         timestamp = datetime.now().strftime('%Y%m%d_%Hh%Mm%Ss')
@@ -152,24 +168,25 @@ def predict():
         filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
         file.save(filepath)
         
-        # Read original image
-        original_img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
-        
-        # Preprocess
-        img_input, original_shape = preprocess_image(filepath)
+        # Preprocess MHA file (returns normalized input and original slice)
+        img_input, original_slice = preprocess_image(filepath)
         
         # Predict
         print("Making prediction...")
         prediction = model.predict(img_input, verbose=0)
         
-        # Postprocess
-        mask = postprocess_mask(prediction[0], original_shape)
+        # Postprocess mask (returns 160x160 binary mask)
+        mask = postprocess_mask(prediction[0], original_slice.shape)
+        
+        # Normalize original slice for display (0-255)
+        original_display = ((original_slice - original_slice.min()) / 
+                           (original_slice.max() - original_slice.min() + 1e-8) * 255).astype(np.uint8)
         
         # Create overlay
-        overlay = create_overlay(original_img, mask)
+        overlay = create_overlay(original_display, mask)
         
         # Convert to base64
-        original_base64 = img_to_base64(original_img)
+        original_base64 = img_to_base64(original_display)
         mask_base64 = img_to_base64(mask)
         overlay_base64 = img_to_base64(overlay)
         
@@ -196,17 +213,6 @@ def predict():
         traceback.print_exc()
         return jsonify({'error': str(e)}), 500
 
-@app.route('/example')
-def example():
-    """Get example image"""
-    example_path = 'image/20251012_09h06m52s_grim.png'
-    if os.path.exists(example_path):
-        with open(example_path, 'rb') as f:
-            img_data = f.read()
-        img_base64 = base64.b64encode(img_data).decode('utf-8')
-        return jsonify({'image': f"data:image/png;base64,{img_base64}"})
-    return jsonify({'error': 'Example image not found'}), 404
-
 if __name__ == '__main__':
     print("\n" + "="*60)
     print("🧠 Brain Tumor Segmentation App")