Update app.py

app.py

@@ -42,7 +42,17 @@ try:
     print("Successfully imported albumentations")
 except ImportError:
     albumentations_available = False
-    print("Warning: albumentations not available, will use OpenCV for transforms")
+    print("Warning: albumentations not available, will try to install it")
+    import subprocess
+    try:
+        subprocess.check_call([
+            "pip", "install", "albumentations"
+        ])
+        import albumentations as A
+        albumentations_available = True
+        print("Successfully installed and imported albumentations")
+    except:
+        print("Failed to install albumentations, will use OpenCV for transforms")
 
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -123,8 +133,10 @@ 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)
@@ -133,6 +145,7 @@ 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:
@@ -159,6 +172,8 @@ 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):
@@ -167,35 +182,59 @@ 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)):
     """
     Preprocess an image for inference
     """
-    # If image is a TIFF path, read it
-    if isinstance(image, str) and (image.lower().endswith('.tif') or image.lower().endswith('.tiff')):
-        image = read_tiff_image(image)
-        if image is None:
-            return None, None
-    # Convert to numpy array if PIL image
+    # If image is already a numpy array, use it directly
+    if isinstance(image, np.ndarray):
+        # Ensure RGB format
+        if len(image.shape) == 2:  # Grayscale
+            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+        elif image.shape[2] == 4:  # RGBA
+            image = image[:, :, :3]
+        
+        # Make a copy for display
+        display_image = image.copy()
+        
+        # Normalize to [0, 1] if needed
+        if display_image.max() > 1.0:
+            image = image.astype(np.float32) / 255.0
+    
+    # Convert PIL image to numpy
     elif isinstance(image, Image.Image):
+        print("Converting PIL image to numpy array")
         image = np.array(image)
+        
+        # Ensure RGB format
+        if len(image.shape) == 2:  # Grayscale
+            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+        elif image.shape[2] == 4:  # RGBA
+            image = image[:, :, :3]
+        
+        # Make a copy for display
+        display_image = image.copy()
+        
+        # Normalize to [0, 1]
+        image = image.astype(np.float32) / 255.0
+    else:
+        print(f"Unsupported image type: {type(image)}")
+        return None, None
     
-    # Ensure RGB format
-    if len(image.shape) == 2:  # Grayscale
-        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
-    elif image.shape[2] == 4:  # RGBA
-        image = image[:, :, :3]
-    
-    # Make a copy for display
-    display_image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.copy()
+    print(f"Image shape: {image.shape}, min: {image.min()}, max: {image.max()}")
     
     # Resize image to the target size
     if albumentations_available:
@@ -216,141 +255,146 @@ def preprocess_image(image, target_size=(128, 128)):
     
     return image_tensor, display_image
 
-def preprocess_uploaded_file(file_obj, target_size=(128, 128)):
+def save_temp_file(file_obj, suffix='.tif'):
     """
-    Process an uploaded file (could be image or TIFF)
+    Save uploaded file to a temporary file
     """
     try:
-        # Save to a temporary file
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as temp_file:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
             temp_path = temp_file.name
-            # Write the content to the temp file
-            if hasattr(file_obj, 'read'):
-                # If it's a file-like object
-                temp_file.write(file_obj.read())
-            elif isinstance(file_obj, (str, bytes)):
-                # If it's a string path or bytes
-                if isinstance(file_obj, str):
-                    with open(file_obj, 'rb') as f:
-                        temp_file.write(f.read())
-                else:
-                    temp_file.write(file_obj)
-        
-        # Check if it's a TIFF file
-        if temp_path.lower().endswith(('.tif', '.tiff')):
-            # Read as TIFF
-            image = read_tiff_image(temp_path)
-            if image is None:
-                os.unlink(temp_path)
-                return None, None
-            
-            # Make a copy for display
-            display_image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.copy()
             
-            # Resize/preprocess
-            if albumentations_available:
-                aug = A.Compose([
-                    A.PadIfNeeded(min_height=target_size[0], min_width=target_size[1], 
-                                 border_mode=cv2.BORDER_CONSTANT, value=0),
-                    A.CenterCrop(height=target_size[0], width=target_size[1])
-                ])
-                augmented = aug(image=image)
-                image_resized = augmented['image']
+            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
+                with open(file_obj, 'rb') as f:
+                    temp_file.write(f.read())
+                print(f"Copied file from {file_obj} to {temp_path}")
             else:
