Equityone commited on
Commit
14caaf8
·
verified ·
1 Parent(s): f9a4c33

Update image_processor.py

Browse files
Files changed (1) hide show
  1. image_processor.py +250 -31
image_processor.py CHANGED
@@ -1,40 +1,259 @@
1
- import cv2
2
- import numpy as np
3
  from PIL import Image, ImageEnhance
4
- from typing import Optional, Dict, Any
 
 
 
 
5
  import logging
 
 
 
6
 
7
- class ImageProcessor:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8
  def __init__(self):
9
- self.cache = {}
10
- logging.info("Image processor initialized")
 
 
 
 
11
 
12
- def enhance_image(self, image: Image.Image, params: Dict[str, Any]) -> Image.Image:
13
- """Amélioration optimisée de la qualité d'image"""
14
  try:
15
- # Paramètres de base
16
- sharpness = params.get('detail_level', 7) / 5
17
- contrast = params.get('contrast', 5) / 5
18
- saturation = params.get('saturation', 5) / 5
19
-
20
- # Application des améliorations
21
- # 1. Netteté
22
- if sharpness != 1:
23
- enhancer = ImageEnhance.Sharpness(image)
24
- image = enhancer.enhance(sharpness)
25
-
26
- # 2. Contraste
27
- if contrast != 1:
28
- enhancer = ImageEnhance.Contrast(image)
29
- image = enhancer.enhance(contrast)
30
-
31
- # 3. Saturation
32
- if saturation != 1:
33
- enhancer = ImageEnhance.Color(image)
34
- image = enhancer.enhance(saturation)
35
-
 
 
 
 
36
  return image
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
37
 
38
  except Exception as e:
39
- logging.error(f"Erreur lors de l'amélioration de l'image: {str(e)}")
40
- return image
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import os
3
  from PIL import Image, ImageEnhance
4
+ import requests
5
+ import io
6
+ import gc
7
+ import json
8
+ from typing import Tuple, Optional, Dict, Any
9
  import logging
10
+ import numpy as np
11
+ import cv2
12
+ from dotenv import load_dotenv
13
 
14
+ # Configuration du logging
15
+ logging.basicConfig(level=logging.DEBUG,
16
+ format='%(asctime)s - %(levelname)s - %(message)s')
17
+ logger = logging.getLogger(__name__)
18
+
19
+ # Chargement des variables d'environnement
20
+ load_dotenv()
21
+
22
+ # Styles artistiques enrichis
23
+ ART_STYLES = {
24
+ "Ultra Réaliste": {
25
+ "prompt_prefix": "ultra realistic photograph, stunning photorealistic quality, unreal engine 5 quality, cinema quality, masterpiece, perfect composition, award winning photography, professional lighting, 8k UHD",
26
+ "negative_prompt": "artificial, digital art, illustration, painting, drawing, artistic, cartoon, anime, unreal, fake, low quality, blurry, soft, deformed",
27
+ "quality_boost": 1.2
28
+ },
29
+ "Photoréaliste": {
30
+ "prompt_prefix": "hyperrealistic studio photograph, extremely detailed, professional photography, perfect lighting, high-end camera, 8k uhd",
31
+ "negative_prompt": "artificial, illustration, painting, animated, cartoon, artistic",
32
+ "quality_boost": 1.1
33
+ },
34
+ "Art Moderne": {
35
+ "prompt_prefix": "modern art style, professional design, contemporary aesthetic, trending artwork, perfect composition",
36
+ "negative_prompt": "old style, vintage, traditional, amateur, low quality",
37
+ "quality_boost": 1.0
38
+ },
39
+ "Minimaliste": {
40
+ "prompt_prefix": "minimalist design, clean composition, elegant simplicity, refined aesthetic",
41
+ "negative_prompt": "complex, cluttered, busy, ornate, detailed",
42
+ "quality_boost": 1.0
43
+ }
44
+ }
45
+
46
+ class ImageGenerator:
47
  def __init__(self):
48
+ self.API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
49
+ token = os.getenv('HUGGINGFACE_TOKEN')
50
+ if not token:
51
+ logger.error("HUGGINGFACE_TOKEN non trouvé!")
52
+ self.headers = {"Authorization": f"Bearer {token}"}
53
+ logger.info("ImageGenerator initialisé")
54
 
55
+ def _enhance_image(self, image: Image.Image, params: Dict[str, Any]) -> Image.Image:
56
+ """Amélioration avancée de la qualité d'image"""
57
  try:
58
+ # Conversion en CV2 pour traitement avancé
59
+ cv2_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
60
+
61
+ # Débruitage adaptatif
62
+ if params.get("quality", 35) > 40:
63
+ cv2_image = cv2.fastNlMeansDenoisingColored(cv2_image, None, 10, 10, 7, 21)
64
+
65
+ # Amélioration des détails
66
+ kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
67
+ cv2_image = cv2.filter2D(cv2_image, -1, kernel)
68
+
69
+ # Reconversion en PIL
70
+ image = Image.fromarray(cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB))
71
+
72
+ # Ajustements fins avec PIL
73
+ enhancers = [
74
+ (ImageEnhance.Sharpness, params.get("detail_level", 7) / 5),
75
+ (ImageEnhance.Contrast, params.get("contrast", 5) / 5),
76
+ (ImageEnhance.Color, params.get("saturation", 5) / 5)
77
+ ]
78
+
79
+ for enhancer_class, factor in enhancers:
80
+ if factor != 1.0:
81
+ image = enhancer_class(image).enhance(factor)
82
+
83
  return image
