Update app.py

app.py

@@ -1,54 +1,50 @@
-from flask import Flask, request, jsonify,render_template, send_file
+from flask import Flask, request, jsonify, render_template, send_file
 from flask_cors import CORS
-from PIL import Image, ImageDraw
+from PIL import Image
 import io
 import json
 import os
 import uuid
 import google.generativeai as genai
-
+import cv2
+import numpy as np
 
 generation_config = {
-  "temperature": 1,
-  "max_output_tokens": 8192,
+    "temperature": 1,
+    "max_output_tokens": 8192,
 }
 
 safety_settings = [
-  {
-    "category": "HARM_CATEGORY_HARASSMENT",
-    "threshold": "BLOCK_NONE"
-  },
-  {
-    "category": "HARM_CATEGORY_HATE_SPEECH",
-    "threshold": "BLOCK_NONE"
-  },
-  {
-    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-    "threshold": "BLOCK_NONE"
-  },
-  {
-    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-    "threshold": "BLOCK_NONE"
-  },
+    {
+        "category": "HARM_CATEGORY_HARASSMENT",
+        "threshold": "BLOCK_NONE"
+    },
+    {
+        "category": "HARM_CATEGORY_HATE_SPEECH",
+        "threshold": "BLOCK_NONE"
+    },
+    {
+        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+        "threshold": "BLOCK_NONE"
+    },
+    {
+        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        "threshold": "BLOCK_NONE"
+    },
 ]
 
 GOOGLE_API_KEY = os.environ.get("TOKEN")
 
-
 # Configuration de l'API Gemini
-genai.configure(api_key=GOOGLE_API_KEY)  # Remplacez par votre clé API
+genai.configure(api_key=GOOGLE_API_KEY)
 
 app = Flask(__name__)
 CORS(app)
 
-
 @app.route('/', methods=['GET'])
 def svt():
     """Renders the SVT page."""
     return render_template("svt.html")
-    
-
-
 
 # Prompt pour la détection d'objets
 DETECTION_PROMPT = "Detect items, with no more than 20 items. Output a json list where each entry contains the 2D bounding box in \"box_2d\" and a text label in \"label\"."
@@ -84,48 +80,56 @@ def analyze_image():
             return jsonify({'error': 'No selected file'}), 400
 
         if file:
-            # Générer un nom de fichier unique pour éviter les conflits
             unique_filename = str(uuid.uuid4()) + os.path.splitext(file.filename)[1]
             filename = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
             file.save(filename)
 
-            # 1. Détection d'objets avec Gemini
-            model = genai.GenerativeModel("gemini-1.5-flash-exp",safety_settings=safety_settings,generation_config=generation_config)
+            # Charger l'image avec OpenCV
+            image = cv2.imread(filename)
+            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convertir en RGB pour l'API Gemini
+
+            # Préparer l'image pour l'API Gemini
             image_part = {
-              "mime_type": "image/jpeg",  # Assurez-vous que cela corresponde au type de votre image
-              "data": open(filename, "rb").read() # Le fichier est lu ici
+                "mime_type": "image/jpeg",
+                "data": cv2.imencode('.jpg', image_rgb)[1].tobytes()
             }
-            response = model.generate_content([DETECTION_PROMPT, image_part])
 
+            # Détection d'objets et dessin des boîtes (avec gestion d'erreur)
+            detection_results = []  # Initialiser une liste vide pour les résultats de détection
             try:
-                detection_results = json.loads(response.text)
-            except json.JSONDecodeError:
-                print(f"Erreur de décodage JSON : {response.text}")
-                return jsonify({'error': 'Erreur lors de la détection des objets (JSON invalide)'}), 500
-
-            # 2. Dessiner les boîtes englobantes
-            image = Image.open(filename)
-            draw = ImageDraw.Draw(image)
-
-            for item in detection_results:
-                box = item['box_2d']
-                label = item['label']
-                draw.rectangle(box, outline=(255, 0, 0), width=2)
-                text_position = (box[0], box[1] - 10)
-                draw.text(text_position, label, fill=(255, 255, 255))
-
-            # Enregistrer l'image avec les boîtes
+                model = genai.GenerativeModel("gemini-1.5-flash-exp", safety_settings=safety_settings, generation_config=generation_config)
+                response = model.generate_content([DETECTION_PROMPT, image_part])
+                cleaned_response_text = response.text.replace('\n', '')
+
+                if cleaned_response_text.startswith("```json"):
+                    cleaned_response_text = cleaned_response_text[7:]
+                if cleaned_response_text.endswith("```"):
+                    cleaned_response_text = cleaned_response_text[:-3]
+
+                detection_results = json.loads(cleaned_response_text)
+
+                # Dessiner les boîtes englobantes avec OpenCV
+                for item in detection_results:
+                    box = item['box_2d']
+                    label = str(item['label'])
+                    x0, y0, x1, y1 = box
+                    cv2.rectangle(image, (x0, y0), (x1, y1), (255, 0, 0), 2)
+                    cv2.putText(image, label, (x0, y0 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
+
+            except Exception as e:
+                print(f"Erreur lors de la détection d'objets ou du dessin des boîtes : {e}")
+
+            # Enregistrer l'image traitée
             output_filename = os.path.join(app.config['UPLOAD_FOLDER'], 'output_' + unique_filename)
-            image.save(output_filename)
+            cv2.imwrite(output_filename, image)
 
-            # 3. Générer la description
-            
+            # Générer la description
+            model = genai.GenerativeModel("gemini-1.5-flash-exp", safety_settings=safety_settings, generation_config=generation_config)
             response = model.generate_content([DESCRIPTION_PROMPT, image_part])
             description = response.text
 
-            # 4. Renvoyer les résultats
             return jsonify({
-                'image_path': '/uploads/' + 'output_' + unique_filename, # Chemin relatif vers l'image
+                'image_path': '/uploads/' + 'output_' + unique_filename,
                 'description': description,
                 'detected_objects': detection_results
             })
@@ -134,7 +138,6 @@ def analyze_image():
         print(f"Une erreur s'est produite : {e}")
         return jsonify({'error': f'Erreur lors du traitement de l\'image : {e}'}), 500
 
-# Servir les fichiers statiques depuis le dossier 'uploads'
 @app.route('/uploads/<filename>')
 def uploaded_file(filename):
     return send_file(os.path.join(app.config['UPLOAD_FOLDER'], filename))