-                image_resized = cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)
-            
-            # Convert to tensor
-            image_tensor = torch.from_numpy(image_resized.transpose(2, 0, 1)).float().unsqueeze(0)
-        else:
-            # Try to open as a regular image
-            try:
-                pil_image = Image.open(temp_path)
-                image_tensor, display_image = preprocess_image(pil_image, target_size)
-            except Exception as e:
-                print(f"Error opening as regular image: {e}")
+                print(f"Unsupported file object type: {type(file_obj)}")
                 os.unlink(temp_path)
-                return None, None
+                return None
+                
+        return temp_path
+    except Exception as e:
+        print(f"Error saving temporary file: {e}")
+        import traceback
+        traceback.print_exc()
+        return None
+
+def process_uploaded_tiff(file_obj):
+    """
+    Process an uploaded TIFF file
+    """
+    try:
+        print(f"Processing uploaded TIFF file: {type(file_obj)}")
+        
+        # Save to a temporary file
+        temp_path = save_temp_file(file_obj)
+        if not temp_path:
+            return None, None
+        
+        # Read the TIFF file
+        image = read_tiff_image(temp_path)
         
         # Clean up
         os.unlink(temp_path)
+        
+        if image is None:
+            return None, None
+        
+        # Make a copy for display
+        display_image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.copy()
+        
+        # Resize/preprocess
+        if albumentations_available:
+            aug = A.Compose([
+                A.PadIfNeeded(min_height=128, min_width=128, 
+                             border_mode=cv2.BORDER_CONSTANT, value=0),
+                A.CenterCrop(height=128, width=128)
+            ])
+            augmented = aug(image=image)
+            image_resized = augmented['image']
+        else:
+            image_resized = cv2.resize(image, (128, 128), interpolation=cv2.INTER_LINEAR)
+        
+        # Convert to tensor
+        image_tensor = torch.from_numpy(image_resized.transpose(2, 0, 1)).float().unsqueeze(0)
+        
         return image_tensor, display_image
     
     except Exception as e:
-        print(f"Error processing uploaded file: {e}")
+        print(f"Error processing uploaded TIFF: {e}")
         import traceback
         traceback.print_exc()
         return None, None
 
-def preprocess_mask(mask_input, target_size=(128, 128)):
+def process_uploaded_mask(file_obj):
     """
-    Preprocess a ground truth mask
+    Process an uploaded mask file
     """
     try:
-        # If mask is a file path or file-like object, process it
-        if isinstance(mask_input, str) or hasattr(mask_input, 'read'):
-            # Save to temp file if it's a file-like object
-            if hasattr(mask_input, 'read'):
-                with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as temp_file:
-                    temp_path = temp_file.name
-                    temp_file.write(mask_input.read())
-                mask_path = temp_path
-            else:
-                mask_path = mask_input
-            
-            # Check if it's a TIFF
-            if mask_path.lower().endswith(('.tif', '.tiff')):
-                mask = read_tiff_mask(mask_path)
-                if mask is None:
-                    if hasattr(mask_input, 'read'):
-                        os.unlink(temp_path)
-                    return None
-            else:
-                # Try as regular image
-                try:
-                    mask = np.array(Image.open(mask_path))
-                    if len(mask.shape) == 3:
-                        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
-                except Exception as e:
-                    print(f"Error opening mask as image: {e}")
-                    if hasattr(mask_input, 'read'):
-                        os.unlink(temp_path)
-                    return None
-            
-            # Clean up if temp file
-            if hasattr(mask_input, 'read'):
-                os.unlink(temp_path)
+        print(f"Processing uploaded mask file: {type(file_obj)}")
         
-        # Convert to numpy array if PIL image
-        elif isinstance(mask_input, Image.Image):
-            mask = np.array(mask_input)
-            if len(mask.shape) == 3:
-                mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
+        # Save to a temporary file
+        temp_path = save_temp_file(file_obj)
+        if not temp_path:
+            return None
         
-        # Already numpy array
+        # Check if it's a TIFF file
+        if temp_path.lower().endswith(('.tif', '.tiff')):
+            mask = read_tiff_mask(temp_path)
         else:
-            mask = mask_input
-            if len(mask.shape) == 3:
-                mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
+            # Try to open as a regular image
+            try:
+                mask_img = Image.open(temp_path)
+                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)
+                return None
+        
+        # Clean up
+        os.unlink(temp_path)
+        
+        if mask is None:
+            return None
         
-        # Resize mask to target size
+        # Resize mask to 128x128
         if albumentations_available:
             aug = A.Compose([
-                A.PadIfNeeded(min_height=target_size[0], min_width=target_size[1], 
+                A.PadIfNeeded(min_height=128, min_width=128, 
                              border_mode=cv2.BORDER_CONSTANT, value=0),
-                A.CenterCrop(height=target_size[0], width=target_size[1])
+                A.CenterCrop(height=128, width=128)
             ])
             augmented = aug(image=mask)
             mask_resized = augmented['image']
         else:
-            mask_resized = cv2.resize(mask, target_size, interpolation=cv2.INTER_NEAREST)
+            mask_resized = cv2.resize(mask, (128, 128), interpolation=cv2.INTER_NEAREST)
         
         # Binarize the mask (0: background, 1: wetland)
-        mask_binary = (mask_resized > 127).astype(np.uint8)
+        mask_binary = (mask_resized > 0).astype(np.uint8)
         
         return mask_binary
     
     except Exception as e:
-        print(f"Error preprocessing mask: {e}")
+        print(f"Error processing uploaded mask: {e}")
         import traceback
         traceback.print_exc()
         return None
@@ -415,27 +459,36 @@ def calculate_metrics(pred_mask, gt_mask):
     
     return metrics
 
-def process_images(input_image=None, input_tiff=None, gt_mask=None):
+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."
         
-        # Process the input image
-        if input_tiff is not None:
+        # Process the input
+        if input_tiff is not None and input_tiff:
             # Process uploaded TIFF file
-            image_tensor, display_image = preprocess_uploaded_file(input_tiff)
-        else:
+            image_tensor, display_image = process_uploaded_tiff(input_tiff)
+            if image_tensor is None:
+                return None, "Failed to process the input TIFF file."
+        elif input_image is not None:
             # Process regular image
             image_tensor, display_image = preprocess_image(input_image)
-        
-        if image_tensor is None or display_image is None:
-            return None, "Failed to process the input image."
+            if image_tensor is None:
+                return None, "Failed to process the input image."
+        else:
+            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."
@@ -444,14 +497,20 @@ def process_images(input_image=None, input_tiff=None, gt_mask=None):
         gt_mask_processed = None
         metrics_text = ""
         
-        if gt_mask is not None:
-            gt_mask_processed = preprocess_mask(gt_mask)
+        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:
@@ -498,6 +557,7 @@ def process_images(input_image=None, input_tiff=None, gt_mask=None):
         result_image = Image.open(buf)
         plt.close(fig)
         
+        print("Processing completed successfully.")
         return result_image, result_text
         
     except Exception as e:
@@ -521,7 +581,8 @@ 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"])
             
-            gt_mask = gr.Image(label="Ground Truth Mask (Optional)", type="numpy")
+            # IMPORTANT CHANGE: Changed ground truth from Image to File for TIFF support
+            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")
         
@@ -546,6 +607,7 @@ with gr.Blocks(title="Wetlands Segmentation from Satellite Imagery") as demo:
     - The model works best with RGB satellite imagery
     - For optimal results, use images with similar characteristics to those used in training
     - The model focuses on identifying wetland regions in natural landscapes
+    - For ground truth masks, both TIFF and standard image formats are supported
     
     **Repository:** [dcrey7/wetlands_segmentation_deeplabsv3plus](https://huggingface.co/dcrey7/wetlands_segmentation_deeplabsv3plus)
     """)
@@ -553,7 +615,7 @@ with gr.Blocks(title="Wetlands Segmentation from Satellite Imagery") as demo:
     # Set up event handlers
     process_btn.click(
         fn=process_images,
-        inputs=[input_image, input_tiff, gt_mask],
+        inputs=[input_image, input_tiff, gt_mask_file],
         outputs=[output_image, output_text]
     )