84
+ except Exception as e:
85
+ logger.error(f"Erreur traitement image: {str(e)}")
86
+ return image
87
+
88
+ def _build_prompt(self, params: Dict[str, Any]) -> str:
89
+ """Construction optimisée du prompt"""
90
+ try:
91
+ style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
92
+
93
+ # Construction du prompt principal
94
+ base_prompt = f"{params['subject']}"
95
+ if params.get('title'):
96
+ base_prompt += f", with text '{params['title']}'"
97
+
98
+ # Ajout des éléments de style
99
+ prompt = f"{base_prompt}, {style_info['prompt_prefix']}"
100
+
101
+ return prompt
102
+ except Exception as e:
103
+ logger.error(f"Erreur prompt: {str(e)}")
104
+ return params['subject']
105
+
106
+ def generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
107
+ try:
108
+ if 'Bearer None' in self.headers['Authorization']:
109
+ return None, "⚠️ Erreur: Token Hugging Face non configuré"
110
+
111
+ # Optimisation des paramètres
112
+ style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
113
+ quality_boost = style_info.get("quality_boost", 1.0)
114
+
115
+ # Préparation de la requête
116
+ prompt = self._build_prompt(params)
117
+ payload = {
118
+ "inputs": prompt,
119
+ "parameters": {
120
+ "negative_prompt": style_info["negative_prompt"],
121
+ "num_inference_steps": min(int(40 * quality_boost), 50),
122
+ "guidance_scale": min(8.0 * quality_boost, 12.0),
123
+ "width": 1024 if params.get("quality", 35) > 40 else 768,
124
+ "height": 1024 if params["orientation"] == "Portrait" else 768
125
+ }
126
+ }
127
+
128
+ response = requests.post(
129
+ self.API_URL,
130
+ headers=self.headers,
131
+ json=payload,
132
+ timeout=45
133
+ )
134
+
135
+ if response.status_code == 200:
136
+ image = Image.open(io.BytesIO(response.content))
137
+ # Application des améliorations de qualité
138
+ enhanced_image = self._enhance_image(image, params)
139
+ return enhanced_image, "✨ Création réussie!"
140
+ else:
141
+ error_msg = f"⚠️ Erreur API {response.status_code}: {response.text}"
142
+ logger.error(error_msg)
143
+ return None, error_msg
144
 
145
  except Exception as e:
146
+ error_msg = f"⚠️ Erreur: {str(e)}"
147
+ logger.exception("Erreur génération:")
148
+ return None, error_msg
149
+ finally:
150
+ gc.collect()
151
+
152
+ def create_interface():
153
+ generator = ImageGenerator()
154
+
155
+ with gr.Blocks(css="style.css") as app:
156
+ gr.HTML("""
157
+ <div class="welcome">
158
+ <h1>🎨 Equity Artisan 3.0</h1>
159
+ <p>Assistant de création d'affiches professionnelles</p>
160
+ </div>
161
+ """)
162
+
163
+ with gr.Column():
164
+ # Contrôles principaux
165
+ with gr.Group():
166
+ gr.Markdown("### 📐 Format et Style")
167
+ with gr.Row():
168
+ format_size = gr.Dropdown(
169
+ choices=["A4", "A3", "A2", "A1"],
170
+ value="A4",
171
+ label="Format"
172
+ )
173
+ orientation = gr.Radio(
174
+ choices=["Portrait", "Paysage"],
175
+ value="Portrait",
176
+ label="Orientation"
177
+ )
178
+ style = gr.Dropdown(
179
+ choices=list(ART_STYLES.keys()),
180
+ value="Art Moderne",
181
+ label="Style artistique"
182
+ )
183
+
184
+ # Description
185
+ with gr.Group():
186
+ gr.Markdown("### 📝 Description")
187
+ subject = gr.Textbox(
188
+ label="Description",
189
+ placeholder="Décrivez votre vision...",
190
+ lines=3
191
+ )
192
+ title = gr.Textbox(
193
+ label="Titre (optionnel)",
194
+ placeholder="Titre à inclure..."
195
+ )
196
+
197
+ # Paramètres avancés
198
+ with gr.Group():
199
+ gr.Markdown("### ⚙️ Paramètres")
200
+ with gr.Row():
201
+ quality = gr.Slider(
202
+ minimum=30,
203
+ maximum=50,
204
+ value=35,
205
+ label="Qualité"
206
+ )
207
+ detail_level = gr.Slider(
208
+ minimum=1,
209
+ maximum=10,
210
+ value=7,
211
+ step=1,
212
+ label="Niveau de Détail"
213
+ )
214
+ creativity = gr.Slider(
215
+ minimum=5,
216
+ maximum=15,
217
+ value=7.5,
218
+ label="Créativité"
219
+ )
220
+
221
+ # Boutons
222
+ with gr.Row():
223
+ generate_btn = gr.Button("✨ Générer", variant="primary")
224
+ clear_btn = gr.Button("🗑️ Effacer")
225
+
226
+ # Résultat
227
+ image_output = gr.Image(label="Résultat")
228
+ status = gr.Textbox(label="Status", interactive=False)
229
+
230
+ def generate(*args):
231
+ params = {
232
+ "format_size": args[0],
233
+ "orientation": args[1],
234
+ "style": args[2],
235
+ "subject": args[3],
236
+ "title": args[4],
237
+ "quality": args[5],
238
+ "detail_level": args[6],
239
+ "creativity": args[7]
240
+ }
241
+ return generator.generate(params)
242
+
243
+ generate_btn.click(
244
+ generate,
245
+ inputs=[format_size, orientation, style, subject, title,
246
+ quality, detail_level, creativity],
247
+ outputs=[image_output, status]
248
+ )
249
+
250
+ clear_btn.click(
251
+ lambda: (None, "🗑️ Image effacée"),
252
+ outputs=[image_output, status]
253
+ )
254
+
255
+ return app
256
+
257
+ if __name__ == "__main__":
258
+ app = create_interface()
259
+ app.launch()