gdrive integration

cloud/storage_gdrive.py

@@ -1,76 +1,123 @@
 # cloud/storage_gdrive.py
 import io
-import os
-import mimetypes
-from typing import Iterable, Optional
+from typing import Iterable, Optional, List
+from dataclasses import dataclass
 
 from google.oauth2 import service_account as gsa
+from google.oauth2.credentials import Credentials as UserCreds
 from googleapiclient.discovery import build
 from googleapiclient.http import MediaIoBaseDownload, MediaIoBaseUpload
 
 from storage import BlobStore, BlobInfo
 
-_SCOPES = ["https://www.googleapis.com/auth/drive"]
+# Read-only scopes (sufficient for listing + downloading images)
+_SCOPES = [
+    "https://www.googleapis.com/auth/drive.readonly",
+    "https://www.googleapis.com/auth/drive.metadata.readonly",
+]
+
+@dataclass
+class _FileRow:
+    id: str
+    name: str
+    size: Optional[int]
+    modified: Optional[str]
+    mime: str
 
 class GDriveStore(BlobStore):
-    def __init__(self, folder_id: str, creds_json_path: str = "", service_account_json: str = ""):
-        if not folder_id:
-            raise ValueError("GDriveStore requires folder_id")
-        self.folder_id = folder_id
+    def __init__(
+        self,
+        folder_id: Optional[str] = None,
+        *,
+        credentials: Optional[UserCreds] = None,          # user OAuth (device flow)
+        service_account_json: Optional[str] = None,       # SA JSON (optional alt)
+    ):
+        self.root_folder_id = folder_id or "root"
 
-        if service_account_json:
-            creds = gsa.Credentials.from_service_account_file(service_account_json, scopes=_SCOPES)
+        if credentials is not None:
+            self.creds = credentials
+        elif service_account_json:
+            self.creds = gsa.Credentials.from_service_account_file(
+                service_account_json, scopes=_SCOPES
+            )
         else:
-            if not creds_json_path:
-                raise ValueError("Provide GDRIVE_CREDENTIALS_JSON (OAuth) or GDRIVE_SERVICE_ACCOUNT_JSON")
-            creds = gsa.Credentials.from_service_account_file(creds_json_path, scopes=_SCOPES)
-
-        self.svc = build("drive", "v3", credentials=creds, cache_discovery=False)
+            raise ValueError(
+                "GDriveStore requires user OAuth credentials or a service_account_json path."
+            )
 
-    def _query(self, extra_q: str = ""):
-        q = f"'{self.folder_id}' in parents and trashed=false"
-        if extra_q:
-            q += f" and {extra_q}"
-        return q
+        self.svc = build("drive", "v3", credentials=self.creds, cache_discovery=False)
 
-    def list(self, prefix: str = "", recursive: bool = False) -> Iterable[BlobInfo]:
+    # ---------------- internal helpers ----------------
+    def _iter_files(self, q: str) -> Iterable[_FileRow]:
         page_token = None
-        name_filter = f"name contains '{prefix}'" if prefix else ""
         while True:
             resp = self.svc.files().list(
-                q=self._query(name_filter),
-                fields="nextPageToken, files(id, name, size, modifiedTime, mimeType)",
+                q=q,
+                fields="nextPageToken, files(id,name,size,modifiedTime,mimeType)",
                 pageToken=page_token,
                 corpora="allDrives",
                 includeItemsFromAllDrives=True,
                 supportsAllDrives=True,
             ).execute()
             for f in resp.get("files", []):
