Repository: mwalmsley/galaxy-datasets (language: python)

Issue:
Update transforms.py
Includes ElasticTransform,  ChannelWiseDropout,  ShiftScaleRotate instead of RandomResizedCrop, and Flip


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index c3358d8..e5c32ec 100644
--- a/.gitignore
+++ b/.gitignore
@@ -151,4 +151,5 @@ internal_urls.py
 # candels table with urls
 hidden_data/*
 
-*.DS_Store
\ No newline at end of file
+*.DS_Store
+.vscode
\ No newline at end of file
diff --git a/README.md b/README.md
index 14315cc..f871e15 100644
--- a/README.md
+++ b/README.md
@@ -1,21 +1,25 @@
-# pytorch-galaxy-datasets
-PyTorch Datasets and PyTorch Lightning Datamodules for loading images and labels from Galaxy Zoo citizen science campaigns.
-
-| Name      | Class | Published | Downloadable | Galaxies
-| ----------- | ----------- | --- | ---- | ---- |
-| Galaxy Zoo 2 | GZ2 | &#x2611; | &#x2611; | ~210k (main sample) |
-| GZ Hubble   | Hubble | &#x2611; | &#x2611; | ~106k (main sample) |
-| GZ CANDELS   | Candels | &#x2611; | &#x2611; | ~50k |
-| GZ DECaLS GZD-5   | DecalsDR5 | &#x2611; | &#x2611; | ~230k |
-| Galaxy Zoo Rings | Rings | &#x2612; | &#x2611; | ~93k |
-| GZ Legacy Survey  | Legs | &#x2612; | z < 0.1 only | ~375k + 8.3m unlabelled |
-| CFHT Tidal* | Tidal | &#x2611; | &#x2611; | 1760 (expert) |
+# galaxy-datasets
 
-Any datasets marked as downloadable but not marked as published are only downloadable internally (for development purposes).
+ML-friendly datasets for major Galaxy Zoo citizen science campaigns.
 
-If a dataset is published but not marked as downloadable (none currently), it means I haven't yet got around to making the download automatic. You can still download it via the paper instructions.
+- PyTorch Datasets and PyTorch Lightning DataModules
+- TensorFlow tf.data.Dataset's
+- Framework-independent (i.e. TensorFlow-friendly) download and augmentation code
 
 You may also be interested in [Galaxy MNIST](https://github.com/mwalmsley/galaxy_mnist) as a simple dataset for teaching/debugging.
+ 
+
+| Name      | Method | PyTorch Dataset | Published | Downloadable | Galaxies
+| ----------- | ----- | ----------- | --- | ---- | ---- |
+| Galaxy Zoo 2 | gz2 | GZ2 | &#x2611; | &#x2611; | ~210k (main sample) |
+| GZ Hubble   | gz_hubble | GZHubble | &#x2611; | &#x2611; | ~106k (main sample) |
+| GZ CANDELS   | gz_candels | GZCandels | &#x2611; | &#x2611; | ~50k |
+| GZ DECaLS GZD-5 | gz_decals_5 | GZDecals5 | &#x2611; | &#x2611; | ~230k (GZD-5 only)|
+| GZ Rings | gz_rings | GZRings | &#x2612; | &#x2611; | ~93k |
+| GZ DESI  | gz_desi | GZDesi | &#x2612; | z < 0.1 only | ~375k + 8.3m unlabelled |
+| CFHT Tidal* | tidal | Tidal | &#x2611; | &#x2611; | 1760 (expert) |
+
+Any datasets marked as downloadable but not marked as published are only downloadable internally (for development purposes).
 
 For each dataset, you must cite/acknowledge the GZ data release paper and the original telescope survey from which the images were derived. See [data.galaxyzoo.org](data.galaxyzoo.org) for the data release paper citations to use.
 
@@ -23,58 +27,86 @@ For each dataset, you must cite/acknowledge the GZ data release paper and the or
 MW reproduced and modified the images in [Walmsley 2019](https://doi.org/10.1093/mnras/sty3232).
 We include it here as a challenging fine-grained morphology classification task with little labelled data.
 
-### Installation
+## Installation
+
+Installing [zoobot](www.github/mwalmsley/zoobot) will automatically install this package as a dependency.
+
+To install directly:
 
-For local development (e.g. adding a new dataset), you can install this by cloning from github, then running `pip install -e .` in the cloned repo root. 
+- `pip install galaxy-datasets[pytorch]` for PyTorch dependencies
+- `pip install galaxy-datasets[tensorflow]` for TensorFlow dependencies
+- `pip install galaxy-datasets[pytorch,tensorflow]` for both
 
-Note that installing `zoobot` will install this package as a dependency (by automatically running `pip install pytorch-galaxy-datasets`). As with any package) pip will install under your `sitepackages` so you won't be able to make changes easily. 
+For local development (e.g. adding a new dataset), you can install this by cloning from github, then running `pip install -e .` in the cloned repo root. This makes changing the code easier than if you don't use the -e, in which case the package is installed under sitepackages.
 
 I suggest either:
-- For development, installing both `zoobot` and `pytorch-galaxy-datasets` via git
+
 - For basic use without changes, installing `zoobot` via pip and allowing pip to manage this dependency
+- For development, installing both `zoobot` and `galaxy-datasets` via git
+
+## Usage
 
-### Usage
+### Framework-Independent
 
-You can load each prepared dataset as a pytorch Dataset like so:
+To download a dataset:
 
-    from pytorch_galaxy_datasets.prepared_datasets import GZ2Dataset
+    from galaxy_datasets import gz2  # or gz_hubble, gz_candels, ...
 
-    gz2_dataset = GZ2Dataset(
-        root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/gz2',
+    catalog, label_cols = gz2(
+        root='your_data_folder/gz2',
         train=True,
-        download=False
+        download=True
     )
-    image, label = gz2_dataset[0]
-    plt.imshow(image)
 
-You will probably want to customise the dataset, selecting a subset of galaxies or labels. Do this with the `{dataset}_setup()` methods.
+This will download the images and train/test catalogs to `root`. Each `catalog` is a pandas DataFrame with the column `file_loc` giving absolute image paths and additional columns `label_cols = ['col_a', 'col_b', ...]` giving the labels (usually, the number of volunteers who gave each answer for each galaxy). If `train=True`, the method returns the train catalog, otherwise, the test catalog.
 
-    from pytorch_galaxy_datasets.prepared_datasets import gz2_setup
+If training Zoobot from scratch, this is all you need. For example, in PyTorch:
 
-    catalog, label_cols = gz2_setup(
-        root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/gz2',
-        train=True,
-        download=False
+    from zoobot.pytorch.training import train_with_pytorch_lightning
+
+    train_with_pytorch_lightning.train_default_zoobot_from_scratch(
+        catalog=catalog,
+        save_dir=save_dir,
+        schema=gz2_schema, # see zoobot/pytorch/examples/minimal_example.py
+        ...
     )
-    adjusted_catalog = gz2_catalog.sample(1000)
 
-You can then customise the catalog and labels before creating a generic GalaxyDataset, which can be used with your own transforms etc. like any other pytorch dataset
+Otherwise, you might like to use the classes in this package to load these catalogs into ML-friendly inputs.
 
-    from pytorch_galaxy_datasets.galaxy_dataset import GalaxyDataset
+### PyTorch
+
+Create a PyTorch Dataset from a catalog like so:
+
+    from galaxy_datasets.pytorch.galaxy_dataset import GalaxyDataset  # generic Dataset for galaxies
 
     dataset = GalaxyDataset(
-        label_cols=['smooth-or-featured_smooth'],
-        catalog=adjusted_catalog,
-        transforms=some_torchvision_transforms_if_you_like
+        catalog=catalog.sample(1000),  # from gz2(...) above
+        label_cols=['smooth-or-featured-gz2_smooth'],
+        transform=optional_transforms_if_you_like
     )
 
-For training models, I recommend using Pytorch Lightning and GalaxyDataModule, which has default transforms for supervised learning.
+Notice how you can adjust the catalog before creating the Dataset. This gives flexibility to try training on e.g. different catalog subsets.
 
-    from pytorch_galaxy_datasets.galaxy_datamodule import GalaxyDataModule
+If you don't want to change anything about the catalog, you can skip the framework-independent download and use a named class from `galaxy_datasets.pytorch`, which takes the same arguments and directly gives a Dataset:
+
+    from galaxy_datasets.pytorch import GZ2
+
+    gz2_dataset = GZ2(
+        root='your_data_folder/gz2',
+        train=True,
+        download=False
+    )
+    image, label = gz2_dataset[0]
+    plt.imshow(image)
+
+You might also find the PyTorch Lightning DataModule under `galaxy_datasets/pytorch/galaxy_datamodule` useful. Zoobot uses this for training and finetuning.
+
+    from galaxy_datasets.pytorch.galaxy_datamodule import GalaxyDataModule
 
     datamodule = GalaxyDataModule(
-        label_cols=['smooth-or-featured_smooth'],
-        catalog=adjusted_catalog
+        label_cols=['smooth-or-featured-gz2_smooth'],
+        catalog=catalog
+        # optional args to specify augmentations
     )
 
     datamodule.prepare_data()
@@ -83,12 +115,34 @@ For training models, I recommend using Pytorch Lightning and GalaxyDataModule, w
         print(images.shape, labels.shape)
         break
 
-You can also get the canonical catalog and label_cols from the Dataset, if you prefer.
+### TensorFlow
+
+To create a tf.data.Dataset from a catalog:
 
-    gz2_catalog = gz2_dataset.catalog
-    gz2_label_cols = gz2_dataset.label_cols
+    import tensorflow as tf
+    from galaxy_datasets.tensorflow.datasets import get_image_dataset, add_transforms_to_dataset
+    from galaxy_datasets.transforms import default_transforms  # same transforms as PyTorch
 
-### Download Notes
+    train_dataset = get_image_dataset(
+        image_paths = catalog['file_loc'],
+        labels=catalog[label_cols].values,
+        requested_img_size=224
+    )
+
+    # specify augmentations
+    transforms = default_transforms()
+
+    # apply augmentations
+    train_dataset = add_transforms_to_dataset(train_dataset, transforms)
+  
+    # batch, shuffle, prefetch for performance
+    train_dataset = train_dataset.shuffle(5000).batch(64).prefetch(tf.data.experimental.AUTOTUNE)
+
+    for images, labels in train_dataset.take(1):
+        print(images.shape, labels.shape)
+        break
+
+## Download Notes
 
 Datasets are downloaded like:
 
@@ -96,4 +150,6 @@ Datasets are downloaded like:
     - images
         - subfolder (except GZ2)
             - image.jpg
-    - {catalog_name(s)}.parquet
\ No newline at end of file
+    - {catalog_name(s)}.parquet
+
+The whole dataset is downloaded regardless of whether `train=True` or `train=False`.
diff --git a/pytorch_galaxy_datasets/__init__.py b/derived_data/.gitkeep
similarity index 100%
rename from pytorch_galaxy_datasets/__init__.py
rename to derived_data/.gitkeep
diff --git a/galaxy_datasets/__init__.py b/galaxy_datasets/__init__.py
new file mode 100644
index 0000000..7cc7406
--- /dev/null
+++ b/galaxy_datasets/__init__.py
@@ -0,0 +1,10 @@
+from galaxy_datasets.shared.gz_candels import gz_candels
+from galaxy_datasets.shared.gz_decals_5 import gz_decals_5
+from galaxy_datasets.shared.gz_hubble import gz_hubble
+from galaxy_datasets.shared.gz2 import gz2
+from galaxy_datasets.shared.tidal import tidal
+
+from galaxy_datasets.check_internal_urls import INTERNAL_URLS_EXIST
+if INTERNAL_URLS_EXIST:
+    from galaxy_datasets.shared.gz_desi import gz_desi
+    from galaxy_datasets.shared.gz_rings import gz_rings
diff --git a/galaxy_datasets/check_internal_urls.py b/galaxy_datasets/check_internal_urls.py
new file mode 100644
index 0000000..8d695d2
--- /dev/null
+++ b/galaxy_datasets/check_internal_urls.py
@@ -0,0 +1,4 @@
+import os
+
+_expected_internal_urls_loc = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'shared/internal_urls.py')
+INTERNAL_URLS_EXIST = os.path.isfile(_expected_internal_urls_loc)
\ No newline at end of file
diff --git a/pytorch_galaxy_datasets/deprecated/decals_dr8.py b/galaxy_datasets/deprecated/decals_dr8.py
similarity index 100%
rename from pytorch_galaxy_datasets/deprecated/decals_dr8.py
rename to galaxy_datasets/deprecated/decals_dr8.py
diff --git a/pytorch_galaxy_datasets/deprecated/gz2_curated.py b/galaxy_datasets/deprecated/gz2_curated.py
similarity index 100%
rename from pytorch_galaxy_datasets/deprecated/gz2_curated.py
rename to galaxy_datasets/deprecated/gz2_curated.py
diff --git a/galaxy_datasets/pytorch/__init__.py b/galaxy_datasets/pytorch/__init__.py
new file mode 100644
index 0000000..838415e
--- /dev/null
+++ b/galaxy_datasets/pytorch/__init__.py
@@ -0,0 +1,7 @@
+from galaxy_datasets.pytorch.datasets import GZCandels, GZDecals5, GZHubble, GZ2, Tidal
+
+try:
+    from galaxy_datasets.pytorch.datasets import GZDesi, GZRings
+except ImportError:
+    # not using logging in case config still required
+    print('GZDESI and GZRings not available from galaxy_datasets.pytorch.datasets - skipping')    
diff --git a/galaxy_datasets/pytorch/datasets.py b/galaxy_datasets/pytorch/datasets.py
new file mode 100644
index 0000000..82b308f
--- /dev/null
+++ b/galaxy_datasets/pytorch/datasets.py
@@ -0,0 +1,87 @@
+from galaxy_datasets.shared import gz2, gz_candels, gz_decals_5, gz_hubble, tidal
+
+from galaxy_datasets.pytorch import galaxy_dataset
+
+
+# TODO could refactor these into same class if needed
+
+class GZCandels(galaxy_dataset.GalaxyDataset):
+    
+    def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
+
+        catalog, label_cols = gz_candels(root, train, download)
+
+        super().__init__(catalog, label_cols, transform, target_transform)
+
+
+class GZDecals5(galaxy_dataset.GalaxyDataset):
+    
+    def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
+
+        catalog, label_cols = gz_decals_5(root, train, download)
+
+        super().__init__(catalog, label_cols, transform, target_transform)
+
+
+
+class GZ2(galaxy_dataset.GalaxyDataset):
+    
+    def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
+
+        catalog, label_cols = gz2(root, train, download)  # no train arg
+
+        super().__init__(catalog, label_cols, transform, target_transform)
+
+
+class GZHubble(galaxy_dataset.GalaxyDataset):
+    
+    def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
+
+        catalog, label_cols = gz_hubble(root, train, download)
+
+        super().__init__(catalog, label_cols, transform, target_transform)
+
+class Tidal(galaxy_dataset.GalaxyDataset):
+    
+    def __init__(self, root, train=True, download=False, transform=None, target_transform=None, label_mode='coarse'):
+
+        catalog, label_cols = tidal(root, train, download, label_mode=label_mode)
+
+        super().__init__(catalog, label_cols, transform, target_transform)
+
+from galaxy_datasets import check_internal_urls
+if check_internal_urls.INTERNAL_URLS_EXIST:
+    from galaxy_datasets.shared import gz_desi, gz_rings
+
+    class GZDesi(galaxy_dataset.GalaxyDataset):
+        
+        def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
+
+            catalog, label_cols = gz_desi(root, train, download)
+
+            super().__init__(catalog, label_cols, transform, target_transform)
+
+    class GZRings(galaxy_dataset.GalaxyDataset):
+        
+        def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
+
+            catalog, label_cols = gz_rings(root, train, download)
+
+            super().__init__(catalog, label_cols, transform, target_transform)
+
+
+
+# temporarily deprecated
+# class Legs(galaxy_dataset.GalaxyDataset):
+    
+#     # based on https://pytorch.org/vision/stable/generated/torchvision.datasets.STL10.html
+#     def __init__(self, root=None, split='train', download=False, transform=None, target_transform=None, train=None):
+#         # train=None is just an exception-raising parameter to avoid confused users using the train=False api
+
+#         catalog, label_cols = legs(root, split, download, train)
+
+#         # paths are not adjusted as cannot be downloaded
+#         # catalog = _temp_adjust_catalog_paths(catalog)
+#         # catalog = adjust_catalog_dtypes(catalog, label_cols)
+
+#         super().__init__(catalog, label_cols, transform, target_transform)
diff --git a/pytorch_galaxy_datasets/galaxy_datamodule.py b/galaxy_datasets/pytorch/galaxy_datamodule.py
similarity index 78%
rename from pytorch_galaxy_datasets/galaxy_datamodule.py
rename to galaxy_datasets/pytorch/galaxy_datamodule.py
index fabaad9..66bc21c 100644
--- a/pytorch_galaxy_datasets/galaxy_datamodule.py
+++ b/galaxy_datasets/pytorch/galaxy_datamodule.py
@@ -1,6 +1,7 @@
 
 from typing import Optional
 import logging
+from functools import partial
 
 import numpy as np
 from sklearn.model_selection import train_test_split
@@ -9,10 +10,8 @@
 import pytorch_lightning as pl
 from torchvision import transforms
 
-import albumentations as A
-from albumentations.pytorch import ToTensorV2
-
-from pytorch_galaxy_datasets import galaxy_dataset
+from galaxy_datasets.pytorch import galaxy_dataset
+from galaxy_datasets.transforms import default_transforms
 
 
 # https://pytorch-lightning.readthedocs.io/en/stable/extensions/datamodules.html
@@ -35,10 +34,10 @@ def __init__(
         predict_catalog=None,
         # augmentation params (sensible supervised defaults)
         greyscale=True,
-        album=False,
-        resize_size=224,
+        # album=False,  # now True always
         crop_scale_bounds=(0.7, 0.8),
         crop_ratio_bounds=(0.9, 1.1),
+        resize_after_crop=224,
         # hardware params
         batch_size=256,  # careful - will affect final performance
         use_memory=False,  # deprecated
@@ -66,7 +65,7 @@ def __init__(
 
         self.batch_size = batch_size
 
-        self.resize_size = resize_size
+        self.resize_after_crop = resize_after_crop
         self.crop_scale_bounds = crop_scale_bounds
         self.crop_ratio_bounds = crop_ratio_bounds
 
@@ -83,14 +82,14 @@ def __init__(
         self.test_fraction = test_fraction
 
         self.greyscale = greyscale
-        self.album = album
+        # self.album = album
 
-        if self.album:
-            logging.info('Using albumentations for augmentations')
-            self.transform_with_album()
-        else:
-            logging.info('Using torchvision for augmentations')
-            self.transform_with_torchvision()
+        # if self.album:
+        logging.info('Using albumentations for augmentations')
+        self.transform_with_album()
+        # else:
+        #     logging.info('Using torchvision for augmentations')
+        #     self.transform_with_torchvision()
 
         self.prefetch_factor = prefetch_factor
         self.dataloader_timeout = 240  # seconds
@@ -99,42 +98,25 @@ def __init__(
         logging.info('Prefetch factor: {}'.format(self.prefetch_factor))
 
     def transform_with_torchvision(self):
-
-        transforms_to_apply = default_torchvision_transforms(self.greyscale, self.resize_size, self.crop_scale_bounds, self.crop_ratio_bounds)
-
-        self.transform = transforms.Compose(transforms_to_apply)
+        raise NotImplementedError('Deprecated in favor of albumentations')
+        # transforms_to_apply = default_torchvision_transforms(self.greyscale, self.resize_size, self.crop_scale_bounds, self.crop_ratio_bounds)
+        # self.transform = transforms.Compose(transforms_to_apply)
 
 
     def transform_with_album(self):
-
-        if self.greyscale:
-            transforms_to_apply = [A.Lambda(name='ToGray', image=ToGray(
-                reduce_channels=True), always_apply=True)]
-        else:
-            transforms_to_apply = []
-
-            transforms_to_apply += [
-                A.ToFloat(),
-                # anything outside of the original image is set to 0.
-                A.Rotate(limit=180, interpolation=1,
-                         always_apply=True, border_mode=0, value=0),
-                A.RandomResizedCrop(
-                    height=self.resize_size,  # after crop resize
-                    width=self.resize_size,
-                    scale=self.crop_scale_bounds,  # crop factor
-                    ratio=self.crop_ratio_bounds,  # crop aspect ratio
-                    interpolation=1,  # This is "INTER_LINEAR" == BILINEAR interpolation. See: https://docs.opencv.org/3.4/da/d54/group__imgproc__transform.html
-                    always_apply=True
-                ),  # new aspect ratio
-                A.VerticalFlip(p=0.5),
-                ToTensorV2()
-            ]
-
-        albumentations_transform = A.Compose(transforms_to_apply)
-
-        # warning - might need a transpose check
+        # gives a transforms = Compose() object
+        transforms_to_apply = default_transforms(
+            crop_scale_bounds=self.crop_scale_bounds,
+            crop_ratio_bounds=self.crop_ratio_bounds,
+            resize_after_crop=self.resize_after_crop,
+            pytorch_greyscale=self.greyscale
+        )
+        # applies that transforms object
         # albumentations expects np array, and returns dict keyed by "image"
-        self.transform = lambda img: albumentations_transform(image=np.array(img))["image"]
+        # transpose changes from BHWC (numpy/TF style) to BCHW (torch style) 
+
+        # cannot use a lambda or define here because must be pickleable for multi-gpu
+        self.transform = partial(do_transform, transforms_to_apply=transforms_to_apply)
 
 
     # only called on main process
@@ -232,6 +214,9 @@ def default_torchvision_transforms(greyscale, resize_size, crop_scale_bounds, cr
     
     return transforms_to_apply
 
+def do_transform(img, transforms_to_apply):
+    return np.transpose(transforms_to_apply(image=np.array(img))["image"], axes=[2, 0, 1]).astype(np.float32)
+
 # torchvision
 class GrayscaleUnweighted(torch.nn.Module):
 
@@ -256,15 +241,3 @@ def __repr__(self):
         return self.__class__.__name__ + '(num_output_channels={0})'.format(self.num_output_channels)
 
 
-# albumentations versuib of GrayscaleUnweighted
-class ToGray():
-
-    def __init__(self, reduce_channels=False):
-        if reduce_channels:
-            self.mean = lambda arr: arr.mean(axis=2, keepdims=True)
-        else:
-            self.mean = lambda arr: arr.mean(
-                axis=2, keepdims=True).repeat(3, axis=2)
-
-    def __call__(self, image, **kwargs):
-        return self.mean(image)
diff --git a/galaxy_datasets/pytorch/galaxy_dataset.py b/galaxy_datasets/pytorch/galaxy_dataset.py
new file mode 100644
index 0000000..71c1755
--- /dev/null
+++ b/galaxy_datasets/pytorch/galaxy_dataset.py
@@ -0,0 +1,122 @@
+from typing import List
+import logging
+
+import numpy as np
+import pandas as pd
+from torch.utils.data import Dataset
+from PIL import Image
+import simplejpeg
+
+
+
+# https://pytorch.org/tutorials/beginner/basics/data_tutorial.html
+class GalaxyDataset(Dataset):
+    def __init__(self, catalog: pd.DataFrame, label_cols=None, transform=None, target_transform=None):
+        """
+        Create custom PyTorch Dataset using catalog of galaxy images
+
+        Catalog and images should already be downloaded.
+        Use methods under galaxy_datasets.prepared_datasets e.g. prepared_datasets.gz2()
+        
+        Note that if you want the canonical datasets (i.e. using the standard train/test split)
+        you can use the galaxy_datasets.pytorch.datasets instead of this class e.g. 
+
+        ```
+        datasets.GZ2Dataset(
+            root=os.path.join(some_dir, 'gz2'),
+            download=True
+        )
+        ```
+        
+        Reads a catalog of jpeg galaxy images
+        if transform is from albumentations, datamodule should know about the transform including 
+        catalog should be split already
+        should have correct image locations under file_loc
+
+        Args:
+            catalog (pd.DataFrame): with images under 'file_loc' and labels under label_cols
+            label_cols (List): columns in catalog with labels
+            transform (callable, optional): See Pytorch Datasets. Defaults to None.
+            target_transform (callable, optional): See Pytorch Datasets. Defaults to None.
+        """
+        # internal catalog has id_str as index for query speed
+        # (dataset.catalog still returns the int-indexed catalog via the catalog method below)
+        # self._catalog = catalog.copy().set_index('id_str', verify_integrity=True)
+        self.catalog = catalog
+        
+        self.label_cols = label_cols
+        self.transform = transform
+        self.target_transform = target_transform
+
+        image_format = self.catalog['file_loc'].iloc[0].split('.')[-1]
+        assert all(self.catalog['file_loc'].str.endswith(image_format)), 'Catalog file extensions not consistent. All image files must use the same extension (e.g. .jpg, .png, etc'
+        # technically would still work, but bad practice
+
+        if image_format == 'png':
+            logging.warning('Loading .png images - this works but is probably slower and more storage-intensive than .jpg images')
+            self.load_image_file = load_png_file
+        elif image_format == 'jpg':
+            self.load_image_file = load_jpg_file
+        else:
+            raise ValueError('File format {} not recognised - should be jpg (preferred) or png'.format(image_format))
+
+
+    def __len__(self) -> int:
+        return len(self.catalog)
+
+
+    def __getitem__(self, idx: int):
+        #the index is id_str so can use that for quick search on 1M+ catalo
+        # galaxy = self._catalog.loc[idx]
+        galaxy = self.catalog.iloc[idx]
+
+        # load the image into memory
+        image_loc = galaxy['file_loc']
+        try:
+            image = self.load_image_file(image_loc)
+            # HWC PIL image
+            # logging.info((image.shape, torch.max(image), image.dtype, label))  # always 0-255 uint8
+        except Exception as e:
+            logging.critical('Cannot load {}'.format(image_loc))
+            raise e
+    
+        if self.transform:
+            image = self.transform(image)
+            # now a CHW tensor, which torchvision wants
+            # logging.info((image.shape, torch.max(image), image.dtype, label))  #  should be 0-1 float
+        # else:
+        #     raise ValueError('Please specify a transform')
+        #     # TODO could use some kind of default transform for torch-friendly image, if needed
+
+        if self.label_cols is None:
+            return image
+        else:
+            # load the labels. If no self.label_cols, will 
+            label = get_galaxy_label(galaxy, self.label_cols)
+
+            if self.target_transform:
+                label = self.target_transform(label)
+
+            
+            return image, label
+
+
+def load_jpg_file(loc):
+    with open(loc, 'rb') as f:
+        return Image.fromarray(decode_jpeg(f.read()))  # HWC PIL image via simplejpeg
+
+def load_png_file(loc):
+    return Image.open(loc)
+
+# def load_encoded_jpeg(loc: str):
+#     with open(loc, "rb") as f:
+#         return f.read()  # bytes, not yet decoded
+
+
+def decode_jpeg(encoded_bytes):
+    return simplejpeg.decode_jpeg(encoded_bytes, fastdct=True, fastupsample=True)
+
+
+def get_galaxy_label(galaxy: pd.Series, label_cols: List) -> np.ndarray:
+    # no longer casts to int64, user now responsible in df. If dtype is mixed, will try to infer with infer_objects
+    return galaxy[label_cols].infer_objects().values.squeeze()  # squeeze for if there's one label_col
diff --git a/galaxy_datasets/shared/__init__.py b/galaxy_datasets/shared/__init__.py
new file mode 100644
index 0000000..8ec6c66
--- /dev/null
+++ b/galaxy_datasets/shared/__init__.py
@@ -0,0 +1,10 @@
+from galaxy_datasets.shared.gz2 import gz2
+from galaxy_datasets.shared.gz_candels import gz_candels
+from galaxy_datasets.shared.gz_decals_5 import gz_decals_5
+from galaxy_datasets.shared.gz_hubble import gz_hubble
+from galaxy_datasets.shared.tidal import tidal
+
+from galaxy_datasets import check_internal_urls
+if check_internal_urls.INTERNAL_URLS_EXIST:
+    from galaxy_datasets.shared.gz_desi import gz_desi
+    from galaxy_datasets.shared.gz_rings import gz_rings
diff --git a/pytorch_galaxy_datasets/prepared_datasets/legs.py b/galaxy_datasets/shared/desi_low_z.py
similarity index 65%
rename from pytorch_galaxy_datasets/prepared_datasets/legs.py
rename to galaxy_datasets/shared/desi_low_z.py
index 8bdbb07..e08a0c8 100644
--- a/pytorch_galaxy_datasets/prepared_datasets/legs.py
+++ b/galaxy_datasets/shared/desi_low_z.py
@@ -2,29 +2,24 @@
 import os
 
 import pandas as pd
-from pytorch_galaxy_datasets import galaxy_datamodule, galaxy_dataset, download_utils
-from pytorch_galaxy_datasets.prepared_datasets import internal_urls
+
+from galaxy_datasets.shared import download_utils
+from galaxy_datasets.check_internal_urls import INTERNAL_URLS_EXIST
+if not INTERNAL_URLS_EXIST:
+    raise FileNotFoundError
+from galaxy_datasets.shared import internal_urls
 
 # TODO could eventually refactor this out of Zoobot as well
 from zoobot.shared import label_metadata
 
+"""
+Downloads DESI galaxies below z<0.1
+Includes GZ DESI labels where available (identical split to )
 
-class LegsDataset(galaxy_dataset.GalaxyDataset):
-    
-    # based on https://pytorch.org/vision/stable/generated/torchvision.datasets.STL10.html
-    def __init__(self, root=None, split='train', download=False, transform=None, target_transform=None, train=None):
-        # train=None is just an exception-raising parameter to avoid confused users using the train=False api
-
-        catalog, label_cols = legs_setup(root, split, download, train)
-
-        # paths are not adjusted as cannot be downloaded
-        # catalog = _temp_adjust_catalog_paths(catalog)
-        # catalog = adjust_catalog_dtypes(catalog, label_cols)
-
-        super().__init__(catalog, label_cols, transform, target_transform)
+Intended for active learning experiments simulated using GZ DESI
+"""
 
-
-def legs_setup(root=None, split='train', download=False, train=None):
+def desi_low_z(root=None, split='train', download=False, train=None):
 
     if train is not None:
         raise ValueError("This dataset has unlabelled data: use split='train', 'test', 'unlabelled' or 'train+unlabelled' rather than train=False etc")
@@ -35,10 +30,15 @@ def legs_setup(root=None, split='train', download=False, train=None):
         # 'Legacy Survey cannot be downloaded - ignoring root {}'.format(root)
         # TODO update for non-manchester users with a manual copy?
 
+
+    
+
+
+
     # resources = (
-    #     (internal_urls.legs_train_catalog, 'bae2906e337bd114af013d02f3782473'),
-    #     (internal_urls.legs_test_catalog, '20919fe512ee8ce4d267790e519fcbf8'),
-    #     (internal_urls.legs_unlabelled_catalog, 'fbf287990add34d2249f584325bc9dca'),
+    #     (internal_urls.legs_train_catalog, '2364a714a1339f020587d374c1838418'),
+    #     (internal_urls.legs_test_catalog, '647bd46f53a06f13eb4df25311ccb9d3'),
+    #     (internal_urls.legs_unlabelled_catalog, '37c9f07b5c058f55d515d2df08ff132a'),
     #     # and the images, split into 8gb chunks
     #     (internal_urls.legs_chunk_00, 'd6ca1051b3dd48cfc5c7f0535b403b2d'),
     #     (internal_urls.legs_chunk_01, 'f258ab647cee076ca66288e25f4a778d'),
@@ -50,7 +50,7 @@ def legs_setup(root=None, split='train', download=False, train=None):
     #     (internal_urls.legs_chunk_07, '583d92b917bd70670d7860e3836cb4a4')
     # )
     resources = (
-        (internal_urls.legs_train_catalog, 'bae2906e337bd114af013d02f3782473'),
+        (internal_urls.legs_train_catalog, '2364a714a1339f020587d374c1838418'),
         (internal_urls.legs_test_catalog, None),
         (internal_urls.legs_unlabelled_catalog, None),
         # and the images, split into 8gb chunks
@@ -105,32 +105,3 @@ def legs_setup(root=None, split='train', download=False, train=None):
 
 
     return catalog, label_cols
-
-
-
-
-if __name__ == '__main__':
-
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    legs_dataset = LegsDataset(
-        root='whatever',
-        train=True,
-        download=False  # will fail except on galahad
-    )
-    legs_catalog = legs_dataset.catalog
-    adjusted_catalog = legs_catalog.sample(1000)
-
-    # user will probably tweak and use images/catalog directly for generic galaxy catalog datamodule
-    # (which makes its own generic datasets internally)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=label_metadata.decals_all_campaigns_label_cols,
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-        
diff --git a/pytorch_galaxy_datasets/download_utils.py b/galaxy_datasets/shared/download_utils.py
similarity index 72%
rename from pytorch_galaxy_datasets/download_utils.py
rename to galaxy_datasets/shared/download_utils.py
index 39f0676..9c81690 100644
--- a/pytorch_galaxy_datasets/download_utils.py
+++ b/galaxy_datasets/shared/download_utils.py
@@ -2,6 +2,7 @@
 import logging
 
 from urllib.error import URLError
+# TODO this is the only pytorch dependency, could swap out perhaps
 from torchvision.datasets.utils import download_and_extract_archive, download_url, check_integrity
 
 
@@ -9,12 +10,15 @@ class DatasetDownloader():
     # responsible for downloading a prespecified set of images/catalogs to a directory
     # supports GalaxyDataset via composition
 
-    def __init__(self, root, resources, images_to_spotcheck=None, image_dirname='images'):
-        # image_dirname should always be imagess; is not properly generalised to update the extract location
+    def __init__(self, root, resources, images_to_spotcheck=None, image_dirname='images', archive_includes_subdir=True):
+        # image_dirname should always be images; is not properly generalised to update the extract location
         self.root = root
         self.image_dir = os.path.join(self.root, image_dirname)
+        if not os.path.exists(self.image_dir):
+            os.mkdir(self.image_dir)
         self.resources = resources
         self.images_to_spotcheck = images_to_spotcheck
+        self.archive_includes_subdir = archive_includes_subdir
 
     def download(self) -> None:
         """Download the data if it doesn't exist already."""
@@ -30,9 +34,14 @@ def download(self) -> None:
             try:
                 logging.info(f"Downloading {url}")
                 if url.endswith('.tar.gz') or url.endswith('.zip'):
+                    if self.archive_includes_subdir:
+                        download_root = self.root  # archive unpacks to subdir by itself
+                    else:  # need to download and extract directly into subdir
+                        download_root = self.image_dir
                     download_and_extract_archive(
-                        url, download_root=self.root, filename=filename, md5=md5)
+                        url, download_root=download_root, filename=filename, md5=md5)
                 else:  # don't try to extract archive, just download
+                    # TODO is this used?
                     download_url(url, root=self.root, filename=filename, md5=md5)
             except URLError as error:
                 logging.info(f"Failed to download (trying next):\n{error}")
diff --git a/pytorch_galaxy_datasets/prepared_datasets/gz2.py b/galaxy_datasets/shared/gz2.py
similarity index 51%
rename from pytorch_galaxy_datasets/prepared_datasets/gz2.py
rename to galaxy_datasets/shared/gz2.py
index f426577..bd576c2 100644
--- a/pytorch_galaxy_datasets/prepared_datasets/gz2.py
+++ b/galaxy_datasets/shared/gz2.py
@@ -1,23 +1,14 @@
 import os
 
 import pandas as pd
-from pytorch_galaxy_datasets import galaxy_datamodule, galaxy_dataset, download_utils
 
 # TODO could eventually refactor this out of Zoobot as well
 from zoobot.shared import label_metadata
 
+from galaxy_datasets.shared import download_utils
 
-# TODO add train flag? currently full dataset
-class GZ2Dataset(galaxy_dataset.GalaxyDataset):
-    
-    def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
 
-        catalog, label_cols = gz2_setup(root, train, download)  # no train arg
-
-        super().__init__(catalog, label_cols, transform, target_transform)
-
-
-def gz2_setup(root, train, download):
+def gz2(root, train, download):
     resources = [
         ('https://dl.dropboxusercontent.com/s/vu77e3sh2s5c250/gz2_train_catalog.parquet', 'f489c9ec7dcf8d99f728bd00ee00b1d0'),  # the train catalog
         ('https://dl.dropboxusercontent.com/s/t8eh6f3oupndpl3/gz2_test_catalog.parquet', '8b2d74c93d35f86cc34f1d058b3b220b'),  # the test catalog
@@ -44,32 +35,3 @@ def gz2_setup(root, train, download):
 
     label_cols = label_metadata.gz2_ortho_label_cols  # default, but you can ignore
     return catalog, label_cols
-
-
-if __name__ == '__main__':
-
-    # can use this directly, e.g. to visualise the dataset
-    gz2_dataset = GZ2Dataset(
-        root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/gz2',
-        download=False
-    )
-    gz2_catalog = gz2_dataset.catalog  # and can get the catalog/cols for tweaking
-    gz2_label_cols = gz2_dataset.label_cols
-
-    # or, can use the setup method directly
-    catalog, label_cols = gz2_setup(
-        root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/gz2',
-        download=False
-    )
-    adjusted_catalog = gz2_catalog.sample(1000)
-
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=['smooth-or-featured_smooth'],
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
\ No newline at end of file
diff --git a/pytorch_galaxy_datasets/prepared_datasets/candels.py b/galaxy_datasets/shared/gz_candels.py
similarity index 76%
rename from pytorch_galaxy_datasets/prepared_datasets/candels.py
rename to galaxy_datasets/shared/gz_candels.py
index c314b66..0cf59b5 100644
--- a/pytorch_galaxy_datasets/prepared_datasets/candels.py
+++ b/galaxy_datasets/shared/gz_candels.py
@@ -2,23 +2,16 @@
 import os
 
 import pandas as pd
-from pytorch_galaxy_datasets import galaxy_dataset, galaxy_datamodule, download_utils
 
 # this one has label_metadata contained in this repo, not zoobot. Others will follow, perhaps.
 from zoobot.shared import label_metadata
 
+from galaxy_datasets.shared import download_utils
 
 
-class CandelsDataset(galaxy_dataset.GalaxyDataset):
-    
-    def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
 
-        catalog, label_cols = candels_setup(root, train, download)
 
-        super().__init__(catalog, label_cols, transform, target_transform)
-
-
-def candels_setup(root, train, download):
+def gz_candels(root, train, download):
     resources = [
         ('https://dl.dropboxusercontent.com/s/cnjvdinnhh1r1md/candels_ortho_train_catalog.parquet', '90593d1bab79a608cf0e645d6fd8e741'),  # train catalog
         ('https://dl.dropboxusercontent.com/s/y83v1gktw72hs0f/candels_ortho_test_catalog.parquet', '1062993dd8df09684b335678ab3fa8e3'),  # test catalog
@@ -91,28 +84,3 @@ def candels_setup(root, train, download):
 }
 
 candels_ortho_questions, candels_ortho_label_cols = label_metadata.extract_questions_and_label_cols(candels_ortho_pairs)
-
-
-if __name__ == '__main__':
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    candels_label_cols, candels_catalog = candels_setup(
-        root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/candels',
-        train=True,
-        download=False
-    )
-    
-    # user will probably tweak and use images/catalog directly for generic galaxy catalog datamodule
-    # (which makes its own generic datasets internally)
-    adjusted_catalog = candels_catalog.sample(1000)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=candels_label_cols,
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-    
\ No newline at end of file
diff --git a/pytorch_galaxy_datasets/prepared_datasets/dr5.py b/galaxy_datasets/shared/gz_decals_5.py
similarity index 51%
rename from pytorch_galaxy_datasets/prepared_datasets/dr5.py
rename to galaxy_datasets/shared/gz_decals_5.py
index fea509f..eabbdd9 100644
--- a/pytorch_galaxy_datasets/prepared_datasets/dr5.py
+++ b/galaxy_datasets/shared/gz_decals_5.py
@@ -1,24 +1,15 @@
 import os
 
 import pandas as pd
-from pytorch_galaxy_datasets import galaxy_datamodule, galaxy_dataset, download_utils
 
 # TODO could eventually refactor this out of Zoobot as well
 from zoobot.shared import label_metadata
 
-# DR8 will be basically the same
-
-
-class DecalsDR5Dataset(galaxy_dataset.GalaxyDataset):
-    
-    def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
-
-        catalog, label_cols = decals_dr5_setup(root, train, download)
-
-        super().__init__(catalog, label_cols, transform, target_transform)
+from galaxy_datasets.shared import download_utils
 
+# DR8 will be basically the same
 
-def decals_dr5_setup(root, train, download):
+def gz_decals_5(root, train, download):
     resources = [
         ('https://dl.dropboxusercontent.com/s/1tuehajonhgv8a2/decals_dr5_ortho_train_catalog.parquet', 'a0cd74edc073fdff068370f6eefeb802'),  # train catalog
         ('https://dl.dropboxusercontent.com/s/3vo6hjlnbqgzuxz/decals_dr5_ortho_test_catalog.parquet', '55820e3712b22e587f6971e4b6c73dfe'),  # test catalog
@@ -54,40 +45,3 @@ def decals_dr5_setup(root, train, download):
 #     for answer_col in label_cols:
 #         catalog[answer_col] = catalog[answer_col].astype(int)
 #     return catalog
-
-
-if __name__ == '__main__':
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    dr5_dataset = DecalsDR5Dataset(
-        root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/decals_dr5',
-        train=True,
-        download=False
-    )
-    dr5_catalog = dr5_dataset.catalog
-    adjusted_catalog = dr5_catalog.sample(1000)
-
-    # user will probably tweak and use images/catalog directly for generic galaxy catalog datamodule
-    # (which makes its own generic datasets internally)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=label_metadata.decals_dr5_ortho_label_cols,
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-        
-
-# class DECaLSDR5DataModule(galaxy_datamodule.GalaxyDataModule):
-
-#     def __init__(self, *args, **kwargs) -> None:
-#         """
-#         Currently identical to GalaxyDataModule - see that description
-#         """
-#         super().__init__(*args, dataset_class=DECaLSDR5Dataset, **kwargs)
-
-#     def prepare_data(self):
-#         DECaLSDR5Dataset(self.root, download=True)
diff --git a/galaxy_datasets/shared/gz_desi.py b/galaxy_datasets/shared/gz_desi.py
new file mode 100644
index 0000000..82506da
--- /dev/null
+++ b/galaxy_datasets/shared/gz_desi.py
@@ -0,0 +1,58 @@
+import os
+
+import pandas as pd
+
+# TODO could eventually refactor this out of Zoobot as well
+from zoobot.shared import label_metadata
+
+from galaxy_datasets.shared import download_utils
+from galaxy_datasets.check_internal_urls import INTERNAL_URLS_EXIST
+if not INTERNAL_URLS_EXIST:
+    raise FileNotFoundError
+from galaxy_datasets.shared import internal_urls
+
+"""
+This downloads the *labelled* GZ DESI galaxies, at all redshifts, with no division between data releases (except label_cols)
+useful to conveniently download galaxies for GZ DESI itself (see gz-decals-classifiers, though this also uses catalogs directly)
+or for other large-scale supervised learning problems or for pretraining ahead of active learning
+does not include unlabelled galaxies (for e.g. contrastive learning, active learning)
+
+Added to make GZ DESI easily downloadable, basically
+"""
+
+def gz_desi(root, train, download):
+    resources = [
+        (internal_urls.gz_desi_train_catalog, '50997398a10886dafbe556d071ffbc68'),
+        (internal_urls.gz_desi_test_catalog, '386f13daf123f439ffb8266b65eac474'),
+        (internal_urls.gz_desi_images_chunk_00, '1d52898581bfb8f08acd13fe77c69901'),
+        (internal_urls.gz_desi_images_chunk_01, 'ada3fe9a94258075bbe62d42e6550440'),
+        (internal_urls.gz_desi_images_chunk_02, 'e187ef2d77ea606244c8631599a31ca1'),
+        (internal_urls.gz_desi_images_chunk_03, '1a649787bd2ecfe3efcb1856a3627548'),
+        (internal_urls.gz_desi_images_chunk_04, 'a7734ef953f68aac7787b8ea9caf4633'),
+        (internal_urls.gz_desi_images_chunk_05, '960b25c4f6fec9967625696bbfa72d07'),
+        (internal_urls.gz_desi_images_chunk_06, 'bfa6fa0b64e49da72053b4f63468afa0'),
+        (internal_urls.gz_desi_images_chunk_07, '1c1b59cc72ef86db765f5f7189c06260'),
+        (internal_urls.gz_desi_images_chunk_08, 'abc32a75cb5eea4d425085abb6ecd6e8')
+    ]
+    images_to_spotcheck = []  # TODO
+
+    downloader = download_utils.DatasetDownloader(root, resources, images_to_spotcheck, archive_includes_subdir=False)
+    if download is True:
+        downloader.download()
+
+
+    # useful_columns = label_cols + ['subfolder', 'filename']
+    if train:
+        train_catalog_loc = os.path.join(root, 'gz_desi_train_catalog.parquet')
+        catalog = pd.read_parquet(train_catalog_loc)
+    else:
+        test_catalog_loc = os.path.join(root, 'gz_desi_test_catalog.parquet')
+        catalog = pd.read_parquet(test_catalog_loc)
+
+    # removed 'root' from here as downloader.image_dir already includes root
+    catalog['file_loc'] = catalog.apply(lambda x: os.path.join(downloader.image_dir, x['subfolder'], x['filename']), axis=1)
+
+    # default, but not actually used here
+    label_cols = label_metadata.decals_all_campaigns_ortho_label_cols
+    return catalog, label_cols
+
diff --git a/pytorch_galaxy_datasets/prepared_datasets/hubble.py b/galaxy_datasets/shared/gz_hubble.py
similarity index 75%
rename from pytorch_galaxy_datasets/prepared_datasets/hubble.py
rename to galaxy_datasets/shared/gz_hubble.py
index d46ad2a..0394ed6 100644
--- a/pytorch_galaxy_datasets/prepared_datasets/hubble.py
+++ b/galaxy_datasets/shared/gz_hubble.py
@@ -1,24 +1,16 @@
 
+from cProfile import label
 import os
 
 import pandas as pd
-from pytorch_galaxy_datasets import galaxy_dataset, galaxy_datamodule, download_utils
 
 # this one has label_metadata contained in this repo, not zoobot. Others will follow, perhaps.
 from zoobot.shared import label_metadata
 
+from galaxy_datasets.shared import download_utils
 
 
-class HubbleDataset(galaxy_dataset.GalaxyDataset):
-    
-    def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
-
-        catalog, label_cols = hubble_setup(root, train, download)
-
-        super().__init__(catalog, label_cols, transform, target_transform)
-
-
-def hubble_setup(root, train, download):
+def gz_hubble(root, train, download):
     resources = [
         ('https://dl.dropboxusercontent.com/s/xnktj9hq6xig0a7/hubble_ortho_train_catalog.parquet', 'c6cb821f7ebefb583dc74488cf7bfc5f'),  # train catalog
         ('https://dl.dropboxusercontent.com/s/1g9lwih9944sys8/hubble_ortho_test_catalog.parquet', '05e01ed822b34400f32977280eebec87'),  # test catalog
@@ -55,8 +47,6 @@ def hubble_setup(root, train, download):
     # 'clump-configuration': ['_straight-line', '_chain', '_cluster-or-irregular', '_spiral'],
     # 'one-clump-brightest': ['_yes', '_no'],
     # 'brightest-clump-central': ['_yes', '_no'],
-    # 'galaxy-symmetrical': ['_yes', '_no'],
-    # 'clumps-embedded-larger-object': ['_yes', '_no'],
     'disk-edge-on': ['_yes', '_no'],
     'bulge-shape': ['_rounded', '_boxy', '_none'],
     'bar': ['_yes', '_no'],
@@ -97,24 +87,12 @@ def hubble_setup(root, train, download):
 
 if __name__ == '__main__':
 
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    hubble_label_cols, hubble_catalog = hubble_setup(
-        root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/hubble',
-        train=True,
-        download=False
-    )
-    
-    # user will probably tweak and use images/catalog directly for generic galaxy catalog datamodule
-    # (which makes its own generic datasets internally)
-    adjusted_catalog = hubble_catalog.sample(1000)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=hubble_label_cols,
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-    
\ No newline at end of file
+    hubble_cols = [x.replace('-hubble', '') for x in hubble_ortho_label_cols]
+    # print('\n'.join(cols))
+
+    decals_cols = [x for x in label_metadata.decals_label_cols]
+    # print('\n'.join(cols))
+
+    all_cols = list(set(hubble_cols).union(decals_cols))
+    all_cols.sort()
+    print('\n'.join(all_cols))
diff --git a/galaxy_datasets/shared/gz_rings.py b/galaxy_datasets/shared/gz_rings.py
new file mode 100644
index 0000000..c91a5de
--- /dev/null
+++ b/galaxy_datasets/shared/gz_rings.py
@@ -0,0 +1,42 @@
+import os
+
+import pandas as pd
+
+from galaxy_datasets.shared import download_utils
+from galaxy_datasets.check_internal_urls import INTERNAL_URLS_EXIST
+if not INTERNAL_URLS_EXIST:
+    raise FileNotFoundError
+from galaxy_datasets.shared import internal_urls
+
+# TODO could eventually refactor this out of Zoobot as well
+from zoobot.shared import label_metadata
+
+
+
+
+def gz_rings(root, train, download):
+    resources = [
+        (internal_urls.rings_train_catalog, 'e2fb6b2bca45cd7f1c58f5b4089a5976'),  # train catalog
+        (internal_urls.rings_test_catalog, '6e3f362a6e19ecd02675eaa48f6727f0'),  # test catalog
+        (internal_urls.rings_images, 'd0950250436a05ce88de747e6af825b6')  # the images
+    ]
+    images_to_spotcheck = []
+
+    downloader = download_utils.DatasetDownloader(root, resources, images_to_spotcheck)
+    if download is True:
+        downloader.download()
+
+    # useful_columns = label_cols + ['subfolder', 'filename']
+    useful_columns = None
+    if train:
+        train_catalog_loc = os.path.join(root, 'rings_train_catalog.parquet')
+        catalog = pd.read_parquet(train_catalog_loc, columns=useful_columns)
+    else:
+        test_catalog_loc = os.path.join(root, 'rings_test_catalog.parquet')
+        catalog = pd.read_parquet(test_catalog_loc, columns=useful_columns)
+
+    catalog['file_loc'] = catalog.apply(lambda x: os.path.join(downloader.image_dir, x['subfolder'], x['filename']), axis=1)
+
+    label_cols = label_metadata.rings_label_cols
+
+    return catalog, label_cols
diff --git a/pytorch_galaxy_datasets/prepared_datasets/tidal.py b/galaxy_datasets/shared/tidal.py
similarity index 54%
rename from pytorch_galaxy_datasets/prepared_datasets/tidal.py
rename to galaxy_datasets/shared/tidal.py
index b1c9b8c..742af34 100644
--- a/pytorch_galaxy_datasets/prepared_datasets/tidal.py
+++ b/galaxy_datasets/shared/tidal.py
@@ -1,22 +1,11 @@
 import os
 
 import pandas as pd
-from pytorch_galaxy_datasets import galaxy_datamodule, galaxy_dataset, download_utils
 
-# TODO could eventually refactor this out of Zoobot as well
-from zoobot.shared import label_metadata
+from galaxy_datasets.shared import download_utils
 
 
-class TidalDataset(galaxy_dataset.GalaxyDataset):
-    
-    def __init__(self, root, train=True, download=False, transform=None, target_transform=None, label_mode='coarse'):
-
-        catalog, label_cols = tidal_setup(root, train, download, label_mode=label_mode)
-
-        super().__init__(catalog, label_cols, transform, target_transform)
-
-
-def tidal_setup(root, train, download, label_mode='coarse'):
+def tidal(root, train, download, label_mode='coarse'):
     resources = [
         ('https://dl.dropboxusercontent.com/s/jq6dyc8q87h92qq/tidal_train_catalog.parquet', '39def8527823f6d4f332fbc209b15a32'),  # train catalog
         ('https://dl.dropboxusercontent.com/s/h8m7j7o3o11mt05/tidal_test_catalog.parquet', '12a8c6a23ddc8e7ad3f5bd597931bbe3'),  # test catalog
@@ -50,27 +39,3 @@ def tidal_setup(root, train, download, label_mode='coarse'):
 # defined outside for use elsewhere
 coarse_tidal_label_cols = ['coarse_tidal_label']
 finegrained_tidal_label_cols = ['finegrained_tidal_label']
-
-if __name__ == '__main__':
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    tidal_catalog, label_cols = tidal_setup(
-        root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/tidal',
-        train=True,
-        download=False,
-    )
-    adjusted_catalog = tidal_catalog.sample(1000)
-
-    # user will probably tweak and use images/catalog directly for generic galaxy catalog datamodule
-    # (which makes its own generic datasets internally)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=label_cols,
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-        
\ No newline at end of file
diff --git a/galaxy_datasets/tensorflow/__init__.py b/galaxy_datasets/tensorflow/__init__.py
new file mode 100644
index 0000000..f7a74b1
--- /dev/null
+++ b/galaxy_datasets/tensorflow/__init__.py
@@ -0,0 +1,1 @@
+from galaxy_datasets.tensorflow.datasets import get_image_dataset, add_transforms_to_dataset
\ No newline at end of file
diff --git a/galaxy_datasets/tensorflow/datasets.py b/galaxy_datasets/tensorflow/datasets.py
new file mode 100644
index 0000000..f54deda
--- /dev/null
+++ b/galaxy_datasets/tensorflow/datasets.py
@@ -0,0 +1,199 @@
+import os
+from functools import partial
+from typing import List
+import logging
+
+import tensorflow as tf
+
+
+def get_image_dataset(
+    image_paths: List,
+    labels=None,
+    check_valid_paths=True,
+    # new preprocessing options now here
+    requested_img_size=None,
+    greyscale=False,
+    permute_channels=False,
+    include_id_str=False) -> tf.data.Dataset:
+    """
+    Load images in a folder as a tf.data dataset
+    Supports jpeg (note the e) and png
+
+    Args:
+        image_paths (list): list of image paths to load
+        requested_img_size (int or None): e.g. 256 for 256x256x3 image. Assumed square. Will resize if size on disk != this. No resizing if None.
+        labels (list or None): If not None, include labels in dataset (see Returns). Must be equal length to image_paths. Defaults to None.
+        greyscale (bool): 
+        permute_channels (bool):
+
+    Raises:
+        FileNotFoundError: at least one path does not match an existing file
+
+    Returns:
+        tf.data.Dataset: yielding batches with 'image' key for image array, 'id_str' for the image path, and 'label' if ``labels`` was provided.
+    """
+    
+    assert len(image_paths) > 0
+    assert isinstance(image_paths[0], str)
+    logging.info('Image paths to load as dataset: {}'.format(len(image_paths)))
+
+    if check_valid_paths:
+        logging.info('Checking if all paths are valid')
+        missing_paths = [path for path in image_paths if not os.path.isfile(path)]
+        if missing_paths:
+            raise FileNotFoundError(f'Missing {len(missing_paths)} images e.g. {missing_paths[0]}')
+        logging.info('All paths exist')
+    else:
+        logging.warning('Skipping valid path check')
+
+    file_format = image_paths[0].split('.')[-1]  # assume all same format
+
+    # load a single image to check the shape
+    test_image = load_image_as_element(image_paths[0])['image']
+    size_on_disk = test_image.numpy().shape[0]  # x dimension (XYC convention, not yet batched)
+    if requested_img_size is None:
+        logging.info('No specific image size requested on load - not resizing (until augmentations')
+        resize_size = None
+    elif size_on_disk == requested_img_size:
+        logging.info('Image size on disk matches requested_img_size of {}, skipping resizing'.format(requested_img_size))  # x dimension of first image, first y index, first channel
+        resize_size = None
+    else:
+        logging.warning('Resizing images from disk size {} to requested size {}'.format(size_on_disk, requested_img_size))
+        resize_size = requested_img_size
+
+    path_ds = tf.data.Dataset.from_tensor_slices([str(path) for path in image_paths])
+
+    # will yield elements of {'image': image, 'id_str': path}
+    image_ds = path_ds.map(lambda x: load_image_as_element(x, mode=file_format), num_parallel_calls=tf.data.AUTOTUNE)  # keep determinstic though
+
+    # applies preprocessing transforms to x['image']: greyscaling/permuting, resizing
+    # TODO could join these into the above function for speed, if needed
+    image_ds = image_ds.map(lambda x: preprocess_image_in_element(x, resize_size=resize_size, greyscale=greyscale, permute_channels=permute_channels), num_parallel_calls=tf.data.AUTOTUNE)
+
+    if labels is not None:
+        assert len(labels) == len(image_paths)
+
+        # if isinstance(labels[0], dict):
+        #     # assume list of dicts, each representing one datapoint e.g. [{'feat_a': 1, 'feat_b': 2}, {'feat_a': 12, 'feat_b': 42}]
+        #     # reshape to columnlike dict e.g. {'feat_a': [1, 12], 'feat_b: [2, 42]} because that's what tf supports
+        #     # (could pass this directly, but inputs tend to be easier to handle row-wise for e.g. shuffling etc)
+        #     label_ds = tf.data.Dataset.from_tensor_slices(pd.DataFrame.from_dict(labels).to_dict(orient="list"))
+        # else:
+            # make it a dict anyway for consistency, keyed by 'label' instead of e.g. 'feat_a', 'feat_b'
+        label_ds = tf.data.Dataset.from_tensor_slices({'label': labels})
+
+        # label_dict is {'label': (256)} or {'feat_a': (256), 'feat_b': (256)}
+        # image_dict is {'id_str': some_id 'image': (image)}
+        # merge the two dicts to create {'id_str': ..., 'image': ..., 'feat_a': ..., 'feat_b': ...}
+        image_ds = tf.data.Dataset.zip((image_ds, label_ds)).map(lambda image_dict, label_dict: {**image_dict, **label_dict})
+        # now yields {'image': , 'id_str': , 'label': } batched dicts
+
+    # now choose how to unpack
+    # (yes, I could probably do this faster/in fewer steps, but this is very readable)
+    if include_id_str:
+        if labels is not None:
+            image_ds = image_ds.map(lambda x: (x['image'], x['label'], x['id_str']))
+        else:
+            image_ds = image_ds.map(lambda x: (x['image'], x['id_str']))
+    else:
+        assert labels is not None
+        image_ds = image_ds.map(lambda x: (x['image'], x['label']))
+
+    return image_ds
+
+
+# https://stackoverflow.com/questions/62544528/tensorflow-decodejpeg-expected-image-jpeg-png-or-gif-got-unknown-format-st?rq=1
+def load_image_as_element(loc, mode='png'):
+    """
+    Load an image file from disk to memory.
+    *Recently changed to return 0-1 floats not 0-255 floats*
+
+    Args:
+        loc (str): Path to image on disk. Includes format e.g. .png.
+        mode (str, optional): Image format. Defaults to 'png'.
+
+    Raises:
+        ValueError: mode is neither png nor jpeg.
+
+    Returns:
+        dict: like {'image': float32 np.ndarray from 0. to 1., 'id_str': ``loc``}
+    """
+    # values will be 0-255, does not normalise. Happens in preprocessing instead.
+    # specify mode explicitly to avoid graph tracing issues
+    image = tf.io.read_file(loc)
+    if mode == 'png':
+        image = tf.image.decode_png(image)
+    # rename jpg to jpeg or validation checks in decode_jpg will fail
+    # might fail, let's see
+    elif (mode == 'jpg') or (mode == 'jpeg'):  
+        # TODO also allow jpg
+        image = tf.image.decode_jpeg(image)
+    else:
+        raise ValueError(f'Image filetype mode {mode} not recognised')
+
+    converted_image = tf.cast(image, tf.float32) / 255.  # floats in 0-1 range
+
+    return {'image': converted_image, 'id_str': loc}  # using the file paths as identifiers
+
+
+def preprocess_image_in_element(element, resize_size=None, greyscale=False, permute_channels=False):
+    # now applies all preprocessing
+
+    # element['image'] must be 0-1 floats, not 0-255 ints, or clipping will ruin
+    image, id_str = element['image'], element['id_str']  # unpack from dict
+
+    if greyscale:
+        # new channel dimension of 1
+        image = tf.reduce_mean(input_tensor=image, axis=-1, keepdims=True)
+        # assert images.shape[1] == input_size
+        # assert images.shape[2] == input_size
+        # assert images.shape[3] == 1
+        # tf.summary.image('b_make_greyscale', images)
+    else:
+        if permute_channels:
+            image = tf.map_fn(permute_channels, image)  # map to each image in element
+        else:
+            image = tf.identity(image)
+            
+    if resize_size is not None:
+        image = resize_image_with_tf(image , size=resize_size)   # initial size = after resize from image on disk (e.g. 424 for GZ pngs) but before crop/zoom
+        image = tf.clip_by_value(image, 0., 1.)  # resizing can cause slight change in min/max
+
+    return {'image': image, 'id_str': id_str}  # pack back into dict
+
+
+def resize_image_with_tf(image, size):
+    # May cause values outside 0-255 margins
+    # allows for image dimension, or no image dimension
+    # May be slow. Ideally, resize the images beforehand on disk (or save as TFRecord, see make_shards.py and tfrecord_datasets.py)
+    return tf.image.resize(image, (size, size), method=tf.image.ResizeMethod.LANCZOS3, antialias=True)
+
+
+def permute_channels(im):
+    assert tf.shape(im)[-1] == 3
+    # tf.random.shuffle shuffles 0th dimension, so need to temporarily swap channel to 0th, shuffle, and swap back
+    return tf.transpose(tf.random.shuffle(tf.transpose(im, perm=[2, 1, 0])), perm=[2, 1, 0])
+
+
+def add_transforms_to_dataset(dataset, transforms):
+    # only works if dataset yields either (images, labels) or (images, id_str), not (images, label_str, id_str)
+    # TODO make work with extra element - a little fiddly due to unpacking in graph mode, though
+
+    # use closure to pass in transforms (which cannot be passed via inp as not tensor-like)
+    def process_data(images, labels):
+        aug_imgs = tf.numpy_function(func=aug_fn, inp=[images], Tout=tf.float32)
+        return aug_imgs, labels
+        # return (aug_imgs) + args  # equivalent to aug_imgs, *args, but works prior to Python 3.8
+
+    def aug_fn(image):
+        # albumentations.transforms expects input like image=x
+        # and returns {'image': image} dict
+        return transforms(image=image)["image"]
+
+    # TODO specify output_size
+
+    return dataset.map(
+        partial(process_data),
+        num_parallel_calls=tf.data.experimental.AUTOTUNE
+    )
+
diff --git a/galaxy_datasets/transforms.py b/galaxy_datasets/transforms.py
new file mode 100644
index 0000000..85463c5
--- /dev/null
+++ b/galaxy_datasets/transforms.py
@@ -0,0 +1,156 @@
+import typing
+
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+def default_transforms(
+    crop_scale_bounds=(0.7, 0.8),
+    crop_ratio_bounds=(0.9, 1.1),
+    resize_after_crop=224, 
+    pytorch_greyscale=False
+    ) -> typing.Dict[str, typing.Any]:
+    if pytorch_greyscale:
+        transforms_to_apply = [A.Lambda(name='ToGray', image=ToGray(
+            reduce_channels=True), always_apply=True)]
+    else:
+        transforms_to_apply = []
+
+    transforms_to_apply += [
+        # A.ToFloat(),
+        # anything outside of the original image is set to 0.
+        A.Rotate(limit=180, interpolation=1,
+                    always_apply=True, border_mode=0, value=0),
+        A.RandomResizedCrop(
+            height=resize_after_crop,  # after crop resize
+            width=resize_after_crop,
+            scale=crop_scale_bounds,  # crop factor
+            ratio=crop_ratio_bounds,  # crop aspect ratio
+            interpolation=1,  # This is "INTER_LINEAR" == BILINEAR interpolation. See: https://docs.opencv.org/3.4/da/d54/group__imgproc__transform.html
+            always_apply=True
+        ),  # new aspect ratio
+        A.VerticalFlip(p=0.5),
+        # ToTensorV2() TODO may be needed with pytorch arg
+    ]
+
+    return A.Compose(transforms_to_apply)
+
+
+def astroaugmentation_transforms(
+    resize_size: int,
+    shift_limit: float,
+    scale_limit: float,
+    rotate_limit: float,
+    p_channelwise_dropout: float,
+    p_elastic=0,
+    elastic_alpha=1,
+    elastic_sigma=1,
+    p_flip=0.5,
+    channelwise_dropout_max_fraction=0.2,
+    channelwise_dropout_min_width=10,
+    channelwise_dropout_min_height=10,
+    channelwise_dropout_max_holes=100,
+    pytorch_greyscale=False,
+) -> typing.Dict[str, typing.Any]:
+    """
+    Astrophysically-inspired image transforms from AstroAugmentations (Bowles in prep.)
+    Intended to mimic common distortions to optical telescope images.
+    
+    Sequentially:
+    - ElasticTransform (optional) from https://ieeexplore.ieee.org/document/1227801
+    - ChannelWiseDropout (good for Decals) from AstroAugmentations
+    - ShiftScaleRotate instead of RandomResizedCrop for more control
+    - Flip along x or y axis
+
+    Args:
+        resize_size (int): resolution
+        shift_limit (float): max relative shift factor. 0.3 can have the galaxy pretty much
+                            at the edge *without scaling*
+        scale_limit (float):  max relative max zoom *in* factor
+        rotate_limit (float): max rotate angle in degrees
+        p_channelwise_dropout (float): _description_
+        p_elastic (int, optional): how much to scale the random field. Defaults to 0.
+        elastic_alpha (int, optional): std. of the Gauss. kernel blurring the random field. Defaults to 1.
+        elastic_sigma (int, optional):  max area fraction to be affected. Defaults to 1.
+        p_flip (float, optional): _description_. Defaults to 0.5.
+        channelwise_dropout_max_fraction (float, optional): _description_. Defaults to 0.2.
+        channelwise_dropout_min_width (int, optional):  minimum width of the hole in pixels. Defaults to 10.
+        channelwise_dropout_min_height (int, optional):  minimum height of the hole in pixels. Defaults to 10.
+        channelwise_dropout_max_holes (int, optional): _description_. Defaults to 100.
+        pytorch_greyscale (bool, optional): _description_. Defaults to False.
+
+    Returns:
+        typing.Dict[str, typing.Any]: _description_
+    """
+    # wrapped in Try/Except to avoid making AstroAugmentations a package requirement for only this optional transform
+    try:
+        from AstroAugmentations import image_domain
+    except ImportError:
+        raise ImportError(
+            'Trying to use astroaugmentation_transforms but AstroAugmentations is not installed\n \
+            Please install via git. See instructions at https://github.com/mb010/AstroAugmentations#quick-start'
+        )
+
+    if pytorch_greyscale:
+        transforms_to_apply = [
+            A.Lambda(
+                name="ToGray", image=ToGray(reduce_channels=True), always_apply=True
+            )
+        ]
+    else:
+        transforms_to_apply = []
+    transforms_to_apply += [
+        A.LongestMaxSize(max_size=resize_size),
+        A.ElasticTransform(
+            alpha=elastic_alpha,
+            sigma=elastic_sigma,
+            alpha_affine=0,
+            interpolation=1,
+            border_mode=1,
+            value=0,
+            p=p_elastic,
+        ),
+        A.Lambda(
+            name="MissingData",
+            image=image_domain.optical.ChannelWiseDropout(
+                max_fraction=channelwise_dropout_max_fraction,
+                min_width=channelwise_dropout_min_width,
+                min_height=channelwise_dropout_min_height,
+                max_holes=channelwise_dropout_max_holes,
+                channelwise_application=True,
+            ),
+            p=p_channelwise_dropout,
+        ),
+        A.ShiftScaleRotate(
+            shift_limit=shift_limit,
+            scale_limit=(0, scale_limit),
+            rotate_limit=rotate_limit,
+            interpolation=2,
+            border_mode=0,
+            p=1,
+        ),
+        A.Flip(p=p_flip),
+        ToTensorV2(),
+    ]
+
+    return A.Compose(transforms_to_apply)
+
+
+# albumentations versuib of GrayscaleUnweighted
+class ToGray():
+    # will do nothing if already greyscale
+
+    def __init__(self, reduce_channels=False):
+        if reduce_channels:
+            self.forward = to_single_greyscale_channel
+        else:
+            self.forward = to_triple_greyscale_channel
+            
+    def __call__(self, image, **kwargs):
+        return self.forward(image)
+
+def to_single_greyscale_channel(img):
+    return img.mean(axis=2, keepdims=True)
+
+def to_triple_greyscale_channel(img):
+    return img.mean(axis=2, keepdims=True).repeat(3, axis=2)
\ No newline at end of file
diff --git a/legs_low_z_copy.sh b/legs_low_z_copy.sh
deleted file mode 100644
index 4f9d1b7..0000000
--- a/legs_low_z_copy.sh
+++ /dev/null
@@ -1,64 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=copy                    # Job name
-#SBATCH --output=copy_%A.log 
-#SBATCH --mem=0                                     # "reserve all the available memory on each node assigned to the job"
-#SBATCH --no-requeue                                    # Do not resubmit a failed job
-#SBATCH --time=23:00:00                                # Time limit hrs:min:sec
-#SBATCH --exclusive   # only one task per node
-#SBATCH --ntasks 1
-#SBATCH --cpus-per-task=24
-pwd; hostname; date
-
-# srun head -n 10 /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt | xargs -I X -n 1 -P 24 echo X
-
-# srun head -n 10 /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt | xargs -I X -n 1 -P 24 -t -d " " rsync --relative /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/./X /share/nas2/walml/galaxy_zoo/decals/dr8/low_z_jpg
-
-# rsync --version
-# xargs --version
-
-# xargs -a /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt -I X -n 1 -P 24 -t rsync --relative /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/./X /share/nas2/walml/galaxy_zoo/decals/dr8/low_z_jpg
-
-# xargs -a /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt -I X -n 1 -P 24 mkdir -p /share/nas2/walml/galaxy_zoo/decals/dr8/low_z_jpg/X
-# xargs -a /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt -I X -n 1 -P 24 cp /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/X /share/nas2/walml/galaxy_zoo/decals/dr8/low_z_jpg/X
-
-# xargs -a /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt -I X -n 1 -P 24 cp /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/X /share/nas2/walml/galaxy_zoo/decals/dr8/low_z_jpg/X
-
-# timestamp() {
-#   "%T" # current time
-# }
-
-# xargs -a /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt -n 50000 -P 1 -t -i tar -czvf /share/nas2/walml/galaxy_zoo/decals/dr8/archive_{timestamp}.tar.gz {}
-
-# tar -cf -a --files-from /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt /share/nas2/walml/galaxy_zoo/decals/dr8/archive_test.tar.gz
-
-# head -n 10 /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt | tar -cvf --auto-compress  /share/nas2/walml/galaxy_zoo/decals/dr8/archive_test.tar.gz 
-
-# tar -cvfz --directory /share/nas2/walml/galaxy_zoo/decals/dr8/jpg /share/nas2/walml/galaxy_zoo/decals/dr8/archive_test.tar.gz 42832/42832_4476.jpg
-
-
-# tar -cvfz /share/nas2/walml/galaxy_zoo/decals/dr8/archive_test.tar.gz /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/42832/42832_4476.jpg
-
-# tar -czvf /share/nas2/walml/galaxy_zoo/decals/dr8/archive_test.tar.gz jpg/42832/42832_4476.jpg jpg/42832/42832_4476.jpg
-
-# tar -czvf /share/nas2/walml/checks_debug.tar.gz checks_old/100000_100499.parquet checks_old/6452500_6452999.parquet
-
-# random not re-triggered
-# head -n 10 /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt | xargs -n 2 -P 1 -t -I @ tar -czvf /share/nas2/walml/galaxy_zoo/decals/dr8/archive_$RANDOM.tar.gz @
-
-# head -n 10 /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt | split -l 2
-
-# ls low_z_jpg_* | xargs -P 1 -t -i tar -czvf -T {} /share/nas2/walml/galaxy_zoo/decals/dr8/{}.tar.gz
-
-# tar -czvf /share/nas2/walml/galaxy_zoo/decals/dr8/low_z_jpg_00.tar.gz
-# tar -czvf /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/temp/debug.tar.gz -T temp/low_z_jpg_00 
-
-
-
-# from jpg dir
-cd /share/nas2/walml/galaxy_zoo/decals/dr8/jpg
-
-# split main file into 200k chunks, named like temp/low_z_jpg_* where * is an int
-split -d -l 200000 /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt temp/low_z_jpg_chunk_
-
-# get the paths to those files (temp/...), and then use xargs to start 24 processes, each using tar to read the paths in a chunkfile (-T) and save to path.tar.gz
-ls temp/low_z_jpg_chunk_* | xargs -n 1 -P 24 -i tar -czvf {}_archive.tar.gz -T {} 
\ No newline at end of file
diff --git a/notebooks/legs_split_dr8_only.ipynb b/notebooks/legs_split_dr8_only.ipynb
deleted file mode 100644
index b83f124..0000000
--- a/notebooks/legs_split_dr8_only.ipynb
+++ /dev/null
@@ -1,540 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "\n",
-    "import pandas as pd\n",
-    "\n",
-    "from sklearn.model_selection import train_test_split\n",
-    "\n",
-    "from zoobot.shared.label_metadata import decals_all_campaigns_ortho_label_cols\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "dr12 = pd.read_parquet('/home/walml/repos/gz-decals-classifiers/data/catalogs/training_catalogs/dr12_ortho_v2_labelled_catalog.parquet')\n",
-    "dr5 = pd.read_parquet('/home/walml/repos/gz-decals-classifiers/data/catalogs/training_catalogs/dr5_ortho_v2_labelled_catalog.parquet')\n",
-    "dr8 = pd.read_parquet('/home/walml/repos/gz-decals-classifiers/data/catalogs/training_catalogs/dr8_ortho_v2_labelled_catalog.parquet')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>dr8_id</th>\n",
-       "      <th>iauname</th>\n",
-       "      <th>redshift</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>5206072</th>\n",
-       "      <td>430531_1076</td>\n",
-       "      <td>None</td>\n",
-       "      <td>0.212848</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4462610</th>\n",
-       "      <td>482154_2597</td>\n",
-       "      <td>None</td>\n",
-       "      <td>0.269002</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>8641681</th>\n",
-       "      <td>476140_1345</td>\n",
-       "      <td>None</td>\n",
-       "      <td>0.118060</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>7054936</th>\n",
-       "      <td>225838_2938</td>\n",
-       "      <td>None</td>\n",
-       "      <td>0.120925</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>8229307</th>\n",
-       "      <td>474717_394</td>\n",
-       "      <td>J133229.55+253853.4</td>\n",
-       "      <td>0.077330</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "              dr8_id              iauname  redshift\n",
-       "5206072  430531_1076                 None  0.212848\n",
-       "4462610  482154_2597                 None  0.269002\n",
-       "8641681  476140_1345                 None  0.118060\n",
-       "7054936  225838_2938                 None  0.120925\n",
-       "8229307   474717_394  J133229.55+253853.4  0.077330"
-      ]
-     },
-     "execution_count": 3,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "megacatalog = pd.read_parquet('/home/walml/repos/decals-rings/results/rings_megacatalog.parquet', columns=['dr8_id', 'iauname', 'redshift'])\n",
-    "megacatalog.sample(5)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "85730\n",
-      "82387\n"
-     ]
-    }
-   ],
-   "source": [
-    "print(len(dr12))\n",
-    "dr12 = pd.merge(dr12, megacatalog, on='iauname', how='inner')\n",
-    "print(len(dr12))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "228059\n",
-      "222754\n"
-     ]
-    }
-   ],
-   "source": [
-    "print(len(dr5))\n",
-    "dr5 = pd.merge(dr5, megacatalog, on='iauname', how='inner')\n",
-    "print(len(dr5))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# replace with cross-matched\n",
-    "dr12['id_str'] = dr12['dr8_id']\n",
-    "dr5['id_str'] = dr5['dr8_id']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "365770"
-      ]
-     },
-     "execution_count": 7,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "labelled = pd.concat([dr12, dr5, dr8], axis=0)\n",
-    "len(labelled)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now make the unlabelled catalog from the remainder"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:2: SettingWithCopyWarning: \n",
-      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
-      "Try using .loc[row_indexer,col_indexer] = value instead\n",
-      "\n",
-      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
-      "  \n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "8323600"
-      ]
-     },
-     "execution_count": 8,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "unlabelled = megacatalog[~megacatalog['dr8_id'].isin(labelled['dr8_id'])]\n",
-    "unlabelled['id_str'] = unlabelled['dr8_id']\n",
-    "len(unlabelled)  # mostly high-z, some redshift filtering would be wise"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:2: SettingWithCopyWarning: \n",
-      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
-      "Try using .loc[row_indexer,col_indexer] = value instead\n",
-      "\n",
-      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
-      "  \n",
-      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:2: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
-      "  \n"
-     ]
-    }
-   ],
-   "source": [
-    "for col in decals_all_campaigns_ortho_label_cols:\n",
-    "    unlabelled[col] = 0\n",
-    "unlabelled = unlabelled.copy().sample(len(unlabelled), random_state=42)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:2: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
-      "  \n",
-      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:3: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
-      "  This is separate from the ipykernel package so we can avoid doing imports until\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "0    /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/34...\n",
-      "0    /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/33...\n",
-      "0    /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/17...\n",
-      "Name: file_loc, dtype: object\n",
-      "/share/nas2/walml/galaxy_zoo/decals/dr8/jpg/43543/43543_92.jpg\n"
-     ]
-    }
-   ],
-   "source": [
-    "for df in [labelled, unlabelled]:  # dr8 already has but it points to the png\n",
-    "    df['brickid'] = df['id_str'].apply(lambda x: x.split('_')[0])\n",
-    "    df['objid'] = df['id_str'].apply(lambda x: x.split('_')[1])\n",
-    "    df['file_loc'] = '/share/nas2/walml/galaxy_zoo/decals/dr8/jpg/' + df['brickid'] + '/' + df['id_str'] + '.jpg'\n",
-    "    print(df['file_loc'][0])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "array(['dr8_id', 'iauname', 'redshift', 'id_str',\n",
-       "       'smooth-or-featured-dr12_smooth',\n",
-       "       'smooth-or-featured-dr12_featured-or-disk',\n",
-       "       'smooth-or-featured-dr12_artifact', 'disk-edge-on-dr12_yes',\n",
-       "       'disk-edge-on-dr12_no', 'has-spiral-arms-dr12_yes',\n",
-       "       'has-spiral-arms-dr12_no', 'bar-dr12_yes', 'bar-dr12_no',\n",
-       "       'bulge-size-dr12_dominant', 'bulge-size-dr12_obvious',\n",
-       "       'bulge-size-dr12_none', 'how-rounded-dr12_completely',\n",
-       "       'how-rounded-dr12_in-between', 'how-rounded-dr12_cigar-shaped',\n",
-       "       'edge-on-bulge-dr12_boxy', 'edge-on-bulge-dr12_none',\n",
-       "       'edge-on-bulge-dr12_rounded', 'spiral-winding-dr12_tight',\n",
-       "       'spiral-winding-dr12_medium', 'spiral-winding-dr12_loose',\n",
-       "       'spiral-arm-count-dr12_1', 'spiral-arm-count-dr12_2',\n",
-       "       'spiral-arm-count-dr12_3', 'spiral-arm-count-dr12_4',\n",
-       "       'spiral-arm-count-dr12_more-than-4', 'merging-dr12_neither',\n",
-       "       'merging-dr12_tidal-debris', 'merging-dr12_both',\n",
-       "       'merging-dr12_merger', 'smooth-or-featured-dr5_smooth',\n",
-       "       'smooth-or-featured-dr5_featured-or-disk',\n",
-       "       'smooth-or-featured-dr5_artifact', 'disk-edge-on-dr5_yes',\n",
-       "       'disk-edge-on-dr5_no', 'has-spiral-arms-dr5_yes',\n",
-       "       'has-spiral-arms-dr5_no', 'bar-dr5_strong', 'bar-dr5_weak',\n",
-       "       'bar-dr5_no', 'bulge-size-dr5_dominant', 'bulge-size-dr5_large',\n",
-       "       'bulge-size-dr5_moderate', 'bulge-size-dr5_small',\n",
-       "       'bulge-size-dr5_none', 'how-rounded-dr5_round',\n",
-       "       'how-rounded-dr5_in-between', 'how-rounded-dr5_cigar-shaped',\n",
-       "       'edge-on-bulge-dr5_boxy', 'edge-on-bulge-dr5_none',\n",
-       "       'edge-on-bulge-dr5_rounded', 'spiral-winding-dr5_tight',\n",
-       "       'spiral-winding-dr5_medium', 'spiral-winding-dr5_loose',\n",
-       "       'spiral-arm-count-dr5_1', 'spiral-arm-count-dr5_2',\n",
-       "       'spiral-arm-count-dr5_3', 'spiral-arm-count-dr5_4',\n",
-       "       'spiral-arm-count-dr5_more-than-4',\n",
-       "       'spiral-arm-count-dr5_cant-tell', 'merging-dr5_none',\n",
-       "       'merging-dr5_minor-disturbance', 'merging-dr5_major-disturbance',\n",
-       "       'merging-dr5_merger', 'smooth-or-featured-dr8_smooth',\n",
-       "       'smooth-or-featured-dr8_featured-or-disk',\n",
-       "       'smooth-or-featured-dr8_artifact', 'disk-edge-on-dr8_yes',\n",
-       "       'disk-edge-on-dr8_no', 'has-spiral-arms-dr8_yes',\n",
-       "       'has-spiral-arms-dr8_no', 'bar-dr8_strong', 'bar-dr8_weak',\n",
-       "       'bar-dr8_no', 'bulge-size-dr8_dominant', 'bulge-size-dr8_large',\n",
-       "       'bulge-size-dr8_moderate', 'bulge-size-dr8_small',\n",
-       "       'bulge-size-dr8_none', 'how-rounded-dr8_round',\n",
-       "       'how-rounded-dr8_in-between', 'how-rounded-dr8_cigar-shaped',\n",
-       "       'edge-on-bulge-dr8_boxy', 'edge-on-bulge-dr8_none',\n",
-       "       'edge-on-bulge-dr8_rounded', 'spiral-winding-dr8_tight',\n",
-       "       'spiral-winding-dr8_medium', 'spiral-winding-dr8_loose',\n",
-       "       'spiral-arm-count-dr8_1', 'spiral-arm-count-dr8_2',\n",
-       "       'spiral-arm-count-dr8_3', 'spiral-arm-count-dr8_4',\n",
-       "       'spiral-arm-count-dr8_more-than-4',\n",
-       "       'spiral-arm-count-dr8_cant-tell', 'merging-dr8_none',\n",
-       "       'merging-dr8_minor-disturbance', 'merging-dr8_major-disturbance',\n",
-       "       'merging-dr8_merger', 'brickid', 'objid', 'file_loc'], dtype=object)"
-      ]
-     },
-     "execution_count": 11,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "df.columns.values"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import numpy as np"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "False"
-      ]
-     },
-     "execution_count": 18,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "np.any(labelled[decals_all_campaigns_ortho_label_cols].isna())"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "True"
-      ]
-     },
-     "execution_count": 23,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "any(labelled[decals_all_campaigns_ortho_label_cols].isna())"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 22,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "False"
-      ]
-     },
-     "execution_count": 22,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "any(labelled[decals_all_campaigns_ortho_label_cols].isna().values.max(axis=1))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "0.0"
-      ]
-     },
-     "execution_count": 17,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "np.any(labelled[decals_all_campaigns_ortho_label_cols].isna(), axis=1).mean()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_catalog, test_catalog = train_test_split(labelled, test_size=0.2, random_state=42)  # note that because we dropped a few rows which couldn't be cross-matched, this is not the same split as the original_images version"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# train_catalog.to_parquet('/home/walml/repos/pytorch-galaxy-datasets/roots/legs/legs_all_campaigns_ortho_dr8_only_train_catalog.parquet', index=False)\n",
-    "# test_catalog.to_parquet('/home/walml/repos/pytorch-galaxy-datasets/roots/legs/legs_all_campaigns_ortho_dr8_only_test_catalog.parquet', index=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# unlabelled.to_parquet('/home/walml/repos/pytorch-galaxy-datasets/roots/legs/legs_all_campaigns_ortho_dr8_only_unlabelled_catalog.parquet', index=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
-   "outputs": [
-    {
-     "ename": "FileNotFoundError",
-     "evalue": "[Errno 2] No such file or directory: '/home/walml/repos/pytorch-galaxy-datasets/roots/legs/legs_all_campaigns_ortho_dr8_only_train_catalog.parquet'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
-      "\u001b[0;32m/tmp/ipykernel_13362/2121439838.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      7\u001b[0m ]:\n\u001b[1;32m      8\u001b[0m     \u001b[0;31m# print hash\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m     \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'rb'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m         \u001b[0mmd5_checksum\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhashlib\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmd5\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhexdigest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
-      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: '/home/walml/repos/pytorch-galaxy-datasets/roots/legs/legs_all_campaigns_ortho_dr8_only_train_catalog.parquet'"
-     ]
-    }
-   ],
-   "source": [
-    "import hashlib\n",
-    "\n",
-    "for loc in [\n",
-    "    '/home/walml/repos/pytorch-galaxy-datasets/roots/legs/legs_all_campaigns_ortho_dr8_only_train_catalog.parquet',\n",
-    "    '/home/walml/repos/pytorch-galaxy-datasets/roots/legs/legs_all_campaigns_ortho_dr8_only_test_catalog.parquet',\n",
-    "    '/home/walml/repos/pytorch-galaxy-datasets/roots/legs/legs_all_campaigns_ortho_dr8_only_unlabelled_catalog.parquet'\n",
-    "]:\n",
-    "    # print hash\n",
-    "    with open(loc, 'rb') as f:\n",
-    "        md5_checksum = hashlib.md5(f.read()).hexdigest()\n",
-    "\n",
-    "    print(md5_checksum)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "interpreter": {
-   "hash": "5ec7511d1fc7a72aba078c8815c4bcacd5404aac7a6282065ad4abec249876bb"
-  },
-  "kernelspec": {
-   "display_name": "Python 3.7.11 ('zoobot')",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.7.11"
-  },
-  "orig_nbformat": 4
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/pytorch_galaxy_datasets/galaxy_dataset.py b/pytorch_galaxy_datasets/galaxy_dataset.py
deleted file mode 100644
index b3fde48..0000000
--- a/pytorch_galaxy_datasets/galaxy_dataset.py
+++ /dev/null
@@ -1,130 +0,0 @@
-import pandas as pd
-import numpy as np
-import torch
-from torch.utils.data import Dataset
-from PIL import Image
-import simplejpeg
-
-
-import torch
-import os
-import logging
-
-import numpy as np
-import pandas as pd
-
-
-# https://pytorch.org/tutorials/beginner/basics/data_tutorial.html
-class GalaxyDataset(Dataset):
-    def __init__(self, catalog: pd.DataFrame, label_cols, transform=None, target_transform=None):
-        # downloaded dataset where everything is all prepared, e.g. from prepared_datasets.gz2.gz2_setup()
-        # reads a catalog of jpeg images, assumed of galaxies
-        # if transform is from albumentations, datamodule should know about the transform including 
-        # catalog should be split already
-        # should have correct image locations under file_loc
-
-        # internal catalog has id_str as index for query speed
-        # (dataset.catalog still returns the int-indexed catalog via the catalog method below)
-        # self._catalog = catalog.copy().set_index('id_str', verify_integrity=True)
-        self.catalog = catalog
-        
-        self.label_cols = label_cols
-        self.transform = transform
-        self.target_transform = target_transform
-
-    def __len__(self):
-        return len(self.catalog)
-
-    # @property
-    # def catalog(self):
-    #     return self._catalog.reset_index(drop=False)
-
-    def __getitem__(self, idx):
-        #the index is id_str so can use that for quick search on 1M+ catalo
-        # galaxy = self._catalog.loc[idx]
-        galaxy = self.catalog.iloc[idx]
-        # option A
-        # img_path = galaxy['file_loc']
-        # image = read_image(img_path) # PIL under the hood: Returns CHW Tensor.
-        # option B - tiny bit faster when CPU-limited
-        with open(galaxy['file_loc'], 'rb') as f:
-            # image = torch.from_numpy(decode_jpeg(f.read()).transpose(2, 0, 1))  # CHW tensor
-            try:
-                image = Image.fromarray(decode_jpeg(f.read()))  # HWC PIL image via simplejpeg
-            except Exception as e:
-                logging.critical('Cannot load {}'.format(galaxy['file_loc']))
-                raise e
-        label = get_galaxy_label(galaxy, self.label_cols)
-
-        # logging.info((image.shape, torch.max(image), image.dtype, label))  # always 0-255 uint8
-
-        if self.transform:
-            # a CHW tensor, which torchvision wants. May change to PIL image.
-            image = self.transform(image)
-
-        if self.target_transform:
-            label = self.target_transform(label)
-
-        # logging.info((image.shape, torch.max(image), image.dtype, label))  #  should be 0-1 float
-        return image, label
-
-
-def load_encoded_jpeg(loc):
-    with open(loc, "rb") as f:
-        return f.read()  # bytes, not yet decoded
-
-
-def decode_jpeg(encoded_bytes):
-    return simplejpeg.decode_jpeg(encoded_bytes, fastdct=True, fastupsample=True)
-
-
-def get_galaxy_label(galaxy, label_cols):
-    # no longer casts to int64, user now responsible in df. If dtype is mixed, will try to infer with infer_objects
-    return galaxy[label_cols].infer_objects().values.squeeze()  # squeeze for if there's one label_col
-
-
-
-
-# # https://pytorch.org/vision/stable/_modules/torchvision/datasets/mnist.html for download process inspiration
-# class GalaxyDatasetFactory():
-
-#     def __init__(self, root, catalog=None, label_cols=None,  transform=None, target_transform=None) -> None:
-
-#         # can use target_transform to turn counts into regression or even classification
-#         # will need another step to drop rows, in DataModule probably
-
-#         self.root = root
-
-#         catalog, label_cols = self.adjust_catalog_and_labels(root, catalog, label_cols)
-
-#         super().__init__(catalog=catalog, label_cols=label_cols,
-#                          transform=transform, target_transform=target_transform)
-
-
-#     def adjust_catalog_and_labels(self, root, catalog, label_cols):
-#         if not self._check_exists():
-#             raise RuntimeError(
-#                 "Dataset not found. You can use download=True to download it")
-
-#         if catalog is None:
-#             logging.info('Loading dataset with default (unsplit) catalog')
-#             catalog = pd.read_parquet(os.path.join(root, self.default_catalog_loc))
-#             catalog['file_loc'] = catalog['filename'].apply(
-#                 lambda x: os.path.join(self.image_dir, x))
-#         else:
-#             logging.info(
-#                 'Overriding default catalog with user-provided catalog (length {})'.format(len(catalog)))
-#             assert isinstance(catalog, pd.DataFrame)
-#             # will always check label_cols as well, below
-#             assert 'file_loc' in catalog.columns.values
-
-#         if label_cols is None:
-#             logging.info('Loading dataset with default label columns')
-#             label_cols = self.default_label_cols
-#         else:
-#             logging.info('User provided GZ2 dataset with custom label cols {}'.format(label_cols))
-#                 # label_cols_present = [col in catalog.columns.values for col in label_cols]
-#             missing_label_cols = set(label_cols) - set(catalog.columns.values)
-#             if len(missing_label_cols) > 0:
-#                 raise KeyError(f'User asked for label columns not present in catalog:\n{missing_label_cols} not in \n{catalog.columns.values}')
-#         return catalog,label_cols
diff --git a/pytorch_galaxy_datasets/prepared_datasets/__init__.py b/pytorch_galaxy_datasets/prepared_datasets/__init__.py
deleted file mode 100644
index c346fe1..0000000
--- a/pytorch_galaxy_datasets/prepared_datasets/__init__.py
+++ /dev/null
@@ -1,12 +0,0 @@
-from pytorch_galaxy_datasets.prepared_datasets.gz2 import GZ2Dataset, gz2_setup
-from pytorch_galaxy_datasets.prepared_datasets.dr5 import DecalsDR5Dataset, decals_dr5_setup
-
-import os
-# i.e. if internal_urls.py is available in the folder of this script
-if os.path.isfile(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'internal_urls.py')):  
-    # import the internal datasets
-    from pytorch_galaxy_datasets.prepared_datasets.legs import LegsDataset, legs_setup
-    from pytorch_galaxy_datasets.prepared_datasets.rings import RingsDataset, rings_setup
-
-# allows for imports like
-# from pytorch_galaxy_datasets.prepared_datasets import GZ2Dataset
\ No newline at end of file
diff --git a/pytorch_galaxy_datasets/prepared_datasets/rings.py b/pytorch_galaxy_datasets/prepared_datasets/rings.py
deleted file mode 100644
index 572e673..0000000
--- a/pytorch_galaxy_datasets/prepared_datasets/rings.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import os
-
-import pandas as pd
-from pytorch_galaxy_datasets import galaxy_datamodule, galaxy_dataset, download_utils
-
-from pytorch_galaxy_datasets.prepared_datasets import internal_urls
-
-# TODO could eventually refactor this out of Zoobot as well
-from zoobot.shared import label_metadata
-
-
-class RingsDataset(galaxy_dataset.GalaxyDataset):
-    
-    def __init__(self, root, train=True, download=False, transform=None, target_transform=None):
-
-        catalog, label_cols = rings_setup(root, train, download)
-
-        super().__init__(catalog, label_cols, transform, target_transform)
-
-
-def rings_setup(root, train, download):
-    resources = [
-        (internal_urls.rings_train_catalog, 'e2fb6b2bca45cd7f1c58f5b4089a5976'),  # train catalog
-        (internal_urls.rings_test_catalog, '6e3f362a6e19ecd02675eaa48f6727f0'),  # test catalog
-        (internal_urls.rings_images, 'd0950250436a05ce88de747e6af825b6')  # the images
-    ]
-    images_to_spotcheck = []
-
-    downloader = download_utils.DatasetDownloader(root, resources, images_to_spotcheck)
-    if download is True:
-        downloader.download()
-
-    # useful_columns = label_cols + ['subfolder', 'filename']
-    useful_columns = None
-    if train:
-        train_catalog_loc = os.path.join(root, 'rings_train_catalog.parquet')
-        catalog = pd.read_parquet(train_catalog_loc, columns=useful_columns)
-    else:
-        test_catalog_loc = os.path.join(root, 'rings_test_catalog.parquet')
-        catalog = pd.read_parquet(test_catalog_loc, columns=useful_columns)
-
-    catalog['file_loc'] = catalog.apply(lambda x: os.path.join(downloader.image_dir, x['subfolder'], x['filename']), axis=1)
-
-    label_cols = label_metadata.rings_label_cols
-
-    return catalog, label_cols
-
-
-if __name__ == '__main__':
-
-    # personal stuff to pick my own root, comment out or adapt to your machine
-    if os.path.isdir('/Users/walml'):
-        root = '/Users/walml/repos/pytorch-galaxy-datasets/roots/rings'
-    elif os.path.sidir('/nvme1/scratch'):
-        root = '/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/rings'
-    else:
-        raise ValueError
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    rings_datset = RingsDataset(
-        root=root,
-        train=True,
-        download=True
-    )
-    rings_catalog = rings_datset.catalog
-    adjusted_catalog = rings_catalog.sample(1000)
-
-    # user will probably tweak and use images/catalog directly for generic galaxy catalog datamodule
-    # (which makes its own generic datasets internally)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=['ring_fraction'],  # counts and totals also available
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-        
\ No newline at end of file
diff --git a/replication/check_all_images_downloaded.py b/replication/check_all_images_downloaded.py
new file mode 100644
index 0000000..838d740
--- /dev/null
+++ b/replication/check_all_images_downloaded.py
@@ -0,0 +1,13 @@
+import os
+
+from galaxy_datasets.pytorch.datasets import GZDesiDataset
+
+if __name__ == '__main__':
+
+    # or any other dataset
+    dataset = GZDesiDataset(root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/gz_desi', train=True, download=False)
+    for loc in dataset.catalog['file_loc']:
+        if not os.path.isfile(loc):
+            raise FileNotFoundError(loc)
+
+    print('All images successfully located')
\ No newline at end of file
diff --git a/notebooks/legs_checksums.py b/replication/get_chunk_checksums.py
similarity index 52%
rename from notebooks/legs_checksums.py
rename to replication/get_chunk_checksums.py
index 5ecc758..ba161b6 100644
--- a/notebooks/legs_checksums.py
+++ b/replication/get_chunk_checksums.py
@@ -1,10 +1,14 @@
 import os
 import hashlib
+import glob
 
 
 if __name__ == '__main__':    
 
-    for loc in ['/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/legs/low_z_jpg_chunk_0{}_archive.tar.gz'.format(n) for n in range(8)]:
+    root_to_checksum = 'gz_desi'
+    locs_to_checksum = glob.glob(f'/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/{root_to_checksum}/*_chunk_0*.tar.gz')
+
+    for loc in locs_to_checksum:
         # print hash
         with open(loc, 'rb') as f:
             md5_checksum = hashlib.md5(f.read()).hexdigest()
diff --git a/replication/list_of_images_to_chunks.sh b/replication/list_of_images_to_chunks.sh
new file mode 100644
index 0000000..76829fc
--- /dev/null
+++ b/replication/list_of_images_to_chunks.sh
@@ -0,0 +1,41 @@
+#!/bin/bash
+#SBATCH --job-name=copy                    # Job name
+#SBATCH --output=copy_%A.log 
+#SBATCH --mem=0                                     # "reserve all the available memory on each node assigned to the job"
+#SBATCH --no-requeue                                    # Do not resubmit a failed job
+#SBATCH --time=23:00:00                                # Time limit hrs:min:sec
+#SBATCH --exclusive   # only one task per node
+#SBATCH --ntasks 1
+#SBATCH --cpus-per-task=24
+pwd; hostname; date
+
+# this script takes a list of images (already saved from elsewhere e.g.) and groups those images into .tar.gz "chunks" for easy upload/download
+
+# from jpg dir
+cd /share/nas2/walml/galaxy_zoo/decals/dr8/jpg
+
+# gz desi
+
+# make sure all_files list is up-to-date (run locally)
+# rsync -avz /nvme1/scratch/walml/repos/pytorch-galaxy-datasets/derived_data/gz_desi_all_files.txt walml@galahad.ast.man.ac.uk:/share/nas2/walml/repos/pytorch-galaxy-datasets/derived_data
+
+# split main file into 50k chunks, named like temp/gz_desi_* where * is an int
+split -d -l 50000 /share/nas2/walml/repos/pytorch-galaxy-datasets/derived_data/gz_desi_all_files.txt gz_desi_chunks/gz_desi_chunk_
+
+# get the paths to those files (temp/...), and then use xargs to start 24 processes, each using tar to read the paths in a chunkfile (-T) and save to path.tar.gz
+ls gz_desi_chunks/gz_desi_chunk_* | xargs -n 1 -P 24 -i tar -czvf {}_archive.tar.gz -T {} 
+
+# then rsync them back as usual
+
+
+
+
+
+
+# low z legs (deprecated for now)
+
+# # split main file into 200k chunks, named like temp/low_z_jpg_* where * is an int
+# split -d -l 200000 /share/nas2/walml/repos/pytorch-galaxy-datasets/notebooks/temp_legs_z_below_0p1_all_files.txt low_z_jpg_chunks/low_z_jpg_chunk_
+
+# # get the paths to those files (low_z_jpg_chunks/...), and then use xargs to start 24 processes, each using tar to read the paths in a chunkfile (-T) and save to path.tar.gz
+# ls low_z_jpg_chunks/low_z_jpg_chunk_* | xargs -n 1 -P 24 -i tar -czvf {}_archive.tar.gz -T {} 
diff --git a/replication/make_list_of_images.py b/replication/make_list_of_images.py
new file mode 100644
index 0000000..3eb1fbf
--- /dev/null
+++ b/replication/make_list_of_images.py
@@ -0,0 +1,57 @@
+import chunk
+import pandas as pd
+
+
+if __name__ == '__main__':
+
+    # TODO dr8_id vs id_str...
+    train = pd.read_parquet('/home/walml/repos/pytorch-galaxy-datasets/roots/gz_desi/gz_desi_train_catalog.parquet', columns=['id_str', 'file_loc', 'brickid'])
+    test = pd.read_parquet('/home/walml/repos/pytorch-galaxy-datasets/roots/gz_desi/gz_desi_test_catalog.parquet', columns=['id_str', 'file_loc', 'brickid'])
+    df = pd.concat([train, test], axis=0).reset_index()
+
+
+
+    # df['id_str'] = df['dr8_id']
+    # df['brickid'] = df['id_str'].apply(lambda x: x.split('_')[0])
+    # df['objid'] = df['id_str'].apply(lambda x: x.split('_')[1])
+
+    # change file_locs for rsync to be relative to rsync root
+    df['file_loc'] = df['brickid'] + '/' + df['id_str'] + '.jpg'
+
+    print(df)
+
+    # print(df['file_loc'][0])
+    print('Galaxies: {}'.format(len(df)))
+
+    all_files = list(df['file_loc'])
+    # print(all_files[:3])
+
+    with open('/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/derived_data/gz_desi_all_files.txt', 'w') as f:
+        f.writelines([x + '\n' for x in all_files])
+    # next, used in legs_low_z_copy.sh
+
+
+
+    # all_files_base_path = '/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/derived_data/gz_desi_chunks'
+
+    # # rsync is really slow, probably super-linearly, at building file list for large transfers
+    # # split into many chunks instead
+    # chunk_size = 50000
+    # file_n = 0
+    # chunk_paths = []
+    # while file_n < len(all_files):
+    #     this_chunk_loc = all_files_base_path + '_{}.txt'.format(file_n)
+
+    #     these_files = all_files[file_n:file_n+chunk_size]
+
+    #     with open(this_chunk_loc, 'w') as f:
+    #         f.write('\n'.join(these_files))
+
+    #     chunk_paths.append(this_chunk_loc)
+    #     file_n += chunk_size
+
+    # print(' \n'.join(chunk_paths))
+
+    # print(f'\nrsync --dry-run --files-from {chunk_paths[0]} walml@galahad.ast.man.ac.uk:/share/nas2/walml/galaxy_zoo/decals/dr8/jpg  /home/walml/repos/pytorch-galaxy-datasets/derived_data/images')
+
+    # # print(f'rsync -v --files-from {chunk_paths[0]} walml@galahad.ast.man.ac.uk:/share/nas2/walml/galaxy_zoo/decals/dr8/jpg  /home/walml/repos/pytorch-galaxy-datasets/derived_data/legs_z_below_0p1_to_upload')
\ No newline at end of file
diff --git a/notebooks/dr5_rsync.ipynb b/replication/notebooks/deprecated/dr5_rsync.ipynb
similarity index 100%
rename from notebooks/dr5_rsync.ipynb
rename to replication/notebooks/deprecated/dr5_rsync.ipynb
diff --git a/notebooks/legs_low_z_rsync.py b/replication/notebooks/deprecated/legs_low_z_rsync.py
similarity index 100%
rename from notebooks/legs_low_z_rsync.py
rename to replication/notebooks/deprecated/legs_low_z_rsync.py
diff --git a/notebooks/legs_split_original_images.ipynb b/replication/notebooks/deprecated/legs_split_original_images.ipynb
similarity index 100%
rename from notebooks/legs_split_original_images.ipynb
rename to replication/notebooks/deprecated/legs_split_original_images.ipynb
diff --git a/notebooks/download_and_split_candels.ipynb b/replication/notebooks/download_and_split_candels.ipynb
similarity index 100%
rename from notebooks/download_and_split_candels.ipynb
rename to replication/notebooks/download_and_split_candels.ipynb
diff --git a/notebooks/download_and_split_hubble.ipynb b/replication/notebooks/download_and_split_hubble.ipynb
similarity index 100%
rename from notebooks/download_and_split_hubble.ipynb
rename to replication/notebooks/download_and_split_hubble.ipynb
diff --git a/notebooks/dr5_split.ipynb b/replication/notebooks/dr5_split.ipynb
similarity index 100%
rename from notebooks/dr5_split.ipynb
rename to replication/notebooks/dr5_split.ipynb
diff --git a/notebooks/gz2_split.ipynb b/replication/notebooks/gz2_split.ipynb
similarity index 100%
rename from notebooks/gz2_split.ipynb
rename to replication/notebooks/gz2_split.ipynb
diff --git a/replication/notebooks/legs_split_dr8_only.ipynb b/replication/notebooks/legs_split_dr8_only.ipynb
new file mode 100644
index 0000000..3186c30
--- /dev/null
+++ b/replication/notebooks/legs_split_dr8_only.ipynb
@@ -0,0 +1,1052 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "\n",
+    "from sklearn.model_selection import train_test_split\n",
+    "\n",
+    "from zoobot.shared.label_metadata import decals_all_campaigns_ortho_label_cols\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dr12 = pd.read_parquet('/home/walml/repos/gz-decals-classifiers/data/catalogs/training_catalogs/dr12_ortho_v5_labelled_catalog.parquet')\n",
+    "dr5 = pd.read_parquet('/home/walml/repos/gz-decals-classifiers/data/catalogs/training_catalogs/dr5_ortho_v5_labelled_catalog.parquet')\n",
+    "dr8 = pd.read_parquet('/home/walml/repos/gz-decals-classifiers/data/catalogs/training_catalogs/dr8_ortho_v5_labelled_catalog.parquet')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>iauname</th>\n",
+       "      <th>in_gzd_c</th>\n",
+       "      <th>data_release</th>\n",
+       "      <th>smooth-or-featured-dr12_total-votes</th>\n",
+       "      <th>smooth-or-featured-dr12_smooth</th>\n",
+       "      <th>smooth-or-featured-dr12_smooth_fraction</th>\n",
+       "      <th>smooth-or-featured-dr12_featured-or-disk</th>\n",
+       "      <th>smooth-or-featured-dr12_featured-or-disk_fraction</th>\n",
+       "      <th>smooth-or-featured-dr12_artifact</th>\n",
+       "      <th>smooth-or-featured-dr12_artifact_fraction</th>\n",
+       "      <th>...</th>\n",
+       "      <th>spiral-arm-count-dr5_3_fraction</th>\n",
+       "      <th>spiral-arm-count-dr5_4_fraction</th>\n",
+       "      <th>spiral-arm-count-dr5_more-than-4_fraction</th>\n",
+       "      <th>spiral-arm-count-dr5_cant-tell_fraction</th>\n",
+       "      <th>merging-dr5_none_fraction</th>\n",
+       "      <th>merging-dr5_minor-disturbance_fraction</th>\n",
+       "      <th>merging-dr5_major-disturbance_fraction</th>\n",
+       "      <th>merging-dr5_merger_fraction</th>\n",
+       "      <th>dr_with_most_votes</th>\n",
+       "      <th>id_str</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>J161835.06+094825.4</td>\n",
+       "      <td>False</td>\n",
+       "      <td>DR1</td>\n",
+       "      <td>47</td>\n",
+       "      <td>29.0</td>\n",
+       "      <td>0.617021</td>\n",
+       "      <td>8.0</td>\n",
+       "      <td>0.170213</td>\n",
+       "      <td>10.0</td>\n",
+       "      <td>0.212766</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>387274_3520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>J171558.33+294918.8</td>\n",
+       "      <td>False</td>\n",
+       "      <td>DR2</td>\n",
+       "      <td>36</td>\n",
+       "      <td>22.0</td>\n",
+       "      <td>0.611111</td>\n",
+       "      <td>4.0</td>\n",
+       "      <td>0.111111</td>\n",
+       "      <td>10.0</td>\n",
+       "      <td>0.277778</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>495330_3526</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>J152604.39+140215.0</td>\n",
+       "      <td>True</td>\n",
+       "      <td>DR2</td>\n",
+       "      <td>38</td>\n",
+       "      <td>28.0</td>\n",
+       "      <td>0.736842</td>\n",
+       "      <td>8.0</td>\n",
+       "      <td>0.210526</td>\n",
+       "      <td>2.0</td>\n",
+       "      <td>0.052632</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>0.800000</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.200000</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>411197_3292</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>J074034.26+240659.4</td>\n",
+       "      <td>False</td>\n",
+       "      <td>DR2</td>\n",
+       "      <td>38</td>\n",
+       "      <td>17.0</td>\n",
+       "      <td>0.447368</td>\n",
+       "      <td>14.0</td>\n",
+       "      <td>0.368421</td>\n",
+       "      <td>7.0</td>\n",
+       "      <td>0.184211</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>465235_5856</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>J161921.36+131433.9</td>\n",
+       "      <td>False</td>\n",
+       "      <td>DR2</td>\n",
+       "      <td>35</td>\n",
+       "      <td>29.0</td>\n",
+       "      <td>0.828571</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>0.028571</td>\n",
+       "      <td>5.0</td>\n",
+       "      <td>0.142857</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>407046_2961</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>82355</th>\n",
+       "      <td>J082618.08+111227.6</td>\n",
+       "      <td>False</td>\n",
+       "      <td>DR2</td>\n",
+       "      <td>37</td>\n",
+       "      <td>28.0</td>\n",
+       "      <td>0.756757</td>\n",
+       "      <td>2.0</td>\n",
+       "      <td>0.054054</td>\n",
+       "      <td>7.0</td>\n",
+       "      <td>0.189189</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>395315_1473</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>82356</th>\n",
+       "      <td>J225716.71+000240.9</td>\n",
+       "      <td>False</td>\n",
+       "      <td>DR1</td>\n",
+       "      <td>44</td>\n",
+       "      <td>35.0</td>\n",
+       "      <td>0.795455</td>\n",
+       "      <td>5.0</td>\n",
+       "      <td>0.113636</td>\n",
+       "      <td>4.0</td>\n",
+       "      <td>0.090909</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>331745_4989</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>82357</th>\n",
+       "      <td>J154806.17+142227.1</td>\n",
+       "      <td>True</td>\n",
+       "      <td>DR2</td>\n",
+       "      <td>29</td>\n",
+       "      <td>10.0</td>\n",
+       "      <td>0.344828</td>\n",
+       "      <td>17.0</td>\n",
+       "      <td>0.586207</td>\n",
+       "      <td>2.0</td>\n",
+       "      <td>0.068966</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>0.833333</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>0.166667</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>412616_3932</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>82358</th>\n",
+       "      <td>J001510.54-010434.1</td>\n",
+       "      <td>False</td>\n",
+       "      <td>DR1</td>\n",
+       "      <td>49</td>\n",
+       "      <td>39.0</td>\n",
+       "      <td>0.795918</td>\n",
+       "      <td>4.0</td>\n",
+       "      <td>0.081633</td>\n",
+       "      <td>6.0</td>\n",
+       "      <td>0.122449</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>324623_1219</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>82359</th>\n",
+       "      <td>J114404.49+202855.9</td>\n",
+       "      <td>False</td>\n",
+       "      <td>DR2</td>\n",
+       "      <td>38</td>\n",
+       "      <td>21.0</td>\n",
+       "      <td>0.552632</td>\n",
+       "      <td>14.0</td>\n",
+       "      <td>0.368421</td>\n",
+       "      <td>3.0</td>\n",
+       "      <td>0.078947</td>\n",
+       "      <td>...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>dr12</td>\n",
+       "      <td>446774_1876</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>82360 rows × 204 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                   iauname  in_gzd_c data_release  \\\n",
+       "0      J161835.06+094825.4     False          DR1   \n",
+       "1      J171558.33+294918.8     False          DR2   \n",
+       "2      J152604.39+140215.0      True          DR2   \n",
+       "3      J074034.26+240659.4     False          DR2   \n",
+       "4      J161921.36+131433.9     False          DR2   \n",
+       "...                    ...       ...          ...   \n",
+       "82355  J082618.08+111227.6     False          DR2   \n",
+       "82356  J225716.71+000240.9     False          DR1   \n",
+       "82357  J154806.17+142227.1      True          DR2   \n",
+       "82358  J001510.54-010434.1     False          DR1   \n",
+       "82359  J114404.49+202855.9     False          DR2   \n",
+       "\n",
+       "       smooth-or-featured-dr12_total-votes  smooth-or-featured-dr12_smooth  \\\n",
+       "0                                       47                            29.0   \n",
+       "1                                       36                            22.0   \n",
+       "2                                       38                            28.0   \n",
+       "3                                       38                            17.0   \n",
+       "4                                       35                            29.0   \n",
+       "...                                    ...                             ...   \n",
+       "82355                                   37                            28.0   \n",
+       "82356                                   44                            35.0   \n",
+       "82357                                   29                            10.0   \n",
+       "82358                                   49                            39.0   \n",
+       "82359                                   38                            21.0   \n",
+       "\n",
+       "       smooth-or-featured-dr12_smooth_fraction  \\\n",
+       "0                                     0.617021   \n",
+       "1                                     0.611111   \n",
+       "2                                     0.736842   \n",
+       "3                                     0.447368   \n",
+       "4                                     0.828571   \n",
+       "...                                        ...   \n",
+       "82355                                 0.756757   \n",
+       "82356                                 0.795455   \n",
+       "82357                                 0.344828   \n",
+       "82358                                 0.795918   \n",
+       "82359                                 0.552632   \n",
+       "\n",
+       "       smooth-or-featured-dr12_featured-or-disk  \\\n",
+       "0                                           8.0   \n",
+       "1                                           4.0   \n",
+       "2                                           8.0   \n",
+       "3                                          14.0   \n",
+       "4                                           1.0   \n",
+       "...                                         ...   \n",
+       "82355                                       2.0   \n",
+       "82356                                       5.0   \n",
+       "82357                                      17.0   \n",
+       "82358                                       4.0   \n",
+       "82359                                      14.0   \n",
+       "\n",
+       "       smooth-or-featured-dr12_featured-or-disk_fraction  \\\n",
+       "0                                               0.170213   \n",
+       "1                                               0.111111   \n",
+       "2                                               0.210526   \n",
+       "3                                               0.368421   \n",
+       "4                                               0.028571   \n",
+       "...                                                  ...   \n",
+       "82355                                           0.054054   \n",
+       "82356                                           0.113636   \n",
+       "82357                                           0.586207   \n",
+       "82358                                           0.081633   \n",
+       "82359                                           0.368421   \n",
+       "\n",
+       "       smooth-or-featured-dr12_artifact  \\\n",
+       "0                                  10.0   \n",
+       "1                                  10.0   \n",
+       "2                                   2.0   \n",
+       "3                                   7.0   \n",
+       "4                                   5.0   \n",
+       "...                                 ...   \n",
+       "82355                               7.0   \n",
+       "82356                               4.0   \n",
+       "82357                               2.0   \n",
+       "82358                               6.0   \n",
+       "82359                               3.0   \n",
+       "\n",
+       "       smooth-or-featured-dr12_artifact_fraction  ...  \\\n",
+       "0                                       0.212766  ...   \n",
+       "1                                       0.277778  ...   \n",
+       "2                                       0.052632  ...   \n",
+       "3                                       0.184211  ...   \n",
+       "4                                       0.142857  ...   \n",
+       "...                                          ...  ...   \n",
+       "82355                                   0.189189  ...   \n",
+       "82356                                   0.090909  ...   \n",
+       "82357                                   0.068966  ...   \n",
+       "82358                                   0.122449  ...   \n",
+       "82359                                   0.078947  ...   \n",
+       "\n",
+       "       spiral-arm-count-dr5_3_fraction  spiral-arm-count-dr5_4_fraction  \\\n",
+       "0                                  NaN                              NaN   \n",
+       "1                                  NaN                              NaN   \n",
+       "2                                  NaN                              NaN   \n",
+       "3                                  NaN                              NaN   \n",
+       "4                                  NaN                              NaN   \n",
+       "...                                ...                              ...   \n",
+       "82355                              NaN                              NaN   \n",
+       "82356                              NaN                              NaN   \n",
+       "82357                              NaN                              NaN   \n",
+       "82358                              NaN                              NaN   \n",
+       "82359                              NaN                              NaN   \n",
+       "\n",
+       "       spiral-arm-count-dr5_more-than-4_fraction  \\\n",
+       "0                                            NaN   \n",
+       "1                                            NaN   \n",
+       "2                                            NaN   \n",
+       "3                                            NaN   \n",
+       "4                                            NaN   \n",
+       "...                                          ...   \n",
+       "82355                                        NaN   \n",
+       "82356                                        NaN   \n",
+       "82357                                        NaN   \n",
+       "82358                                        NaN   \n",
+       "82359                                        NaN   \n",
+       "\n",
+       "       spiral-arm-count-dr5_cant-tell_fraction  merging-dr5_none_fraction  \\\n",
+       "0                                          NaN                        NaN   \n",
+       "1                                          NaN                        NaN   \n",
+       "2                                          NaN                   0.800000   \n",
+       "3                                          NaN                        NaN   \n",
+       "4                                          NaN                        NaN   \n",
+       "...                                        ...                        ...   \n",
+       "82355                                      NaN                        NaN   \n",
+       "82356                                      NaN                        NaN   \n",
+       "82357                                      NaN                   0.833333   \n",
+       "82358                                      NaN                        NaN   \n",
+       "82359                                      NaN                        NaN   \n",
+       "\n",
+       "       merging-dr5_minor-disturbance_fraction  \\\n",
+       "0                                         NaN   \n",
+       "1                                         NaN   \n",
+       "2                                         0.0   \n",
+       "3                                         NaN   \n",
+       "4                                         NaN   \n",
+       "...                                       ...   \n",
+       "82355                                     NaN   \n",
+       "82356                                     NaN   \n",
+       "82357                                     0.0   \n",
+       "82358                                     NaN   \n",
+       "82359                                     NaN   \n",
+       "\n",
+       "       merging-dr5_major-disturbance_fraction  merging-dr5_merger_fraction  \\\n",
+       "0                                         NaN                          NaN   \n",
+       "1                                         NaN                          NaN   \n",
+       "2                                         0.0                     0.200000   \n",
+       "3                                         NaN                          NaN   \n",
+       "4                                         NaN                          NaN   \n",
+       "...                                       ...                          ...   \n",
+       "82355                                     NaN                          NaN   \n",
+       "82356                                     NaN                          NaN   \n",
+       "82357                                     0.0                     0.166667   \n",
+       "82358                                     NaN                          NaN   \n",
+       "82359                                     NaN                          NaN   \n",
+       "\n",
+       "       dr_with_most_votes       id_str  \n",
+       "0                    dr12  387274_3520  \n",
+       "1                    dr12  495330_3526  \n",
+       "2                    dr12  411197_3292  \n",
+       "3                    dr12  465235_5856  \n",
+       "4                    dr12  407046_2961  \n",
+       "...                   ...          ...  \n",
+       "82355                dr12  395315_1473  \n",
+       "82356                dr12  331745_4989  \n",
+       "82357                dr12  412616_3932  \n",
+       "82358                dr12  324623_1219  \n",
+       "82359                dr12  446774_1876  \n",
+       "\n",
+       "[82360 rows x 204 columns]"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dr12"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>dr8_id</th>\n",
+       "      <th>iauname</th>\n",
+       "      <th>redshift</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>6323738</th>\n",
+       "      <td>650362_3400</td>\n",
+       "      <td>None</td>\n",
+       "      <td>0.163198</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32059</th>\n",
+       "      <td>111374_5632</td>\n",
+       "      <td>None</td>\n",
+       "      <td>0.204382</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7068380</th>\n",
+       "      <td>564595_5357</td>\n",
+       "      <td>None</td>\n",
+       "      <td>0.198830</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1046919</th>\n",
+       "      <td>588287_2869</td>\n",
+       "      <td>None</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1495212</th>\n",
+       "      <td>534627_1502</td>\n",
+       "      <td>None</td>\n",
+       "      <td>0.191561</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "              dr8_id iauname  redshift\n",
+       "6323738  650362_3400    None  0.163198\n",
+       "32059    111374_5632    None  0.204382\n",
+       "7068380  564595_5357    None  0.198830\n",
+       "1046919  588287_2869    None       NaN\n",
+       "1495212  534627_1502    None  0.191561"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "megacatalog = pd.read_parquet('/home/walml/repos/decals-rings/results/rings_megacatalog.parquet', columns=['dr8_id', 'iauname', 'redshift'])\n",
+    "megacatalog.sample(5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "82360\n",
+      "82360\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(len(dr12))\n",
+    "dr12 = pd.merge(dr12, megacatalog, on='iauname', how='inner')\n",
+    "print(len(dr12))\n",
+    "# 85730\n",
+    "# 82387"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "222754\n",
+      "222754\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(len(dr5))\n",
+    "dr5 = pd.merge(dr5, megacatalog, on='iauname', how='inner')\n",
+    "print(len(dr5))\n",
+    "# 228059\n",
+    "# 222754"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "401250"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "labelled = pd.concat([dr12, dr5, dr8], axis=0)\n",
+    "len(labelled)\n",
+    "# 365770"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now make the unlabelled catalog from the remainder"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:2: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  \n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "8593234"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "unlabelled = megacatalog[~megacatalog['dr8_id'].isin(labelled['dr8_id'])]\n",
+    "unlabelled['id_str'] = unlabelled['dr8_id']\n",
+    "len(unlabelled)  # mostly high-z, some redshift filtering would be wise\n",
+    "# 8323600"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:2: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  \n",
+      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:2: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
+      "  \n"
+     ]
+    }
+   ],
+   "source": [
+    "for col in decals_all_campaigns_ortho_label_cols:\n",
+    "    unlabelled[col] = 0\n",
+    "unlabelled = unlabelled.copy().sample(len(unlabelled), random_state=42)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:3: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
+      "  This is separate from the ipykernel package so we can avoid doing imports until\n",
+      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:4: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
+      "  after removing the cwd from sys.path.\n",
+      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
+      "  \"\"\"\n",
+      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:6: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
+      "  \n",
+      "/nvme1/scratch/miniconda3/envs/zoobot/lib/python3.7/site-packages/ipykernel_launcher.py:7: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
+      "  import sys\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0    /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/38...\n",
+      "0    /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/30...\n",
+      "0    /share/nas2/walml/galaxy_zoo/decals/dr8/jpg/10...\n",
+      "Name: file_loc, dtype: object 0    387274_3520.jpg\n",
+      "0    302281_3952.jpg\n",
+      "0    107510_4476.jpg\n",
+      "Name: filename, dtype: object\n",
+      "/share/nas2/walml/galaxy_zoo/decals/dr8/jpg/43543/43543_92.jpg 43543_92.jpg\n"
+     ]
+    }
+   ],
+   "source": [
+    "for df in [labelled, unlabelled]:  # dr8 already has but it points to the png\n",
+    "    df['dr8_id'] = df['id_str'] # ensure consistency\n",
+    "    df['brickid'] = df['id_str'].apply(lambda x: x.split('_')[0])\n",
+    "    df['objid'] = df['id_str'].apply(lambda x: x.split('_')[1])\n",
+    "    df['subfolder'] = df['brickid']\n",
+    "    df['filename'] = df['id_str'] + '.jpg'\n",
+    "    df['file_loc'] = '/share/nas2/walml/galaxy_zoo/decals/dr8/jpg/' + df['brickid'] + '/' + df['filename']\n",
+    "\n",
+    "    print(df['file_loc'][0], df['filename'][0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "False"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.any(labelled[decals_all_campaigns_ortho_label_cols].isna())\n",
+    "# any(labelled[decals_all_campaigns_ortho_label_cols].isna()) doesn't work right"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.0"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.any(labelled[decals_all_campaigns_ortho_label_cols].isna(), axis=1).mean()\n",
+    "# any(labelled[decals_all_campaigns_ortho_label_cols].isna().values.max(axis=1)) doesn't work right"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_catalog, test_catalog = train_test_split(labelled, test_size=0.2, random_state=42)  # note that because we dropped a few rows which couldn't be cross-matched, this is not the same split as the original_images version"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_catalog_loc = '/home/walml/repos/pytorch-galaxy-datasets/roots/gz_desi/gz_desi_train_catalog.parquet'\n",
+    "test_catalog_loc = '/home/walml/repos/pytorch-galaxy-datasets/roots/gz_desi/gz_desi_test_catalog.parquet'\n",
+    "unlabelled_catalog_loc = '/home/walml/repos/pytorch-galaxy-datasets/roots/desi_unlabelled/desi_unlabelled_catalog.parquet'\n",
+    "\n",
+    "train_catalog = train_catalog.reset_index(drop=True)\n",
+    "test_catalog = test_catalog.reset_index(drop=True)\n",
+    "unlabelled = unlabelled.reset_index(drop=True)\n",
+    "\n",
+    "train_catalog.to_parquet(train_catalog_loc, index=False)\n",
+    "test_catalog.to_parquet(test_catalog_loc, index=False)\n",
+    "\n",
+    "# this will not be uploaded as images, wayyyy to big\n",
+    "unlabelled.to_parquet(unlabelled_catalog_loc, index=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "50997398a10886dafbe556d071ffbc68\n",
+      "386f13daf123f439ffb8266b65eac474\n",
+      "6ca344a997661e1e58a3c2128d4ff6d8\n"
+     ]
+    }
+   ],
+   "source": [
+    "import hashlib\n",
+    "\n",
+    "for loc in [\n",
+    "    train_catalog_loc,\n",
+    "    test_catalog_loc,\n",
+    "    unlabelled_catalog_loc\n",
+    "]:\n",
+    "    # print hash\n",
+    "    with open(loc, 'rb') as f:\n",
+    "        md5_checksum = hashlib.md5(f.read()).hexdigest()\n",
+    "\n",
+    "    print(md5_checksum)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>filename</th>\n",
+       "      <th>subfolder</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>330243_908.jpg</td>\n",
+       "      <td>330243</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>473526_238.jpg</td>\n",
+       "      <td>473526</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>335382_2814.jpg</td>\n",
+       "      <td>335382</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>495307_503.jpg</td>\n",
+       "      <td>495307</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>378617_1005.jpg</td>\n",
+       "      <td>378617</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>320995</th>\n",
+       "      <td>367213_749.jpg</td>\n",
+       "      <td>367213</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>320996</th>\n",
+       "      <td>311438_6735.jpg</td>\n",
+       "      <td>311438</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>320997</th>\n",
+       "      <td>507311_2341.jpg</td>\n",
+       "      <td>507311</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>320998</th>\n",
+       "      <td>321115_1721.jpg</td>\n",
+       "      <td>321115</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>320999</th>\n",
+       "      <td>379987_2914.jpg</td>\n",
+       "      <td>379987</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>321000 rows × 2 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "               filename subfolder\n",
+       "0        330243_908.jpg    330243\n",
+       "1        473526_238.jpg    473526\n",
+       "2       335382_2814.jpg    335382\n",
+       "3        495307_503.jpg    495307\n",
+       "4       378617_1005.jpg    378617\n",
+       "...                 ...       ...\n",
+       "320995   367213_749.jpg    367213\n",
+       "320996  311438_6735.jpg    311438\n",
+       "320997  507311_2341.jpg    507311\n",
+       "320998  321115_1721.jpg    321115\n",
+       "320999  379987_2914.jpg    379987\n",
+       "\n",
+       "[321000 rows x 2 columns]"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "train_catalog[['filename', 'subfolder']]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# allow user to filter by redshift as needed"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "5ec7511d1fc7a72aba078c8815c4bcacd5404aac7a6282065ad4abec249876bb"
+  },
+  "kernelspec": {
+   "display_name": "Python 3.7.11 ('zoobot')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.11"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/rings_split.ipynb b/replication/notebooks/rings_split.ipynb
similarity index 100%
rename from notebooks/rings_split.ipynb
rename to replication/notebooks/rings_split.ipynb
diff --git a/notebooks/tidal_split.ipynb b/replication/notebooks/tidal_split.ipynb
similarity index 100%
rename from notebooks/tidal_split.ipynb
rename to replication/notebooks/tidal_split.ipynb
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index fd37c94..0000000
--- a/requirements.txt
+++ /dev/null
@@ -1,6 +0,0 @@
-torch >= 1.10
-torchvision
-simplejpeg
-albumentations
-numpy
-pandas
diff --git a/setup.py b/setup.py
index 874ec43..41be381 100644
--- a/setup.py
+++ b/setup.py
@@ -4,14 +4,14 @@
     long_description = fh.read()
 
 setuptools.setup(
-    name="pytorch_galaxy_datasets",
-    version="0.0.1",
+    name="galaxy-datasets",
+    version="0.0.2",
     author="Mike Walmsley",
     author_email="walmsleymk1@gmail.com",
-    description="Galaxy Zoo datasets for PyTorch/Lightning",
+    description="Galaxy Zoo datasets for PyTorch/TensorFlow",
     long_description=long_description,
     long_description_content_type="text/markdown",
-    url="https://github.com/mwalmsley/pytorch-galaxy-datasets",
+    url="https://github.com/mwalmsley/galaxy-datasets",
     classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: GNU General Public License (GPL)",
@@ -20,5 +20,22 @@
         "Environment :: GPU :: NVIDIA CUDA"
     ],
     packages=setuptools.find_packages(),
-    python_requires=">=3.7"
+    python_requires=">=3.7",
+    extras_require={
+        # these match zoobot's reqs
+        # TODO since zoobot depends on this, can now remove from zoobot reqs if desired?
+        'pytorch': [
+            'torch == 1.10.1',
+            'torchvision == 0.11.2',
+            'torchaudio == 0.10.1',
+            'pytorch-lightning==1.6.5',  # 1.7 requires protobuf version incompatible with tensorflow/tensorboard. Otherwise works.
+            'simplejpeg',
+            'albumentations'
+        ],
+        'tensorflow': [
+            'tensorflow >= 2.10.0',
+            'protobuf <= 3.19'  # tensorflow incompatible above this (usually resolved by pip automatically)
+        ]
+        # if you need both, specify both
+    }
 )


Test patch:
diff --git a/tests/pytorch/test_datasets.py b/tests/pytorch/test_datasets.py
new file mode 100644
index 0000000..d550d4a
--- /dev/null
+++ b/tests/pytorch/test_datasets.py
@@ -0,0 +1,93 @@
+import pytest
+
+import os
+
+from galaxy_datasets.pytorch import datasets, galaxy_datamodule
+
+
+def base_root_dir():
+    if os.path.isdir('/nvme1/scratch'):
+        return '/nvme1/scratch/walml/repos/galaxy-datasets/roots'
+    elif os.path.isdir('/share/nas2'):
+        return '/share/nas2/walml/repos/_data'
+    else:
+        raise FileNotFoundError
+
+
+def tidal_dataset():
+    return datasets.Tidal(
+        root=os.path.join(base_root_dir(), 'tidal'),
+        train=True,
+        download=False
+    )
+
+
+def gz_decals_dataset():
+    return datasets.GZDecals5(
+        root=os.path.join(base_root_dir(), 'gz_decals'),
+        train=True,
+        download=False
+    )
+
+def gz_desi_dataset():
+    return datasets.GZDesi(
+        root=os.path.join(base_root_dir(), 'gz_desi'),
+        train=True,
+        download=False
+    )
+
+
+def gz2_dataset():
+    return datasets.GZ2(
+        root=os.path.join(base_root_dir(), 'gz2'),
+        download=False
+    )
+
+
+def gz_rings_dataset():
+    return datasets.GZRings(
+        root=os.path.join(base_root_dir(), 'gz_rings'),
+        train=True,
+        download=False
+    )
+
+
+def gz_hubble_dataset():
+    return datasets.GZHubble(
+        root=os.path.join(base_root_dir(), 'gz_hubble'),
+        download=False
+    )
+
+
+def gz_candels_dataset():
+    return datasets.GZCandels(
+        root=os.path.join(base_root_dir(), 'gz_candels'),
+        download=False
+    )
+
+# https://docs.pytest.org/en/6.2.x/fixture.html#using-marks-with-parametrized-fixtures
+# pytest.param(gz_desi_dataset, marks=pytest.mark.skip)
+@pytest.fixture(params=[gz2_dataset, gz_candels_dataset, gz_decals_dataset, gz_hubble_dataset, gz_rings_dataset, tidal_dataset, gz_desi_dataset])
+def dataset(request):
+    return request.param()
+
+
+def test_dataset(dataset):
+
+    catalog = dataset.catalog
+    label_cols = dataset.label_cols
+
+    adjusted_catalog = catalog.sample(1000)
+
+    # user will probably tweak and use images/catalog directly for generic galaxy catalog datamodule
+    # (which makes its own generic datasets internally)
+    datamodule = galaxy_datamodule.GalaxyDataModule(
+        label_cols=label_cols,
+        catalog=adjusted_catalog
+    )
+
+    datamodule.prepare_data()
+    datamodule.setup()
+    for images, labels in datamodule.train_dataloader():
+        print(images.shape, labels.shape)
+        break
diff --git a/tests/readme/test_readme.py b/tests/readme/test_readme.py
new file mode 100644
index 0000000..757985a
--- /dev/null
+++ b/tests/readme/test_readme.py
@@ -0,0 +1,91 @@
+import py
+import pytest
+
+import os
+
+@pytest.fixture()
+def example_root():
+    if os.path.isdir('/share/nas2'):
+        return '/share/nas2/walml/repos/_data/gz2'
+    else:
+        return '/nvme1/scratch/walml/repos/galaxy-datasets/roots/gz2'
+
+@pytest.fixture()  # also a test
+def get_catalog(example_root):
+    from galaxy_datasets import gz2  # or gz_hubble, gz_candels, ...
+
+    catalog, label_cols = gz2(
+        root=example_root,
+        train=True,
+        download=True
+    )
+    return catalog, label_cols
+
+def test_pytorch_dataset_custom(get_catalog):
+    from galaxy_datasets.pytorch.galaxy_dataset import GalaxyDataset  # generic Dataset for galaxies
+
+    catalog, _ = get_catalog
+
+    dataset = GalaxyDataset(
+        catalog=catalog.sample(1000),  # from gz2(...) above
+        label_cols=['smooth-or-featured-gz2_smooth']
+    )
+
+def test_pytorch_dataset_canonical(example_root):
+
+    from galaxy_datasets.pytorch import GZ2
+
+    gz2_dataset = GZ2(
+        root=example_root,
+        train=True,
+        download=False
+    )
+    image, label = gz2_dataset[0]
+
+
+def test_pytorch_datamodule_custom(get_catalog):
+
+    catalog, _ = get_catalog
+    
+    from galaxy_datasets.pytorch.galaxy_datamodule import GalaxyDataModule
+
+    datamodule = GalaxyDataModule(
+        label_cols=['smooth-or-featured-gz2_smooth'],
+        catalog=catalog
+        # optional args to specify augmentations
+    )
+
+    datamodule.prepare_data()
+    datamodule.setup()
+    for images, labels in datamodule.train_dataloader():
+        print(images.shape, labels.shape)
+        break
+
+# currently has protobuf version error locally
+@pytest.mark.skip
+def test_tensorflow_dataset(get_catalog):
+
+    import tensorflow as tf
+    from galaxy_datasets.tensorflow.datasets import get_image_dataset, add_transforms_to_dataset
+    from galaxy_datasets.transforms import default_transforms  # same transforms as PyTorch
+
+    catalog, label_cols = get_catalog
+
+    train_dataset = get_image_dataset(
+        image_paths = catalog['file_loc'],
+        labels=catalog[label_cols].values,
+        requested_img_size=224
+    )
+
+    # specify augmentations
+    transforms = default_transforms()
+
+    # apply augmentations
+    train_dataset = add_transforms_to_dataset(train_dataset, transforms)
+
+    # batch, shuffle, prefetch for performance
+    train_dataset = train_dataset.shuffle(5000).batch(64).prefetch(tf.data.experimental.AUTOTUNE)
+
+    for images, labels in train_dataset.take(1):
+        print(images.shape, labels.shape)
+        break
\ No newline at end of file
diff --git a/tests/shared/test_downloads.py b/tests/shared/test_downloads.py
new file mode 100644
index 0000000..b2a1424
--- /dev/null
+++ b/tests/shared/test_downloads.py
@@ -0,0 +1,33 @@
+import pytest
+
+import os
+import numpy as np
+import pandas as pd
+from galaxy_datasets.shared import gz_candels, gz_hubble, gz_decals_5, gz2, gz_desi, gz_rings, tidal
+
+# https://docs.pytest.org/en/6.2.x/fixture.html#using-marks-with-parametrized-fixtures
+# in order of dataset size
+# pytest.param(gz_desi.setup, marks=pytest.mark.skip)
+@pytest.fixture(params=[tidal, gz_rings, gz_candels, gz_hubble, gz2, gz_decals_5, gz_desi])
+# @pytest.fixture(params=[gz_desi.setup])
+def setup_method(request):
+    return request.param  # param is the func, param() calls the func (which breaks here as no args)
+
+def test_download_dataset(setup_method, tmp_path):
+    catalog, label_cols = setup_method(
+        train=True,
+        root=tmp_path,
+        download=True
+    )
+
+    assert isinstance(catalog, pd.DataFrame)
+    assert 'file_loc' in catalog.columns.values
+    assert len(catalog) > 0
+    
+    assert isinstance(label_cols, list)
+    assert len(label_cols) > 0
+    for col in label_cols:
+        assert col in catalog.columns.values
+
+    for loc in catalog.sample(10, random_state=42)['file_loc']:
+        assert os.path.isfile(loc)
diff --git a/tests/tensorflow/augmentations.py b/tests/tensorflow/augmentations.py
new file mode 100644
index 0000000..4604349
--- /dev/null
+++ b/tests/tensorflow/augmentations.py
@@ -0,0 +1,84 @@
+
+# from albumentations import Compose
+
+# purely so we can import more easily from e.g. zoobot
+# (needs only galaxy_datasets.tensorflow.augmentations)
+from galaxy_datasets.transforms import default_transforms  
+
+import matplotlib.pyplot as plt
+
+from galaxy_datasets.tensorflow import galaxy_dataset
+from galaxy_datasets.shared import tidal
+from galaxy_datasets.transforms import default_transforms
+
+from albumentations import (
+    Compose, HorizontalFlip,
+    Rotate
+)
+
+if __name__ == '__main__':
+
+
+    catalog, label_cols = tidal.setup(
+        root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/tidal',
+        train=True,
+        download=False)
+
+    dataset = galaxy_dataset.get_image_dataset(
+        image_paths=catalog['file_loc'],
+        file_format='jpg',
+        requested_img_size=64,
+        labels=catalog['coarse_tidal_label'],
+        greyscale=True,
+        permute_channels=False,
+        include_id_str=False
+    )
+   
+    for images, labels in dataset.take(1):
+        print(images.shape)
+        print(labels)
+
+    # transforms = Compose([
+    #     Rotate(limit=90),
+    #     HorizontalFlip()
+    # ])
+    transforms = default_transforms(
+        greyscale=True,
+        crop_scale_bounds=(0.7, 0.8),
+        crop_ratio_bounds=(0.9, 1.1),
+        resize_after_crop=48
+    )
+    # dataset = add_augmentations_to_dataset(dataset, transforms)
+
+    # for image, label in dataset.take(1):
+    #     print(image.shape)
+    #     print(label)
+
+    # # and then apply batch and then prefetch, as needed
+    # dataset = dataset.shuffle(5000).batch(16).prefetch(buffer_size=AUTOTUNE)
+
+    # for images, labels in dataset.take(1):
+    #     print(images.shape)
+    #     print(labels)
+
+    # plt.imshow(images[0])
+    # plt.show()
+
+
+    n_repeats = 6
+    n_galaxies = 4
+
+    fig, rows = plt.subplots(ncols=n_galaxies, nrows=n_repeats)
+
+    for n_repeat in range(n_repeats):
+
+        for row in rows:
+            augmented = add_transforms_to_dataset(dataset, transforms)
+
+            augmented_subset = augmented.take(n_galaxies)
+            for n_galaxy, (image, _) in enumerate(augmented_subset):
+                row[n_galaxy].imshow(image)
+                row[n_galaxy].axis('off')
+
+    plt.tight_layout()
+    plt.show()
diff --git a/tests/test_datamodules.py b/tests/test_datamodules.py
deleted file mode 100644
index 39d5505..0000000
--- a/tests/test_datamodules.py
+++ /dev/null
@@ -1,122 +0,0 @@
-import pytest
-
-import os
-
-from pytorch_galaxy_datasets import galaxy_datamodule
-from pytorch_galaxy_datasets.prepared_datasets import candels, dr5, gz2, legs, rings, tidal
-
-@pytest.fixture
-def base_root_dir():
-    if os.path.isdir('/nvme1/scratch'):
-        return '/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots'
-    elif os.path.isdir('/share/nas2'):
-        return '/share/nas2/walml/repos/_data'
-    else:
-        raise FileNotFoundError
-
-
-def test_tidal_datamodule(base_root_dir):
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    catalog, label_cols = tidal.tidal_setup(
-        root=os.path.join(base_root_dir, 'tidal'),
-        train=True,
-        download=False
-    )
-    adjusted_catalog = catalog.sample(1000)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=label_cols,  # counts and totals also available
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-        
-
-
-def test_rings_datamodule(base_root_dir):
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    catalog, label_cols = rings.rings_setup(
-        root=os.path.join(base_root_dir, 'rings'),
-        train=True,
-        download=False
-    )
-    adjusted_catalog = catalog.sample(1000)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=label_cols,  # counts and totals also available
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-        
-
-def test_candels_datamodule(base_root_dir):
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    catalog, label_cols = candels.candels_setup(
-        root=os.path.join(base_root_dir, 'candels'),
-        train=True,
-        download=False
-    )
-    adjusted_catalog = catalog.sample(1000)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=label_cols,  # counts and totals also available
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-        
-
-def test_dr5_datamodule(base_root_dir):
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    catalog, label_cols = dr5.decals_dr5_setup(
-        root=os.path.join(base_root_dir, 'decals_dr5'),
-        train=True,
-        download=False
-    )
-    adjusted_catalog = catalog.sample(1000)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=label_cols,  # counts and totals also available
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
-
-
-@pytest.mark.skipif(not os.path.isdir('/share/nas2/walml'), reason="Data only exists on Galahad")
-def test_legs_datamodule():
-
-    # first download is basically just a convenient way to get the images and canonical catalogs
-    catalog, label_cols = legs.legs_setup(
-        # root=os.path.join(base_root_dir, '/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/legs',
-        train=True,
-        download=False
-    )
-    adjusted_catalog = catalog.sample(1000)
-    datamodule = galaxy_datamodule.GalaxyDataModule(
-        label_cols=label_cols,  # counts and totals also available
-        catalog=adjusted_catalog
-    )
-
-    datamodule.prepare_data()
-    datamodule.setup()
-    for images, labels in datamodule.train_dataloader():
-        print(images.shape, labels.shape)
-        break
diff --git a/tests/test_downloads.py b/tests/test_downloads.py
deleted file mode 100644
index 41bdcc8..0000000
--- a/tests/test_downloads.py
+++ /dev/null
@@ -1,72 +0,0 @@
-import pytest
-
-import os
-import numpy as np
-from pytorch_galaxy_datasets.prepared_datasets import candels, dr5, gz2, legs, rings, tidal
-
-# first as smallest
-def test_download_tidal_dataset(tmp_path):
-    dataset = tidal.TidalDataset(
-        train=True,
-        root=tmp_path,
-        download=True
-    )
-
-    for image, label in dataset:
-        print(np.array(image).shape, label)
-        break
-
-
-def test_download_candels_dataset(tmp_path):
-    dataset = candels.CandelsDataset(
-        root=tmp_path,
-        download=True
-    )
-    for image, label in dataset:
-        print(np.array(image).shape, label)
-        break
-
-
-def test_download_dr5_dataset(tmp_path):
-    dataset = dr5.DecalsDR5Dataset(
-        root=tmp_path,
-        download=True
-    )
-    for image, label in dataset:
-        print(np.array(image).shape, label)
-        break
-
-
-def test_download_gz2_dataset(tmp_path):
-    dataset = gz2.GZ2Dataset(
-        root=tmp_path,
-        download=True
-    )
-    for image, label in dataset:
-        print(np.array(image).shape, label)
-        break
-
-
-def test_download_rings_dataset(tmp_path):
-
-    dataset = rings.RingsDataset(
-        # root='/nvme1/scratch/walml/repos/pytorch-galaxy-datasets/roots/rings_root/temp',
-        root=tmp_path,
-        download=True
-    )
-    for image, label in dataset:
-        print(np.array(image).shape, label)
-        break
-
-
-@pytest.mark.skipif(not os.path.isdir('/share/nas2/walml'), reason="Data only exists on Galahad")
-def test_download_legs_dataset(tmp_path):
-    dataset = legs.LegsDataset(
-        root=tmp_path,
-        download=True,
-        train='labelled'
-    )
-    for image, label in dataset:
-        print(np.array(image).shape, label)
-        break
-