Switch to Gemini + update requirements

app.py

@@ -1,6 +1,7 @@
 import os
 import time
 from datetime import datetime
+from typing import List
 
 from fastapi import FastAPI, UploadFile, Form
 from fastapi.responses import FileResponse, JSONResponse
@@ -13,9 +14,13 @@ import uvicorn
 
 from dotenv import load_dotenv
 
-# Load environment variables
+# Load environment variables from .env
 load_dotenv()
-print("Gemini API Key Loaded:", os.getenv("GEMINI_API_KEY") is not None)
+
+# Confirm API key is loaded
+print("Gemini API Key Loaded:", os.getenv("GOOGLE_API_KEY") is not None)
+
+
 
 # ---------------------------
 # CONFIG
@@ -29,7 +34,7 @@ LIFETIME = 24 * 60 * 60  # 24 hours
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(RESULTS_DIR, exist_ok=True)
 
-# Gemini setup
+# Gemini API key
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 model = genai.GenerativeModel("gemini-1.5-flash")
 
@@ -48,55 +53,84 @@ def check_size(filepath):
         os.remove(filepath)
         raise ValueError(f"File too large! Max {MAX_SIZE_MB}MB allowed.")
 
-def process_image(input_img, bg_choice, bg_upload=None, brand_color="#FFFFFF", logo_upload=None, logo_opacity=80, logo_position="bottom-right"):
-    """Main image processor with custom bg, brand color, and logo support"""
+def replace_background(input_path, bg_choice):
+    """Replace background with selected file"""
+    check_size(input_path)
+    input_img = Image.open(input_path).convert("RGBA")
+    fg = remove(input_img)
+
+    bg_path = os.path.join(BG_DIR, bg_choice)
+    bg = Image.open(bg_path).convert("RGBA").resize(fg.size)
+
+    result = Image.alpha_composite(bg, fg)
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    result_path = os.path.join(RESULTS_DIR, f"result_{timestamp}.png")
+    result.save(result_path)
+    cleanup_old_files(RESULTS_DIR)
+    return result_path
+
+def process_image(input_img, bg_choice, bg_upload, logo_upload, logo_transparency, logo_position, brand_color):
     if input_img is None:
         return []
 
     temp_path = os.path.join(UPLOAD_DIR, f"upload_{int(time.time())}.png")
     input_img.save(temp_path)
 
-    fg = remove(Image.open(temp_path).convert("RGBA"))
-
-    # --- Background selection ---
-    if bg_choice and bg_choice in os.listdir(BG_DIR):
-        bg = Image.open(os.path.join(BG_DIR, bg_choice)).convert("RGBA").resize(fg.size)
-    elif bg_choice == "Custom Upload" and bg_upload is not None:
-        bg = bg_upload.convert("RGBA").resize(fg.size)
-    elif bg_choice == "Solid Brand Color" and brand_color:
-        bg = Image.new("RGBA", fg.size, brand_color)
+    # Background selection
+    if bg_upload is not None:
+        bg = bg_upload.convert("RGBA").resize(input_img.size)
     else:
-        bg = Image.new("RGBA", fg.size, (255, 255, 255, 255))
+        bg_path = os.path.join(BG_DIR, bg_choice)
+        bg = Image.open(bg_path).convert("RGBA").resize(input_img.size)
+
+    # Foreground (removed background)
+    fg = remove(input_img.convert("RGBA"))
 
+    # Merge main photo with background
     result = Image.alpha_composite(bg, fg)
 
-    # --- Logo overlay ---
+    # If logo uploaded
     if logo_upload is not None:
         logo = logo_upload.convert("RGBA")
-        logo_size = int(result.width * 0.15)
-        logo.thumbnail((logo_size, logo_size))
-        alpha = logo.split()[3].point(lambda p: int(p * (logo_opacity / 100)))
+        # Resize logo (20% of image width)
+        scale = result.width // 5
+        logo.thumbnail((scale, scale))
+        # Apply transparency
+        alpha = logo.split()[3].point(lambda p: p * (logo_transparency / 100))
         logo.putalpha(alpha)
 
