Spaces:
Sleeping
Sleeping
File size: 6,420 Bytes
4687353 | 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 | from __future__ import annotations
import json
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, List, Tuple
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
import torch
from PIL import Image, UnidentifiedImageError
from torch.utils.data import DataLoader, Dataset
def load_json(path: Path) -> Any:
with path.open("r", encoding="utf-8") as handle:
return json.load(handle)
class ChairAttributeDataset(Dataset):
def __init__(
self,
split_path: Path,
label_vocab_path: Path,
repo_root: Path,
image_cache_dir: Path,
) -> None:
payload = load_json(split_path)
if "samples" not in payload:
raise RuntimeError(f"Invalid split file format: {split_path}")
self.samples: List[Dict[str, Any]] = payload["samples"]
self.label_vocab: Dict[str, int] = load_json(label_vocab_path)
self.num_labels = len(self.label_vocab)
self.repo_root = repo_root
self.image_cache_dir = image_cache_dir
self.image_cache_dir.mkdir(parents=True, exist_ok=True)
self.corrupt_image_count = 0
def __len__(self) -> int:
return len(self.samples)
def _resolve_image_path(self, row: Dict[str, Any]) -> Path:
local_rel = row.get("image_path", "")
if local_rel:
local = self.repo_root / local_rel
if local.exists():
return local
image_id = str(row.get("image_id"))
return self.image_cache_dir / f"{image_id}.jpg"
def _candidate_urls(self, row: Dict[str, Any]) -> List[str]:
primary = str(row.get("image_url") or "").strip()
if not primary:
return []
candidates = [primary]
if "/VG_100K/" in primary:
candidates.append(primary.replace("/VG_100K/", "/VG_100K_2/"))
elif "/VG_100K_2/" in primary:
candidates.append(primary.replace("/VG_100K_2/", "/VG_100K/"))
deduped: List[str] = []
seen = set()
for url in candidates:
if url in seen:
continue
seen.add(url)
deduped.append(url)
return deduped
def _download_if_missing(self, row: Dict[str, Any], path: Path, force: bool = False) -> None:
if path.exists() and not force:
return
urls = self._candidate_urls(row)
if not urls:
raise RuntimeError(f"Missing image_url for image_id={row.get('image_id')}")
last_error: Exception | None = None
for url in urls:
try:
request = Request(url, headers={"User-Agent": "Mozilla/5.0"})
with urlopen(request, timeout=30) as response:
image_bytes = response.read()
# Validate bytes are actually decodable image data before caching.
with Image.open(BytesIO(image_bytes)) as probe:
probe.verify()
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("wb") as handle:
handle.write(image_bytes)
return
except (HTTPError, URLError, TimeoutError, OSError, UnidentifiedImageError) as exc:
last_error = exc
continue
raise RuntimeError(
f"Failed to download valid image for image_id={row.get('image_id')} from candidate URLs."
) from last_error
def _open_local_image(self, path: Path) -> Image.Image:
with Image.open(path) as image:
return image.convert("RGB")
def _load_image(self, row: Dict[str, Any]) -> Image.Image:
image_path = self._resolve_image_path(row)
self._download_if_missing(row, image_path)
try:
return self._open_local_image(image_path)
except (UnidentifiedImageError, OSError):
self.corrupt_image_count += 1
if image_path.exists():
image_path.unlink(missing_ok=True)
# Retry once by re-downloading. Some VG image URLs occasionally produce
# partial/bad files; this makes loading resilient instead of crashing workers.
try:
self._download_if_missing(row, image_path, force=True)
return self._open_local_image(image_path)
except Exception:
return Image.new("RGB", (224, 224), color=(0, 0, 0))
def _labels_to_multihot(self, labels: List[str]) -> torch.Tensor:
target = torch.zeros(self.num_labels, dtype=torch.float32)
for label in labels:
if label in self.label_vocab:
target[self.label_vocab[label]] = 1.0
return target
def __getitem__(self, idx: int) -> Dict[str, Any]:
row = self.samples[idx]
image = self._load_image(row)
labels = row.get("attributes_norm", [])
target = self._labels_to_multihot(labels)
return {
"image": image,
"target": target,
"image_id": int(row["image_id"]),
"labels": labels,
}
def collate_samples(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
return {
"images": [item["image"] for item in batch],
"targets": torch.stack([item["target"] for item in batch], dim=0),
"image_ids": [item["image_id"] for item in batch],
"labels": [item["labels"] for item in batch],
}
def build_dataloader(
dataset: Dataset,
batch_size: int,
shuffle: bool,
num_workers: int,
pin_memory: bool,
) -> DataLoader:
return DataLoader(
dataset,
batch_size=batch_size,
shuffle=shuffle,
num_workers=num_workers,
pin_memory=pin_memory,
collate_fn=collate_samples,
)
def class_pos_weights(label_frequencies: Dict[str, int], label_vocab: Dict[str, int]) -> torch.Tensor:
counts = torch.ones(len(label_vocab), dtype=torch.float32)
for label, idx in label_vocab.items():
counts[idx] = float(max(1, label_frequencies.get(label, 1)))
weights = counts.max() / counts
return weights
def split_paths(processed_root: Path, split_dir_name: str = "splits") -> Dict[str, Path]:
split_dir = processed_root / split_dir_name
return {
"train": split_dir / "train.json",
"val": split_dir / "val.json",
"test": split_dir / "test.json",
}
|