Spaces:
Runtime error
Runtime error
File size: 7,918 Bytes
621b022 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 | 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)
|