-        if logo_position == "top-left":
-            pos = (20, 20)
-        elif logo_position == "top-right":
-            pos = (result.width - logo.width - 20, 20)
-        elif logo_position == "bottom-left":
-            pos = (20, result.height - logo.height - 20)
-        elif logo_position == "center":
-            pos = ((result.width - logo.width) // 2, (result.height - logo.height) // 2)
-        else:  # default bottom-right
-            pos = (result.width - logo.width - 20, result.height - logo.height - 20)
-
-        result.alpha_composite(logo, pos)
-
+        # Position logo
+        pos_map = {
+            "Top-Left": (10, 10),
+            "Top-Right": (result.width - logo.width - 10, 10),
+            "Bottom-Left": (10, result.height - logo.height - 10),
+            "Bottom-Right": (result.width - logo.width - 10, result.height - logo.height - 10),
+            "Center": ((result.width - logo.width) // 2, (result.height - logo.height) // 2),
+        }
+        result.paste(logo, pos_map[logo_position], logo)
+
+    # Apply brand colour tint
+    if brand_color:
+        tint = Image.new("RGBA", result.size, brand_color + "20")  # ~12% opacity overlay
+        result = Image.alpha_composite(result, tint)
+
+    # Save result
     timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
     result_path = os.path.join(RESULTS_DIR, f"result_{timestamp}.png")
     result.save(result_path)
     cleanup_old_files(RESULTS_DIR)
+    return [result_path]
 
+
+'''
+def process_image(input_img, bg_choice):
+    if input_img is None:
+        return []
+    temp_path = os.path.join(UPLOAD_DIR, f"upload_{int(time.time())}.png")
+    input_img.save(temp_path)
+    result_path = replace_background(temp_path, bg_choice)
     return [result_path]
+    '''
 
 def generate_caption(prompt="Promote my product"):
     try:
@@ -115,9 +149,10 @@ def generate_caption(prompt="Promote my product"):
 # ---------------------------
 app = FastAPI(title="SnapLift API")
 
+# Allow CORS for mobile app access
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # Change later for security
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -129,8 +164,9 @@ async def process_image_api(file: UploadFile, bg_choice: str = Form(...)):
         input_path = os.path.join(UPLOAD_DIR, file.filename)
         with open(input_path, "wb") as f:
             f.write(await file.read())
-        result_path = process_image(Image.open(input_path), bg_choice)
-        return FileResponse(result_path[0])
+
+        result_path = replace_background(input_path, bg_choice)
+        return FileResponse(result_path)
     except Exception as e:
         return JSONResponse(content={"error": str(e)}, status_code=400)
 
@@ -146,50 +182,74 @@ with gr.Blocks(css="footer {display:none !important}") as demo:
     gr.Markdown("# ✨ SnapLift – AI Social Media Booster")
     gr.Markdown("Upload your product photo, replace background, and auto-generate marketing captions + hashtags!")
 
-    # --- Image Editor ---
+    
+        # Image Editor Tab
     with gr.Tab("📸 Image Editor"):
         with gr.Row():
             input_img = gr.Image(type="pil", label="Upload Main Photo")
 
             with gr.Column():
                 bg_choices = gr.Dropdown(
-                    choices=os.listdir(BG_DIR) + ["Custom Upload", "Solid Brand Color"],
+                    choices=os.listdir(BG_DIR),
                     value=os.listdir(BG_DIR)[0] if os.listdir(BG_DIR) else None,
                     label="Choose Background"
                 )
-                bg_preview = gr.Image(label="Background Preview", type="filepath")
-                bg_upload = gr.Image(type="pil", label="Upload Custom Background")
-                brand_color = gr.ColorPicker(label="Pick Brand Colour", value="#FFFFFF")
-
-        logo_upload = gr.Image(type="pil", label="Upload Brand Logo (Optional)")
-        logo_opacity = gr.Slider(minimum=0, maximum=100, value=80, step=5, label="Logo Transparency (%)")
-        logo_position = gr.Radio(
-            choices=["top-left", "top-right", "bottom-left", "bottom-right", "center"],
-            value="bottom-right",
-            label="Logo Position"
-        )
-
-        def update_preview(bg_choice):
-            return os.path.join(BG_DIR, bg_choice) if bg_choice in os.listdir(BG_DIR) else None
+                bg_upload = gr.Image(type="pil", label="Or Upload Your Background (Optional)")
 
-        bg_choices.change(fn=update_preview, inputs=bg_choices, outputs=bg_preview)
+        # Logo + Branding Controls
+        with gr.Row():
+            logo_upload = gr.Image(type="pil", label="Upload Logo (Optional)")
+            logo_transparency = gr.Slider(0, 100, value=70, label="Logo Transparency (%)")
+            logo_position = gr.Dropdown(
+                ["Top-Left", "Top-Right", "Bottom-Left", "Bottom-Right", "Center"],
+                value="Bottom-Right",
+                label="Logo Position"
+            )
+            brand_color = gr.ColorPicker(label="Brand Colour")
 
         btn = gr.Button("✨ Generate New Photo")
         output_imgs = gr.Gallery(label="Generated Image", elem_id="gallery", columns=1, rows=1)
+
         btn.click(
             fn=process_image,
-            inputs=[input_img, bg_choices, bg_upload, brand_color, logo_upload, logo_opacity, logo_position],
+            inputs=[input_img, bg_choices, bg_upload, logo_upload, logo_transparency, logo_position, brand_color],
             outputs=output_imgs
         )
 
-    # --- Caption Generator ---
-    with gr.Tab("✍️ Caption Generator"):
-        prompt = gr.Textbox(label="Enter product/promotion text", value="Promote my skincare product")
-        btn2 = gr.Button("💡 Suggest Captions + Hashtags")
-        caption_box = gr.Textbox(label="Suggested Posts (multi-platform)", lines=12)
-        btn2.click(fn=generate_caption, inputs=[prompt], outputs=[caption_box])
 
-# ---------------------------
+
+    '''
+        # Image Editor Tab
+    with gr.Tab("📸 Image Editor"):
+        with gr.Row():
+            input_img = gr.Image(type="pil", label="Upload Photo")
+
+            with gr.Column():
+                bg_choices = gr.Dropdown(
+                    choices=os.listdir(BG_DIR),
+                    value=os.listdir(BG_DIR)[0] if os.listdir(BG_DIR) else None,
+                    label="Choose Background"
+                )
+                bg_preview = gr.Image(
+                    label="Background Preview",
+                    type="filepath"
+                )
+
+        # Update preview whenever background changes
+        def update_preview(bg_choice):
+            return os.path.join(BG_DIR, bg_choice) if bg_choice else None
+
+        bg_choices.change(fn=update_preview, inputs=bg_choices, outputs=bg_preview)
+
+        btn = gr.Button("✨ Generate New Photo")
+        output_imgs = gr.Gallery(label="Generated Image", elem_id="gallery", columns=1, rows=1)
+        btn.click(fn=process_image, inputs=[input_img, bg_choices], outputs=output_imgs)
+
+    
+    '''
+    
+
+    
 # START SERVER
 # ---------------------------
 if __name__ == "__main__":
@@ -200,3 +260,7 @@ if __name__ == "__main__":
 
     threading.Thread(target=run_gradio).start()
     uvicorn.run(app, host="0.0.0.0", port=8000)
+
+threading.Thread(target=start_fastapi, daemon=True).start()
+
+demo.launch(server_name="0.0.0.0", server_port=7860)