-                yield BlobInfo(
-                    f["name"],
-                    size=int(f.get("size", 0)) if f.get("size") else None,
+                yield _FileRow(
+                    id=f["id"],
+                    name=f["name"],
+                    size=int(f["size"]) if f.get("size") else None,
                     modified=f.get("modifiedTime"),
-                    is_dir=(f["mimeType"] == "application/vnd.google-apps.folder"),
+                    mime=f.get("mimeType", ""),
                 )
             page_token = resp.get("nextPageToken")
             if not page_token:
                 break
 
-    def _find_id(self, name: str) -> Optional[str]:
-        resp = self.svc.files().list(
-            q=self._query(f"name='{name}'"),
-            fields="files(id, name)",
-            corpora="allDrives",
-            includeItemsFromAllDrives=True,
-            supportsAllDrives=True,
-        ).execute()
-        files = resp.get("files", [])
-        return files[0]["id"] if files else None
+    def _children_of(self, folder_id: str, name_filter: str = "") -> Iterable[_FileRow]:
+        q = f"'{folder_id}' in parents and trashed=false"
+        if name_filter:
+            # Drive query uses single quotes; no special characters expected in 'contains'
+            q += f" and name contains '{name_filter}'"
+        yield from self._iter_files(q)
+
+    def _resolve_path(self, prefix_path: str) -> str:
+        """
+        Resolve a 'folder1/folder2' path starting from self.root_folder_id and return the final folder ID.
+        If any component is missing, returns self.root_folder_id (so listing won't crash).
+        """
+        if not prefix_path or prefix_path.strip() == "":
+            return self.root_folder_id
+        cur = self.root_folder_id
+        parts: List[str] = [p for p in prefix_path.split("/") if p]
+        for seg in parts:
+            # find a child folder with this name
+            q = f"'{cur}' in parents and trashed=false and mimeType='application/vnd.google-apps.folder' and name='{seg}'"
+            match = next(self._iter_files(q), None)
+            if not match:
+                # folder segment not found; fall back to current (best effort)
+                return cur
+            cur = match.id
+        return cur
+
+    def _extract_id(self, key: str) -> str:
+        # Accept "id|name" or bare "id"
+        return key.split("|", 1)[0]
+
+    # ---------------- BlobStore implementation ----------------
+    def list(self, prefix: str = "", recursive: bool = False) -> Iterable[BlobInfo]:
+        # Treat prefix as a folder path (like S3). Final path component is a folder.
+        folder_id = self._resolve_path(prefix)
+        for f in self._children_of(folder_id):
+            is_dir = (f.mime == "application/vnd.google-apps.folder")
+            # Use "id|name" so the UI can display a friendly name but we keep a unique key
+            key = f"{f.id}|{f.name}"
+            yield BlobInfo(
+                key=key,
+                size=f.size,
+                modified=f.modified,
+                is_dir=is_dir,
+            )
 
     def read_bytes(self, key: str) -> bytes:
-        fid = self._find_id(key)
-        if not fid:
-            raise FileNotFoundError(key)
+        fid = self._extract_id(key)
         req = self.svc.files().get_media(fileId=fid, supportsAllDrives=True)
         buf = io.BytesIO()
         downloader = MediaIoBaseDownload(buf, req)
@@ -80,26 +127,26 @@ class GDriveStore(BlobStore):
         return buf.getvalue()
 
     def write_bytes(self, key: str, data: bytes, content_type: Optional[str] = None) -> None:
+        # Writes are not used by your app, but kept for interface parity.
+        fid = self._extract_id(key)
         media = MediaIoBaseUpload(io.BytesIO(data), mimetype=content_type or "application/octet-stream", resumable=False)
-        fid = self._find_id(key)
-        meta = {"name": key, "parents": [self.folder_id]}
-        if fid:
+        # If fid refers to an existing file, update; otherwise create in root.
+        try:
             self.svc.files().update(fileId=fid, media_body=media, supportsAllDrives=True).execute()
-        else:
+        except Exception:
+            meta = {"name": key.split("|",1)[-1], "parents": [self.root_folder_id]}
             self.svc.files().create(body=meta, media_body=media, supportsAllDrives=True, fields="id").execute()
 
     def head(self, key: str) -> BlobInfo:
-        fid = self._find_id(key)
-        if not fid:
-            raise FileNotFoundError(key)
+        fid = self._extract_id(key)
         f = self.svc.files().get(
             fileId=fid,
             fields="id,name,size,modifiedTime,mimeType",
-            supportsAllDrives=True
+            supportsAllDrives=True,
         ).execute()
         return BlobInfo(
-            f["name"],
+            key=f"{f['id']}|{f['name']}",
             size=int(f.get("size", 0)) if f.get("size") else None,
             modified=f.get("modifiedTime"),
-            is_dir=(f["mimeType"] == "application/vnd.google-apps.folder"),
+            is_dir=(f.get("mimeType") == "application/vnd.google-apps.folder"),
         )

infer.py

@@ -13,7 +13,7 @@ from torchcam.methods import GradCAM
 from torchcam.utils import overlay_mask
 
 # Track last resolution for Diagnostics
-LAST_WEIGHT_SOURCE = None     # one of: local | s3 | gdrive | hf_manifest | hf_env
+LAST_WEIGHT_SOURCE = None     # one of: local | s3 | gdrive | hf_manifest 
 LAST_WEIGHT_DETAIL = ""       # key/repo/filename etc
 LAST_WEIGHT_PATH = ""         # final local path used
 
@@ -241,14 +241,14 @@ def _resolve_ckpt(ckpt_path: str, model_name: str, models_root: str) -> str:
 
     # Order: depends on mode
     if mode == "off":
-        order = ["local", "hf_manifest", "hf_env"]
+        order = ["local", "hf_manifest"]
     elif mode == "prefer_cloud":
-        order = ["local", "cloud", "cloud_auto", "hf_manifest", "hf_env"]
+        order = ["local", "cloud", "cloud_auto", "hf_manifest"]
     elif mode == "prefer_hf":
-        order = ["local", "hf_manifest", "hf_env", "cloud", "cloud_auto"]
+        order = ["local", "hf_manifest", "cloud", "cloud_auto"]
     else:  # auto (per-allow-list)
-        order = ["local", "hf_manifest", "hf_env", "cloud", "cloud_auto"] if (model_name not in (os.getenv("CLOUD_WEIGHTS_ALLOW",""))) else \
-                ["local", "cloud", "cloud_auto", "hf_manifest", "hf_env"]
+        order = ["local", "cloud", "cloud_auto", "hf_manifest"] if _cloud_mode_for(model_name) == "prefer_cloud" \
+                else ["local", "hf_manifest", "cloud", "cloud_auto"]
 
     # Try sources in decided order
     for src in order:
@@ -287,21 +287,10 @@ def _resolve_ckpt(ckpt_path: str, model_name: str, models_root: str) -> str:
                 except Exception as e:
                     _log(f"Manifest read error: {e}")
 
-        if src == "hf_env":
-            repo = os.getenv("HF_MODEL_REPO")
-            if repo:
-                fname = os.getenv("HF_MODEL_FILE", os.path.basename(ckpt_path) or "best.pth")
-                _log(f"Downloading from env → repo={repo} file={fname}")
-                path = hf_hub_download(repo_id=repo, filename=fname, token=os.getenv("HF_TOKEN"))
-                LAST_WEIGHT_SOURCE, LAST_WEIGHT_DETAIL, LAST_WEIGHT_PATH = "hf_env", f"{repo}:{fname}", path
-                _log(f"Downloaded to cache: {path}")
-                return path
-
     # If we got here, nothing worked
     parts = [f"local={'yes' if avail_local else 'no'}",
              f"cloud={'yes' if cloud_exists else 'no'} key={cloud_key or '-'} err={cloud_err or '-'}",
-             f"hf_manifest={'present' if (Path(models_root)/'manifest.json').exists() else 'absent'}",
-             f"hf_env={'yes' if os.getenv('HF_MODEL_REPO') else 'no'}"]
+             f"hf_manifest={'present' if (Path(models_root)/'manifest.json').exists() else 'absent'}"]
     raise FileNotFoundError("Could not resolve weights. Probes: " + ", ".join(parts))
 
 #--- main loading function---

