app.py

from flask import Flask, render_template, request, jsonify
import google.generativeai as genai
import os
import tempfile
from dotenv import load_dotenv
from google.generativeai.types import Content, Part

# Load environment variables
load_dotenv()

# Configure Flask app
app = Flask(__name__)

# Configure Gemini API
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

# Route for home page
@app.route("/")
def home():
    return render_template("index.html")

# Route for processing images
@app.route("/process", methods=["POST"])
def process_image():
    try:
        # Get data from request
        data = request.json
        image_data = data.get("image")
        object_type = data.get("objectType")

        if not image_data or not object_type:
            return jsonify({"success": False, "message": "Invalid input data"})

        # Decode base64 image data
        import base64
        image_bytes = base64.b64decode(image_data.split(",")[1])

        # Create temporary directory
        temp_dir = tempfile.mkdtemp()
        upload_path = os.path.join(temp_dir, "input_image.png")
        with open(upload_path, "wb") as f:
            f.write(image_bytes)

        # Construct Gemini prompt
        gemini_prompt = f"Remove the {object_type} from the image and fill the area naturally."

        # Create request content
        contents = [
            Content(
                role="user",
                parts=[
                    Part.from_text(gemini_prompt),
                    Part.from_data(mime_type="image/png", data=image_bytes)
                ]
            )
        ]

        # Generate content with Gemini
        model = genai.GenerativeModel('gemini-2.0-flash-exp-image-generation')
        response = model.generate_content(
            contents=contents,
            generation_config={
                "temperature": 1,
                "top_p": 0.95,
                "top_k": 40,
                "max_output_tokens": 8192,
            },
            safety_settings={
                "HARM_CATEGORY_CIVIC_INTEGRITY": "BLOCK_NONE"
            }
        )

        # Process response
        if response.candidates:
            for candidate in response.candidates:
                for part in candidate.content.parts:
                    if hasattr(part, 'inline_data'):
                        output_path = os.path.join(temp_dir, "result.png")
                        with open(output_path, "wb") as f:
                            f.write(part.inline_data.data)
                        return jsonify({
                            "success": True,
                            "resultPath": f"/static/results/result.png"
                        })

        return jsonify({"success": False, "message": "No valid image data found in response"})

    except Exception as e:
        return jsonify({"success": False, "message": str(e)})

# Run the app
if __name__ == "__main__":
    app.run(debug=True)