Update src\data\plant.py

src/data/plant.py

@@ -0,0 +1,75 @@
+"""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"]