Upload dataset_ivf.py with huggingface_hub

dataset_ivf.py

@@ -0,0 +1,75 @@
+# dataset_ivf.py
+import numpy as np, pandas as pd, torch 
+from torch.utils.data import Dataset
+from PIL import Image, ImageFile
+import os
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+class IVFSequenceDataset(Dataset):
+    def __init__(self, df, resize=500, norm="minmax01"):
+        #self.df = pd.read_csv(index_csv)
+        self.df = df
+        self.resize = resize
+        self.norm = norm
+
+    def _read_gray(self, path):
+        img = Image.open(path)
+        if img is None:
+            raise FileNotFoundError(path)
+        # Resize if needed
+        if self.resize is not None:
+            img = img.resize((self.resize, self.resize), Image.BILINEAR)
+        return np.array(img, dtype="float32")
+
+    def _normalize_video(self, vol):
+        if self.norm == "zscore":
+            m, s = vol.mean(), vol.std() + 1e-6
+            vol = (vol - m) / s
+        elif self.norm == "minmax01":
+            lo, hi = np.percentile(vol, 1), np.percentile(vol, 99)
+            vol = (vol - lo) / (hi - lo + 1e-6)
+            vol = np.clip(vol, 0, 1)
+        return vol
+
+    def __getitem__(self, idx):
+        row = self.df.iloc[idx]
+        if pd.isna(row["embryo_paths"]) or pd.isna(row["empty_well_paths"]) or pd.isna(row["sample_paths"]):
+            print(f"Row {idx} has missing path data: ", row.to_string(index = False))
+            raise ValueError(f"Row {idx} has missing path data")
+
+        embryo_paths = row["embryo_paths"].split("|")
+        empty_well_paths = row["empty_well_paths"].split("|")
+        sample_paths = row["sample_paths"].split("|")
+
+        embryo_frames = [self._read_gray(p) for p in embryo_paths]
+        embryo_vol = np.stack(embryo_frames, axis=0)  
+        embryo_vol = self._normalize_video(embryo_vol)
+        embryo_vol = embryo_vol[:,None, :, :] 
+
+        empty_well_frames = [self._read_gray(p) for p in empty_well_paths]
+        empty_well_vol = np.stack(empty_well_frames, axis=0)  
+        empty_well_vol = self._normalize_video(empty_well_vol)
+        empty_well_vol = empty_well_vol[:,None, :, :]        
+
+        sample_frames = [self._read_gray(p) for p in sample_paths]
+        sample_vol = np.stack(sample_frames, axis=0)  
+        sample_vol = self._normalize_video(sample_vol)
+        sample_vol = sample_vol[:,None, :, :]
+        return torch.from_numpy(embryo_vol), torch.from_numpy(empty_well_vol),  torch.from_numpy(sample_vol) 
+
+    def __len__(self):
+        return len(self.df)
+
+if __name__ == "__main__":
+
+    ds = IVFSequenceDataset(pd.read_csv(os.path.abspath("index.csv")), resize=128, norm="minmax01")
+    total_size = len(ds)
+
+    train_size = int(0.85 * total_size)
+    val_size = total_size - train_size
+
+    generator = torch.Generator().manual_seed(42)
+    _, val_set = torch.utils.data.random_split(ds, [train_size, val_size], generator=generator)
+    val_df = ds.df.iloc[val_set.indices]
+    pd.set_option('display.max_rows', None)
+    print(val_df[['cell_id']])
+