requirements.txt

@@ -13,4 +13,5 @@ google-api-python-client>=2.137.0
 google-auth>=2.34.0
 google-auth-httplib2>=0.2.0
 google-auth-oauthlib>=1.2.0
+requests>=2.31
 matplotlib>=3.8

storage.py

@@ -1,8 +1,7 @@
 # storage.py
 from __future__ import annotations
 
-import io
-import os
+import io, tempfile, os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Iterable, Optional
@@ -42,19 +41,31 @@ class BlobStore:
 
 # ---------- Local FS baseline ----------
 class LocalStore(BlobStore):
-    def __init__(self, root: Path | None = None):
-        import os
-        from pathlib import Path as _P
-        import tempfile
-        # 1) APP_DATA_DIR wins if set
+    def __init__(self, root: str | Path | None = None):
+        """
+        Always land on a writable directory.
+        Priority:
+          1) explicit root (if provided)
+          2) APP_DATA_DIR (if set)
+          3) /tmp/label_assistant  (always writable on Spaces/containers)
+        """
+        # choose base
         if root:
-            self.root = _P(root)
+            base = Path(root)
         elif os.getenv("APP_DATA_DIR"):
-            self.root = _P(os.getenv("APP_DATA_DIR"))
+            base = Path(os.getenv("APP_DATA_DIR"))
         else:
-            # 2) Always-writable on HF Spaces / most containers
-            self.root = _P(tempfile.gettempdir()) / "label_assistant"
-        self.root.mkdir(parents=True, exist_ok=True)
+            base = Path(tempfile.gettempdir()) / "label_assistant"
+
+        # ensure it exists; if it fails, force /tmp
+        try:
+            base.mkdir(parents=True, exist_ok=True)
+        except Exception:
+            base = Path(tempfile.gettempdir()) / "label_assistant"
+            base.mkdir(parents=True, exist_ok=True)
+
+        self.root = base
+        print(f"[LocalStore] using {self.root}", flush=True)
 
     def _p(self, key: str) -> Path:
         return self.root / key

ui/sidebar.py

@@ -1,10 +1,15 @@
 # ui/sidebar.py
-import os, io, time
+import os, io, time, json, requests
 from pathlib import Path
 import streamlit as st
 from PIL import Image
-from storage import get_store_from_env, LocalStore
+from storage import LocalStore
 from cloud.storage_s3 import S3Store
+from google.oauth2.credentials import Credentials
+from cloud.storage_gdrive import GDriveStore 
+from googleapiclient.discovery import build
+from google_auth_oauthlib.flow import Flow
+from streamlit.components.v1 import html
 
 # ------------- small logger -------------
 def _log_s3(msg: str):
@@ -12,7 +17,7 @@ def _log_s3(msg: str):
     st.session_state.setdefault("aws_logs", [])
     st.session_state.aws_logs.append(f"[{ts}] {msg}")
 
-def _logs_panel():
+def _s3_logs_panel():
     with st.sidebar.expander("S3 logs", expanded=False):
         logs = st.session_state.get("aws_logs", [])
         if not logs:
@@ -25,22 +30,57 @@ def _logs_panel():
 def build_sidebar_datasource(BASE_DIR: Path):
     st.sidebar.markdown("---")
     st.sidebar.subheader("Data source")
-    src_choice = st.sidebar.radio("Load from", ["Local", "AWS S3", "Google Drive"], index=0, horizontal=False)
+    src_choice = st.sidebar.radio(
+        "Load from",
+        ["Local", "AWS S3", "Google Drive"],
+        index=0,
+        horizontal=False
+    )
 
     if src_choice == "Local":
+        os.environ.setdefault("APP_DATA_DIR", "/tmp/label_assistant")  # never /app
         store = LocalStore()
-        _logs_panel()
         return store, src_choice
 
     if src_choice == "AWS S3":
-        store = _s3_connect_ui()  # includes connected banner + logs + browser
+        store = _s3_connect_ui()                 # includes logs + browser
         return store, src_choice
 
-    # Google Drive – use your existing get_store_from_env path
-    os.environ.setdefault("BLOB_BACKEND", "gdrive")
-    store = get_store_from_env("gdrive")
-    _logs_panel()
-    return store, src_choice
+    if src_choice == "Google Drive":
+        with st.sidebar:
+            # (0) If a redirect URL was staged, perform it immediately (same tab).
+            redir = st.session_state.pop("_oauth_redirect", None)
+            if redir:
+                html(f'<script>window.location.replace("{redir}");</script>', height=0)
+                st.stop()
+
+            # (1) Handle callback first (if we just came back from Google)
+            creds = _get_pkce_creds_from_session() or _handle_google_callback_pkce()
+
+            # (2) Not signed in yet → start PKCE and force a reliable redirect
+            if not creds:
+                # Generate the auth URL and show it as a clickable link
+                url = _start_google_oauth_pkce()
+                if url:
+                    st.link_button("Sign in with Google", url, use_container_width=True)
+                st.caption("After granting access, you'll be redirected back here automatically.")
+                os.environ.setdefault("APP_DATA_DIR", "/tmp/label_assistant")
+                return LocalStore(), src_choice
+
+            # (3) Show whoami + sign-out once signed in
+            email, who = _drive_whoami(creds)
+            st.success(f"Signed in as **{who}** · {email}")
+            if st.button("Sign out of Google", use_container_width=True):
+                for k in ("g_creds", "g_state", "g_client_config", "_oauth_redirect"):
+                    st.session_state.pop(k, None)
+                st.rerun()
+
+        # (4) Use Drive store once signed in
+        store = GDriveStore(credentials=creds)
+        _gdrive_browser_ui(store)
+        return store, src_choice
+
+    return LocalStore(), "Local"
 
 
 # ------------- S3 connection UI -------------
@@ -109,9 +149,9 @@ def _s3_connect_ui():
         _remote_browser_ui(store)
     else:
         st.info("Enter your S3 details and click **Connect**.")
-    _logs_panel()
+    _s3_logs_panel()
     # When not connected, return harmless LocalStore so the rest of the app doesn't break
