Update app.py

app.py

@@ -2,56 +2,69 @@ 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):
-    img = Image.open(image)
+    # 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
-    
-    # Crear un directorio seguro para almacenar las imágenes optimizadas
     output_dir = "/tmp/optimized_images"
     os.makedirs(output_dir, exist_ok=True)
     
-    # Generar archivos temporales en un directorio seguro
-    lossless_output_path = os.path.join(output_dir, "lossless.png")
-    lossy_output_path = os.path.join(output_dir, "lossy.jpg")
-    reduced_output_path = os.path.join(output_dir, "reduced_resolution.jpg")
-    webp_lossy_output_path = os.path.join(output_dir, "lossy.webp")
-
-    # 1. Compresión sin pérdida (PNG)
-    img.save(lossless_output_path, format="PNG", optimize=png_optimize)
-    lossless_size = os.path.getsize(lossless_output_path) / 1024
-    lossless_diff = original_size - lossless_size
-    lossless_img = Image.open(lossless_output_path)
+    results = []
     
-    # 2. Compresión con pérdida (JPEG)
-    img.save(lossy_output_path, format="JPEG", quality=jpeg_quality, optimize=True)
-    lossy_size = os.path.getsize(lossy_output_path) / 1024
-    lossy_diff = original_size - lossy_size
-    lossy_img = Image.open(lossy_output_path)
-    
-    # 3. Reducción de resolución (JPEG)
-    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
-    reduced_diff = original_size - reduced_size
-    reduced_img = Image.open(reduced_output_path)
+    # 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))
 
-    # 4. Compresión con WebP (con pérdida)
-    img.save(webp_lossy_output_path, format="WEBP", quality=webp_quality, optimize=True)
-    webp_lossy_size = os.path.getsize(webp_lossy_output_path) / 1024
-    webp_lossy_diff = original_size - webp_lossy_size
-    webp_lossy_img = Image.open(webp_lossy_output_path)
+    # 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))
     
-    # Retornar las imágenes optimizadas, tamaños y rutas para descarga
-    return [
-        lossless_img, f"Sin pérdida: {lossless_size:.2f} KB \n(diferencia: {-lossless_diff:.2f} KB)", lossless_output_path,
-        lossy_img, f"Con pérdida: {lossy_size:.2f} KB \n(diferencia: {-lossy_diff:.2f} KB)", lossy_output_path,
-        reduced_img, f"Reducción de resolución: {reduced_size:.2f} KB \n(diferencia: {-reduced_diff:.2f} KB)", reduced_output_path,
-        webp_lossy_img, f"WebP con pérdida: {webp_lossy_size:.2f} KB \n(diferencia: {-webp_lossy_diff:.2f} KB)", webp_lossy_output_path
-    ]
+    return results
 
 # Función para aplicar un modelo seleccionado desde Hugging Face
 def apply_model(image, model_name):
@@ -60,7 +73,7 @@ def apply_model(image, model_name):
 
 with gr.Blocks() as demo:
     with gr.Tab("Optimización Tradicional"):
-        image_input = gr.File(label="Sube tu imagen", file_types=['image'])
+        image_input = gr.File(label="Sube tu imagen", file_types=['image', 'heic'])
         optimize_button = gr.Button("Optimizar")
 
         with gr.Row():
@@ -101,7 +114,7 @@ with gr.Blocks() as demo:
         )
 
     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'])
+        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