mcp_server.py

import gradio as gr
from src.utils.change_format import change_format
from src.utils.remove_background import remove_background
from src.utils.generate_image import generate_image
from src.utils.apply_filter import apply_filter
from src.utils.add_text import add_text_to_image
from src.utils.watermark import add_watermark, remove_watermark
from src.utils.describe import describe_image
from src.utils.compress import compress_image
import base64
from PIL import Image
import io
import requests

def image_to_base64(image):
    if image is None:
        return None
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode()

def base64_to_image(base64_str):
    if not base64_str:
        return None
    
    # Remove data URI prefix if present (e.g., "data:image/png;base64,")
    if isinstance(base64_str, str) and "base64," in base64_str:
        base64_str = base64_str.split("base64,", 1)[1]
    
    try:
        # Strip any whitespace that might be in the base64 string
        if isinstance(base64_str, str):
            base64_str = base64_str.strip()
        
        # Decode the base64 data
        image_data = base64.b64decode(base64_str)
        
        # Check if we have data
        if not image_data:
            print("Decoded base64 data is empty")
            return None
        
        # Attempt to open the image
        image = Image.open(io.BytesIO(image_data))
        
        # Convert the image to ensure it's valid
        return image.copy()
        
    except base64.binascii.Error as e:
        print(f"Base64 decoding error: {str(e)}")
        if isinstance(base64_str, str):
            preview = base64_str[:30] + "..." if len(base64_str) > 30 else base64_str
            print(f"Base64 preview: {preview}")
        return None
        
    except Exception as e:
        print(f"Error converting base64 to image: {str(e)}")
        
        # Print preview of the base64 string for debugging
        if isinstance(base64_str, str):
            preview = base64_str[:30] + "..." if len(base64_str) > 30 else base64_str
            print(f"Base64 preview: {preview}")
        
        # Additional debug information
        if 'image_data' in locals() and image_data:
            try:
                magic_bytes = image_data[:12].hex()
                print(f"First 12 bytes: {magic_bytes}")
            except:
                pass
        
        return None
    
def url_to_base64(url):
    response = requests.get(url)
    return base64.b64encode(response.content).decode()

def gradio_remove_background(image):
    if image is None:
        return None
    base64_img = image_to_base64(image)
    result = remove_background(f"data:image/png;base64,{base64_img}")
    
    # Check if the result is directly a base64 string or has an image_data key
    if isinstance(result, str):
        return base64_to_image(result)
    elif isinstance(result, dict) and "image_data" in result:
        # If image_data contains a data URI prefix
        if isinstance(result["image_data"], str) and result["image_data"].startswith("data:"):
            # The response already contains the full data URI
            return base64_to_image(result["image_data"])
        else:
            # Try to process it as a regular base64 string
            try:
                return base64_to_image(result["image_data"])
            except Exception as e:
                print(f"Error processing image data: {e}")
                return None
    else:
        print(f"Unexpected response format from remove_background: {type(result)}")
        return None

def gradio_describe_image(image):
    if image is None:
        return "No image provided"
    base64_img = image_to_base64(image)
    return describe_image(base64_img)

def gradio_change_format(image, format_type):
    if image is None:
        return None
    base64_img = image_to_base64(image)
    result = change_format(base64_img, format_type)
    return base64_to_image(result)

def gradio_generate_image(prompt, width=512, height=512):
    result = generate_image(prompt, width, height)
    return base64_to_image(result["b64"])

def gradio_apply_filter(image, filter_type):
    if image is None:
        return None
    base64_img = image_to_base64(image)
    result = apply_filter(base64_img, filter_type)
    return base64_to_image(result)

def gradio_add_text(image, text, x=50, y=50, font_size=20, color="white"):
    if image is None:
        return None
    base64_img = image_to_base64(image)
    result = add_text_to_image(base64_img, text, x, y, font_size, color)
    return base64_to_image(result)

def gradio_add_watermark(image, watermark_text, opacity=0.5):
    if image is None:
        return None
    base64_img = image_to_base64(image)
    result = add_watermark(base64_img, watermark_text, opacity)
    return base64_to_image(result)

def gradio_remove_watermark(image):
    if image is None:
        return None
    base64_img = image_to_base64(image)
    result = remove_watermark(base64_img)
    return base64_to_image(result)

def gradio_compress_image(image, quality=80):
    if image is None:
        return None
    base64_img = image_to_base64(image)
    result = compress_image(base64_img, quality)
    return base64_to_image(result)

def create_gradio_interface():
    with gr.Blocks(title="Image Processing Service", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🖼️ Image Processing Service")
        gr.Markdown("Servicio completo de procesamiento de imágenes")
        
        with gr.Tabs():
            with gr.Tab("🎨 Generate Image"):
                with gr.Row():
                    prompt_input = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate")
                    with gr.Column():
                        width_input = gr.Slider(256, 1024, 512, label="Width")
                        height_input = gr.Slider(256, 1024, 512, label="Height")
                generate_btn = gr.Button("Generate", variant="primary")
                generated_output = gr.Image(label="Generated Image")
                
                generate_btn.click(
                    gradio_generate_image,
                    [prompt_input, width_input, height_input],
                    generated_output
                )
            
            with gr.Tab("🔍 Describe Image"):
                with gr.Row():
                    describe_input = gr.Image(label="Upload Image", type="pil")
                    description_output = gr.Textbox(label="Description", lines=4)
                
                describe_input.change(gradio_describe_image, describe_input, description_output)
            
            with gr.Tab("✂️ Remove Background"):
                with gr.Row():
                    bg_input = gr.Image(label="Upload Image", type="pil")
                    bg_output = gr.Image(label="Background Removed")
                
                bg_input.change(gradio_remove_background, bg_input, bg_output)
            
            with gr.Tab("🎭 Apply Filters"):
                with gr.Row():
                    filter_input = gr.Image(label="Upload Image", type="pil")
                    with gr.Column():
                        filter_type = gr.Dropdown(
                            ["blur", "sharpen", "vintage", "black_white", "sepia"],
                            label="Filter Type",
                            value="blur"
                        )
                        filter_output = gr.Image(label="Filtered Image")
                
                filter_input.change(gradio_apply_filter, [filter_input, filter_type], filter_output)
                filter_type.change(gradio_apply_filter, [filter_input, filter_type], filter_output)
            
            with gr.Tab("📝 Add Text"):
                with gr.Row():
                    text_input = gr.Image(label="Upload Image", type="pil")
                    with gr.Column():
                        text_content = gr.Textbox(label="Text", placeholder="Enter text to add")
                        with gr.Row():
                            text_x = gr.Number(label="X Position", value=50)
                            text_y = gr.Number(label="Y Position", value=50)
                        with gr.Row():
                            font_size = gr.Slider(10, 100, 20, label="Font Size")
                            text_color = gr.ColorPicker(label="Color", value="#FFFFFF")
                        text_output = gr.Image(label="Image with Text")
                
                inputs = [text_input, text_content, text_x, text_y, font_size, text_color]
                for inp in inputs:
                    inp.change(gradio_add_text, inputs, text_output)
            
            with gr.Tab("💧 Watermark"):
                with gr.Tabs():
                    with gr.Tab("Add Watermark"):
                        with gr.Row():
                            watermark_input = gr.Image(label="Upload Image", type="pil")
                            with gr.Column():
                                watermark_text = gr.Textbox(label="Watermark Text")
                                watermark_opacity = gr.Slider(0.1, 1.0, 0.5, label="Opacity")
                                watermark_output = gr.Image(label="Watermarked Image")
                        
                        inputs = [watermark_input, watermark_text, watermark_opacity]
                        for inp in inputs:
                            inp.change(gradio_add_watermark, inputs, watermark_output)
                    
                    with gr.Tab("Remove Watermark"):
                        with gr.Row():
                            unwatermark_input = gr.Image(label="Upload Image", type="pil")
                            unwatermark_output = gr.Image(label="Watermark Removed")
                        
                        unwatermark_input.change(gradio_remove_watermark, unwatermark_input, unwatermark_output)
            
            with gr.Tab("🗜️ Compress"):
                with gr.Row():
                    compress_input = gr.Image(label="Upload Image", type="pil")
                    with gr.Column():
                        quality_slider = gr.Slider(10, 100, 80, label="Quality %")
                        compress_output = gr.Image(label="Compressed Image")
                
                compress_input.change(gradio_compress_image, [compress_input, quality_slider], compress_output)
                quality_slider.change(gradio_compress_image, [compress_input, quality_slider], compress_output)
            
            with gr.Tab("🔄 Change Format"):
                with gr.Row():
                    format_input = gr.Image(label="Upload Image", type="pil")
                    with gr.Column():
                        format_type = gr.Dropdown(
                            ["PNG", "JPEG", "WEBP", "BMP"],
                            label="Output Format",
                            value="PNG"
                        )
                        format_output = gr.Image(label="Converted Image")
                
                format_input.change(gradio_change_format, [format_input, format_type], format_output)
                format_type.change(gradio_change_format, [format_input, format_type], format_output)
        
        gr.Markdown("---")
        gr.Markdown("💡 **Status**: Active | Procesamiento de imágenes en tiempo real")
    
    return demo

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860, 
        share=True,
        show_error=True
    )