File size: 3,153 Bytes
37960d3 | 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 | """Plant CSV dataset for data/wa_plants/manifest.csv (observation-separated).
Reuses RSNA train_mor pattern: manifest + ImageNet norm, 336px, Aug.
"""
from __future__ import annotations
import csv
from pathlib import Path
import random
from PIL import Image
import warnings
# faster decode: raise limit (100Mpx originals) + draft shrink for large JPEGs
Image.MAX_IMAGE_PIXELS = 300_000_000
warnings.filterwarnings("ignore", category=Image.DecompressionBombWarning)
import torch
from torch.utils.data import Dataset
import torchvision.transforms as T
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)
class PlantDataset(Dataset):
def __init__(self, manifest: Path, split: str, img_size: int = 336, augment: bool = False):
self.split = split
self.img_size = img_size
rows = []
# species -> idx map built from manifest train split (500 classes)
with open(manifest, newline="", encoding="utf-8") as f:
for r in csv.DictReader(f):
if r["split"] == split and r["status"] in ("downloaded","skip_exists"):
rows.append(r)
# stable class order sorted
species = sorted({r["species"] for r in rows})
self.species_to_idx = {s:i for i,s in enumerate(species)}
self.rows = rows
# class counts for balanced sampling
self.augment = augment
tfms = []
if augment:
tfms = [
T.RandomResizedCrop(img_size, scale=(0.7,1.0)),
T.RandomHorizontalFlip(),
T.ColorJitter(0.2,0.2,0.2,0.05),
T.ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD),
]
else:
tfms = [T.Resize(int(img_size*1.14)), T.CenterCrop(img_size), T.ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)]
self.tf = T.Compose(tfms)
def __len__(self): return len(self.rows)
def __getitem__(self, i):
r = self.rows[i]
p = Path(r["path"])
# manifest stores absolute win path; if not found try relative to data/wa_plants
if not p.exists():
# try finding under train/val subfolders by gbifID
base = Path(__file__).resolve().parents[2] / "data" / "wa_plants"
for split in ("train","val"):
cand = base / split / r["species"].replace(" ","_").replace("/","_")[:120] / f"{r['gbifID']}.jpg"
if cand.exists():
p = cand; break
try:
im = Image.open(p)
# fast draft for huge JPEGs (8× shrink before full decode) - no effect on small images
try:
if max(im.size) > 1024:
# draft uses libjpeg shrink 1/2/4/8
im.draft("RGB", (768, 768))
except:
pass
im = im.convert("RGB")
except Exception:
im = Image.new("RGB", (self.img_size, self.img_size))
# species_to_idx may have been remapped after init (train vs val) - lookup safely
y = self.species_to_idx.get(r["species"], 0)
x = self.tf(im)
return x, y, r["gbifID"]
|