Update app.py

app.py

@@ -2,37 +2,61 @@ import gradio as gr
 import torch
 from PIL import Image
 import numpy as np
+import subprocess
+import sys
+import io
+import matplotlib.pyplot as plt
+import traceback
+
+# --- 0. PATCH IN PLACE: Temporarily fix the mmseg version check ---
+# This code block forces the required mmcv-full to be installed correctly.
+
+try:
+    print("INFO: Attempting to install pre-built mmcv-full...")
+    # This installs the mmcv-full wheel pre-built for PyTorch 1.13, which includes the necessary _ext modules.
+    subprocess.check_call([
+        sys.executable, '-m', 'pip', 'install', 
+        'mmcv-full==1.7.1', 
+        '-f', 'https://download.openmmlab.com/mmcv/dist/cpu/torch1.13/index.html'
+    ])
+    print("INFO: Successfully installed pre-built mmcv-full.")
+except subprocess.CalledProcessError as e:
+    print(f"FATAL ERROR: Failed to install pre-built mmcv-full via subprocess. {e}")
+    # Exit if critical dependency fails to install
+    sys.exit(1)
 
 # --- 1. Load Custom Model Utilities ---
-# NOTE: These imports MUST match the files you copied from the GitHub repo.
-# Example imports - adjust these if the model files are deeper in subfolders!
-from mmseg.apis import init_segmentor, inference_segmentor # Core MMSeg functions
-from mmseg.datasets import build_dataloader, build_dataset # Utilities
+# These imports rely on the code being copied correctly and the mmcv patch working.
+try:
+    # These imports should now work because mmcv is installed correctly
+    from mmseg.apis import init_segmentor, inference_segmentor 
+except Exception as e:
+    print(f"FATAL ERROR: Failed to import mmseg utilities: {e}")
+    sys.exit(1)
 
 
 # --- 2. CONFIGURATION ---
-# Define the paths for the files you placed in the repository
+# Ensure these paths match your file names and structure
 WEIGHTS_PATH = "R50_ReLeM.pth"
-CONFIG_FILE = "configs/foodnet/SETR_Naive_768x768_80k_base_RM.py" # Replace with actual config file from the repo
+CONFIG_FILE = "configs/foodnet/SETR_Naive_768x768_80k_base_RM.py" 
+DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
 
 # --- 3. Model Loading Function ---
 @torch.no_grad()
 def load_relem_model():
     """Initializes the segmentation model and loads the pre-trained weights."""
     try:
-        # 1. Initialize the segmentor using MMSegmentation's utility
-        # This requires the config file and the checkpoint path
         model = init_segmentor(
-            CONFIG_FILE, 
-            checkpoint=WEIGHTS_PATH, 
-            device='cuda:0' if torch.cuda.is_available() else 'cpu'
+            CONFIG_FILE,             
+            checkpoint=WEIGHTS_PATH,             
+            device=DEVICE
         )
         model.eval()
-        print("ReLeM Model loaded successfully!")
+        print(f"ReLeM Model loaded successfully onto {DEVICE}!")
         return model
     except Exception as e:
-        print(f"Error loading model: {e}")
-        # Return a flag if loading fails
+        print(f"CRITICAL ERROR: Model failed to load weights or config: {e}")
+        traceback.print_exc()
         return None
 
 # Load the model once when the Space starts
@@ -40,61 +64,51 @@ RELEM_MODEL = load_relem_model()
 
 
 # --- 4. Inference Function for Gradio ---
-# --- 4. Inference Function for Gradio (REVISED) ---
-# --- 4. Inference Function for Gradio (ROBUST LOGGING) ---
 def segment_food(input_image: Image.Image):
-    """Takes a PIL Image and returns a segmentation mask image."""
-    
+    """Takes a PIL Image, runs inference, and returns a colorful segmentation mask."""
+        
     if RELEM_MODEL is None:
-        # If model failed to load at startup, this prints the error
-        print("RUNTIME ERROR: RELEM_MODEL is None, failing inference.")
-        return "Error: Model failed to load at startup. Check full build logs."
+        return "Error: Model failed to load at startup. Check build logs."
 
     try:
-        # Step 1: Save input image temporarily
+        # Step 1: Save input image temporarily (Required by mmseg's inference pipeline)
         temp_path = "/tmp/input_img.png"
         input_image.save(temp_path)
-        print(f"INFO: Saved input image to {temp_path}")
         
-        # Step 2: Run Inference (This is where the memory/config crash occurs)
+        # Step 2: Run Inference (Produces the raw class ID map)
+        # **This is the point where an OOM (Out of Memory) crash usually happens**
         result = inference_segmentor(RELEM_MODEL, temp_path)
-        print("INFO: Inference completed successfully.")
-
+        
         # Step 3: Post-process the result into a COLORFUL image
         seg_mask_array = result[0]
         
-        # --- ROBUST COLOR MASK CREATION ---
-        # We use matplotlib to create a visible, colored mask from the raw ID array
-        import matplotlib.pyplot as plt
-        
-        fig, ax = plt.subplots(figsize=(input_image.width / 100, input_image.height / 100)) # Sizing helps prevent memory spikes
-        ax.imshow(seg_mask_array, cmap='nipy_spectral') # Use a distinct color map
+        # --- MATPLOTLIB VISUALIZATION (Robust Color Mask) ---
+        # Create a new figure to plot the mask with distinct colors
+        fig, ax = plt.subplots(figsize=(8, 8)) # Use a moderate size
+        ax.imshow(seg_mask_array, cmap='nipy_spectral', interpolation='nearest') # Use colorful colormap
         ax.axis('off')
         
-        # Save the figure to a buffer (in memory)
-        import io
+        # Save the figure to an in-memory buffer
         buf = io.BytesIO()
         plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
-        plt.close(fig) # Close the figure immediately to free up memory
+        plt.close(fig) # Free memory
         buf.seek(0)
         
         # Return the saved image buffer as a PIL Image
-        print("INFO: Successfully created color mask.")
         return Image.open(buf)
     
     except Exception as e:
-        # **THIS CATCHES THE CRASH AND PRINTS IT TO THE LOGS**
+        # Catches memory errors or other runtime crashes
         print(f"RUNTIME CRASH: Inference failed with error: {e}")
-        import traceback
         traceback.print_exc()
-        return f"Inference failed at runtime: {e}. Check logs for traceback."
+        return f"Inference failed at runtime. Error: {e}. Try a smaller image."
 
 # --- 5. GRADIO INTERFACE ---
 gr.Interface(
     fn=segment_food,
     inputs=gr.Image(type="pil", label="Upload Food Image"),
     outputs=gr.Image(type="pil", label="ReLeM Segmentation Mask"),
-    title="ReLeM (FoodSeg103) Segmentation Demo",
-    description="Custom deployment of the ReLeM PyTorch model. **NOTE:** Model loading requires the full code/config structure from the GitHub repo.",
+    title="ReLeM (FoodSeg103) Deployment Final Attempt",
+    description="Custom deployment of the ReLeM PyTorch model. Check logs for deployment status.",
     allow_flagging="never"
 ).launch()