Update app.py

app.py

@@ -5,6 +5,7 @@ import json
 import base64
 import logging
 import uuid
+import time
 from typing import List, Dict, Any, Tuple
 
 from flask import Flask, request, jsonify
@@ -33,7 +34,9 @@ log = logging.getLogger("wardrobe-server")
 
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
 if not GEMINI_API_KEY:
-    log.warning("GEMINI_API_KEY not set — gemini calls will fail.")
+    log.warning("GEMINI_API_KEY not set — gemini calls will fail (but fallback still works).")
+
+client = genai.Client(api_key=GEMINI_API_KEY)
 
 # Firebase config (read service account JSON from env)
 FIREBASE_ADMIN_JSON = os.getenv("FIREBASE_ADMIN_JSON", "").strip()
@@ -42,8 +45,6 @@ FIREBASE_STORAGE_BUCKET = os.getenv("FIREBASE_STORAGE_BUCKET", "").strip()  # op
 if FIREBASE_ADMIN_JSON and not FIREBASE_ADMIN_AVAILABLE:
     log.warning("FIREBASE_ADMIN_JSON provided but firebase-admin SDK is not installed. Install firebase-admin.")
 
-client = genai.Client(api_key=GEMINI_API_KEY)
-
 app = Flask(__name__)
 CORS(app)
 
@@ -80,9 +81,10 @@ def init_firebase_admin_if_needed():
         log.exception("Failed to initialize firebase admin: %s", e)
         raise
 
-def upload_b64_to_firebase(base64_str: str, path: str, content_type="image/jpeg") -> str:
+def upload_b64_to_firebase(base64_str: str, path: str, content_type="image/jpeg", metadata: dict = None) -> str:
     """
     Upload base64 string to Firebase Storage at `path` (e.g. detected/uid/item.jpg).
+    Optionally attach metadata dict (custom metadata).
     Returns a public URL when possible, otherwise returns gs://<bucket>/<path>.
     """
     if not FIREBASE_ADMIN_JSON:
@@ -91,7 +93,6 @@ def upload_b64_to_firebase(base64_str: str, path: str, content_type="image/jpeg"
     if not FIREBASE_ADMIN_AVAILABLE:
         raise RuntimeError("firebase-admin not available")
 
-    # decode base64 (strip data:... prefix if present)
     raw = base64_str
     if raw.startswith("data:"):
         raw = raw.split(",", 1)[1]
@@ -99,9 +100,17 @@ def upload_b64_to_firebase(base64_str: str, path: str, content_type="image/jpeg"
     data = base64.b64decode(raw)
 
     try:
-        bucket = fb_storage.bucket()  # uses default bucket from app options
+        bucket = fb_storage.bucket()
         blob = bucket.blob(path)
         blob.upload_from_string(data, content_type=content_type)
+        # attach metadata if provided
+        if metadata:
+            try:
+                # google-cloud-storage uses blob.metadata (dict)
+                blob.metadata = metadata
+                blob.patch()
+            except Exception as me:
+                log.warning("Failed to patch metadata for %s: %s", path, me)
         try:
             blob.make_public()
             return blob.public_url
@@ -123,6 +132,7 @@ def read_image_bytes(file_storage) -> Tuple[np.ndarray, int, int, bytes]:
     try:
         img = ImageOps.exif_transpose(img)
     except Exception:
+        # ignore if EXIF not present or transpose fails
         pass
     img = img.convert("RGB")
     w, h = img.size
@@ -136,6 +146,7 @@ def crop_and_b64(bgr_img: np.ndarray, x: int, y: int, w: int, h: int, max_side=5
     crop = bgr_img[y:y2, x:x2]
     if crop.size == 0:
         return ""
+    # resize if too large
     max_dim = max(crop.shape[0], crop.shape[1])
     if max_dim > max_side:
         scale = max_side / max_dim
@@ -174,6 +185,7 @@ def fallback_contour_crops(bgr_img, max_items=8) -> List[Dict[str, Any]]:
             "thumbnail_b64": b64,
             "source": "fallback"
         })
+    # if still none, split into grid
     if not items:
         h_half, w_half = h_img//2, w_img//2
         rects = [
@@ -200,6 +212,7 @@ def process_image():
         return jsonify({"error": "missing photo"}), 400
     file = request.files["photo"]
 
+    # optional uid from form fields (client can supply for grouping)
     uid = (request.form.get("uid") or request.args.get("uid") or "anon").strip() or "anon"
 
     try:
@@ -208,6 +221,10 @@ def process_image():
         log.error("invalid image: %s", e)
         return jsonify({"error": "invalid image"}), 400
 
+    # generate a per-request session id used to mark temporary uploads
+    session_id = str(uuid.uuid4())
+
+    # Build a prompt instructing Gemini to detect garments and return normalized bbox list
     user_prompt = (
         "You are an assistant that extracts clothing detections from a single image. "
         "Return a JSON object with a single key 'items' which is an array. Each item must have: "
@@ -265,14 +282,16 @@ def process_image():
             log.warning("Could not parse Gemini JSON: %s", e)
             parsed = None
 
-        items_out = []
+        items_out: List[Dict[str, Any]] = []
         if parsed and isinstance(parsed.get("items"), list) and len(parsed["items"])>0:
             for it in parsed["items"]:
                 try:
                     label = str(it.get("label","unknown"))[:48]
                     bbox = it.get("bbox",{})
-                    nx = float(bbox.get("x",0)); ny = float(bbox.get("y",0))
-                    nw = float(bbox.get("w",0)); nh = float(bbox.get("h",0))
+                    nx = float(bbox.get("x",0))
+                    ny = float(bbox.get("y",0))
+                    nw = float(bbox.get("w",0))
+                    nh = float(bbox.get("h",0))
                     nx = max(0.0, min(1.0, nx)); ny = max(0.0,min(1.0,ny))
                     nw = max(0.0, min(1.0, nw)); nh = max(0.0, min(1.0, nh))
                     px = int(nx * img_w); py = int(ny * img_h)
@@ -296,7 +315,7 @@ def process_image():
             log.info("Gemini returned no items or parse failed — using fallback contour crops.")
             items_out = fallback_contour_crops(bgr_img, max_items=8)
 
-        # === AUTO-UPLOAD all detection thumbnails to Firebase Storage (so client doesn't upload) ===
+        # === AUTO-UPLOAD all detection thumbnails to Firebase Storage (temporary, marked by session_id) ===
         if FIREBASE_ADMIN_JSON and FIREBASE_ADMIN_AVAILABLE:
             safe_uid = "".join(ch for ch in uid if ch.isalnum() or ch in ("-", "_")) or "anon"
             for itm in items_out:
@@ -306,55 +325,53 @@ def process_image():
                 item_id = itm.get("id") or str(uuid.uuid4())
                 path = f"detected/{safe_uid}/{item_id}.jpg"
                 try:
-                    url = upload_b64_to_firebase(b64, path, content_type="image/jpeg")
+                    meta = {
+                        "tmp": "true",
+                        "session_id": session_id,
+                        "uploaded_by": safe_uid,
+                        "uploaded_at": str(int(time.time()))
+                    }
+                    url = upload_b64_to_firebase(b64, path, content_type="image/jpeg", metadata=meta)
                     itm["thumbnail_url"] = url
-                    # keep thumbnail_b64 for potential debugging but remove to keep payload small:
+                    # remove raw base64 to keep response small
                     itm.pop("thumbnail_b64", None)
-                    log.debug("Auto-uploaded thumbnail for %s -> %s", item_id, url)
+                    # add session marker to response item
+                    itm["_session_id"] = session_id
+                    log.debug("Auto-uploaded thumbnail for %s -> %s (session=%s)", item_id, url, session_id)
                 except Exception as up_e:
                     log.warning("Auto-upload failed for %s: %s", item_id, up_e)
-                    # keep thumbnail_b64 as fallback (client can still display base64)
+                    # leave thumbnail_b64 for client fallback display
         else:
             if not FIREBASE_ADMIN_JSON:
                 log.info("FIREBASE_ADMIN_JSON not set; skipping server-side thumbnail upload.")
             else:
                 log.info("Firebase admin SDK not available; skipping server-side thumbnail upload.")
 
-        return jsonify({"ok": True, "items": items_out, "debug": {"raw_model_text": raw_text[:1600]}}), 200
+        return jsonify({"ok": True, "items": items_out, "session_id": session_id, "debug": {"raw_model_text": raw_text[:1600]}}), 200
 
     except Exception as ex:
         log.exception("Processing error: %s", ex)
         try:
             items_out = fallback_contour_crops(bgr_img, max_items=8)
-            # attempt auto-upload on fallback results as well
-            if FIREBASE_ADMIN_JSON and FIREBASE_ADMIN_AVAILABLE:
-                safe_uid = "".join(ch for ch in uid if ch.isalnum() or ch in ("-", "_")) or "anon"
-                for itm in items_out:
-                    b64 = itm.get("thumbnail_b64")
-                    if not b64:
-                        continue
-                    item_id = itm.get("id") or str(uuid.uuid4())
-                    path = f"detected/{safe_uid}/{item_id}.jpg"
-                    try:
-                        url = upload_b64_to_firebase(b64, path, content_type="image/jpeg")
-                        itm["thumbnail_url"] = url
-                        itm.pop("thumbnail_b64", None)
-                    except Exception as up_e:
-                        log.warning("Failed to upload fallback thumbnail for %s: %s", item_id, up_e)
-            return jsonify({"ok": True, "items": items_out, "debug": {"error": str(ex)}}), 200
+            return jsonify({"ok": True, "items": items_out, "session_id": session_id, "debug": {"error": str(ex)}}), 200
         except Exception as e2:
             log.exception("Fallback also failed: %s", e2)
             return jsonify({"error": "internal failure", "detail": str(e2)}), 500
 
-# ---------- New endpoint: finalize (keep selected -> delete rest) ----------
+# ---------- Finalize endpoint: keep selected and delete only session's temp files ----------
 @app.route("/finalize_detections", methods=["POST"])
 def finalize_detections():
     """
     Body JSON:
-    { "uid": "user123", "keep_ids": ["id1","id2",...] }
+    { "uid": "user123", "keep_ids": ["id1","id2",...], "session_id": "<session id from /process>" }
 
-    Server will delete all detected/<uid>/* files whose id is NOT in keep_ids.
-    Returns kept (id -> thumbnail_url) and deleted ids.
+    Server will delete only detected/<uid>/* files whose:
+      - metadata.tmp == "true"
+      - metadata.session_id == session_id
+      - item_id NOT in keep_ids
+
+    Returns:
+      { ok: True, kept: [...], deleted: [...], errors: [...] }
     """
     try:
         body = request.get_json(force=True)
@@ -363,12 +380,21 @@ def finalize_detections():
 
     uid = (body.get("uid") or request.args.get("uid") or "anon").strip() or "anon"
     keep_ids = set(body.get("keep_ids") or [])
+    session_id = (body.get("session_id") or request.args.get("session_id") or "").strip()
+
+    if not session_id:
+        return jsonify({"error": "session_id required for finalize to avoid unsafe deletes"}), 400
 
     if not FIREBASE_ADMIN_JSON or not FIREBASE_ADMIN_AVAILABLE:
         return jsonify({"error": "firebase admin not configured"}), 500
 
-    init_firebase_admin_if_needed()
-    bucket = fb_storage.bucket()
+    try:
+        init_firebase_admin_if_needed()
+        bucket = fb_storage.bucket()
+    except Exception as e:
+        log.exception("Firebase init error in finalize: %s", e)
+        return jsonify({"error": "firebase admin init failed", "detail": str(e)}), 500
+
     safe_uid = "".join(ch for ch in uid if ch.isalnum() or ch in ("-", "_")) or "anon"
     prefix = f"detected/{safe_uid}/"
 
@@ -380,9 +406,20 @@ def finalize_detections():
         blobs = list(bucket.list_blobs(prefix=prefix))
         for blob in blobs:
             try:
-                name = blob.name  # detected/<safe_uid>/<itemid>.jpg
+                name = blob.name  # e.g. "detected/<safe_uid>/<itemid>.jpg"
                 fname = name.split("/")[-1]
+                # only accept files of form <id>.<ext>
+                if "." not in fname:
+                    # skip unexpected filenames
+                    continue
                 item_id = fname.rsplit(".", 1)[0]
+
+                md = blob.metadata or {}
+                # only consider temporary files that match this session id
+                if str(md.get("session_id", "")) != session_id or str(md.get("tmp", "")).lower() not in ("true", "1", "yes"):
+                    # skip (not part of this session / not temporary)
+                    continue
+
                 if item_id in keep_ids:
                     # ensure public URL available if possible
                     try:
@@ -392,7 +429,7 @@ def finalize_detections():
                         url = f"gs://{bucket.name}/{name}"
                     kept.append({"id": item_id, "thumbnail_url": url})
                 else:
-                    # delete unkept
+                    # delete unkept (safe: only temporary and session-matching files)
                     try:
                         blob.delete()
                         deleted.append(item_id)