Spaces:
Runtime error
Runtime error
| 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" | |
| 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) | |