Spaces:
Sleeping
Sleeping
File size: 2,941 Bytes
e5c238d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
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) |