Update app.py

app.py

@@ -12,7 +12,6 @@ from io import BytesIO
 import urllib.request
 import tempfile
 import rasterio
-from rasterio.plot import reshape_as_image
 import warnings
 warnings.filterwarnings("ignore")
 
@@ -69,7 +68,7 @@ if smp_available:
     )
 else:
     # Fallback to a simple model that won't actually work but allows the UI to load
-    print("Warning: Using a placeholder model that won't produce correct predictions.")
+    print("Warning: Using a placeholder model that won't produce valid predictions.")
     from torch import nn
     class PlaceholderModel(nn.Module):
         def __init__(self):
@@ -133,10 +132,8 @@ def read_tiff_image(tiff_path):
     This matches your training data loading approach
     """
     try:
-        print(f"Reading TIFF image from: {tiff_path}")
         # Read the image using rasterio (get RGB channels)
         with rasterio.open(tiff_path) as src:
-            print(f"TIFF opened successfully. Number of bands: {src.count}")
             # Check if we have enough bands
             if src.count >= 3:
                 red = src.read(1)
@@ -145,7 +142,6 @@ def read_tiff_image(tiff_path):
                 
                 # Stack to create RGB image
                 image = np.dstack((red, green, blue)).astype(np.float32)
-                print(f"RGB image created, shape: {image.shape}, min: {image.min()}, max: {image.max()}")
                 
                 # Normalize to [0, 1]
                 if image.max() > 0:
@@ -154,7 +150,6 @@ def read_tiff_image(tiff_path):
                 return image
             else:
                 # If less than 3 bands, handle accordingly
-                print(f"Warning: TIFF file has only {src.count} bands, RGB expected")
                 bands = [src.read(i+1) for i in range(src.count)]
                 # If only one band, duplicate to create RGB
                 if len(bands) == 1:
@@ -172,8 +167,6 @@ def read_tiff_image(tiff_path):
                 return image
     except Exception as e:
         print(f"Error reading TIFF file: {e}")
-        import traceback
-        traceback.print_exc()
         return None
 
 def read_tiff_mask(mask_path):
@@ -182,17 +175,12 @@ def read_tiff_mask(mask_path):
     This matches your training data loading approach
     """
     try:
-        print(f"Reading TIFF mask from: {mask_path}")
         # Read mask
         with rasterio.open(mask_path) as src:
-            print(f"Mask TIFF opened successfully. Number of bands: {src.count}")
             mask = src.read(1).astype(np.uint8)
-            print(f"Mask shape: {mask.shape}, min: {mask.min()}, max: {mask.max()}, unique values: {np.unique(mask)}")
         return mask
     except Exception as e:
         print(f"Error reading mask file: {e}")
-        import traceback
-        traceback.print_exc()
         return None
 
 def preprocess_image(image, target_size=(128, 128)):
@@ -216,7 +204,6 @@ def preprocess_image(image, target_size=(128, 128)):
     
     # Convert PIL image to numpy
     elif isinstance(image, Image.Image):
-        print("Converting PIL image to numpy array")
         image = np.array(image)
         
         # Ensure RGB format
@@ -234,8 +221,6 @@ def preprocess_image(image, target_size=(128, 128)):
         print(f"Unsupported image type: {type(image)}")
         return None, None
     
-    print(f"Image shape: {image.shape}, min: {image.min()}, max: {image.max()}")
-    
     # Resize image to the target size
     if albumentations_available:
         # Use albumentations to match training preprocessing
@@ -255,59 +240,54 @@ def preprocess_image(image, target_size=(128, 128)):
     
     return image_tensor, display_image
 
-def save_temp_file(file_obj, suffix='.tif'):
-    """
-    Save uploaded file to a temporary file
-    """
+def extract_file_content(file_obj):
+    """Extract content from the file object, handling different types"""
     try:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
-            temp_path = temp_file.name
-            
-            if hasattr(file_obj, 'name'):
-                # If it has a name attribute, it's likely a FileStorage object
-                file_obj.save(temp_path)
-                print(f"Saved file from FileStorage object to {temp_path}")
-            elif hasattr(file_obj, 'read'):
-                # If it's a file-like object
-                content = file_obj.read()
-                if isinstance(content, str):  # It's text, not binary
-                    content = content.encode('utf-8')
-                temp_file.write(content)
-                print(f"Wrote {len(content)} bytes to {temp_path}")
-            elif isinstance(file_obj, bytes):
-                # If it's already bytes
-                temp_file.write(file_obj)
-                print(f"Wrote {len(file_obj)} bytes to {temp_path}")
-            elif isinstance(file_obj, str):
-                # If it's a path
+        if hasattr(file_obj, 'name') and isinstance(file_obj, str):
+            # Handle Gradio's NamedString
+            content = file_obj
+            if os.path.exists(content):
+                # It's a path
+                with open(content, 'rb') as f:
+                    return f.read()
+            else:
+                # It's content
+                return content.encode('latin1')
+        elif hasattr(file_obj, 'read'):
+            # File-like object
+            return file_obj.read()
+        elif isinstance(file_obj, bytes):
+            # Already bytes
+            return file_obj
+        elif isinstance(file_obj, str):
+            # String path
+            if os.path.exists(file_obj):
                 with open(file_obj, 'rb') as f:
-                    temp_file.write(f.read())
-                print(f"Copied file from {file_obj} to {temp_path}")
+                    return f.read()
             else:
-                print(f"Unsupported file object type: {type(file_obj)}")
-                os.unlink(temp_path)
-                return None
-                
-        return temp_path
+                return file_obj.encode('utf-8')
+        else:
+            print(f"Unsupported file object type: {type(file_obj)}")
+            return None
     except Exception as e:
-        print(f"Error saving temporary file: {e}")
-        import traceback
-        traceback.print_exc()
+        print(f"Error extracting file content: {e}")
         return None
 
 def process_uploaded_tiff(file_obj):
-    """
-    Process an uploaded TIFF file
-    """
+    """Process an uploaded TIFF file"""
     try:
-        print(f"Processing uploaded TIFF file: {type(file_obj)}")
+        # Get file content
+        file_content = extract_file_content(file_obj)
+        if file_content is None:
+            print("Failed to extract file content")
+            return None, None
         
         # Save to a temporary file
-        temp_path = save_temp_file(file_obj)
-        if not temp_path:
-            return None, None
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as temp_file:
+            temp_path = temp_file.name
+            temp_file.write(file_content)
         
-        # Read the TIFF file
+        # Read as TIFF
         image = read_tiff_image(temp_path)
         
         # Clean up
@@ -343,16 +323,24 @@ def process_uploaded_tiff(file_obj):
         return None, None
 
 def process_uploaded_mask(file_obj):
-    """
-    Process an uploaded mask file
-    """
+    """Process an uploaded mask file"""
     try:
-        print(f"Processing uploaded mask file: {type(file_obj)}")
+        # Get file content
+        file_content = extract_file_content(file_obj)
+        if file_content is None:
+            return None
         
         # Save to a temporary file
-        temp_path = save_temp_file(file_obj)
-        if not temp_path:
-            return None
+        # Determine suffix based on file name if available
+        suffix = '.tif'
+        if hasattr(file_obj, 'name'):
+            file_name = getattr(file_obj, 'name')
+            if isinstance(file_name, str) and '.' in file_name:
+                suffix = '.' + file_name.split('.')[-1].lower()
+        
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
+            temp_path = temp_file.name
+            temp_file.write(file_content)
         
         # Check if it's a TIFF file
         if temp_path.lower().endswith(('.tif', '.tiff')):
@@ -364,7 +352,6 @@ def process_uploaded_mask(file_obj):
                 mask = np.array(mask_img)
                 if len(mask.shape) == 3:
                     mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
-                print(f"Opened mask as regular image, shape: {mask.shape}")
             except Exception as e:
                 print(f"Error opening mask as regular image: {e}")
                 os.unlink(temp_path)
@@ -420,8 +407,6 @@ def predict_segmentation(image_tensor):
         return pred
     except Exception as e:
         print(f"Error during prediction: {e}")
-        import traceback
-        traceback.print_exc()
         return None
 
 def calculate_metrics(pred_mask, gt_mask):
@@ -464,11 +449,6 @@ def process_images(input_image=None, input_tiff=None, gt_mask_file=None):
     Process input images and generate predictions
     """
     try:
-        print("\n---- Starting new processing request ----")
-        print(f"Input image type: {type(input_image) if input_image is not None else None}")
-        print(f"Input TIFF type: {type(input_tiff) if input_tiff is not None else None}")
-        print(f"Ground truth mask type: {type(gt_mask_file) if gt_mask_file is not None else None}")
-        
         # Check if we have input
         if input_image is None and input_tiff is None:
             return None, "Please upload an image or TIFF file."
@@ -488,7 +468,6 @@ def process_images(input_image=None, input_tiff=None, gt_mask_file=None):
             return None, "No valid input provided."
         
         # Get prediction
-        print("Running model prediction...")
         pred_mask = predict_segmentation(image_tensor)
         if pred_mask is None:
             return None, "Failed to generate prediction."
@@ -498,19 +477,13 @@ def process_images(input_image=None, input_tiff=None, gt_mask_file=None):
         metrics_text = ""
         
         if gt_mask_file is not None and gt_mask_file:
-            print("Processing ground truth mask...")
             gt_mask_processed = process_uploaded_mask(gt_mask_file)
             
             if gt_mask_processed is not None:
-                print("Calculating metrics...")
                 metrics = calculate_metrics(pred_mask, gt_mask_processed)
                 metrics_text = "\n".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
-                print(f"Metrics calculated: {metrics}")
-            else:
-                print("Failed to process ground truth mask.")
         
         # Create visualization
-        print("Creating visualization...")
         fig = plt.figure(figsize=(12, 6))
         
         if gt_mask_processed is not None:
@@ -549,15 +522,14 @@ def process_images(input_image=None, input_tiff=None, gt_mask_file=None):
         if metrics_text:
             result_text += f"\n\nEvaluation Metrics:\n{metrics_text}"
         
-        # Convert figure to image
-        buf = BytesIO()
+        # Convert figure to image for display
         plt.tight_layout()
-        plt.savefig(buf, format='png', dpi=150)
+        buf = BytesIO()
+        plt.savefig(buf, format='png')
         buf.seek(0)
         result_image = Image.open(buf)
         plt.close(fig)
         
-        print("Processing completed successfully.")
         return result_image, result_text
         
     except Exception as e:
@@ -581,7 +553,7 @@ with gr.Blocks(title="Wetlands Segmentation from Satellite Imagery") as demo:
             with gr.Tab("Upload TIFF"):
                 input_tiff = gr.File(label="Upload TIFF File", file_types=[".tif", ".tiff"])
             
-            # IMPORTANT CHANGE: Changed ground truth from Image to File for TIFF support
+            # Ground truth mask as file upload
             gt_mask_file = gr.File(label="Ground Truth Mask (Optional)", file_types=[".tif", ".tiff", ".png", ".jpg", ".jpeg"])
             
             process_btn = gr.Button("Analyze Image", variant="primary")