app.py

import gradio as gr
import cv2
import numpy as np
import requests
from app.image_processing import ImageProcessor
from app.utils import ImageUtils
import logging
import os
from tempfile import NamedTemporaryFile
import json

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def remove_background_bria(image: np.ndarray) -> np.ndarray:
    """Remove background using Bria API and preserve transparency."""
    logger.info("Starting Bria background removal process...")

    # Validate API token early
    api_token = os.getenv('BRIA_API_TOKEN')
    if not api_token:
        logger.error("BRIA_API_TOKEN not found in environment variables")
        return image
    logger.info("Found BRIA_API_TOKEN in environment variables")

    # Create temporary file
    temp_file = None
    processed_temp_file = None
    try:
        # Log input image details
        logger.info(f"Input image shape: {image.shape}")
        logger.info(f"Input image dtype: {image.dtype}")

        # Save numpy array as temporary image file
        temp_file = NamedTemporaryFile(suffix='.png', delete=False)
        logger.info(f"Creating temporary file: {temp_file.name}")

        # Convert RGB to BGR for OpenCV
        bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        logger.info("Converted image from RGB to BGR")

        # Save image
        success = cv2.imwrite(temp_file.name, bgr_image)
        if not success:
            logger.error("Failed to save temporary image file")
            return image
        logger.info("Successfully saved temporary image file")

        # Verify file exists and get size
        file_size = os.path.getsize(temp_file.name)
        logger.info(f"Temporary file size: {file_size} bytes")

        # Prepare API request
        url = "https://engine.prod.bria-api.com/v1/background/remove"
        headers = {'api_token': api_token}
        files = [('file', ('image.png', open(temp_file.name, 'rb'), 'image/png'))]

        logger.info("Sending request to Bria API...")
        response = requests.post(url, headers=headers, files=files)

        # Log response details
        logger.info(f"API Response Status Code: {response.status_code}")
        logger.info(f"API Response Headers: {response.headers}")

        if response.status_code != 200:
            logger.error(f"Bria API error: {response.status_code}")
            logger.error(f"Response content: {response.text}")
            return image

        # Parse API response
        try:
            response_data = response.json()
            logger.info(f"API Response JSON: {response_data}")

            result_url = response_data.get('result_url')
            if not result_url:
                logger.error("No result_url in API response")
                return image

            logger.info(f"Downloading result from: {result_url}")
            
            # Download and process the image
            processed_temp_file = NamedTemporaryFile(suffix='.png', delete=False)
            with requests.get(result_url, stream=True) as r:
                r.raise_for_status()
                with open(processed_temp_file.name, 'wb') as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        f.write(chunk)
            
            # Read the processed image with transparency
            processed_img = cv2.imread(processed_temp_file.name, cv2.IMREAD_UNCHANGED)
            
            if processed_img is None:
                logger.error("Failed to load processed image")
                return image

            logger.info(f"Processed image shape: {processed_img.shape}")

            # Ensure the image has an alpha channel
            if len(processed_img.shape) == 3 and processed_img.shape[2] == 3:
                processed_img = cv2.cvtColor(processed_img, cv2.COLOR_BGR2BGRA)
                processed_img[:, :, 3] = 255  # Set full opacity

            # Convert from BGRA to RGBA for Gradio
            if processed_img.shape[2] == 4:
                processed_img = cv2.cvtColor(processed_img, cv2.COLOR_BGRA2RGBA)
                logger.info("Converted BGRA to RGBA")
            else:
                processed_img = cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB)
                logger.info("Converted BGR to RGB")

            return processed_img

        except Exception as e:
            logger.error(f"Error processing API response: {str(e)}")
            logger.exception("Full traceback:")
            return image

    except Exception as e:
        logger.error(f"Error in Bria background removal: {str(e)}")
        logger.exception("Full traceback:")
        return image
    
    finally:
        # Clean up temporary files
        try:
            if temp_file and os.path.exists(temp_file.name):
                os.unlink(temp_file.name)
            if processed_temp_file and os.path.exists(processed_temp_file.name):
                os.unlink(processed_temp_file.name)
        except Exception as e:
            logger.error(f"Error cleaning up temporary files: {str(e)}")


def cleanup_on_shutdown():
    """Cleanup function to be called on application shutdown"""
    logger.info("Cleaning up temporary files...")
    try:
        upload_dir, output_dir = ImageUtils.setup_directories()
        ImageUtils.cleanup_old_files(upload_dir)
        ImageUtils.cleanup_old_files(output_dir)
        logger.info("Temporary files cleaned up successfully")
    except Exception as e:
        logger.error(f"Error during cleanup: {str(e)}")

def process_image_pipeline(
    input_image: np.ndarray,
    bg_method: str,
    filter_type: str,
    bg_color_r: int,
    bg_color_g: int,
    bg_color_b: int,
    brightness: float,
    contrast: float,
    saturation: float,
    rotation: int,
    flip_horizontal: bool,
    flip_vertical: bool
) -> np.ndarray:
    """Process image with selected parameters"""
    try:
        # Validate input image
        if not ImageUtils.validate_image_array(input_image):
            logger.error("Invalid input image")
            return input_image
            
        # Resize image if needed
        input_image = ImageUtils.resize_image(input_image)
            
        # Create background color tuple
        bg_color = (bg_color_r, bg_color_g, bg_color_b, 255)
        
        try:
            # For advanced (Bria) method, apply background removal first
            if bg_method == "advanced":
                # Apply adjustments to input image
                adjusted_image = ImageProcessor.adjust_image(
                    input_image,
                    brightness=brightness,
                    contrast=contrast,
                    saturation=saturation,
                    rotation=rotation,
                    flip_horizontal=flip_horizontal,
                    flip_vertical=flip_vertical
                )
                
                # Remove background using Bria
                processed = remove_background_bria(adjusted_image)
                
                # Apply filter if needed
                if filter_type != "none":
                    processed = ImageProcessor.apply_image_filter(
                        processed,
                        filter_type=filter_type
                    )
                
                # Ensure we have the correct shape for RGBA
                if len(processed.shape) == 3 and processed.shape[2] == 4:
                    return processed
                else:
                    logger.error(f"Unexpected processed image shape: {processed.shape}")
                    return input_image
                    
            else:
                # For other methods, follow the original pipeline
                # Apply adjustments first
                input_image = ImageProcessor.adjust_image(
                    input_image,
                    brightness=brightness,
                    contrast=contrast,
                    saturation=saturation,
                    rotation=rotation,
                    flip_horizontal=flip_horizontal,
                    flip_vertical=flip_vertical
                )
                
                # Then remove background
                processed = ImageProcessor.remove_background(
                    input_image,
                    method=bg_method,
                    bg_color=bg_color
                )
                
                # Finally apply filter
                processed = ImageProcessor.apply_image_filter(
                    processed,
                    filter_type=filter_type
                )
                
                return processed
            
        except Exception as e:
            logger.error(f"Processing error: {str(e)}")
            logger.exception("Full traceback:")
            return input_image
            
    except Exception as e:
        logger.error(f"Pipeline error: {str(e)}")
        logger.exception("Full traceback:")
        return input_image

def create_gradio_interface():
    """Create and configure Gradio interface with support for RGBA images."""
    # Ensure storage directories exist
    upload_dir, output_dir = ImageUtils.setup_directories()
    
    # Define input components
    input_image = gr.Image(
        label="Upload Image",
        type="numpy",
        sources="upload"
    )
    
    # Add advanced method to background removal choices
    bg_methods = list(ImageProcessor.BG_REMOVAL_METHODS.keys())
    bg_methods.append("advanced")
    
    bg_method = gr.Dropdown(
        choices=bg_methods,
        value="basic",
        label="Background Removal Method",
        info="Choose method for removing image background. 'Advanced' uses Bria's RMBG 2.0 model."
    )
    
    # Define available filters from the image processor's filter dictionary
    filter_type = gr.Dropdown(
        choices=["none", "grayscale", "sepia", "blur", "sharpen", 
                "edge_detect", "emboss", "sketch", "watercolor", "invert"],
        value="none",
        label="Filter Type",
        info="Select a filter to apply to the image"
    )
    
    with gr.Row():
        bg_color_r = gr.Slider(
            minimum=0,
            maximum=255,
            value=173,
            step=1,
            label="Background Red",
            info="Red component of background color"
        )
        bg_color_g = gr.Slider(
            minimum=0,
            maximum=255,
            value=216,
            step=1,
            label="Background Green",
            info="Green component of background color"
        )
        bg_color_b = gr.Slider(
            minimum=0,
            maximum=255,
            value=230,
            step=1,
            label="Background Blue",
            info="Blue component of background color"
        )
    
    with gr.Row():
        brightness = gr.Slider(
            minimum=0.0,
            maximum=2.0,
            value=1.0,
            step=0.1,
            label="Brightness"
        )
        contrast = gr.Slider(
            minimum=0.0,
            maximum=2.0,
            value=1.0,
            step=0.1,
            label="Contrast"
        )
        saturation = gr.Slider(
            minimum=0.0,
            maximum=2.0,
            value=1.0,
            step=0.1,
            label="Saturation"
        )
    
    with gr.Row():
        rotation = gr.Slider(
            minimum=-180,
            maximum=180,
            value=0,
            step=90,
            label="Rotation"
        )
        flip_horizontal = gr.Checkbox(
            label="Flip Horizontal",
            value=False
        )
        flip_vertical = gr.Checkbox(
            label="Flip Vertical",
            value=False
        )

    # Define output with PNG format
    output_image = gr.Image(
        label="Processed Image",
        type="numpy",
        format="png",
        show_download_button=True,  # Enable download button
        height=500  # Set a reasonable display height
    )
    
    # Create interface
    iface = gr.Interface(
        fn=process_image_pipeline,
        inputs=[
            input_image,
            bg_method,
            filter_type,
            bg_color_r,
            bg_color_g,
            bg_color_b,
            brightness,
            contrast,
            saturation,
            rotation,
            flip_horizontal,
            flip_vertical
        ],
        outputs=output_image,
        title="Advanced Image Processing Tool",
        description="""
        Upload an image to remove its background and apply professional filters.
        
        Features:
        - Multiple background removal methods including Bria's RMBG 2.0 model
        - Various image filters
        - Custom background color selection
        - Automatic image validation and resizing
        - Support for high-resolution images
        - Download processed images
        
        Note: The 'Advanced' background removal method requires a valid Bria API token.
        """,
        cache_examples=True,
        theme=gr.themes.Soft()
    )
    
    return iface

if __name__ == "__main__":
    try:
        # Clean up old files before starting
        upload_dir, output_dir = ImageUtils.setup_directories()
        ImageUtils.cleanup_old_files(upload_dir)
        ImageUtils.cleanup_old_files(output_dir)
        
        # Create and launch the interface
        iface = create_gradio_interface()
        import atexit
        atexit.register(cleanup_on_shutdown)
        iface.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False
        )
    except Exception as e:
        logger.error(f"Application startup error: {e}")