Update app.py

app.py

@@ -1,38 +1,59 @@
-from flask import Flask, render_template, request, send_file
 import os
+import base64
 import mimetypes
+from flask import Flask, render_template, request, jsonify
+from werkzeug.utils import secure_filename
 from google import genai
 from google.genai import types
-from io import BytesIO
 
+# Initialize Flask app
 app = Flask(__name__)
 
-# Initialize Gemini client
-
+# Set your Gemini API key (or set GEMINI_API_KEY in your environment)
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 client = genai.Client(api_key=GEMINI_API_KEY)
 
-def save_binary_file(file_name, data):
-    """Save binary data to a file."""
-    with open(file_name, "wb") as f:
-        f.write(data)
-
-def generate_gemini_output(user_input, image):
+# Create necessary directories
+UPLOAD_FOLDER = 'uploads'
+RESULT_FOLDER = os.path.join('static')
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(RESULT_FOLDER, exist_ok=True)
+
+def generate_gemini_output(object_type, image_data_url):
+    """
+    Generate output from Gemini by removing the object specified.
+    The image_data_url is expected to be a base64 data URL.
+    """
     model = "gemini-2.0-flash-exp-image-generation"
-
-    # Upload the image to Gemini
     files = []
-    if image:
-        uploaded_file = client.files.upload(file=image)
+
+    # Decode the image data from the data URL
+    if image_data_url:
+        try:
+            header, encoded = image_data_url.split(',', 1)
+        except ValueError:
+            raise ValueError("Invalid image data")
+        binary_data = base64.b64decode(encoded)
+        # Determine file extension from header
+        if "png" in header.lower():
+            ext = ".png"
+        else:
+            ext = ".jpg"
+        temp_filename = secure_filename("temp_image" + ext)
+        temp_filepath = os.path.join(UPLOAD_FOLDER, temp_filename)
+        with open(temp_filepath, "wb") as f:
+            f.write(binary_data)
+        # Upload file to Gemini
+        uploaded_file = client.files.upload(file=temp_filepath)
         files.append(uploaded_file)
 
-    # Prepare the input content
+    # Prepare content parts for Gemini
     parts = []
     if files:
         parts.append(types.Part.from_uri(file_uri=files[0].uri, mime_type=files[0].mime_type))
-
-    # Incorporate Gemini magic: modify the prompt to instruct removal of the user-specified element.
-    if user_input:
-        magic_prompt = f"Remove {user_input} from the image"
+    if object_type:
+        # Create Gemini magic prompt for object removal
+        magic_prompt = f"Remove {object_type} from the image"
         parts.append(types.Part.from_text(text=magic_prompt))
 
     contents = [types.Content(role="user", parts=parts)]
@@ -52,6 +73,7 @@ def generate_gemini_output(user_input, image):
     result_text = None
     result_image = None
 
+    # Stream the output from Gemini API
     for chunk in client.models.generate_content_stream(
         model=model,
         contents=contents,
@@ -59,29 +81,47 @@ def generate_gemini_output(user_input, image):
     ):
         if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
             continue
-
         part = chunk.candidates[0].content.parts[0]
-
         if part.inline_data:
-            file_name = "generated_output"
             file_extension = mimetypes.guess_extension(part.inline_data.mime_type) or ".png"
-            file_path = os.path.join("static", f"{file_name}{file_extension}")
-            save_binary_file(file_path, part.inline_data.data)
-            result_image = file_path
+            output_filename = secure_filename("generated_output" + file_extension)
+            result_image_path = os.path.join(RESULT_FOLDER, output_filename)
+            with open(result_image_path, "wb") as f:
+                f.write(part.inline_data.data)
+            result_image = result_image_path  # Path relative to static folder
         else:
             result_text = part.text
 
     return result_text, result_image
 
-@app.route("/", methods=["GET", "POST"])
+@app.route("/")
 def index():
-    result_text = None
-    result_image = None
-    if request.method == "POST":
-        user_input = request.form.get("user_input")
-        image = request.files.get("image_input")
-        result_text, result_image = generate_gemini_output(user_input, image)
-    return render_template("index.html", result_text=result_text, result_image=result_image)
+    # Render the front-end HTML (which includes inline CSS and JS)
+    return render_template("index.html")
+
+@app.route("/process", methods=["POST"])
+def process():
+    try:
+        # Expecting JSON with keys "image" (base64 data URL) and "objectType"
+        data = request.get_json(force=True)
+        image_data = data.get("image")
+        object_type = data.get("objectType", "").strip()
+        if not image_data or not object_type:
+            return jsonify({"success": False, "message": "Missing image data or object type."}), 400
+
+        # Generate Gemini output
+        result_text, result_image = generate_gemini_output(object_type, image_data)
+        if not result_image:
+            return jsonify({"success": False, "message": "Failed to generate image."}), 500
+
+        # Create a URL to serve the image from the static folder.
+        # Assuming your static folder is served at '/static'
+        image_url = f"/static/{os.path.basename(result_image)}"
+
+        return jsonify({"success": True, "resultPath": image_url, "resultText": result_text})
+    except Exception as e:
+        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+    # Run the app; debug can be set to False in production.
+    app.run(host="0.0.0.0", port=int(os.getenv("PORT", 5000)), debug=True)