Update app.py

app.py

@@ -44,34 +44,44 @@ RELEM_MODEL = load_relem_model()
 
 
 # --- 4. Inference Function for Gradio ---
+# --- 4. Inference Function for Gradio (REVISED) ---
 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."
 
     try:
-        # Use MMSegmentation's inference pipeline 
-        # The input is usually a filepath, so we need to save and then load
-        
-        # 1. Save input image temporarily
+        # Step 1: Save input image temporarily
         temp_path = "/tmp/input_img.png"
         input_image.save(temp_path)
         
-        # 2. Run Inference
+        # Step 2: Run Inference (Produces the class ID map)
         result = inference_segmentor(RELEM_MODEL, temp_path)
         
-        # 3. Post-process the result (usually a numpy array) into a color mask image
-        # The result is a segmentation map (array of class IDs).
-        # We use a simple utility to convert the ID map to a visible color mask.
-        seg_mask_array = result[0] 
-        color_mask = Image.fromarray(seg_mask_array.astype(np.uint8)).convert("L") 
-        # NOTE: Full color mapping requires the class labels/palette, which you must also copy from the repo.
+        # 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:
+
+        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.
         
-        return color_mask 
+        # 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")
     
     except Exception as e:
-        return f"Inference failed: {e}"
+        # This will catch any error during inference and tell you what went wrong.
+        return f"Inference failed at runtime: {e}"
 
 # --- 5. GRADIO INTERFACE ---
 gr.Interface(