Update app.py

app.py

@@ -1,90 +1,95 @@
-from flask import Flask, render_template, request, jsonify
-import google.generativeai as genai
 import os
-import tempfile
+import io
 import base64
-import logging
-from dotenv import load_dotenv
+import tempfile
+from flask import Flask, render_template, request, jsonify
+from google import genai
+from google.genai import types
+from PIL import Image
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
+# Configure Gemini API key using an environment variable
+genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
-# Load environment variables
-load_dotenv()
+# Initialize Gemini client
+client = genai.Client()
 
+# Initialize Flask app
 app = Flask(__name__)
-genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
-@app.route("/")
-def home():
-    return render_template("index.html")
+def save_image(image_data):
+    """Save the image from a base64 string to a temporary file and return its path."""
+    # image_data is expected to be in the format "data:image/png;base64,...."
+    header, encoded = image_data.split(',', 1)
+    image_bytes = base64.b64decode(encoded)
+    image = Image.open(io.BytesIO(image_bytes))
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    image.save(temp_file, "PNG")
+    return temp_file.name
 
-@app.route("/process", methods=["POST"])
-def process_image():
-    try:
-        data = request.json
-        image_data = data.get("image")
-        object_type = data.get("objectType", "").strip()
+def remove_object_from_image(image_path, object_type):
+    """Use Gemini API to remove a specified object from the image."""
+    # Upload the image file to Gemini
+    uploaded_file = client.files.upload(file=image_path)
+    
+    # Prepare the input parts:
+    # 1. The image file
+    parts = [types.Part.from_uri(file_uri=uploaded_file.uri, mime_type="image/png")]
+    # 2. The Gemini magic text instructing removal
+    if object_type:
+        parts.append(types.Part.from_text(text=f"Remove {object_type} from the image"))
+    
+    contents = [types.Content(role="user", parts=parts)]
+    generate_content_config = types.GenerateContentConfig(
+        temperature=1,
+        top_p=0.95,
+        top_k=40,
+        max_output_tokens=8192,
+        response_modalities=["image", "text"],
+        safety_settings=[types.SafetySetting(category="HARM_CATEGORY_CIVIC_INTEGRITY", threshold="OFF")],
+        response_mime_type="text/plain",
+    )
 
-        # Validate inputs
-        if not image_data or not object_type:
-            return jsonify({"success": False, "message": "Missing required parameters"})
+    result_image = None
 
-        # Decode image
-        try:
-            header, encoded = image_data.split(",", 1)
-            image_bytes = base64.b64decode(encoded)
-        except Exception as e:
-            logging.error(f"Image decoding failed: {str(e)}")
-            return jsonify({"success": False, "message": "Invalid image data"})
+    # Process the generation stream from Gemini
+    for chunk in client.models.generate_content_stream(
+        model="gemini-2.0-flash-exp-image-generation",
+        contents=contents,
+        config=generate_content_config,
+    ):
+        if chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts:
+            part = chunk.candidates[0].content.parts[0]
+            if part.inline_data:
+                # Save the generated binary image data
+                file_name = "generated_output.png"
+                with open(file_name, "wb") as f:
+                    f.write(part.inline_data.data)
+                result_image = file_name
 
-        # Temporary files
-        with tempfile.TemporaryDirectory() as temp_dir:
-            input_path = os.path.join(temp_dir, "input.png")
-            with open(input_path, "wb") as f:
-                f.write(image_bytes)
+    return result_image
 
-            # Configure model with safety settings
-            model = genai.GenerativeModel('gemini-2.0-flash-exp-image-generation')
-            prompt = f"Remove {object_type} from image"
-            
-            response = model.generate_content(
-                [prompt, genai.upload_file(input_path)],
-                generation_config={
-                    "temperature": 0.9,
-                    "top_p": 0.95,
-                    "top_k": 32,
-                    "max_output_tokens": 4096,
-                },
-                safety_settings={
-                    "HARM_CATEGORY_CIVIC_INTEGRITY": "BLOCK_NONE",
-                    "HARM_CATEGORY_HARASSMENT": "BLOCK_NONE",
-                    "HARM_CATEGORY_HATE_SPEECH": "BLOCK_NONE",
-                    "HARM_CATEGORY_SEXUALLY_EXPLICIT": "BLOCK_NONE",
-                    "HARM_CATEGORY_DANGEROUS_CONTENT": "BLOCK_NONE"
-                }
-            )
+@app.route('/')
+def index():
+    """Render the main page."""
+    return render_template('index.html')
 
-            # Process response
-            output_path = os.path.join(temp_dir, "result.png")
-            for chunk in response:
-                if chunk.candidates:
-                    for part in chunk.candidates[0].content.parts:
-                        if hasattr(part, 'inline_data'):
-                            with open(output_path, "wb") as f:
-                                f.write(part.inline_data.data)
-                            return jsonify({
-                                "success": True,
-                                "resultPath": output_path
-                            })
-                        elif hasattr(part, 'text'):
-                            logging.info(f"Text response: {part.text}")
+@app.route('/process', methods=['POST'])
+def process_image():
+    """Handle image processing via POST request."""
+    data = request.get_json()
+    image_data = data['image']
+    object_type = data['objectType']
 
-            return jsonify({"success": False, "message": "No valid output generated"})
+    # Save the uploaded image locally
+    image_path = save_image(image_data)
 
+    try:
+        # Use Gemini to remove the object from the image
+        result_image = remove_object_from_image(image_path, object_type)
+        return jsonify({'success': True, 'resultPath': result_image})
     except Exception as e:
-        logging.error(f"Processing error: {str(e)}")
-        return jsonify({"success": False, "message": f"Processing error: {str(e)}"})
+        return jsonify({'success': False, 'message': str(e)})
 
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+if __name__ == '__main__':
+    # For local testing; in production, your hosting provider will manage the server.
+    app.run(debug=True)