Update app.py

app.py

@@ -38,14 +38,17 @@ def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']
 
 def preprocess_image(image_path):
-    """Preprocess image for U-Net model (256x256 grayscale)"""
+    """Preprocess image for U-Net model (256x256 RGB)"""
     try:
-        # Read image
-        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
+        # Read image as RGB
+        img = cv2.imread(image_path)
         
         if img is None:
             raise ValueError("Could not read image")
         
+        # Convert BGR to RGB
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        
         print(f"Original image shape: {img.shape}")
         
         # Resize to 256x256
@@ -54,9 +57,8 @@ def preprocess_image(image_path):
         # Normalize to [0, 1]
         img_normalized = img_resized / 255.0
         
-        # Add dimensions: (1, 256, 256, 1)
+        # Add batch dimension: (1, 256, 256, 3)
         img_input = np.expand_dims(img_normalized, axis=0)
-        img_input = np.expand_dims(img_input, axis=-1)
         
         print(f"Model input shape: {img_input.shape}")
         
@@ -80,23 +82,16 @@ def create_colored_mask(prediction, original_shape):
 
 def create_overlay(original, mask, alpha=0.5):
     """Create overlay of original image and segmentation mask"""
-    # Convert grayscale to RGB
-    if len(original.shape) == 2:
-        original_rgb = cv2.cvtColor(original, cv2.COLOR_GRAY2RGB)
-    else:
-        original_rgb = original
-    
+    # Original is already RGB
     # Blend
-    overlay = cv2.addWeighted(original_rgb, 1 - alpha, mask, alpha, 0)
+    overlay = cv2.addWeighted(original, 1 - alpha, mask, alpha, 0)
     
     return overlay
 
 def img_to_base64(img):
     """Convert numpy image to base64 string"""
-    if len(img.shape) == 2:
-        img_pil = Image.fromarray(img)
-    else:
-        img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+    # Images are already in RGB format
+    img_pil = Image.fromarray(img.astype('uint8'))
     
     buf = BytesIO()
     img_pil.save(buf, format='PNG')