Update app.py

app.py

@@ -10,6 +10,19 @@ import io
 import gdown
 from transformers import TFSegformerForSemanticSegmentation
 
+
+try:
+    # Limit GPU memory growth
+    gpus = tf.config.experimental.list_physical_devices('GPU')
+    if gpus:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        st.sidebar.success(f"GPU available: {len(gpus)} device(s)")
+    else:
+        st.sidebar.warning("No GPU detected, using CPU")
+except Exception as e:
+    st.sidebar.error(f"GPU config error: {e}")
+
 st.set_page_config(
     page_title="Pet Segmentation with SegFormer",
     page_icon="🐶",
@@ -32,11 +45,12 @@ NUM_CLASSES = len(ID2LABEL)
 def download_model_from_drive():
     # Create a models directory
     os.makedirs("models", exist_ok=True)
-    model_path = "models/best_model"
+    model_path = "models/tf_model.h5"
     
     if not os.path.exists(model_path):
-        # Fixed Google Drive URL format for gdown
-        url = "https://drive.google.com/file/d/1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3/view?usp=sharing"
+        # Extract the file ID from the sharing URL
+        file_id = "1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3"
+        url = f"https://drive.google.com/uc?id={file_id}"
         try:
             gdown.download(url, model_path, quiet=False)
             st.success("Model downloaded successfully from Google Drive.")
@@ -126,8 +140,18 @@ def preprocess_image(image):
     # Store original image for display
     original_img = img_array.copy()
     
-    # Resize to target size
-    img_resized = tf.image.resize(img_array, (IMAGE_SIZE, IMAGE_SIZE))
+    # Resize to target size with preserve_aspect_ratio=False
+    img_resized = tf.image.resize(
+        img_array, 
+        (IMAGE_SIZE, IMAGE_SIZE),
+        method='bilinear',
+        preserve_aspect_ratio=False,  # Ensure exact dimensions
+        antialias=True
+    )
+    
+    # Verify dimensions with assertion
+    tf.debugging.assert_equal(tf.shape(img_resized)[0:2], [IMAGE_SIZE, IMAGE_SIZE], 
+                             message="Image dimensions don't match expected size")
     
     # Normalize
     img_normalized = normalize_image(img_resized)
@@ -150,8 +174,21 @@ def create_mask(pred_mask):
     Returns:
         Processed mask (2D array)
     """
+    # Take argmax along the class dimension (axis=1 for batch data)
     pred_mask = tf.math.argmax(pred_mask, axis=1)
+    
+    # Remove batch dimension and convert to numpy
     pred_mask = tf.squeeze(pred_mask)
+    
+    # Resize to match original image size if needed
+    if pred_mask.shape[0] != IMAGE_SIZE or pred_mask.shape[1] != IMAGE_SIZE:
+        pred_mask = tf.image.resize(
+            tf.expand_dims(pred_mask, axis=-1),
+            (IMAGE_SIZE, IMAGE_SIZE),
+            method='nearest'
+        )
+        pred_mask = tf.squeeze(pred_mask)
+    
     return pred_mask.numpy()
 
 def colorize_mask(mask):
@@ -308,21 +345,46 @@ def main():
             
             # Preprocess and predict
             with st.spinner("Generating segmentation mask..."):
-                # Preprocess the image
-                img_tensor, original_img = preprocess_image(image)
-                
-                # Make prediction
-                outputs = model(pixel_values=img_tensor, training=False)
-                logits = outputs.logits
-                
-                # Create visualization mask
-                mask = create_mask(logits)
-                
-                # Colorize the mask
-                colorized_mask = colorize_mask(mask)
-                
-                # Create overlay
-                overlay = create_overlay(original_img, colorized_mask, alpha=overlay_opacity)
+                try:
+                    # Preprocess the image
+                    img_tensor, original_img = preprocess_image(image)
+                    
+                    # Print shape to debug
+                    st.write(f"DEBUG - Input tensor shape: {img_tensor.shape}")
+                    
+                    # Make prediction with error handling
+                    try:
+                        outputs = model(pixel_values=img_tensor, training=False)
+                        logits = outputs.logits
+                        
+                        # Create visualization mask
+                        mask = create_mask(logits)
+                        
+                        # Colorize the mask
+                        colorized_mask = colorize_mask(mask)
+                        
+                        # Create overlay
+                        overlay = create_overlay(original_img, colorized_mask, alpha=overlay_opacity)
+                    except Exception as inference_error:
+                        st.error(f"Inference error: {inference_error}")
+                        st.write("Trying alternative approach...")
+                        
+                        # Alternative: resize to exactly 512x512 with crop_or_pad
+                        img_resized = tf.image.resize_with_crop_or_pad(
+                            original_img, IMAGE_SIZE, IMAGE_SIZE
+                        )
+                        img_normalized = normalize_image(img_resized)
+                        img_transposed = tf.transpose(img_normalized, (2, 0, 1))
+                        img_tensor = tf.expand_dims(img_transposed, axis=0)
+                        
+                        outputs = model(pixel_values=img_tensor, training=False)
+                        logits = outputs.logits
+                        mask = create_mask(logits)
+                        colorized_mask = colorize_mask(mask)
+                        overlay = create_overlay(original_img, colorized_mask, alpha=overlay_opacity)
+                except Exception as e:
+                    st.error(f"Failed to process image: {e}")
+                    st.stop()
             
             # Display results
             with col2: