Enhance image handling and cleanup in background processing

- Added logic to retain references to original images for proper cleanup after processing in cli.py and core.py.
- Implemented image closing to ensure file handles are released, particularly on Windows, preventing potential resource leaks.
- Updated test cases to utilize context managers for image opening, ensuring proper closure and resource management across various test files.

Author: kocabiyik

packages/python/src/withoutbg/cli.py

@@ -172,6 +172,7 @@ def progress_callback(progress: float) -> None:
 
     # Save result
     save_kwargs: dict[str, Any] = {}
+    original_result = result  # Keep reference to original for cleanup
     if format.lower() == "jpg":
         # Convert RGBA to RGB for JPEG
         if result.mode == "RGBA":
@@ -191,6 +192,12 @@ def progress_callback(progress: float) -> None:
         format=pil_format.get(format.lower(), format.upper()),
         **save_kwargs,
     )
+    
+    # Close the result image to ensure file handles are released on Windows
+    result.close()
+    # Also close the original if we created a new background
+    if result is not original_result:
+        original_result.close()
 
 
 def _process_batch(
@@ -244,6 +251,7 @@ def _process_batch(
 
                 # Save result
                 save_kwargs: dict[str, Any] = {}
+                original_result = result  # Keep reference to original for cleanup
                 if format.lower() == "jpg":
                     if result.mode == "RGBA":
                         background = Image.new("RGB", result.size, (255, 255, 255))
@@ -267,6 +275,12 @@ def _process_batch(
                     format=pil_format.get(format.lower(), format.upper()),
                     **save_kwargs,
                 )
+                
+                # Close the result image to ensure file handles are released on Windows
+                result.close()
+                # Also close the original if we created a new background
+                if result is not original_result:
+                    original_result.close()
 
             except Exception as e:
                 if verbose:

packages/python/src/withoutbg/core.py

@@ -194,6 +194,7 @@ def remove_background_batch(
 
             if output_path:
                 result.save(output_path)
+                # Note: Keep result in memory for return, don't close it yet
 
             results.append(result)
 
@@ -295,6 +296,7 @@ def remove_background_batch(
 
             if output_path:
                 result.save(output_path)
+                # Note: Keep result in memory for return, don't close it yet
 
             results.append(result)
 

packages/python/tests/conftest.py

@@ -74,8 +74,19 @@ def expected_outputs_dir(test_data_dir):
 @pytest.fixture
 def temp_dir():
     """Create a temporary directory for test outputs."""
+    import gc
+    import platform
+    import shutil
+    import time
+
     with tempfile.TemporaryDirectory() as temp_dir:
         yield Path(temp_dir)
+        
+        # Windows-specific cleanup: force garbage collection and wait briefly
+        # to ensure PIL releases file handles before cleanup
+        if platform.system() == "Windows":
+            gc.collect()
+            time.sleep(0.1)  # Small delay to allow file handles to be released
 
 
 @pytest.fixture

packages/python/tests/integration/test_real_image_processing.py

@@ -140,9 +140,9 @@ def test_ice_cream_processing(
         # Check against expected output with IoU metric
         expected_path = expected_outputs_dir / "test-ice-cream.png"
         if expected_path.exists():
-            expected = Image.open(expected_path)
-            # Use IoU for interpretable alpha channel comparison
-            assert_alpha_iou(result, expected, min_iou=0.97)
+            with Image.open(expected_path) as expected:
+                # Use IoU for interpretable alpha channel comparison
+                assert_alpha_iou(result, expected, min_iou=0.97)
 
     def test_core_api_with_real_model(self, test_images_dir):
         """Test core API WithoutBG class with real processing."""
@@ -154,17 +154,16 @@ def test_core_api_with_real_model(self, test_images_dir):
         assert_alpha_channel_valid(result)
 
         # Result should preserve original dimensions
-        original = Image.open(input_path)
-        assert result.size == original.size
+        with Image.open(input_path) as original:
+            assert result.size == original.size
 
     def test_different_input_formats(self, real_opensource_model, test_images_dir):
         """Test processing with different input formats."""
         input_path = test_images_dir / "test-ice-cream.png"
-        original_image = Image.open(input_path)
-
-        # Test with PIL Image
-        result_pil = real_opensource_model.remove_background(original_image)
-        assert_alpha_channel_valid(result_pil)
+        with Image.open(input_path) as original_image:
+            # Test with PIL Image
+            result_pil = real_opensource_model.remove_background(original_image)
+            assert_alpha_channel_valid(result_pil)
 
         # Test with file path (string)
         result_path = real_opensource_model.remove_background(str(input_path))
@@ -261,9 +260,8 @@ def test_processing_consistency(self, real_opensource_model, test_images_dir):
     def test_pipeline_stages_integration(self, real_opensource_model, test_images_dir):
         """Test that the 3-stage pipeline produces reasonable outputs."""
         input_path = test_images_dir / "test-ice-cream.png"
-        original = Image.open(input_path)
-
-        result = real_opensource_model.remove_background(original)
+        with Image.open(input_path) as original:
+            result = real_opensource_model.remove_background(original)
 
         # Validate the result has expected characteristics for ice cream image
         result_array = np.array(result)
@@ -291,12 +289,11 @@ def test_alpha_iou_metric(
         result = real_opensource_model.remove_background(input_path)
 
         if expected_path.exists():
-            expected = Image.open(expected_path)
-
-            # Calculate IoU (should be very high since expected was generated
-            # with same model)
-            iou = calculate_alpha_iou(result, expected)
-            print(f"Alpha IoU score: {iou:.4f}")
+            with Image.open(expected_path) as expected:
+                # Calculate IoU (should be very high since expected was generated
+                # with same model)
+                iou = calculate_alpha_iou(result, expected)
+                print(f"Alpha IoU score: {iou:.4f}")
 
             # Should be perfect since we generated expected output with same model
             assert iou >= 0.99, f"IoU should be near-perfect: {iou:.4f}"

packages/python/tests/performance/test_batch_performance.py

@@ -40,12 +40,13 @@ def _mock_remove_bg(input_image, **kwargs):
 
         # Load image to get size
         if isinstance(input_image, str):
-            img = Image.open(input_image)
+            with Image.open(input_image) as img:
+                size = img.size
         else:
-            img = input_image
+            size = input_image.size
 
         # Return mock result
-        return Image.new("RGBA", img.size, color=(100, 150, 200, 128))
+        return Image.new("RGBA", size, color=(100, 150, 200, 128))
 
     return _mock_remove_bg
 
@@ -280,8 +281,9 @@ def mock_processing_with_errors(input_image, **kwargs):
 
             # Simulate normal processing for valid files
             time.sleep(0.01)
-            img = Image.open(input_image)
-            return Image.new("RGBA", img.size, color=(100, 150, 200, 128))
+            with Image.open(input_image) as img:
+                size = img.size
+            return Image.new("RGBA", size, color=(100, 150, 200, 128))
 
         try:
             model = WithoutBG.opensource()

packages/python/tests/test_api.py

@@ -39,12 +39,16 @@ def mock_alpha_path(self):
     @pytest.fixture
     def test_image(self, test_image_path):
         """Load test image."""
-        return Image.open(test_image_path)
+        img = Image.open(test_image_path)
+        yield img
+        img.close()
 
     @pytest.fixture
     def mock_alpha_image(self, mock_alpha_path):
         """Load mock alpha channel image."""
-        return Image.open(mock_alpha_path)
+        img = Image.open(mock_alpha_path)
+        yield img
+        img.close()
 
     def _create_mock_alpha_response(self, mock_alpha_image):
         """Create mock API response using real alpha channel image."""

packages/python/tests/test_cli.py

@@ -304,12 +304,12 @@ def test_output_format_options(
             assert output_path.exists()
 
             # Verify the file was saved with correct PIL format
-            saved_image = Image.open(output_path)
-            expected_pil_formats = {"png": "PNG", "jpg": "JPEG", "webp": "WEBP"}
-            expected_format = expected_pil_formats[fmt]
-            assert (
-                saved_image.format == expected_format
-            ), f"Expected {expected_format}, got {saved_image.format}"
+            with Image.open(output_path) as saved_image:
+                expected_pil_formats = {"png": "PNG", "jpg": "JPEG", "webp": "WEBP"}
+                expected_format = expected_pil_formats[fmt]
+                assert (
+                    saved_image.format == expected_format
+                ), f"Expected {expected_format}, got {saved_image.format}"
 
     @patch("src.withoutbg.cli.WithoutBG")
     def test_jpeg_quality_setting(
@@ -604,8 +604,8 @@ def test_rgba_to_jpg_conversion(self, mock_withoutbg_class, temp_dir):
         assert output_path.exists()
 
         # Verify the saved image is RGB (not RGBA)
-        saved_image = Image.open(output_path)
-        assert saved_image.mode == "RGB"
+        with Image.open(output_path) as saved_image:
+            assert saved_image.mode == "RGB"
 
     def test_directory_as_input_without_batch_flag(self, temp_dir):
         """Test providing directory as input without --batch flag."""
@@ -650,9 +650,9 @@ def create_test_image(path, size=(100, 100), color=(255, 0, 0), mode="RGB"):
         assert image_path.exists()
 
         # Verify image properties
-        image = Image.open(image_path)
-        assert image.size == (100, 100)
-        assert image.mode == "RGB"
+        with Image.open(image_path) as image:
+            assert image.size == (100, 100)
+            assert image.mode == "RGB"
 
     def test_create_corrupted_file(self, temp_dir):
         """Test utility for creating corrupted image files."""
@@ -679,30 +679,29 @@ def verify_output_properties(
             """Verify properties of output image file."""
             assert output_path.exists(), f"Output file {output_path} does not exist"
 
-            image = Image.open(output_path)
-
-            if expected_format:
-                assert image.format.lower() == expected_format.lower()
-            if expected_mode:
-                assert image.mode == expected_mode
-            if expected_size:
-                assert image.size == expected_size
-
-            return image
+            with Image.open(output_path) as image:
+                if expected_format:
+                    assert image.format.lower() == expected_format.lower()
+                if expected_mode:
+                    assert image.mode == expected_mode
+                if expected_size:
+                    assert image.size == expected_size
+                
+                return image.size  # Return size instead of image object
 
         # Create test output file
         test_image = Image.new("RGBA", (200, 150), color=(255, 0, 0, 128))
         output_path = temp_dir / "output.png"
         test_image.save(output_path)
 
         # Test verification
-        verified_image = verify_output_properties(
+        verified_size = verify_output_properties(
             output_path,
             expected_format="PNG",
             expected_mode="RGBA",
             expected_size=(200, 150),
         )
-        assert verified_image.size == (200, 150)
+        assert verified_size == (200, 150)
 
 
 class TestCLIPerformance: