from __future__ import annotations import hashlib import shutil import zipfile from dataclasses import dataclass from pathlib import Path from typing import Iterable, Iterator, List, Optional, Sequence, Tuple from urllib.parse import parse_qs, urlparse SUPPORTED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"} PROJECT_ROOT = Path(__file__).resolve().parent.parent DRIVE_CACHE_ROOT = PROJECT_ROOT / ".cache" / "xdl-colitis-demo" / "drive-downloads" @dataclass(frozen=True) class DriveDownloadProgress: message: str current: int total: int paths: List[Path] filenames: List[str] done: bool = False error: str = "" def iter_image_paths(folder: Path) -> List[Path]: return sorted( p for p in folder.rglob("*") if p.is_file() and p.suffix.lower() in SUPPORTED_IMAGE_EXTENSIONS ) def normalize_uploaded_paths(uploaded_files: Optional[Sequence[str] | str]) -> List[Path]: if uploaded_files is None: return [] if isinstance(uploaded_files, str): raw_paths: Iterable[str] = [uploaded_files] else: raw_paths = uploaded_files collected: List[Path] = [] for raw_path in raw_paths: if not raw_path: continue path = Path(str(raw_path)) if path.is_dir(): collected.extend(iter_image_paths(path)) elif path.is_file() and path.suffix.lower() in SUPPORTED_IMAGE_EXTENSIONS: collected.append(path) return sorted({p.resolve() for p in collected}) def parse_google_drive_link(link: str) -> Tuple[Optional[str], Optional[str]]: parsed = urlparse(link.strip()) host = parsed.netloc.lower() if "drive.google.com" not in host and "docs.google.com" not in host: return None, None query_id = parse_qs(parsed.query).get("id", [None])[0] if query_id: return "file", query_id parts = [part for part in parsed.path.split("/") if part] if "folders" in parts: idx = parts.index("folders") if idx + 1 < len(parts): return "folder", parts[idx + 1] if "file" in parts and "d" in parts: idx = parts.index("d") if idx + 1 < len(parts): return "file", parts[idx + 1] return None, None def is_google_drive_link(link: str) -> bool: kind, drive_id = parse_google_drive_link(link) return bool(kind and drive_id) def _cache_dir_for_link(link: str, cache_root: Path) -> Path: cache_key = hashlib.sha256(link.strip().encode("utf-8")).hexdigest()[:16] return cache_root / cache_key def _extract_zip_files(folder: Path) -> None: for zip_path in sorted(folder.rglob("*.zip")): extract_dir = zip_path.with_suffix("") extract_dir.mkdir(parents=True, exist_ok=True) with zipfile.ZipFile(zip_path) as archive: for member in archive.infolist(): target = (extract_dir / member.filename).resolve() if not target.is_relative_to(extract_dir.resolve()): raise ValueError(f"Unsafe path in ZIP archive: {member.filename}") archive.extractall(extract_dir) def download_google_drive_folder_images( link: str, cache_root: Path = DRIVE_CACHE_ROOT, ) -> Iterator[DriveDownloadProgress]: cleaned_link = link.strip() if not cleaned_link: yield DriveDownloadProgress("", 0, 0, [], [], done=True) return kind, drive_id = parse_google_drive_link(cleaned_link) if kind != "folder" or not drive_id: yield DriveDownloadProgress( "Google Drive link must be a public folder link. File or ZIP links are not supported here.", 0, 0, [], [], done=True, error="Google Drive link must be a public folder link. File or ZIP links are not supported here.", ) return cache_dir = _cache_dir_for_link(cleaned_link, cache_root) complete_marker = cache_dir / ".download_complete" if complete_marker.exists(): cached_paths = iter_image_paths(cache_dir) if cached_paths: yield DriveDownloadProgress( f"Using cached Google Drive folder ({len(cached_paths)} images found).", len(cached_paths), len(cached_paths), cached_paths, [path.name for path in cached_paths], done=True, ) return if cache_dir.exists(): shutil.rmtree(cache_dir) cache_dir.mkdir(parents=True, exist_ok=True) try: import gdown yield DriveDownloadProgress("Retrieving Google Drive folder contents...", 0, 0, [], []) drive_files = gdown.download_folder( id=drive_id, output=str(cache_dir), quiet=True, use_cookies=False, skip_download=True, ) image_files = [ item for item in drive_files if Path(str(item.local_path)).suffix.lower() in SUPPORTED_IMAGE_EXTENSIONS ] filenames = [Path(str(item.path)).name or Path(str(item.local_path)).name for item in image_files] total = len(image_files) if total == 0: message = "Google Drive folder did not contain any supported image files." yield DriveDownloadProgress(message, 0, 0, [], [], done=True, error=message) return yield DriveDownloadProgress( f"Google Drive folder contains {total} supported images.", 0, total, iter_image_paths(cache_dir), filenames, ) for index, item in enumerate(image_files, start=1): local_path = Path(str(item.local_path)) local_path.parent.mkdir(parents=True, exist_ok=True) display_name = Path(str(item.path)).name or local_path.name yield DriveDownloadProgress( f"Downloading {index} / {total}: {display_name}", index - 1, total, iter_image_paths(cache_dir), filenames, ) gdown.download( id=str(item.id), output=str(local_path), quiet=True, use_cookies=False, ) yield DriveDownloadProgress( f"Downloaded {index} / {total}: {display_name}", index, total, iter_image_paths(cache_dir), filenames, ) image_paths = iter_image_paths(cache_dir) if not image_paths: message = "Google Drive download did not contain any supported image files." yield DriveDownloadProgress(message, 0, total, [], filenames, done=True, error=message) return complete_marker.write_text("ok", encoding="utf-8") yield DriveDownloadProgress( f"Google Drive download complete. {len(image_paths)} images found.", total, total, image_paths, [path.name for path in image_paths] or filenames, done=True, ) except Exception as exc: if cache_dir.exists(): shutil.rmtree(cache_dir) message = f"Google Drive download failed: {exc}" yield DriveDownloadProgress(message, 0, 0, [], [], done=True, error=message) def download_google_drive_images(link: str, cache_root: Path = DRIVE_CACHE_ROOT) -> Tuple[List[Path], str]: final_paths: List[Path] = [] final_error = "" for progress in download_google_drive_folder_images(link, cache_root=cache_root): final_paths = progress.paths final_error = progress.error return final_paths, final_error def clean_drive_downloads(cache_root: Path = DRIVE_CACHE_ROOT) -> None: if cache_root.exists(): shutil.rmtree(cache_root)