Fixed OpenMP Fork Errors and OpenCV Resize Error

gradio_app.py

@@ -4,6 +4,12 @@ os.environ["OMP_NUM_THREADS"] = "1"  # Limit OpenMP to single thread
 os.environ["PYTHONFAULTHANDLER"] = "1"  # Better error reporting
 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"  # Avoid Intel MKL errors
 os.environ["OMP_THREAD_LIMIT"] = "1"  # Limit total number of OpenMP threads
+os.environ["MKL_NUM_THREADS"] = "1"  # Intel MKL threading
+os.environ["NUMEXPR_NUM_THREADS"] = "1"  # NumExpr threading
+os.environ["OPENBLAS_NUM_THREADS"] = "1"  # OpenBLAS threading
+os.environ["VECLIB_MAXIMUM_THREADS"] = "1"  # Apple vecLib threading
+os.environ["OMP_DISPLAY_ENV"] = "FALSE"  # Suppress OpenMP environment display
+os.environ["KMP_WARNINGS"] = "FALSE"  # Suppress KMP warnings
 
 # Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
 # except for the third-party components listed below.

hy3dshape/hy3dshape/preprocessors.py

@@ -88,17 +88,54 @@ class ImageProcessorV2:
         return result, mask
 
     def load_image(self, image, border_ratio=0.15, to_tensor=True):
+        # Handle different input types
         if isinstance(image, str):
+            # Load from file path
             image = cv2.imread(image, cv2.IMREAD_UNCHANGED)
+            if image is None:
+                raise ValueError(f"Could not load image from path: {image}")
             image, mask = self.recenter(image, border_ratio=border_ratio)
             image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         elif isinstance(image, Image.Image):
+            # Convert PIL Image to numpy array
             image = image.convert("RGBA")
             image = np.asarray(image)
             image, mask = self.recenter(image, border_ratio=border_ratio)
-
-        image = cv2.resize(image, (self.size, self.size), interpolation=cv2.INTER_CUBIC)
-        mask = cv2.resize(mask, (self.size, self.size), interpolation=cv2.INTER_NEAREST)
+        elif isinstance(image, np.ndarray):
+            # Handle numpy array input
+            if image.size == 0:
+                raise ValueError("Input image array is empty")
+            image, mask = self.recenter(image, border_ratio=border_ratio)
+        else:
+            # Handle any other type by trying to convert to numpy array
+            try:
+                if hasattr(image, 'convert'):
+                    # Assume it's a PIL-like image
+                    image = image.convert("RGBA")
+                    image = np.asarray(image)
+                else:
+                    # Try direct conversion to numpy array
+                    image = np.asarray(image)
+                
+                if image.size == 0:
+                    raise ValueError("Converted image array is empty")
+                image, mask = self.recenter(image, border_ratio=border_ratio)
+            except Exception as e:
+                raise ValueError(f"Could not process image input of type {type(image)}: {e}")
+
+        # Validate that we have valid arrays before resizing
+        if not isinstance(image, np.ndarray) or image.size == 0:
+            raise ValueError("Image processing failed - invalid image array")
+        if not isinstance(mask, np.ndarray) or mask.size == 0:
+            raise ValueError("Image processing failed - invalid mask array")
+
+        # Resize with error handling
+        try:
+            image = cv2.resize(image, (self.size, self.size), interpolation=cv2.INTER_CUBIC)
+            mask = cv2.resize(mask, (self.size, self.size), interpolation=cv2.INTER_NEAREST)
+        except cv2.error as e:
+            raise ValueError(f"OpenCV resize failed: {e}. Image shape: {image.shape if hasattr(image, 'shape') else 'unknown'}, Mask shape: {mask.shape if hasattr(mask, 'shape') else 'unknown'}")
+        
         mask = mask[..., np.newaxis]
 
         if to_tensor:

hy3dshape/hy3dshape/rembg.py

@@ -73,40 +73,78 @@ class BackgroundRemover():
     def __init__(self):
         # Set multiprocessing start method to 'spawn' to avoid OpenMP fork issues
         # Only set if not already set
-        if multiprocessing.get_start_method(allow_none=True) != 'spawn':
+        current_method = multiprocessing.get_start_method(allow_none=True)
+        if current_method != 'spawn':
             try:
                 multiprocessing.set_start_method('spawn', force=True)
-            except RuntimeError:
+                print(f"Set multiprocessing start method to 'spawn' (was: {current_method})")
+            except RuntimeError as e:
                 # If already set and not 'spawn', we'll have to work with what we have
-                print("Warning: Could not set multiprocessing start method to 'spawn'")
+                print(f"Warning: Could not set multiprocessing start method to 'spawn': {e}")
+                print(f"Current method: {multiprocessing.get_start_method()}")
         
         # We don't initialize the session here anymore as it will be created in the subprocess
         self.session = None
+        self.max_retries = 2  # Maximum number of retries for background removal
 
     def __call__(self, image: Image.Image):
+        if not isinstance(image, Image.Image):
+            print(f"Warning: Expected PIL Image, got {type(image)}. Attempting conversion...")
+            try:
+                if hasattr(image, 'convert'):
+                    image = image.convert('RGBA')
+                else:
+                    # Try to create PIL Image from input
+                    image = Image.fromarray(np.asarray(image))
+            except Exception as e:
+                print(f"Could not convert input to PIL Image: {e}")
+                return image
+        
         try:
+            # Ensure image is in RGBA mode
+            if image.mode != 'RGBA':
+                image = image.convert('RGBA')
+            
             # Convert image to bytes for passing to subprocess
             img_bytes = io.BytesIO()
             image.save(img_bytes, format='PNG')
             img_bytes = img_bytes.getvalue()
             
             # Create a process context
-            ctx = multiprocessing.get_context('spawn')
+            try:
+                ctx = multiprocessing.get_context('spawn')
+            except RuntimeError:
+                # Fallback to default context if spawn is not available
+                ctx = multiprocessing
             
             # Try with alpha matting first (high quality)
-            with ctx.Pool(processes=1) as pool:
-                result = pool.apply(_remove_bg_in_process, args=(img_bytes, True))
-            
-            if result is not None:
-                return Image.open(io.BytesIO(result))
+            for attempt in range(self.max_retries):
+                try:
+                    with ctx.Pool(processes=1) as pool:
+                        result = pool.apply(_remove_bg_in_process, args=(img_bytes, True))
+                    
+                    if result is not None:
+                        return Image.open(io.BytesIO(result))
+                    break  # Exit retry loop if we got a result (even if None)
+                except Exception as e:
+                    print(f"Alpha matting attempt {attempt + 1} failed: {e}")
+                    if attempt == self.max_retries - 1:
+                        break
             
             # If alpha matting fails, try without it
             print("Alpha matting failed, falling back to standard background removal")
-            with ctx.Pool(processes=1) as pool:
-                result = pool.apply(_remove_bg_in_process, args=(img_bytes, False))
-            
-            if result is not None:
-                return Image.open(io.BytesIO(result))
+            for attempt in range(self.max_retries):
+                try:
+                    with ctx.Pool(processes=1) as pool:
+                        result = pool.apply(_remove_bg_in_process, args=(img_bytes, False))
+                    
+                    if result is not None:
+                        return Image.open(io.BytesIO(result))
+                    break  # Exit retry loop if we got a result (even if None)
+                except Exception as e:
+                    print(f"Standard background removal attempt {attempt + 1} failed: {e}")
+                    if attempt == self.max_retries - 1:
+                        break
             
             # If all else fails, return the original image
             print("Background removal completely failed, returning original image")