app.py

import gradio as gr
from PIL import Image
import os
import tempfile
import pyheif
from transformers import pipeline

def open_heic_image(image_path):
    heif_file = pyheif.read(image_path)
    img = Image.frombytes(
        heif_file.mode, 
        heif_file.size, 
        heif_file.data,
        "raw",
        heif_file.mode,
        heif_file.stride,
    )
    return img

def optimize_image(image, png_optimize, jpeg_quality, jpeg_resolution, webp_quality):
    # Manejar archivos HEIC
    if image.name.lower().endswith(".heic"):
        img = open_heic_image(image.name)
    else:
        img = Image.open(image)

    # Convertir la imagen a RGB si está en RGBA
    if img.mode == 'RGBA':
        img = img.convert('RGB')
    
    original_size = os.path.getsize(image.name) / 1024  # tamaño en KB
    output_dir = "/tmp/optimized_images"
    os.makedirs(output_dir, exist_ok=True)
    
    results = []
    
    # Optimización para PNG
    if png_optimize:
        lossless_output_path = os.path.join(output_dir, "lossless.png")
        img.save(lossless_output_path, format="PNG", optimize=True)
        lossless_size = os.path.getsize(lossless_output_path) / 1024
        results.append(("PNG", lossless_output_path, lossless_size, original_size))

    # Optimización para JPEG
    if jpeg_quality != 50:
        lossy_output_path = os.path.join(output_dir, "lossy.jpg")
        img.save(lossy_output_path, format="JPEG", quality=jpeg_quality, optimize=True)
        lossy_size = os.path.getsize(lossy_output_path) / 1024
        results.append(("JPEG", lossy_output_path, lossy_size, original_size))

    # Reducción de resolución (JPEG)
    if jpeg_resolution != 50:
        reduced_output_path = os.path.join(output_dir, "reduced_resolution.jpg")
        new_resolution = (img.width * jpeg_resolution // 100, img.height * jpeg_resolution // 100)
        reduced_img = img.resize(new_resolution, Image.LANCZOS)
        reduced_img.save(reduced_output_path, format="JPEG", quality=jpeg_quality, optimize=True)
        reduced_size = os.path.getsize(reduced_output_path) / 1024
        results.append(("JPEG (resolución reducida)", reduced_output_path, reduced_size, original_size))

    # Optimización para WebP
    if webp_quality != 50:
        webp_lossy_output_path = os.path.join(output_dir, "lossy.webp")
        img.save(webp_lossy_output_path, format="WEBP", quality=webp_quality, optimize=True)
        webp_lossy_size = os.path.getsize(webp_lossy_output_path) / 1024
        results.append(("WebP", webp_lossy_output_path, webp_lossy_size, original_size))
    
    return results

# Función para aplicar un modelo seleccionado desde Hugging Face
def apply_model(image, model_name):
    model_pipeline = pipeline("image-super-resolution", model=model_name)
    return model_pipeline(image)

with gr.Blocks() as demo:
    with gr.Tab("Optimización Tradicional"):
        # Aceptar imágenes y archivos .heic
        image_input = gr.File(label="Sube tu imagen", file_types=['image', '.heic'])
        optimize_button = gr.Button("Optimizar")

        with gr.Row():
            with gr.Column():
                optimized_output1 = gr.Image(label="Optimización sin pérdida")
                png_optimize = gr.Checkbox(label="Optimizar PNG", value=True)
                download_button1 = gr.File(label="Descargar", visible=True)
                optimized_size1 = gr.Text(value="", interactive=False, show_label=False)

            with gr.Column():
                optimized_output2 = gr.Image(label="Optimización con pérdida (JPEG)")
                jpeg_quality = gr.Slider(label="Calidad JPEG", minimum=10, maximum=100, value=50, step=1)
                download_button2 = gr.File(label="Descargar", visible=True)
                optimized_size2 = gr.Text(value="", interactive=False, show_label=False)
                
            with gr.Column():
                optimized_output3 = gr.Image(label="Reducción de resolución (JPEG)")
                jpeg_resolution = gr.Slider(label="Resolución JPEG (%)", minimum=10, maximum=100, value=50, step=1)
                download_button3 = gr.File(label="Descargar", visible=True)
                optimized_size3 = gr.Text(value="", interactive=False, show_label=False)

            with gr.Column():
                optimized_output4 = gr.Image(label="Optimización WebP con pérdida")
                webp_quality = gr.Slider(label="Calidad WebP", minimum=10, maximum=100, value=50, step=1)
                download_button4 = gr.File(label="Descargar", visible=True)
                optimized_size4 = gr.Text(value="", interactive=False, show_label=False)

        # Conectar cada control a la función de optimización
        optimize_button.click(
            fn=optimize_image,
            inputs=[image_input, png_optimize, jpeg_quality, jpeg_resolution, webp_quality],
            outputs=[
                optimized_output1, optimized_size1, download_button1,
                optimized_output2, optimized_size2, download_button2,
                optimized_output3, optimized_size3, download_button3,
                optimized_output4, optimized_size4, download_button4
            ]
        )

    with gr.Tab("Optimización con Modelos de Hugging Face"):
        hf_image_input = gr.File(label="Sube tu imagen para optimización avanzada", file_types=['image', '.heic'])
        model_selector = gr.Dropdown(
            label="Selecciona un modelo",
            choices=["xinntao/Real-ESRGAN", "google/ddpm-cifar10-32", "facebook/ddpm"],  # Añade los modelos disponibles
            value="xinntao/Real-ESRGAN"
        )
        hf_output = gr.Image(label="Resultado")
        hf_button = gr.Button("Aplicar Modelo")

        hf_button.click(
            fn=apply_model,
            inputs=[hf_image_input, model_selector],
            outputs=hf_output
        )

demo.launch()