Update app.py

app.py

@@ -45,43 +45,53 @@ 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."""
     
     if RELEM_MODEL is None:
-        # If model failed to load, return the error message
-        return "Error: Model failed to load. Check logs for details."
+        # 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."
 
     try:
         # Step 1: Save input image temporarily
         temp_path = "/tmp/input_img.png"
         input_image.save(temp_path)
+        print(f"INFO: Saved input image to {temp_path}")
         
-        # Step 2: Run Inference (Produces the class ID map)
+        # Step 2: Run Inference (This is where the memory/config crash occurs)
         result = inference_segmentor(RELEM_MODEL, temp_path)
-        
-        # Step 3: Use the MMSegmentation utility to visualize the result
-        # NOTE: This requires the 'show_result_pyplot' utility and the model's palette.
-        # Since we don't have the utility, we will use a simplified NumPy approach that is less prone to memory errors:
+        print("INFO: Inference completed successfully.")
 
+        # Step 3: Post-process the result into a COLORFUL image
         seg_mask_array = result[0]
         
-        # --- SIMPLE NUMPY VISUALIZATION (Requires model to define a palette) ---
-        # MMSegmentation models usually contain a palette definition.
-        # If the model has a PALETTE property, use it. Otherwise, we convert to a colored PIL Image.
+        # --- 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
+        ax.axis('off')
+        
+        # Save the figure to a buffer (in memory)
+        import io
+        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
+        buf.seek(0)
         
-        # Assuming RELEM_MODEL has a palette defined (common in MMSegmentation models)
-        if hasattr(RELEM_MODEL, 'PALETTE'):
-            # This is a placeholder as the PALETTE logic is complex.
-            # We'll just return the raw NumPy array which Gradio might handle better:
-            return Image.fromarray(seg_mask_array.astype(np.uint8))
-        else:
-            # Fallback to the simple (but likely black/grayscale) conversion
-            return Image.fromarray(seg_mask_array.astype(np.uint8)).convert("L")
+        # Return the saved image buffer as a PIL Image
+        print("INFO: Successfully created color mask.")
+        return Image.open(buf)
     
     except Exception as e:
-        # This will catch any error during inference and tell you what went wrong.
-        return f"Inference failed at runtime: {e}"
+        # **THIS CATCHES THE CRASH AND PRINTS IT TO THE LOGS**
+        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."
 
 # --- 5. GRADIO INTERFACE ---
 gr.Interface(