-    return store if st.session_state.get("aws_store_ok") else LocalStore(Path.cwd()/ "data")
+    return store if st.session_state.get("aws_store_ok") else LocalStore()
 
 # ------------- S3 browser -------------
 def _remote_browser_ui(STORE: S3Store):
@@ -190,6 +230,422 @@ def _remote_browser_ui(STORE: S3Store):
         st.sidebar.info("No images found for this prefix. Try a different filter or check your Prefix above.")
         _log_s3(f"No images found for '{list_prefix}'")
 
+# ---------- Google Drive connection + browser ----------
+# ---------- Google Drive (PKCE OAuth) ----------
+DRIVE_SCOPES = [
+    "https://www.googleapis.com/auth/drive.readonly",
+    "https://www.googleapis.com/auth/drive.metadata.readonly",
+    "https://www.googleapis.com/auth/drive.file",   
+]
+
+def _start_google_oauth_pkce():
+    client_id = os.getenv("GOOGLE_OAUTH_CLIENT_ID")
+    client_secret = os.getenv("GOOGLE_OAUTH_CLIENT_SECRET")
+    redirect_uri = os.getenv("GOOGLE_OAUTH_REDIRECT_URI")
+    if not (client_id and client_secret and redirect_uri):
+        st.error("Missing GOOGLE_OAUTH_CLIENT_ID / GOOGLE_OAUTH_CLIENT_SECRET / GOOGLE_OAUTH_REDIRECT_URI.")
+        return None
+
+    client_config = {
+        "web": {
+            "client_id": client_id,
+            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+            "token_uri": "https://oauth2.googleapis.com/token",
+            "client_secret": client_secret,
+            "redirect_uris": [redirect_uri],
+        }
+    }
+    flow = Flow.from_client_config(client_config, scopes=DRIVE_SCOPES, redirect_uri=redirect_uri)
+    # prompt=consent ensures refresh_token on repeated logins during testing
+    auth_url, state = flow.authorization_url(access_type="offline", include_granted_scopes="true", prompt="consent")
+    st.session_state["g_state"] = state
+    st.session_state["g_client_config"] = client_config
+    return auth_url
+
+def _handle_google_callback_pkce():
+    params = st.query_params
+    if "code" not in params or "state" not in params:
+        return None
+
+    code = params.get("code")
+    state = params.get("state")
+    if isinstance(code, list):  code = code[0]
+    if isinstance(state, list): state = state[0]
+
+    # Enforce state only if we have one saved (same-session case).
+    expected = st.session_state.get("g_state")
+    if expected and state != expected:
+        st.error("OAuth state mismatch. Please try signing in again.")
+        return None
+
+    # Try to get client config from session; if missing (new session), rebuild from ENV.
+    client_config = st.session_state.get("g_client_config")
+    if not client_config:
+        cid  = os.getenv("GOOGLE_OAUTH_CLIENT_ID")
+        csec = os.getenv("GOOGLE_OAUTH_CLIENT_SECRET")
+        ruri = os.getenv("GOOGLE_OAUTH_REDIRECT_URI")
+        if not (cid and csec and ruri):
+            st.error("Missing GOOGLE_OAUTH_CLIENT_ID/SECRET/REDIRECT_URI while handling callback.")
+            return None
+        client_config = {
+            "web": {
+                "client_id": cid,
+                "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+                "token_uri": "https://oauth2.googleapis.com/token",
+                "client_secret": csec,
+                "redirect_uris": [ruri],
+            }
+        }
+
+    redirect_uri = client_config["web"]["redirect_uris"][0]
+    flow = Flow.from_client_config(client_config, scopes=DRIVE_SCOPES, redirect_uri=redirect_uri)
+    try:
+        flow.fetch_token(code=code)
+    except Exception as e:
+        st.error(f"OAuth token exchange failed: {e}")
+        return None
+
+    creds = flow.credentials
+    st.session_state["g_creds"] = {
+        "token": creds.token,
+        "refresh_token": creds.refresh_token,
+        "token_uri": creds.token_uri,
+        "client_id": creds.client_id,
+        "client_secret": creds.client_secret,
+        "scopes": creds.scopes,
+    }
+
+    # Clear query params so we don’t reprocess the callback on every rerun
+    try:
+        st.query_params.clear()
+    except Exception:
+        pass
+    return creds
+
+def _get_pkce_creds_from_session():
+    data = st.session_state.get("g_creds")
+    return Credentials(**data) if data else None
+
+def _gdrive_display_name(k: str) -> str:
+    return k.split("|", 1)[1] if "|" in k else k
+
+def _gdev_start():
+    cid = os.getenv("GOOGLE_OAUTH_CLIENT_ID", "").strip()
+    if not cid:
+        st.error("Missing GOOGLE_OAUTH_CLIENT_ID env. Set it in your Space/local .env.")
+        return
+    if not cid.endswith(".apps.googleusercontent.com"):
+        st.warning("GOOGLE_OAUTH_CLIENT_ID doesn't look like a valid OAuth client ID (…apps.googleusercontent.com). Double-check the value.")
+
+    try:
+        resp = requests.post(
+            "https://oauth2.googleapis.com/device/code",
+            data={"client_id": cid, "scope": DRIVE_SCOPES},
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            timeout=15,
+        )
+    except Exception as e:
+        st.error(f"Google device-code start failed: network error: {e}")
+        return
+
+    if resp.status_code != 200:
+        # Show the exact JSON/text so you know what's wrong
+        try:
+            err = resp.json()
+        except Exception:
+            err = resp.text
+        st.error(f"Device-code start failed: {resp.status_code} {err}")
+        st.session_state["gdev_error"] = err
+        return
+
+    data = resp.json()
+    st.session_state["gdev"] = data
+    st.session_state["gdev_started_at"] = time.time()
+    st.session_state["gcreds"] = None
+
+def _gdev_poll():
+    if "gdev" not in st.session_state:
+        return
+    cid = os.getenv("GOOGLE_OAUTH_CLIENT_ID")
+    csec = os.getenv("GOOGLE_OAUTH_CLIENT_SECRET", "")
+    dev = st.session_state["gdev"]
+    try:
+        resp = requests.post(
+            "https://oauth2.googleapis.com/token",
+            data={
+                "client_id": cid,
+                "client_secret": csec,  # optional for desktop/device, harmless if set
+                "device_code": dev["device_code"],
+                "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+            },
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            timeout=15,
+        )
+    except Exception as e:
+        st.error(f"Google token poll failed: network error: {e}")
+        return
+
+    if resp.status_code == 400:
+        err = (resp.json() or {}).get("error")
+        # expected while user hasn’t approved yet
+        if err in ("authorization_pending", "slow_down"):
+            return
+        st.error(f"Google auth error: {err}")
+        st.session_state["gdev_error"] = resp.text
+        return
+
+    if resp.status_code != 200:
+        st.error(f"Google token poll failed: {resp.status_code} {resp.text}")
+        st.session_state["gdev_error"] = resp.text
+        return
+
+    tok = resp.json()
+    creds = Credentials(
+        token=tok["access_token"],
+        refresh_token=tok.get("refresh_token"),
+        token_uri="https://oauth2.googleapis.com/token",
+        client_id=cid,
+        client_secret=csec or None,
+        scopes=DRIVE_SCOPES.split(),
+    )
+    st.session_state["gcreds"] = {
+        "token": creds.token,
+        "refresh_token": creds.refresh_token,
+        "token_uri": creds.token_uri,
+        "client_id": creds.client_id,
+        "client_secret": creds.client_secret,
+        "scopes": creds.scopes,
+    }
+
+def _drive_whoami(creds):
+    """Return (email, name) for the signed-in account."""
+    try:
+        svc = build("drive", "v3", credentials=creds, cache_discovery=False)
+        me = svc.about().get(fields="user(emailAddress, displayName)").execute()
+        u = me.get("user", {}) or {}
+        return u.get("emailAddress") or "unknown", u.get("displayName") or "unknown"
+    except Exception as e:
+        return None, f"whoami failed: {e}"
+
+def _drive_debug_list_root(creds, folder_id="root"):
+    """Return count and sample names from the current root to verify listing works."""
+    try:
+        svc = build("drive", "v3", credentials=creds, cache_discovery=False)
+        q = f"'{folder_id or 'root'}' in parents and trashed=false"
+        resp = svc.files().list(
+            q=q,
+            fields="files(id,name,mimeType), nextPageToken",
+            corpora="allDrives",
+            includeItemsFromAllDrives=True,
+            supportsAllDrives=True,
+            pageSize=20
+        ).execute()
+        files = resp.get("files", [])
+        names = [f["name"] for f in files[:10]]
+        return len(files), names
+    except Exception as e:
+        return None, f"list failed: {e}"
+
+def _drive_login_block():
+    """UI block for Google Drive sign-in via Device Code."""
+    creds_data = st.session_state.get("gcreds")
+    if creds_data:
+        st.success("Connected to Google Drive.")
+        if st.button("Sign out of Google", use_container_width=True):
+            st.session_state.pop("gcreds", None)
+            st.session_state.pop("gdev", None)
+            st.rerun()
+        return Credentials(**creds_data)
+
+    dev = st.session_state.get("gdev")
+    if not dev:
+        if st.button("Sign in with Google", use_container_width=True):
+            # clear a prior error so it doesn’t linger after a new attempt
+            st.session_state.pop("gdev_error", None)
+            _gdev_start()
+            st.rerun()
+
+        # ← show the last Google response even if device-code init failed
+        ge = st.session_state.get("gdev_error")
+        if ge:
+            st.sidebar.markdown("**Google auth error (last response):**")
+            st.sidebar.code(json.dumps(ge, indent=2) if isinstance(ge, dict) else str(ge))
+
+        st.caption("You’ll use a short code on accounts.google.com/device to grant Drive read-only access.")
+        return None
+
+    # Show code + link and poll
+    st.info(f"Go to **{dev['verification_url']}** and enter this code:\n\n**{dev['user_code']}**")
+    col1, col2 = st.columns([1,1])
+    with col1:
+        if st.button("I’ve approved", use_container_width=True):
+            _gdev_poll()
+            st.rerun()
+    with col2:
+        if st.button("Cancel", use_container_width=True):
+            st.session_state.pop("gdev", None)
+            st.rerun()
+
+    # Optional: auto-poll every few seconds
+    st.caption("Waiting for approval…")
+    _gdev_poll()
+
+    if "gdev_error" in st.session_state:
+        st.sidebar.caption("Google response:")
+        st.sidebar.code(json.dumps(st.session_state["gdev_error"], indent=2) if isinstance(st.session_state["gdev_error"], dict) else str(st.session_state["gdev_error"]))
+
+    return None
+
+def _gdrive_connect_ui():
+    st.sidebar.markdown("**Connect to your Google Drive**")
+    with st.sidebar.expander("Drive connection", expanded=True):
+        # Two modes: Service Account (recommended on Spaces) OR OAuth Device Code.
+        mode = st.radio("Auth mode", ["Service Account (JSON)", "OAuth Device Code"], horizontal=False, key="gd_mode")
+
+        if mode == "Service Account (JSON)":
+            sa_json = st.text_input(
+                "Path to service account JSON (mounted/secret)",
+                value=st.session_state.get("gd_sa_json", os.getenv("GDRIVE_SERVICE_ACCOUNT_JSON","")),
+                key="gd_sa_json"
+            )
+            folder_id = st.text_input(
+                "Root Folder ID (share this folder with your service account email)",
+                value=st.session_state.get("gd_folder", os.getenv("GDRIVE_FOLDER_ID","")),
+                key="gd_folder"
+            )
+
+            c1, c2 = st.columns([1,1])
+            if c1.button("Connect", use_container_width=True):
+                try:
+                    store = GDriveStore(service_account_json=sa_json or None, root_folder_id=folder_id or None)
+                    # probe
+                    _ = next(iter(store.list(prefix="", recursive=False)), None)
+                    st.session_state.gd_store = {
+                        "ok": True, "err": None, "mode": "sa",
+                        "sa_json": sa_json, "folder": folder_id
+                    }
+                    st.success("Connected to Google Drive (Service Account).")
+                    st.rerun()
+                except Exception as e:
+                    st.session_state.gd_store = {"ok": False, "err": str(e)}
+                    st.error(f"Connection failed: {e}")
+
+            if c2.button("Sign out", use_container_width=True):
+                st.session_state.pop("gd_store", None)
+                st.rerun()
+
+        else:  # OAuth Device Code (user account)
+            client_id = st.text_input("OAuth Client ID", value=st.session_state.get("gd_cid",""), key="gd_cid")
+            client_secret = st.text_input("OAuth Client Secret", value=st.session_state.get("gd_csec",""), type="password", key="gd_csec")
+            folder_id = st.text_input("Root Folder ID (optional)", value=st.session_state.get("gd_folder",""), key="gd_folder")
+
+            c1, c2 = st.columns([1,1])
+            if c1.button("Connect", use_container_width=True):
+                try:
+                    store = GDriveStore(
+                        client_id=client_id or None,
+                        client_secret=client_secret or None,
+                        root_folder_id=folder_id or None,
+                        oauth_mode="device"  # GDriveStore should implement device-code flow
+                    )
+                    _ = next(iter(store.list(prefix="", recursive=False)), None)
+                    st.session_state.gd_store = {
+                        "ok": True, "err": None, "mode": "oauth",
+                        "cid": client_id, "csec": client_secret, "folder": folder_id
+                    }
+                    st.success("Connected to Google Drive (OAuth).")
+                    st.rerun()
+                except Exception as e:
+                    st.session_state.gd_store = {"ok": False, "err": str(e)}
+                    st.error(f"Connection failed: {e}")
+
+            if c2.button("Sign out", use_container_width=True):
+                st.session_state.pop("gd_store", None)
+                st.rerun()
+
+    # Connected banner + browser
+    gd = st.session_state.get("gd_store")
+    if gd and gd.get("ok"):
+        st.sidebar.success(f"Connected to Drive (mode={gd.get('mode')}).")
+        if gd["mode"] == "sa":
+            store = GDriveStore(service_account_json=gd.get("sa_json") or None,
+                                root_folder_id=gd.get("folder") or None)
+        else:
+            store = GDriveStore(client_id=gd.get("cid") or None,
+                                client_secret=gd.get("csec") or None,
+                                root_folder_id=gd.get("folder") or None,
+                                oauth_mode="device")
+        _gdrive_browser_ui(store)
+        return store
+
+    st.info("Enter your Drive credentials and click **Connect**.")
+    return LocalStore()
+
+def _gdrive_browser_ui(STORE: GDriveStore):
+    st.sidebar.markdown("---")
+    st.sidebar.subheader("Browse Drive images")
+
+    list_prefix = st.sidebar.text_input("Subfolder filter (path / ‘folder1/folder2’)", value=st.session_state.get("gd_prefix",""))
+    c1, c2 = st.sidebar.columns(2)
+
+    if c1.button("List", use_container_width=True):
+        st.session_state.gd_prefix = list_prefix
+        items = []
+        try:
+            raw = list(STORE.list(prefix=list_prefix, recursive=True))
+            for b in raw:
+                if b.is_dir:
+                    continue
+                if b.key.lower().endswith((".jpg",".jpeg",".png",".bmp",".webp")):
+                    items.append(b)
+        except Exception as e:
+            st.sidebar.error(f"Drive list failed: {e}")
+            items = []
+        st.session_state.gd_list = items
+
+    if c2.button("Load all listed", use_container_width=True):
+        loaded = 0
+        for b in st.session_state.get("gd_list", []):
+            try:
+                data = STORE.read_bytes(b.key)
+                im = Image.open(io.BytesIO(data)).convert("RGB")
+                key = f"gdrive://{b.key}"
+                st.session_state.images[key] = {
+                    "key": key, "name": _gdrive_display_name(b.key), "pil": im,
+                    "boxes": [], "preds": [], "user_labels": [], "actions": [], "canvas_json": None
+                }
+                if st.session_state.active_key is None:
+                    st.session_state.active_key = key
+                loaded += 1
+            except Exception as e:
+                st.sidebar.warning(f"Load failed for {b.key}: {e}")
+        if loaded:
+            st.sidebar.success(f"Loaded {loaded} image(s)")
+            st.rerun()
+
+    remote_items = st.session_state.get("gd_list", [])
+    if remote_items:
+        st.sidebar.caption(f"{len(remote_items)} image(s) listed under '{list_prefix or '(root)'}'")
+        for b in remote_items:
+            cols = st.sidebar.columns([0.7, 0.3])
+            cols[0].caption(f"🖼️ {_gdrive_display_name(b.key)}")
+            if cols[1].button("Load", key=f"gd_load_{b.key}"):
+                try:
+                    data = STORE.read_bytes(b.key)
+                    im = Image.open(io.BytesIO(data)).convert("RGB")
+                    key = f"gdrive://{b.key}"
+                    st.session_state.images[key] = {
+                        "key": key, "name": _gdrive_display_name(b.key), "pil": im,
+                        "boxes": [], "preds": [], "user_labels": [], "actions": [], "canvas_json": None
+                    }
+                    if st.session_state.active_key is None:
+                        st.session_state.active_key = key
+                    st.rerun()
+                except Exception as e:
+                    st.sidebar.error(f"Load failed: {e}")
+    else:
+        st.sidebar.info("No images found. Try another subfolder.")
+
 # ------------- Loaded images sidebar -------------
 def show_loaded_images_sidebar():
     st.sidebar.markdown("---")