Update app.py

app.py

@@ -8,12 +8,17 @@ 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.
+# --- CRITICAL PATCH: Fix for 'container_abcs' not found in torch._six ---
+# This makes older code compatible with PyTorch 1.13.1 by providing the correct import.
+try:
+    from torch._six import container_abcs
+except ImportError:
+    import collections.abc as container_abcs
 
+# --- 0. FORCE INSTALL: Install pre-built mmcv-full for _ext modules ---
 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.
+    # This installs the mmcv-full wheel pre-built for PyTorch 1.13, which includes the necessary compiled _ext modules.
     subprocess.check_call([
         sys.executable, '-m', 'pip', 'install', 
         'mmcv-full==1.7.1', 
@@ -22,21 +27,21 @@ try:
     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)
+    # We allow the code to continue execution but the model will likely fail to load later
+    pass # Continue execution, but model will likely fail to load
+
 
-# --- 1. Load Custom Model Utilities ---
-# These imports rely on the code being copied correctly and the mmcv patch working.
+# --- 1. Load Custom Model Utilities (Must come after mmcv is installed) ---
 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)
+    # Returning None here will trigger the "Error: Model failed to load" message in the app.
+    init_segmentor = None
+    inference_segmentor = None
 
 
 # --- 2. CONFIGURATION ---
-# 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" 
 DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
@@ -45,6 +50,9 @@ DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
 @torch.no_grad()
 def load_relem_model():
     """Initializes the segmentation model and loads the pre-trained weights."""
+    if init_segmentor is None:
+        return None # Skip if imports failed
+
     try:
         model = init_segmentor(
             CONFIG_FILE,             
@@ -59,7 +67,6 @@ def load_relem_model():
         traceback.print_exc()
         return None
 
-# Load the model once when the Space starts
 RELEM_MODEL = load_relem_model()
 
 
@@ -68,7 +75,7 @@ def segment_food(input_image: Image.Image):
     """Takes a PIL Image, runs inference, and returns a colorful segmentation mask."""
         
     if RELEM_MODEL is None:
-        return "Error: Model failed to load at startup. Check build logs."
+        return "Error: Model failed to load at startup. Check build logs for reason."
 
     try:
         # Step 1: Save input image temporarily (Required by mmseg's inference pipeline)
@@ -76,29 +83,25 @@ def segment_food(input_image: Image.Image):
         input_image.save(temp_path)
         
         # 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)
         
         # Step 3: Post-process the result into a COLORFUL image
         seg_mask_array = result[0]
         
         # --- 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
+        fig, ax = plt.subplots(figsize=(8, 8))
+        ax.imshow(seg_mask_array, cmap='nipy_spectral', interpolation='nearest')
         ax.axis('off')
         
         # 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) # Free memory
+        plt.close(fig)
         buf.seek(0)
         
-        # Return the saved image buffer as a PIL Image
         return Image.open(buf)
     
     except Exception as e:
-        # Catches memory errors or other runtime crashes
         print(f"RUNTIME CRASH: Inference failed with error: {e}")
         traceback.print_exc()
         return f"Inference failed at runtime. Error: {e}. Try a smaller image."