Optimized Docker build to fix OOM errors

app.py

@@ -11,12 +11,16 @@ import os
 from pathlib import Path
 from typing import Dict, List
 
+# Ensure OMP_NUM_THREADS is a valid integer to avoid libgomp warnings
+if not os.environ.get("OMP_NUM_THREADS", "").isdigit():
+    os.environ["OMP_NUM_THREADS"] = "1"
+
 import cv2
 import gradio as gr
 import numpy as np
 import torch
 
-from estimator import generate_naive_mask
+from masks import generate_naive_mask
 
 _slimsam_model = None
 _slimsam_processor = None

estimator.py

@@ -13,6 +13,7 @@ import numpy as np
 import torch
 import cv2
 
+from masks import generate_naive_mask
 logger = logging.getLogger(__name__)
 
 # Add FoundationPose to Python path
@@ -34,37 +35,6 @@ except ImportError as e:
     trimesh = None
 
 
-def generate_naive_mask(
-    rgb_image: np.ndarray,
-    min_percentage: float = 1.0,
-    max_percentage: float = 90.0
-) -> tuple[np.ndarray, np.ndarray, float, bool]:
-    """Generate a naive foreground mask using brightness + Otsu thresholding.
-
-    Returns:
-        mask_bool: Boolean mask (H, W)
-        debug_mask: uint8 mask for visualization (H, W)
-        mask_percentage: % of pixels active in mask_bool
-        fallback_full_image: True if the mask was replaced by full-image mask
-    """
-    gray = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2GRAY)
-    _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-
-    kernel = np.ones((5, 5), np.uint8)
-    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # Fill holes
-    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)   # Remove noise
-
-    debug_mask = mask.copy()
-    mask_bool = mask.astype(bool)
-    mask_percentage = (mask_bool.sum() / mask_bool.size) * 100
-
-    fallback_full_image = False
-    if mask_percentage < min_percentage or mask_percentage > max_percentage:
-        fallback_full_image = True
-        mask_bool = np.ones((rgb_image.shape[0], rgb_image.shape[1]), dtype=bool)
-        debug_mask = np.ones((rgb_image.shape[0], rgb_image.shape[1]), dtype=np.uint8) * 255
-
-    return mask_bool, debug_mask, mask_percentage, fallback_full_image
 
 
 class FoundationPoseEstimator: