Update image_processor.py

image_processor.py

@@ -1,40 +1,259 @@
-import cv2
-import numpy as np
+import gradio as gr
+import os
 from PIL import Image, ImageEnhance
-from typing import Optional, Dict, Any
+import requests
+import io
+import gc
+import json
+from typing import Tuple, Optional, Dict, Any
 import logging
+import numpy as np
+import cv2
+from dotenv import load_dotenv
 
-class ImageProcessor:
+# Configuration du logging
+logging.basicConfig(level=logging.DEBUG,
+                   format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Chargement des variables d'environnement
+load_dotenv()
+
+# Styles artistiques enrichis
+ART_STYLES = {
+    "Ultra Réaliste": {
+        "prompt_prefix": "ultra realistic photograph, stunning photorealistic quality, unreal engine 5 quality, cinema quality, masterpiece, perfect composition, award winning photography, professional lighting, 8k UHD",
+        "negative_prompt": "artificial, digital art, illustration, painting, drawing, artistic, cartoon, anime, unreal, fake, low quality, blurry, soft, deformed",
+        "quality_boost": 1.2
+    },
+    "Photoréaliste": {
+        "prompt_prefix": "hyperrealistic studio photograph, extremely detailed, professional photography, perfect lighting, high-end camera, 8k uhd",
+        "negative_prompt": "artificial, illustration, painting, animated, cartoon, artistic",
+        "quality_boost": 1.1
+    },
+    "Art Moderne": {
+        "prompt_prefix": "modern art style, professional design, contemporary aesthetic, trending artwork, perfect composition",
+        "negative_prompt": "old style, vintage, traditional, amateur, low quality",
+        "quality_boost": 1.0
+    },
+    "Minimaliste": {
+        "prompt_prefix": "minimalist design, clean composition, elegant simplicity, refined aesthetic",
+        "negative_prompt": "complex, cluttered, busy, ornate, detailed",
+        "quality_boost": 1.0
+    }
+}
+
+class ImageGenerator:
     def __init__(self):
-        self.cache = {}
-        logging.info("Image processor initialized")
+        self.API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
+        token = os.getenv('HUGGINGFACE_TOKEN')
+        if not token:
+            logger.error("HUGGINGFACE_TOKEN non trouvé!")
+        self.headers = {"Authorization": f"Bearer {token}"}
+        logger.info("ImageGenerator initialisé")
 
-    def enhance_image(self, image: Image.Image, params: Dict[str, Any]) -> Image.Image:
-        """Amélioration optimisée de la qualité d'image"""
+    def _enhance_image(self, image: Image.Image, params: Dict[str, Any]) -> Image.Image:
+        """Amélioration avancée de la qualité d'image"""
         try:
-            # Paramètres de base
-            sharpness = params.get('detail_level', 7) / 5
-            contrast = params.get('contrast', 5) / 5
-            saturation = params.get('saturation', 5) / 5
-
-            # Application des améliorations
-            # 1. Netteté
-            if sharpness != 1:
-                enhancer = ImageEnhance.Sharpness(image)
-                image = enhancer.enhance(sharpness)
-
-            # 2. Contraste
-            if contrast != 1:
-                enhancer = ImageEnhance.Contrast(image)
-                image = enhancer.enhance(contrast)
-
-            # 3. Saturation
-            if saturation != 1:
-                enhancer = ImageEnhance.Color(image)
-                image = enhancer.enhance(saturation)
-
+            # Conversion en CV2 pour traitement avancé
+            cv2_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+            
+            # Débruitage adaptatif
+            if params.get("quality", 35) > 40:
+                cv2_image = cv2.fastNlMeansDenoisingColored(cv2_image, None, 10, 10, 7, 21)
+            
+            # Amélioration des détails
+            kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
+            cv2_image = cv2.filter2D(cv2_image, -1, kernel)
+            
+            # Reconversion en PIL
+            image = Image.fromarray(cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB))
+            
+            # Ajustements fins avec PIL
+            enhancers = [
+                (ImageEnhance.Sharpness, params.get("detail_level", 7) / 5),
+                (ImageEnhance.Contrast, params.get("contrast", 5) / 5),
+                (ImageEnhance.Color, params.get("saturation", 5) / 5)
+            ]
+            
+            for enhancer_class, factor in enhancers:
+                if factor != 1.0:
+                    image = enhancer_class(image).enhance(factor)
+            
             return image
+        except Exception as e:
+            logger.error(f"Erreur traitement image: {str(e)}")
+            return image
+
+    def _build_prompt(self, params: Dict[str, Any]) -> str:
+        """Construction optimisée du prompt"""
+        try:
+            style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
+            
+            # Construction du prompt principal
+            base_prompt = f"{params['subject']}"
+            if params.get('title'):
+                base_prompt += f", with text '{params['title']}'"
+
+            # Ajout des éléments de style
+            prompt = f"{base_prompt}, {style_info['prompt_prefix']}"
+            
+            return prompt
+        except Exception as e:
+            logger.error(f"Erreur prompt: {str(e)}")
+            return params['subject']
+
+    def generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
+        try:
+            if 'Bearer None' in self.headers['Authorization']:
+                return None, "⚠️ Erreur: Token Hugging Face non configuré"
+
+            # Optimisation des paramètres
+            style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
+            quality_boost = style_info.get("quality_boost", 1.0)
+
+            # Préparation de la requête
+            prompt = self._build_prompt(params)
+            payload = {
+                "inputs": prompt,
+                "parameters": {
+                    "negative_prompt": style_info["negative_prompt"],
+                    "num_inference_steps": min(int(40 * quality_boost), 50),
+                    "guidance_scale": min(8.0 * quality_boost, 12.0),
+                    "width": 1024 if params.get("quality", 35) > 40 else 768,
+                    "height": 1024 if params["orientation"] == "Portrait" else 768
+                }
+            }
+
+            response = requests.post(
+                self.API_URL,
+                headers=self.headers,
+                json=payload,
+                timeout=45
+            )
+
+            if response.status_code == 200:
+                image = Image.open(io.BytesIO(response.content))
+                # Application des améliorations de qualité
+                enhanced_image = self._enhance_image(image, params)
+                return enhanced_image, "✨ Création réussie!"
+            else:
+                error_msg = f"⚠️ Erreur API {response.status_code}: {response.text}"
+                logger.error(error_msg)
+                return None, error_msg
 
         except Exception as e:
-            logging.error(f"Erreur lors de l'amélioration de l'image: {str(e)}")
-            return image
+            error_msg = f"⚠️ Erreur: {str(e)}"
+            logger.exception("Erreur génération:")
+            return None, error_msg
+        finally:
+            gc.collect()
+
+def create_interface():
+    generator = ImageGenerator()
+
+    with gr.Blocks(css="style.css") as app:
+        gr.HTML("""
+            <div class="welcome">
+                <h1>🎨 Equity Artisan 3.0</h1>
+                <p>Assistant de création d'affiches professionnelles</p>
+            </div>
+        """)
+
+        with gr.Column():
+            # Contrôles principaux
+            with gr.Group():
+                gr.Markdown("### 📐 Format et Style")
+                with gr.Row():
+                    format_size = gr.Dropdown(
+                        choices=["A4", "A3", "A2", "A1"],
+                        value="A4",
+                        label="Format"
+                    )
+                    orientation = gr.Radio(
+                        choices=["Portrait", "Paysage"],
+                        value="Portrait",
+                        label="Orientation"
+                    )
+                    style = gr.Dropdown(
+                        choices=list(ART_STYLES.keys()),
+                        value="Art Moderne",
+                        label="Style artistique"
+                    )
+
+            # Description
+            with gr.Group():
+                gr.Markdown("### 📝 Description")
+                subject = gr.Textbox(
+                    label="Description",
+                    placeholder="Décrivez votre vision...",
+                    lines=3
+                )
+                title = gr.Textbox(
+                    label="Titre (optionnel)",
+                    placeholder="Titre à inclure..."
+                )
+
+            # Paramètres avancés
+            with gr.Group():
+                gr.Markdown("### ⚙️ Paramètres")
+                with gr.Row():
+                    quality = gr.Slider(
+                        minimum=30,
+                        maximum=50,
+                        value=35,
+                        label="Qualité"
+                    )
+                    detail_level = gr.Slider(
+                        minimum=1,
+                        maximum=10,
+                        value=7,
+                        step=1,
+                        label="Niveau de Détail"
+                    )
+                    creativity = gr.Slider(
+                        minimum=5,
+                        maximum=15,
+                        value=7.5,
+                        label="Créativité"
+                    )
+
+            # Boutons
+            with gr.Row():
+                generate_btn = gr.Button("✨ Générer", variant="primary")
+                clear_btn = gr.Button("🗑️ Effacer")
+
+            # Résultat
+            image_output = gr.Image(label="Résultat")
+            status = gr.Textbox(label="Status", interactive=False)
+
+        def generate(*args):
+            params = {
+                "format_size": args[0],
+                "orientation": args[1],
+                "style": args[2],
+                "subject": args[3],
+                "title": args[4],
+                "quality": args[5],
+                "detail_level": args[6],
+                "creativity": args[7]
+            }
+            return generator.generate(params)
+
+        generate_btn.click(
+            generate,
+            inputs=[format_size, orientation, style, subject, title,
+                   quality, detail_level, creativity],
+            outputs=[image_output, status]
+        )
+
+        clear_btn.click(
+            lambda: (None, "🗑️ Image effacée"),
+            outputs=[image_output, status]
+        )
+
+    return app
+
+if __name__ == "__main__":
+    app = create_interface()
+    app.launch()