hammaster/cutoutai

1import os2import io3import base644import numpy as np5from PIL import Image, ImageDraw6import cutoutai7import logging8 9# Setup logging10logging.basicConfig(level=logging.INFO)11logger = logging.getLogger("TestCutoutAI")12 13def create_test_image(path="test_input.png"):14    """Create a synthetic test image with bubbles and a central object."""15    # 512x512 white background16    img = Image.new("RGB", (512, 512), (240, 240, 240))17    draw = ImageDraw.Draw(img)18 19    # Draw a "subject" (blue circle)20    draw.ellipse([150, 150, 362, 362], fill=(0, 0, 255), outline=(0, 0, 0))21 22    # Draw "bubbles" (small circles)23    draw.ellipse([50, 50, 80, 80], fill=(200, 200, 255, 128), outline=(100, 100, 100))24    draw.ellipse([400, 100, 430, 130], fill=(200, 200, 255, 128), outline=(100, 100, 100))25    draw.ellipse([100, 400, 140, 440], fill=(255, 200, 200, 128), outline=(100, 100, 100))26 27    # Draw some "fine detail" (thin lines)28    draw.line([256, 0, 256, 150], fill=(0, 0, 0), width=1)29 30    img.save(path)31    logger.info(f"Created test image: {path}")32    return path33 34def test_processing():35    """Test the core processing logic."""36    input_path = create_test_image()37 38    # Use 'lite' variant for faster testing if possible,39    # but the prompt asks for BiRefNet quality analysis.40    # Note: Loading the model will take time and requires internet + torch.41    # If we are in a restricted environment, this might fail.42 43    try:44        processor = cutoutai.CutoutAI(model_variant="lite") # Using lite for faster test45 46        logger.info("Running process()...")47        result = processor.process(48            input_path,49            capture_all_elements=True,50            edge_refinement=True,51            edge_radius=2,52            output_format="pil"53        )54 55        output_path = "test_output.png"56        result.save(output_path)57        logger.info(f"Saved result to: {output_path}")58 59        # Check if output is RGBA60        if result.mode == "RGBA":61            logger.info("SUCCESS: Output is in RGBA mode.")62        else:63            logger.error(f"FAILURE: Output mode is {result.mode}, expected RGBA.")64 65    except Exception as e:66        logger.error(f"Error during processing: {e}")67        logger.info("Note: This test requires torch and transformers to be installed and working.")68 69if __name__ == "__main__":70    test_processing()71