Repository: TheCacophonyProject/classifier-pipeline (language: python)

Issue:
Updated dbus call for trap trigger



Discussion / hints:

Fix patch:
diff --git a/.travis.yml b/.travis.yml
index b05a2595..7a282184 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -16,8 +16,7 @@ before_script:
 script:
 - pytest -s
 - "./check-style"
-- python3 load.py -r -c tests/test-config.yaml
-- python3 build.py -c tests/test-config.yaml
+- python3 build.py -c tests/test-config.yaml tests --ext ".cptv"
 - python3 motiontest.py --thermal-config-file tests/config.toml -c tests/test-config.yaml
   tests/clips/possum.cptv
 - sudo mkdir /var/spool/cptv
diff --git a/README.md b/README.md
index 79ff9b85..1f25fe11 100644
--- a/README.md
+++ b/README.md
@@ -7,35 +7,19 @@ The output is a TensorFlow model that can identify thermal video clips of animal
 
 # Scripts
 
-### load.py
-Processes CPTV clips and metadata, extracts the tagged frames (supplied by the metadata) and adds them to the hdf5 database.
-
-The tracking algorithm tries to distinguish between animal tracks and false positives, but is not 100% reliable.  For this reason the output of the tracking algorithm should be checked by hand.
-
 
 ### build.py
 Creates training, validation and testing datasets from database of clips & tracks.
-Datasets contain frames and segments ( e.g. 45 frames).
+Datasets contain frames or segments ( e.g. 45 frames).
 
 Frames (important frames) are calculated by choosing frames with a mass ( the count of pixels that have been deemed an object by track extraction) between the lower and upper quartiles of a tracks mass distribution.
 
-Frames are also checked to see if they are noisy frames.
-This attempts to remove the following types of noisy frames:
-- Frames where it has bad tracking e.g. in the middle of an animal
-- Frames where only part of the animal is shown (maybe leaving the cameras field of vision)
-- Frames where nothing is there but noise
-
-Segments are calculated either by (depending on config):
-
-Choosing random permutations of the important frames.
-- Number of Segments (# of important frames - segment duration) // 9 segments are selected
-
-or by choosing segment duration consecutive frames whose mass is above a certain amount        
+Segments are calculated  by choosing segment duration consecutive frames whose mass is above a certain amount        
 - Number of Segments up to (# of frames - segment duration) // segment-frame-spacing
 
 Datasets are split by camera and location ( to try and remove any bias that may occur from using a camera in multiple sets).
 
-Some labels have low amounts of data so a single camera is split into 2 cameras e.g. Wallabies and Leoparidaes
+Some labels have low amounts of data so they are split by clip_id i.e. wallaby and penguin
 
 ### train.py
 Trains a neural net using a provided test / train / validation dataset.
@@ -64,8 +48,8 @@ Choose the region with greatest mass if any regions exist (these are points of i
 
 Otherwise take the frame with the highest mean pixel value and find the highest mean pixel 64 x64 region
 
-### evaluate.py
-Evaluates the performance of a classify.py run and generates reports.
+### modelevaluate.py
+Evaluates the performance of a model
 
 # Setup
 
@@ -97,15 +81,11 @@ CPTV files can be downloaded using the [cptv-downloader](https://github.com/TheC
 
 First download the CPTV files by running
 
-`python cptv-download.py --user x --password x`
-
-Next load the track files.  This can take some time
-
-`python load.py all -v -p`
+`python cptv-download.py <dir> <user> <password>`
 
 Now we can build the data set
 
-`python build.py`
+`python build.py <dir> --ext ".cptv"`
 
 And finally train the model
 
diff --git a/build.py b/build.py
index ce2d5569..235c87c4 100644
--- a/build.py
+++ b/build.py
@@ -16,18 +16,22 @@
 from config.config import Config
 from ml_tools.dataset import Dataset
 from ml_tools.datasetstructures import Camera
+from ml_tools.tfwriter import create_tf_records
+from ml_tools.irwriter import save_data as save_ir_data
+from ml_tools.thermalwriter import save_data as save_thermal_data
 
-from ml_tools.irwriter import create_tf_records as create_ir_records
-from ml_tools.thermalwriter import create_tf_records as create_thermal_records
 
 import numpy as np
 
-MAX_TEST_TRACKS = 10
-MAX_TEST_SAMPLES = 10
+from pathlib import Path
 
-MIN_SAMPLES = 100
-MIN_TRACKS = 100
-LOW_SAMPLES_LABELS = ["bird/kiwi"]
+MAX_TEST_TRACKS = 1
+MAX_TEST_SAMPLES = 1
+
+MIN_SAMPLES = 1
+MIN_TRACKS = 1
+# LOW_SAMPLES_LABELS = ["bird", "cat", "possum"]
+LOW_SAMPLES_LABELS = []
 
 
 def load_config(config_file):
@@ -36,9 +40,7 @@ def load_config(config_file):
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--rebuild-important", action="count", help="Rebuild important frames"
-    )
+
     parser.add_argument(
         "-m",
         "--min-samples",
@@ -53,6 +55,10 @@ def parse_args():
         type=float,
         help="Percentage of training set to add extra augmentations of",
     )
+    parser.add_argument("--split-file", help="Json file defining a split")
+    parser.add_argument(
+        "--ext", default=".hdf5", help="Extension of files to load .mp4,.cptv,.hdf5"
+    )
 
     parser.add_argument("-c", "--config-file", help="Path to config file to use")
     parser.add_argument("-d", "--date", help="Use clips after this")
@@ -63,19 +69,7 @@ def parse_args():
         default=False,
         help="Use consecutive frames for segments",
     )
-    parser.add_argument(
-        "--bb",
-        "--balance-bins",
-        action="count",
-        help="Balance bins so each track has even percentage of being picked",
-    )
-    parser.add_argument(
-        "--bl",
-        "--balance-labels",
-        action="count",
-        help="Balance labels so that they have are distributed as defined in config",
-    )
-
+    parser.add_argument("data_dir", help="Directory of hdf5 files")
     args = parser.parse_args()
     if args.date:
         if args.date == "None":
@@ -94,20 +88,38 @@ def parse_args():
         MIN_SAMPLES = args.min_samples
         MIN_TRACKS = args.min_samples
         # args.date = datetime.datetime.now() - datetime.timedelta(days=30)
-
+    if args.data_dir is not None:
+        args.data_dir = Path(args.data_dir)
     logging.info("Loading training set up to %s", args.date)
     return args
 
 
+def show_clips_breakdown(dataset):
+    print("Clips breakdown:")
+
+    for label in dataset.labels:
+        print(label)
+        lbl_samples = dataset.samples_by_label.get(label, [])
+        if dataset.label_mapping:
+            actual_labels = {}
+            for s in lbl_samples:
+                actual_labels.setdefault(s.original_label, set()).add(s.clip_id)
+            for actual, clips in actual_labels.items():
+                print("  {:<20} {} clips".format(actual, len(clips)))
+
+        else:
+            clips = [s.clip_id for s in lbl_samples]
+            clips = set(clips)
+            print("  {:<20} {} clips".format(label, len(clips)))
+
+
 def show_tracks_breakdown(dataset):
     print("Tracks breakdown:")
-    samples = dataset.samples
 
     for label in dataset.labels:
         lbl_samples = dataset.samples_by_label.get(label, [])
         tracks = [s.unique_track_id for s in lbl_samples]
         tracks = set(tracks)
-        print("labels are", label)
         print("  {:<20} {} tracks".format(label, len(tracks)))
 
 
@@ -122,13 +134,55 @@ def show_cameras_tracks(dataset):
 def show_cameras_breakdown(dataset):
     print("Cameras breakdown")
     samples_by_camera = {}
-    for sample in dataset.samples:
-        if sample.camera not in samples_by_camera:
-            samples_by_camera[sample.camera] = []
-        samples_by_camera[sample.camera].append(sample)
+    for clip in dataset.clips:
+        if clip.camera not in samples_by_camera:
+            samples_by_camera[clip.camera] = []
+
+        for track in clip.tracks:
+            samples_by_camera[clip.camera].extend(track.samples)
 
     for camera, samples in samples_by_camera.items():
-        print("{:<20} {}".format(camera, len(samples)))
+        print(
+            "Camera: {:<20} {}".format(
+                "None" if camera is None else camera, len(samples)
+            )
+        )
+
+
+def show_bins_breakdown(dataset):
+    print("Bins breakdown")
+    samples_by_bins = {}
+    for clip in dataset.clips:
+        for t in clip.tracks:
+            for sample in t.samples:
+                if sample.bin_id not in samples_by_bins:
+                    samples_by_bins[sample.bin_id] = []
+                samples_by_bins[sample.bin_id].append(sample)
+
+    for bin_id, samples in samples_by_bins.items():
+        print(
+            "Bin Id: {:<20} {}".format(
+                "None" if bin_id is None else bin_id, len(samples)
+            )
+        )
+
+
+def show_stations_breakdown(dataset):
+    print("Stations breakdown")
+    samples_by_station = {}
+    for clip in dataset.clips:
+        if clip.station_id not in samples_by_station:
+            samples_by_station[clip.station_id] = []
+
+        for track in clip.tracks:
+            samples_by_station[clip.station_id].extend(track.samples)
+
+    for station, samples in samples_by_station.items():
+        print(
+            "StationId: {:<20} {}".format(
+                "None" if station is None else station, len(samples)
+            )
+        )
 
 
 def show_samples_breakdown(dataset):
@@ -151,14 +205,14 @@ def print_cameras(train, validation, test):
     print()
 
 
-def print_counts(dataset, train, validation, test):
+def print_counts(train, validation, test):
     print("Counts per class:")
     print("-" * 90)
     print("{:<20} {:<21} {:<21} {:<21}".format("Class", "Train", "Validation", "Test"))
     print("-" * 90)
     print("Samples / Tracks/ Bins/ weight")
     # display the dataset summary
-    for label in dataset.labels:
+    for label in train.labels:
         print(
             "{:<20} {:<20} {:<20} {:<20}".format(
                 label,
@@ -170,98 +224,153 @@ def print_counts(dataset, train, validation, test):
     print()
 
 
-def split_label(dataset, label, existing_test_count=0, max_samples=None):
+# default split is by stationid, but some labels dont have many stations so best to just split by clip
+split_by_clip = ["penguin", "wallaby"]
+
+
+def split_label(
+    dataset,
+    label,
+    counts,
+    train_count,
+    validation_count,
+    test_count,
+    existing_test_count=0,
+    max_samples=None,
+    use_test=True,
+):
     # split a label from dataset such that vlaidation is 15% or MIN_TRACKS
+    # dont split these by location and camera
+    logging.info(
+        "Splitting %s have counts (Tracks/Samples/Bins) %s already have (Tracks/Samples) in train %s validation %s and in test %s",
+        label,
+        counts,
+        train_count,
+        validation_count,
+        test_count,
+    )
+
     samples = dataset.samples_by_label.get(label, [])
-    if max_samples is not None:
-        samples = np.random.choice(
-            samples, min(len(samples), max_samples), replace=False
-        )
-    samples_by_bin = {}
-    total_tracks = set()
-    for sample in samples:
-        total_tracks.add(sample.track_id)
-        if sample.bin_id not in samples_by_bin:
-            samples_by_bin[sample.bin_id] = []
-        samples_by_bin[sample.bin_id].append(sample)
-    total_tracks = len(total_tracks)
-    sample_bins = [sample.bin_id for sample in samples]
+    sample_bins = set([sample.bin_id for sample in samples])
+
+    if counts[2] < 4 or counts[0] < 100:
+        global split_by_clip
+        split_by_clip.append(label)
+
+    if label in split_by_clip:
+        sample_bins = set([sample.clip_id for sample in samples])
+        print(label, " Splitting by clip")
+        samples_by_bin = {}
+        for clip in dataset.clips:
+            if clip.clip_id in sample_bins:
+                samples_by_bin[clip.clip_id] = clip.get_samples()
+    else:
+        samples_by_bin = dataset.samples_by_bin
     if len(sample_bins) == 0:
         return None, None, None
 
-    # sample_bins duplicates
-    sample_bins = list(set(sample_bins))
+    if max_samples is not None:
+        sample_bins = np.random.choice(
+            sample_bins, min(len(sample_bins), max_samples), replace=False
+        )
+
+    sample_count = counts[1]
+    total_tracks = counts[0]
+
+    sample_bins = list(sample_bins)
 
     random.shuffle(sample_bins)
-    train_c = Camera("{}-Train".format(label))
-    validate_c = Camera("{}-Val".format(label))
-    test_c = Camera("{}-Test".format(label))
+    train_c = []
+    validate_c = []
+    test_c = [] if use_test else None
 
     camera_type = "validate"
     add_to = validate_c
     last_index = 0
     label_count = 0
-    total = len(samples)
     min_t = MIN_SAMPLES
 
     if label in LOW_SAMPLES_LABELS:
         min_t = 10
-    num_validate_samples = max(total * 0.15, min_t)
+    num_validate_samples = max(sample_count * 0.15, min_t) - validation_count[1]
     num_test_samples = (
-        min(MAX_TEST_SAMPLES, max(total * 0.05, min_t)) - existing_test_count
+        min(MAX_TEST_SAMPLES, max(sample_count * 0.05, min_t)) - test_count[1]
     )
-    # should have test covered by test set
 
     min_t = MIN_TRACKS
 
     if label in LOW_SAMPLES_LABELS:
         min_t = 10
 
-    num_validate_tracks = max(total_tracks * 0.15, min_t)
+    num_validate_tracks = max(total_tracks * 0.15, min_t) - validation_count[0]
     num_test_tracks = (
-        min(MAX_TEST_TRACKS, max(total_tracks * 0.05, min_t)) - existing_test_count
+        min(MAX_TEST_TRACKS, max(total_tracks * 0.05, min_t)) - test_count[0]
     )
     track_limit = num_validate_tracks
     sample_limit = num_validate_samples
     tracks = set()
-    print(
-        f"{label} - looking for val {num_validate_tracks} tracks out of {total_tracks} tracks and {num_validate_samples} samples from a total of {total} samples  with {num_test_tracks} test tracks and {num_test_samples} test samples"
+    logging.info(
+        f"{label} - looking for val {num_validate_tracks} tracks out of {total_tracks} tracks and {num_validate_samples} samples from a total of {sample_count} samples  with {num_test_tracks} test tracks and {num_test_samples} test samples"
     )
-    for i, sample_bin in enumerate(sample_bins):
-        samples = samples_by_bin[sample_bin]
-        for sample in samples:
-            if sample.label == label:
-                tracks.add(sample.track_id)
-                label_count += 1
-
-            sample.camera = "{}-{}".format(sample.camera, camera_type)
-            add_to.add_sample(sample)
-        samples_by_bin[sample_bin] = []
-        last_index = i
-        track_count = len(tracks)
-        if label_count >= sample_limit and track_count >= track_limit:
-            # 100 more for test
-            if add_to == validate_c:
-                add_to = test_c
-                camera_type = "test"
-                if num_test_samples <= 0:
+    # if sample_limit < 0 and track_limit < 0:
+    #     add_to = test_c
+    #     camera_type = "test"
+    #     sample_limit = num_test_samples
+    #     track_limit = num_test_tracks
+    #     label_count = 0
+    #     tracks = set()
+    if sample_limit > 0 and track_limit > 0:
+        for i, sample_bin in enumerate(sample_bins):
+            samples = samples_by_bin[sample_bin]
+            for sample in samples:
+                if sample.label == label:
+                    tracks.add(sample.track_id)
+                    label_count += 1
+
+                # sample.camera = "{}-{}".format(sample.camera, camera_type)
+                add_to.append(sample)
+
+                if label in split_by_clip:
+                    dataset.remove_sample(sample)
+            if label not in split_by_clip:
+                while len(samples) > 0:
+                    dataset.remove_sample(samples[0])
+            samples_by_bin[sample_bin] = []
+            last_index = i
+            track_count = len(tracks)
+            if label_count >= sample_limit and track_count >= track_limit:
+                # 100 more for test
+                if add_to == validate_c:
+                    add_to = test_c
+                    camera_type = "test"
+                    if num_test_samples <= 0:
+                        break
+                    sample_limit = num_test_samples
+                    track_limit = num_test_tracks
+                    label_count = 0
+                    tracks = set()
+                    if use_test is False:
+                        break
+                else:
                     break
-                sample_limit = num_test_samples
-                track_limit = num_test_tracks
-                label_count = 0
-                tracks = set()
-            else:
-                break
 
-    sample_bins = sample_bins[last_index + 1 :]
+    # sample_bins = sample_bins[last_index + 1 :]
+    # clearining anyway
     camera_type = "train"
     added = 0
     for i, sample_bin in enumerate(sample_bins):
         samples = samples_by_bin[sample_bin]
         for sample in samples:
-            sample.camera = "{}-{}".format(sample.camera, camera_type)
-            train_c.add_sample(sample)
+            # sample.camera = "{}-{}".format(sample.camera, camera_type)
+            train_c.append(sample)
             added += 1
+
+            if label in split_by_clip:
+                dataset.remove_sample(sample)
+        if label not in split_by_clip:
+            while len(samples) > 0:
+                dataset.remove_sample(samples[0])
+
         samples_by_bin[sample_bin] = []
     return train_c, validate_c, test_c
 
@@ -269,29 +378,95 @@ def split_label(dataset, label, existing_test_count=0, max_samples=None):
 def get_test_set_camera(dataset, test_clips, after_date):
     # load test set camera from tst_clip ids and all clips after a date
     test_c = Camera("Test-Set-Camera")
-    test_samples = [
+    print("after date is", after_date)
+    test_clips = [
         sample
-        for sample in dataset.samples
-        if sample.clip_id in test_clips
+        for clip in dataset.clips
+        if clip.clip_id in test_clips
         or after_date is not None
-        and sample.start_time.replace(tzinfo=pytz.utc) > after_date
+        and clip.rec_time.replace(tzinfo=pytz.utc) > after_date
     ]
-    for sample in test_samples:
-        dataset.remove_sample(sample)
-        test_c.add_sample(sample)
+    for clip in test_clips:
+        for track in clip.tracks:
+            for sample in track.samples:
+                dataset.remove_sample(sample)
+                test_c.add_sample(sample)
     return test_c
 
 
-def split_randomly(db_file, dataset, config, args, test_clips=[], balance_bins=True):
+def split_by_file(dataset, config, split_file, base_dir):
+    base_dir = Path(base_dir)
+    with open(split_file, "r") as f:
+        split = json.load(f)
+
+    samples_by_source = {}
+    for s in dataset.samples_by_id.values():
+        samples_by_source.setdefault(s.source_file.name, []).append(s)
+    datasets = []
+    for name in ["train", "validation", "test"]:
+        logging.info("Loading %s", name)
+        split_dataset = Dataset(
+            dataset.dataset_dir, name, config, label_mapping=dataset.label_mapping
+        )
+        if name == "train":
+            split_dataset.enable_augmentation = True
+        elif name == "test":
+            split_dataset.skip_ffc = False
+        split_files = split.get(name, [])
+        for f in split_files:
+            # print("loading", base_dir / f["source"])
+            # continue
+            file = base_dir / f["source"]
+            if file.exists():
+                try:
+                    # can filter crappy segments here, or can do at train stage to have some way of testing different thresholds
+                    split_dataset.load_clip(file, dont_filter_segment=True)
+                except:
+                    logging.error("Could not load %s", file, exc_info=True)
+            else:
+                pass
+                # logging.warn("No source file %s found for %s", f, name)
+        datasets.append(split_dataset)
+
+    print_counts(*datasets)
+    train, val, test = split_randomly(datasets[0], config, None, use_test=False)
+    datasets = [train, val, datasets[2]]
+    return datasets
+
+
+def split_randomly(dataset, config, date, test_clips=[], use_test=True):
     # split data randomly such that a clip is only in one dataset
     # have tried many ways to split i.e. location and cameras found this is simplest
     # and the results are the same
-    train = Dataset(db_file, "train", config)
+    train = Dataset(
+        dataset.dataset_dir,
+        "train",
+        config,
+        label_mapping=dataset.label_mapping,
+        raw=dataset.raw,
+        ext=dataset.ext,
+    )
     train.enable_augmentation = True
-    validation = Dataset(db_file, "validation", config)
-    test = Dataset(db_file, "test", config)
-    test_c = get_test_set_camera(dataset, test_clips, args.date)
-    test_cameras = [test_c]
+    validation = Dataset(
+        dataset.dataset_dir,
+        "validation",
+        config,
+        label_mapping=dataset.label_mapping,
+        raw=dataset.raw,
+        ext=dataset.ext,
+    )
+    test = None
+    if use_test:
+        test_c = get_test_set_camera(dataset, test_clips, date)
+        test_cameras = [test_c]
+        test = Dataset(
+            dataset.dataset_dir,
+            "test",
+            config,
+            label_mapping=dataset.label_mapping,
+            raw=dataset.raw,
+            ext=dataset.ext,
+        )
     validate_cameras = []
     train_cameras = []
     min_label = None
@@ -301,66 +476,224 @@ def split_randomly(db_file, dataset, config, args, test_clips=[], balance_bins=T
             continue
         if min_label is None or label_count < min_label[1]:
             min_label = (label, label_count)
-    for label in dataset.labels:
-        existing_test_count = len(test.samples_by_label.get(label, []))
+    lbl_order = sorted(
+        dataset.labels,
+        key=lambda lbl: len(dataset.samples_by_label.get(lbl, [])),
+    )
+    if "wallaby" in lbl_order:
+        # make sure we do wallaby first so we get a good split
+        lbl_order.remove("wallaby")
+        lbl_order.insert(0, "wallaby")
+    if "pest" in lbl_order:
+        # make sure its last
+        lbl_order.remove("pest")
+        lbl_order.append("pest")
+
+    lbl_counts = {}
+    for lbl in dataset.labels:
+        samples = dataset.samples_by_label.get(lbl, [])
+        tracks = set([s.track_id for s in samples])
+        bins = set([s.bin_id for s in samples])
+
+        lbl_counts[lbl] = (len(tracks), len(samples), len(bins))
+
+    logging.debug("lbl order is %s", lbl_order)
+    train_counts = {}
+    validation_counts = {}
+    test_counts = {}
+    existing_test_count = 0
+    for label in lbl_order:
+        # existing_test_count = len(test.samples_by_label.get(label, []))
         train_c, validate_c, test_c = split_label(
             dataset,
             label,
+            counts=lbl_counts[label],
+            train_count=train_counts.get(label, (0, 0)),
+            validation_count=validation_counts.get(label, (0, 0)),
+            test_count=test_counts.get(label, (0, 0)),
             existing_test_count=existing_test_count,
+            use_test=use_test,
             # max_samples=min_label[1],
         )
         if train_c is not None:
-            train_cameras.append(train_c)
+            add_samples(dataset.labels, train, train_c, train_counts)
         if validate_c is not None:
-            validate_cameras.append(validate_c)
+            add_samples(dataset.labels, validation, validate_c, validation_counts)
         if test_c is not None:
-            test_cameras.append(test_c)
+            add_samples(dataset.labels, test, test_c, test_counts)
+        logging.debug("Train counts %s", train_counts)
+        logging.debug("val counts %s", validation_counts)
 
-    add_camera_samples(dataset.labels, train, train_cameras, balance_bins)
-    add_camera_samples(dataset.labels, validation, validate_cameras, balance_bins)
-    add_camera_samples(dataset.labels, test, test_cameras, balance_bins)
     return train, validation, test
 
 
-def add_camera_samples(
+def add_samples(
     labels,
     dataset,
-    cameras,
-    balance_bins=None,
+    samples,
+    counts,
 ):
-    # add camera tracks to the daaset and calculate segments and bins
-    all_samples = []
-    for label in labels:
-        for camera in cameras:
-            samples = camera.label_to_samples.get(label, {}).values()
-            all_samples.extend(list(samples))
-    dataset.add_samples(all_samples)
-    dataset.balance_bins()
+    by_labels = {}
+    counts = {}
+    for s in samples:
+        if s.label not in by_labels:
+            by_labels[s.label] = []
+        by_labels[s.label].append(s)
+
+    for label, samples in by_labels.items():
+        track_count = len(set([s.track_id for s in samples]))
+        counts[label] = (
+            track_count,
+            len(samples),
+        )
+    dataset.add_samples(samples)
 
 
-def validate_datasets(datasets, test_clips, date):
+def validate_datasets(datasets, test_bins, date):
     # check that clips are only in one dataset
     # that only test set has clips after date
     # that test set is the only dataset with test_clips
-    for dataset in datasets[:2]:
-        for track in dataset.tracks:
-            assert track.start_time < date
+
+    # for dataset in datasets[:2]:
+    #     for track in dataset.tracks:
+    #         assert track.start_time < date
 
     for i, dataset in enumerate(datasets):
-        clips = set([sample.clip_id for sample in dataset.samples])
-        if test_clips is not None and dataset.name != "test":
+        dont_check = set(
+            [
+                sample.bin_id
+                for sample in dataset.samples_by_id.values()
+                if sample.label in split_by_clip
+            ]
+        )
+        bins = set([sample.bin_id for sample in dataset.samples_by_id.values()])
+        clips = set([sample.clip_id for sample in dataset.samples_by_id.values()])
+
+        bins = bins - dont_check
+        if test_bins is not None and dataset.name != "test":
             assert (
-                len(clips.intersection(set(test_clips))) == 0
-            ), "test clips should only be in test set"
-        if len(clips) == 0:
+                len(bins.intersection(set(test_bins))) == 0
+            ), "test bins should only be in test set"
+        if len(bins) == 0:
             continue
         for other in datasets[(i + 1) :]:
             if dataset.name == other.name:
                 continue
-            other_clips = set([sample.clip_id for sample in other.samples])
-            # other_tracks = set([track.track_id for track in other.tracks])
+            dont_check = set(
+                [
+                    sample.bin_id
+                    for sample in other.samples_by_id.values()
+                    if sample.label in split_by_clip
+                ]
+            )
+            other_bins = set([sample.bin_id for sample in other.samples_by_id.values()])
+            other_bins = other_bins - dont_check
+            other_clips = set(
+                [sample.clip_id for sample in other.samples_by_id.values()]
+            )
+
+            intersection = bins.intersection(set(other_bins))
+
+            assert (
+                len(bins.intersection(set(other_bins))) == 0
+            ), "bins should only be in one set"
 
-            assert clips != other_clips, "clips should only be in one set"
+            assert (
+                len(clips.intersection(set(other_clips))) == 0
+            ), "clips should only be in one set"
+
+
+land_birds = [
+    "pukeko",
+    "california quail",
+    "brown quail",
+    "black swan",
+    "quail",
+    "pheasant",
+    "penguin",
+    "duck",
+    "chicken",
+    "rooster",
+]
+
+label_paths_dl = "https://raw.githubusercontent.com/TheCacophonyProject/cacophony-web/main/api/classifications/label_paths.json"
+
+
+def dl_mappings():
+    import requests
+
+    logging.info("Downloading mappings file from %s ", label_paths_dl)
+    response = requests.get(label_paths_dl)
+    response.raise_for_status()
+
+    mapping_content = response.content.decode()
+    print(mapping_content)
+    with open("label_paths.json", "w") as f:
+        f.write(mapping_content)
+    return mapping_content
+
+
+def get_mappings():
+    labels_path = Path("label_paths.json")
+    if not labels_path.exists():
+        label_paths = json.loads(dl_mappings())
+    with open("label_paths.json", "r") as f:
+        label_paths = json.load(f)
+    regroup = {}
+    for l, path in label_paths.items():
+        if l in land_birds:
+            regroup[l] = l
+            continue
+        split_path = path.split(".")
+        if len(split_path) == 1:
+            regroup[l] = l
+        elif path.startswith("all.mammal"):
+            if len(split_path) == 4:
+                regroup[l] = split_path[-2]
+            else:
+                regroup[l] = l
+        else:
+            # print("l", l, " has ", path)
+            parent = split_path[-2]
+            # print("Parent is", parent, path)
+            if parent == "kiwi" or split_path[-1] == "kiwi":
+                regroup[l] = "kiwi"
+            elif parent == "other":
+                regroup[l] = l
+
+            else:
+                if "bird." in path:
+                    regroup[l] = "bird"
+
+                elif len(split_path) > 2:
+                    regroup[l] = split_path[-3]
+                else:
+                    regroup[l] = split_path[-1]
+
+    return regroup
+
+
+def dump_split_ids(datasets, out_file="datasplit.json"):
+    splits = {}
+    logging.info("Wrinting split ids to %s", out_file)
+    for d in datasets:
+        samples_by_source = d.get_samples_by_source()
+        clips = []
+        for source, samples in samples_by_source.items():
+            tags = set([s.label for s in samples])
+            clips.append(
+                {
+                    "clip_id": samples[0].clip_id,
+                    "source": source.name,
+                    "station_id": "{}".format(samples[0].station_id),
+                    "tags": list(tags),
+                }
+            )
+        # clips = set(clips)
+        splits[d.name] = clips
+    with open(out_file, "w") as f:
+        json.dump(splits, f)
+    return
 
 
 def main():
@@ -368,58 +701,87 @@ def main():
     args = parse_args()
     config = load_config(args.config_file)
     logging.info("Building for type %s", config.train.type)
+
     test_clips = config.build.test_clips()
     if test_clips is None:
         test_clips = []
     logging.info("# of test clips are %s", len(test_clips))
-    db_file = os.path.join(config.tracks_folder, "dataset.hdf5")
-    dataset = Dataset(
-        db_file, "dataset", config, consecutive_segments=args.consecutive_segments
+    label_mapping = get_mappings()
+    master_dataset = Dataset(
+        args.data_dir,
+        "dataset",
+        config,
+        consecutive_segments=args.consecutive_segments,
+        label_mapping=label_mapping,
+        raw=False if args.ext == ".hdf5" else True,
+        ext=args.ext,
     )
-
-    tracks_loaded, total_tracks = dataset.load_clips()
-    # return
-    dataset.labels.sort()
-    print(
-        "Loaded {}/{} tracks, found {:.1f}k samples".format(
-            tracks_loaded, total_tracks, len(dataset.samples) / 1000
+    base_dir = Path(config.tracks_folder)
+    record_dir = base_dir / "training-data"
+    record_dir.mkdir(parents=True, exist_ok=True)
+
+    if args.split_file:
+        logging.info("Loading datasets from split file %s", args.split_file)
+        datasets = split_by_file(master_dataset, config, args.split_file, args.data_dir)
+        labels = set()
+        for dataset in datasets:
+            labels.update(dataset.labels)
+        labels = list(labels)
+        labels.sort()
+        for dataset in datasets:
+            dataset.labels = labels
+    else:
+        tracks_loaded, total_tracks = master_dataset.load_clips(
+            dont_filter_segment=True
         )
-    )
-    for key, value in dataset.filtered_stats.items():
-        if value != 0:
-            print("  {} filtered {}".format(key, value))
-
-    print()
-    show_tracks_breakdown(dataset)
-    print()
-    show_samples_breakdown(dataset)
-    print()
-    show_cameras_breakdown(dataset)
-    print()
-    print("Splitting data set into train / validation")
-    datasets = split_randomly(db_file, dataset, config, args, test_clips)
-    validate_datasets(datasets, test_clips, args.date)
 
-    print_counts(dataset, *datasets)
+        master_dataset.labels.sort()
+        print(
+            "Loaded {}/{} tracks, found {:.1f}k samples".format(
+                tracks_loaded, total_tracks, len(master_dataset.clips) / 1000
+            )
+        )
+        for key, value in master_dataset.filtered_stats.items():
+            if value != 0:
+                print("  {} filtered {}".format(key, value))
+
+        print()
+        show_clips_breakdown(master_dataset)
+        print()
+        show_samples_breakdown(master_dataset)
+        print()
+        show_cameras_breakdown(master_dataset)
+        print()
+        show_stations_breakdown(master_dataset)
+        print()
+        show_bins_breakdown(master_dataset)
+        print()
+        print("Splitting data set into train / validation")
+
+        datasets = split_randomly(master_dataset, config, args.date, test_clips)
+        validate_datasets(datasets, test_clips, args.date)
+        dump_split_ids(datasets, record_dir / "datasplit.json")
+
+    print_counts(*datasets)
     print("split data")
-    base_dir = config.tracks_folder
-    record_dir = os.path.join(base_dir, "training-data/")
+
     dataset_counts = {}
-    create_tf_records = create_thermal_records
+    # create_tf_records = create_thermal_records
     if config.train.type == "IR":
         threshold = (
             config.tracking[config.train.type]
             .motion.threshold_for_model(config.train.type)
             .background_thresh
         )
-
-        create_tf_records = create_ir_records
+        # create_tf_records = create_ir_records
     else:
         threshold = None
 
-    train_set = datasets[0]
+    # create pre augmented records if samples are very low
+    # not in use probably needs to be re tested if used
     aug_percent = args.aug_percent
     if aug_percent is not None:
+        train_set = datasets[0]
         for l, samples in train_set.samples_by_label.items():
             track_dic = {}
             for s in samples:
@@ -440,16 +802,46 @@ def main():
                     new_samples.append(new)
                 train_set.add_samples(new_samples)
         print("Count post augmentation")
-        print_counts(dataset, *datasets)
+        print_counts(*datasets)
+
     for dataset in datasets:
         dir = os.path.join(record_dir, dataset.name)
-        create_tf_records(dataset, dir, datasets[0].labels, threshold, num_shards=100)
+        if config.train.type == "IR":
+            create_tf_records(
+                dataset,
+                dir,
+                datasets[0].labels,
+                save_ir_data,
+                num_shards=100,
+                back_thresh=threshold,
+            )
+        else:
+            extra_args = None
+            if args.ext != ".hdf5":
+                extra_args = {
+                    "segment_frame_spacing": master_dataset.segment_spacing * 9,
+                    "segment_width": master_dataset.segment_length,
+                    "segment_type": master_dataset.segment_type,
+                    "segment_min_avg_mass": master_dataset.segment_min_avg_mass,
+                    "max_segments": master_dataset.max_segments,
+                    "dont_filter_segment": True,
+                    "skip_ffc": True,
+                    "tag_precedence": config.load.tag_precedence,
+                }
+            create_tf_records(
+                dataset,
+                dir,
+                datasets[0].labels,
+                save_thermal_data,
+                num_shards=100,
+                num_frames=dataset.segment_length,
+                **extra_args,
+            )
         counts = {}
         for label in dataset.labels:
             count = len(dataset.samples_by_label.get(label, []))
             counts[label] = count
         dataset_counts[dataset.name] = counts
-        # dataset.saveto_numpy(os.path.join(base_dir))
     # dont need dataset anymore just need some meta
     meta_filename = f"{base_dir}/training-meta.json"
     meta_data = {
diff --git a/classify/clipclassifier.py b/classify/clipclassifier.py
index 3637f63a..032dada5 100644
--- a/classify/clipclassifier.py
+++ b/classify/clipclassifier.py
@@ -8,12 +8,11 @@
 
 import cv2
 from classify.trackprediction import Predictions
-from load.clip import Clip
-from load.cliptrackextractor import ClipTrackExtractor, is_affected_by_ffc
+from track.clip import Clip
+from track.cliptrackextractor import ClipTrackExtractor, is_affected_by_ffc
 from ml_tools import tools
 from ml_tools.kerasmodel import KerasModel
-from load.irtrackextractor import IRTrackExtractor
-from ml_tools.preprocess import preprocess_segment
+from track.irtrackextractor import IRTrackExtractor
 from ml_tools.previewer import Previewer
 from track.track import Track
 
diff --git a/config/trackingconfig.py b/config/trackingconfig.py
index a95a8bcf..f6e983f1 100644
--- a/config/trackingconfig.py
+++ b/config/trackingconfig.py
@@ -21,7 +21,7 @@
 
 from .defaultconfig import DefaultConfig, deep_copy_map_if_key_not_exist
 from .trackingmotionconfig import TrackingMotionConfig
-from load.cliptrackextractor import ClipTrackExtractor
+from track.cliptrackextractor import ClipTrackExtractor
 from track.track import RegionTracker
 
 
diff --git a/label_paths.json b/label_paths.json
new file mode 100644
index 00000000..b7ec5264
--- /dev/null
+++ b/label_paths.json
@@ -0,0 +1,536 @@
+{
+    "all": "all",
+    "mammal": "all.mammal",
+    "cat": "all.mammal.cat",
+    "deer": "all.mammal.deer",
+    "dog": "all.mammal.dog",
+    "goat": "all.mammal.goat",
+    "hedgehog": "all.mammal.hedgehog",
+    "human": "all.mammal.human",
+    "leporidae": "all.mammal.leporidae",
+    "hare": "all.mammal.leporidae.hare",
+    "rabbit": "all.mammal.leporidae.rabbit",
+    "mustelid": "all.mammal.mustelid",
+    "ferret": "all.mammal.mustelid.ferret",
+    "stoat": "all.mammal.mustelid.stoat",
+    "weasel": "all.mammal.mustelid.weasel",
+    "pig": "all.mammal.pig",
+    "possum": "all.mammal.possum",
+    "rodent": "all.mammal.rodent",
+    "mouse": "all.mammal.rodent.mouse",
+    "rat": "all.mammal.rodent.rat",
+    "sealion": "all.mammal.sealion",
+    "sheep": "all.mammal.sheep",
+    "wallaby": "all.mammal.wallaby",
+    "other": "all.other",
+    "false-positive": "all.other.falsepositive",
+    "frog": "all.other.frog",
+    "insect": "all.other.insect",
+    "weta": "all.other.insect.weta",
+    "interesting": "all.other.interesting",
+    "lizard": "all.other.lizard",
+    "part": "all.other.part",
+    "pest": "all.other.pest",
+    "poor tracking": "all.other.poor_tracking",
+    "rain": "all.other.rain",
+    "static": "all.other.static",
+    "unidentified": "all.other.unidentified",
+    "vehicle": "all.other.vehicle",
+    "water": "all.other.water",
+    "wind": "all.other.wind",
+    "bird": "all.bird",
+    "albatross": "all.bird.albatross",
+    "antipodean albatross": "all.bird.albatross.antipodean_albatross",
+    "black-footed albatross": "all.bird.albatross.blackfooted_albatross",
+    "laysan albatross": "all.bird.albatross.laysan_albatross",
+    "light-mantled sooty albatross": "all.bird.albatross.lightmantled_sooty_albatross",
+    "northern royal albatross": "all.bird.albatross.northern_royal_albatross",
+    "sooty albatross": "all.bird.albatross.sooty_albatross",
+    "southern royal albatross": "all.bird.albatross.southern_royal_albatross",
+    "wandering albatross": "all.bird.albatross.wandering_albatross",
+    "australasian shoveler": "all.bird.australasian_shoveler",
+    "australian coot": "all.bird.australian_coot",
+    "avocet": "all.bird.avocet",
+    "red-necked avocet": "all.bird.avocet.rednecked_avocet",
+    "bellbird": "all.bird.bellbird",
+    "bittern": "all.bird.bittern",
+    "australasian bittern": "all.bird.bittern.australasian_bittern",
+    "little bittern": "all.bird.bittern.little_bittern",
+    "black-tailed native-hen": "all.bird.blacktailed_nativehen",
+    "blackbird": "all.bird.blackbird",
+    "eurasian blackbird": "all.bird.blackbird.eurasian_blackbird",
+    "booby": "all.bird.booby",
+    "brown booby": "all.bird.booby.brown_booby",
+    "masked booby": "all.bird.booby.masked_booby",
+    "red-footed booby": "all.bird.booby.redfooted_booby",
+    "bulbul": "all.bird.bulbul",
+    "red-vented bulbul": "all.bird.bulbul.redvented_bulbul",
+    "bunting": "all.bird.bunting",
+    "cirl bunting": "all.bird.bunting.cirl_bunting",
+    "chaffinch": "all.bird.chaffinch",
+    "chatham island taiko": "all.bird.chatham_island_taiko",
+    "chestnut-breasted shelduck": "all.bird.chestnutbreasted_shelduck",
+    "chicken": "all.bird.chicken",
+    "chukor": "all.bird.chukor",
+    "cockatoo": "all.bird.cockatoo",
+    "sulphur-crested cockatoo": "all.bird.cockatoo.sulphurcrested_cockatoo",
+    "corncrake": "all.bird.corncrake",
+    "crake": "all.bird.crake",
+    "australian crake": "all.bird.crake.australian_crake",
+    "marsh crake": "all.bird.crake.marsh_crake",
+    "spotless crake": "all.bird.crake.spotless_crake",
+    "crane": "all.bird.crane",
+    "creeper": "all.bird.creeper",
+    "brown creeper": "all.bird.creeper.brown_creeper",
+    "cuckoo": "all.bird.cuckoo",
+    "black-faced cuckoo-shrike": "all.bird.cuckoo.blackfaced_cuckooshrike",
+    "channel-billed cuckoo": "all.bird.cuckoo.channelbilled_cuckoo",
+    "fan-tailed cuckoo": "all.bird.cuckoo.fantailed_cuckoo",
+    "long-tailed cuckoo": "all.bird.cuckoo.longtailed_cuckoo",
+    "oriental cuckoo": "all.bird.cuckoo.oriental_cuckoo",
+    "pallid cuckoo": "all.bird.cuckoo.pallid_cuckoo",
+    "shining cuckoo": "all.bird.cuckoo.shining_cuckoo",
+    "curlew": "all.bird.curlew",
+    "bristle-thighed curlew": "all.bird.curlew.bristlethighed_curlew",
+    "eastern curlew": "all.bird.curlew.eastern_curlew",
+    "dabchick": "all.bird.dabchick",
+    "new zealand dabchick": "all.bird.dabchick.new_zealand_dabchick",
+    "darter": "all.bird.darter",
+    "dollarbird": "all.bird.dollarbird",
+    "dotterel": "all.bird.dotterel",
+    "banded dotterel": "all.bird.dotterel.banded_dotterel",
+    "black-fronted dotterel": "all.bird.dotterel.blackfronted_dotterel",
+    "new zealand dotterel": "all.bird.dotterel.new_zealand_dotterel",
+    "oriental dotterel": "all.bird.dotterel.oriental_dotterel",
+    "red-kneed dotterel": "all.bird.dotterel.redkneed_dotterel",
+    "dove": "all.bird.dove",
+    "barbary dove": "all.bird.dove.barbary_dove",
+    "rose-crowned fruit-dove": "all.bird.dove.rosecrowned_fruitdove",
+    "spotted dove": "all.bird.dove.spotted_dove",
+    "dowitcher": "all.bird.dowitcher",
+    "asiatic dowitcher": "all.bird.dowitcher.asiatic_dowitcher",
+    "duck": "all.bird.duck",
+    "australian white-eyed duck": "all.bird.duck.australian_whiteeyed_duck",
+    "australian wood duck": "all.bird.duck.australian_wood_duck",
+    "blue duck": "all.bird.duck.blue_duck",
+    "grey duck": "all.bird.duck.grey_duck",
+    "muscovy duck": "all.bird.duck.muscovy_duck",
+    "pink-eared duck": "all.bird.duck.pinkeared_duck",
+    "plumed whistling duck": "all.bird.duck.plumed_whistling_duck",
+    "dunlin": "all.bird.dunlin",
+    "dunnock": "all.bird.dunnock",
+    "dusky woodswallow": "all.bird.dusky_woodswallow",
+    "eagle": "all.bird.eagle",
+    "white-bellied sea eagle": "all.bird.eagle.whitebellied_sea_eagle",
+    "egret": "all.bird.egret",
+    "cattle egret": "all.bird.egret.cattle_egret",
+    "little egret": "all.bird.egret.little_egret",
+    "plumed egret": "all.bird.egret.plumed_egret",
+    "falcon": "all.bird.falcon",
+    "new zealand falcon": "all.bird.falcon.new_zealand_falcon",
+    "fantail": "all.bird.fantail",
+    "new zealand fantail": "all.bird.fantail.new_zealand_fantail",
+    "fernbird": "all.bird.fernbird",
+    "flycatcher": "all.bird.flycatcher",
+    "satin flycatcher": "all.bird.flycatcher.satin_flycatcher",
+    "frigatebird": "all.bird.frigatebird",
+    "great frigatebird": "all.bird.frigatebird.great_frigatebird",
+    "lesser frigatebird": "all.bird.frigatebird.lesser_frigatebird",
+    "fulmar": "all.bird.fulmar",
+    "antarctic fulmar": "all.bird.fulmar.antarctic_fulmar",
+    "fulmar prion": "all.bird.fulmar.fulmar_prion",
+    "northern fulmar": "all.bird.fulmar.northern_fulmar",
+    "galah": "all.bird.galah",
+    "gannet": "all.bird.gannet",
+    "australasian gannet": "all.bird.gannet.australasian_gannet",
+    "cape gannet": "all.bird.gannet.cape_gannet",
+    "godwit": "all.bird.godwit",
+    "bar-tailed godwit": "all.bird.godwit.bartailed_godwit",
+    "black-tailed godwit": "all.bird.godwit.blacktailed_godwit",
+    "hudsonian godwit": "all.bird.godwit.hudsonian_godwit",
+    "golden whistler": "all.bird.golden_whistler",
+    "goldfinch": "all.bird.goldfinch",
+    "european goldfinch": "all.bird.goldfinch.european_goldfinch",
+    "goose": "all.bird.goose",
+    "canada goose": "all.bird.goose.canada_goose",
+    "cape barren goose": "all.bird.goose.cape_barren_goose",
+    "greylag goose": "all.bird.goose.greylag_goose",
+    "grebe": "all.bird.grebe",
+    "australasian crested grebe": "all.bird.grebe.australasian_crested_grebe",
+    "australasian little grebe": "all.bird.grebe.australasian_little_grebe",
+    "hoary-headed grebe": "all.bird.grebe.hoaryheaded_grebe",
+    "greenfinch": "all.bird.greenfinch",
+    "european greenfinch": "all.bird.greenfinch.european_greenfinch",
+    "greenshank": "all.bird.greenshank",
+    "common greenshank": "all.bird.greenshank.common_greenshank",
+    "guineafowl": "all.bird.guineafowl",
+    "helmeted guineafowl": "all.bird.guineafowl.helmeted_guineafowl",
+    "gull": "all.bird.gull",
+    "black-billed gull": "all.bird.gull.blackbilled_gull",
+    "franklin's gull": "all.bird.gull.franklins_gull",
+    "laughing gull": "all.bird.gull.laughing_gull",
+    "pacific gull": "all.bird.gull.pacific_gull",
+    "red-billed gull": "all.bird.gull.redbilled_gull",
+    "southern black-backed gull": "all.bird.gull.southern_blackbacked_gull",
+    "harrier": "all.bird.harrier",
+    "swamp harrier": "all.bird.harrier.swamp_harrier",
+    "heron": "all.bird.heron",
+    "grey heron": "all.bird.heron.grey_heron",
+    "nankeen night heron": "all.bird.heron.nankeen_night_heron",
+    "pacific heron": "all.bird.heron.pacific_heron",
+    "reef heron": "all.bird.heron.reef_heron",
+    "white heron": "all.bird.heron.white_heron",
+    "white-faced heron": "all.bird.heron.whitefaced_heron",
+    "ibis": "all.bird.ibis",
+    "glossy ibis": "all.bird.ibis.glossy_ibis",
+    "straw-necked ibis": "all.bird.ibis.strawnecked_ibis",
+    "white ibis": "all.bird.ibis.white_ibis",
+    "kaka": "all.bird.kaka",
+    "kakapo": "all.bird.kakapo",
+    "kea": "all.bird.kea",
+    "kereru": "all.bird.kereru",
+    "kestrel": "all.bird.kestrel",
+    "nankeen kestrel": "all.bird.kestrel.nankeen_kestrel",
+    "kingfisher": "all.bird.kingfisher",
+    "sacred kingfisher": "all.bird.kingfisher.sacred_kingfisher",
+    "kite": "all.bird.kite",
+    "black kite": "all.bird.kite.black_kite",
+    "kiwi": "all.bird.kiwi",
+    "bird/kiwi": "all.bird.kiwi",
+    "great spotted kiwi": "all.bird.kiwi.great_spotted_kiwi",
+    "little spotted kiwi": "all.bird.kiwi.little_spotted_kiwi",
+    "north island brown kiwi": "all.bird.kiwi.north_island_brown_kiwi",
+    "okarito brown kiwi": "all.bird.kiwi.okarito_brown_kiwi",
+    "southern brown kiwi": "all.bird.kiwi.southern_brown_kiwi",
+    "knot": "all.bird.knot",
+    "great knot": "all.bird.knot.great_knot",
+    "lesser knot": "all.bird.knot.lesser_knot",
+    "kokako": "all.bird.kokako",
+    "north island kokako": "all.bird.kokako.north_island_kokako",
+    "south island kokako": "all.bird.kokako.south_island_kokako",
+    "kookaburra": "all.bird.kookaburra",
+    "laughing kookaburra": "all.bird.kookaburra.laughing_kookaburra",
+    "lesser yellowlegs": "all.bird.lesser_yellowlegs",
+    "magpie": "all.bird.magpie",
+    "australian magpie": "all.bird.magpie.australian_magpie",
+    "magpie-lark": "all.bird.magpie.magpielark",
+    "mallard": "all.bird.mallard",
+    "martin": "all.bird.martin",
+    "fairy martin": "all.bird.martin.fairy_martin",
+    "tree martin": "all.bird.martin.tree_martin",
+    "mollymawk": "all.bird.mollymawk",
+    "atlantic yellow-nosed mollymawk": "all.bird.mollymawk.atlantic_yellownosed_mollymawk",
+    "black-browed mollymawk": "all.bird.mollymawk.blackbrowed_mollymawk",
+    "buller's mollymawk": "all.bird.mollymawk.bullers_mollymawk",
+    "campbell black-browed mollymawk": "all.bird.mollymawk.campbell_blackbrowed_mollymawk",
+    "chatham island mollymawk": "all.bird.mollymawk.chatham_island_mollymawk",
+    "grey-headed mollymawk": "all.bird.mollymawk.greyheaded_mollymawk",
+    "indian ocean yellow-nosed mollymawk": "all.bird.mollymawk.indian_ocean_yellownosed_mollymawk",
+    "salvin's mollymawk": "all.bird.mollymawk.salvins_mollymawk",
+    "white-capped mollymawk": "all.bird.mollymawk.whitecapped_mollymawk",
+    "monarch": "all.bird.monarch",
+    "black-faced monarch": "all.bird.monarch.blackfaced_monarch",
+    "moorhen": "all.bird.moorhen",
+    "common moorhen": "all.bird.moorhen.common_moorhen",
+    "dusky moorhen": "all.bird.moorhen.dusky_moorhen",
+    "morepork": "all.bird.morepork",
+    "norfolk morepork": "all.bird.morepork.norfolk_morepork",
+    "myna": "all.bird.myna",
+    "common myna": "all.bird.myna.common_myna",
+    "needletail": "all.bird.needletail",
+    "white-throated needletail": "all.bird.needletail.whitethroated_needletail",
+    "noddy": "all.bird.noddy",
+    "black noddy": "all.bird.noddy.black_noddy",
+    "brown noddy": "all.bird.noddy.brown_noddy",
+    "grey noddy": "all.bird.noddy.grey_noddy",
+    "norfolk gerygone": "all.bird.norfolk_gerygone",
+    "norfolk golden whistler": "all.bird.norfolk_golden_whistler",
+    "norfolk silkeye": "all.bird.norfolk_silkeye",
+    "northern shoveler": "all.bird.northern_shoveler",
+    "owl": "all.bird.owl",
+    "barn owl": "all.bird.owl.barn_owl",
+    "little owl": "all.bird.owl.little_owl",
+    "oystercatcher": "all.bird.oystercatcher",
+    "chatham island oystercatcher": "all.bird.oystercatcher.chatham_island_oystercatcher",
+    "south island pied oystercatcher": "all.bird.oystercatcher.south_island_pied_oystercatcher",
+    "variable oystercatcher": "all.bird.oystercatcher.variable_oystercatcher",
+    "paradise shelduck": "all.bird.paradise_shelduck",
+    "parakeet": "all.bird.parakeet",
+    "antipodes island parakeet": "all.bird.parakeet.antipodes_island_parakeet",
+    "forbes' parakeet": "all.bird.parakeet.forbes_parakeet",
+    "orange-fronted parakeet": "all.bird.parakeet.orangefronted_parakeet",
+    "red-crowned parakeet": "all.bird.parakeet.redcrowned_parakeet",
+    "reischek's parakeet": "all.bird.parakeet.reischeks_parakeet",
+    "rose-ringed parakeet": "all.bird.parakeet.roseringed_parakeet",
+    "yellow-crowned parakeet": "all.bird.parakeet.yellowcrowned_parakeet",
+    "parrot": "all.bird.parrot",
+    "norfolk parrot": "all.bird.parrot.norfolk_parrot",
+    "partridge": "all.bird.partridge",
+    "red-legged partridge": "all.bird.partridge.redlegged_partridge",
+    "peafowl": "all.bird.peafowl",
+    "pelican": "all.bird.pelican",
+    "australian pelican": "all.bird.pelican.australian_pelican",
+    "penguin": "all.bird.penguin",
+    "adelie penguin": "all.bird.penguin.adelie_penguin",
+    "chinstrap penguin": "all.bird.penguin.chinstrap_penguin",
+    "eastern rockhopper penguin": "all.bird.penguin.eastern_rockhopper_penguin",
+    "emperor penguin": "all.bird.penguin.emperor_penguin",
+    "erect-crested penguin": "all.bird.penguin.erectcrested_penguin",
+    "fiordland crested penguin": "all.bird.penguin.fiordland_crested_penguin",
+    "gentoo penguin": "all.bird.penguin.gentoo_penguin",
+    "king penguin": "all.bird.penguin.king_penguin",
+    "little penguin": "all.bird.penguin.little_penguin",
+    "macaroni penguin": "all.bird.penguin.macaroni_penguin",
+    "moseley's rockhopper penguin": "all.bird.penguin.moseleys_rockhopper_penguin",
+    "royal penguin": "all.bird.penguin.royal_penguin",
+    "snares crested penguin": "all.bird.penguin.snares_crested_penguin",
+    "western rockhopper penguin": "all.bird.penguin.western_rockhopper_penguin",
+    "yellow-eyed penguin": "all.bird.penguin.yelloweyed_penguin",
+    "petrel": "all.bird.petrel",
+    "antarctic petrel": "all.bird.petrel.antarctic_petrel",
+    "black petrel": "all.bird.petrel.black_petrel",
+    "black-bellied storm petrel": "all.bird.petrel.blackbellied_storm_petrel",
+    "black-winged petrel": "all.bird.petrel.blackwinged_petrel",
+    "blue petrel": "all.bird.petrel.blue_petrel",
+    "bulwer's petrel": "all.bird.petrel.bulwers_petrel",
+    "cape petrel": "all.bird.petrel.cape_petrel",
+    "chatham petrel": "all.bird.petrel.chatham_petrel",
+    "collared petrel": "all.bird.petrel.collared_petrel",
+    "common diving petrel": "all.bird.petrel.common_diving_petrel",
+    "cook's petrel": "all.bird.petrel.cooks_petrel",
+    "gould's petrel": "all.bird.petrel.goulds_petrel",
+    "grey petrel": "all.bird.petrel.grey_petrel",
+    "grey-backed storm petrel": "all.bird.petrel.greybacked_storm_petrel",
+    "grey-faced petrel": "all.bird.petrel.greyfaced_petrel",
+    "herald petrel": "all.bird.petrel.herald_petrel",
+    "juan fernandez petrel": "all.bird.petrel.juan_fernandez_petrel",
+    "kerguelen petrel": "all.bird.petrel.kerguelen_petrel",
+    "kermadec petrel": "all.bird.petrel.kermadec_petrel",
+    "kermadec storm petrel": "all.bird.petrel.kermadec_storm_petrel",
+    "leach's storm petrel": "all.bird.petrel.leachs_storm_petrel",
+    "mottled petrel": "all.bird.petrel.mottled_petrel",
+    "new zealand storm petrel": "all.bird.petrel.new_zealand_storm_petrel",
+    "northern giant petrel": "all.bird.petrel.northern_giant_petrel",
+    "phoenix petrel": "all.bird.petrel.phoenix_petrel",
+    "providence petrel": "all.bird.petrel.providence_petrel",
+    "pycroft's petrel": "all.bird.petrel.pycrofts_petrel",
+    "snow petrel": "all.bird.petrel.snow_petrel",
+    "soft-plumaged petrel": "all.bird.petrel.softplumaged_petrel",
+    "south georgian diving petrel": "all.bird.petrel.south_georgian_diving_petrel",
+    "southern giant petrel": "all.bird.petrel.southern_giant_petrel",
+    "stejneger's petrel": "all.bird.petrel.stejnegers_petrel",
+    "tahiti petrel": "all.bird.petrel.tahiti_petrel",
+    "westland petrel": "all.bird.petrel.westland_petrel",
+    "white-bellied storm petrel": "all.bird.petrel.whitebellied_storm_petrel",
+    "white-chinned petrel": "all.bird.petrel.whitechinned_petrel",
+    "white-faced storm petrel": "all.bird.petrel.whitefaced_storm_petrel",
+    "white-headed petrel": "all.bird.petrel.whiteheaded_petrel",
+    "white-naped petrel": "all.bird.petrel.whitenaped_petrel",
+    "wilson's storm petrel": "all.bird.petrel.wilsons_storm_petrel",
+    "phalarope": "all.bird.phalarope",
+    "grey phalarope": "all.bird.phalarope.grey_phalarope",
+    "red-necked phalarope": "all.bird.phalarope.rednecked_phalarope",
+    "wilson's phalarope": "all.bird.phalarope.wilsons_phalarope",
+    "pheasant": "all.bird.pheasant",
+    "common pheasant": "all.bird.pheasant.common_pheasant",
+    "pigeon": "all.bird.pigeon",
+    "chatham island pigeon": "all.bird.pigeon.chatham_island_pigeon",
+    "rock pigeon": "all.bird.pigeon.rock_pigeon",
+    "pintail": "all.bird.pintail",
+    "northern pintail": "all.bird.pintail.northern_pintail",
+    "pipit": "all.bird.pipit",
+    "new zealand pipit": "all.bird.pipit.new_zealand_pipit",
+    "plover": "all.bird.plover",
+    "american golden plover": "all.bird.plover.american_golden_plover",
+    "greater sand plover": "all.bird.plover.greater_sand_plover",
+    "grey plover": "all.bird.plover.grey_plover",
+    "lesser sand plover": "all.bird.plover.lesser_sand_plover",
+    "pacific golden plover": "all.bird.plover.pacific_golden_plover",
+    "red-capped plover": "all.bird.plover.redcapped_plover",
+    "semipalmated plover": "all.bird.plover.semipalmated_plover",
+    "shore plover": "all.bird.plover.shore_plover",
+    "spur-winged plover": "all.bird.plover.spurwinged_plover",
+    "pratincole": "all.bird.pratincole",
+    "oriental pratincole": "all.bird.pratincole.oriental_pratincole",
+    "prion": "all.bird.prion",
+    "antarctic prion": "all.bird.prion.antarctic_prion",
+    "broad-billed prion": "all.bird.prion.broadbilled_prion",
+    "fairy prion": "all.bird.prion.fairy_prion",
+    "salvin's prion": "all.bird.prion.salvins_prion",
+    "thin-billed prion": "all.bird.prion.thinbilled_prion",
+    "pukeko": "all.bird.pukeko",
+    "purple swamphens": "all.bird.purple_swamphens",
+    "quail": "all.bird.quail",
+    "brown quail": "all.bird.quail.brown_quail",
+    "california quail": "all.bird.quail.california_quail",
+    "rail": "all.bird.rail",
+    "auckland island rail": "all.bird.rail.auckland_island_rail",
+    "banded rail": "all.bird.rail.banded_rail",
+    "rainbow lorikeet": "all.bird.rainbow_lorikeet",
+    "redpoll": "all.bird.redpoll",
+    "common redpoll": "all.bird.redpoll.common_redpoll",
+    "rifleman": "all.bird.rifleman",
+    "robin": "all.bird.robin",
+    "black robin": "all.bird.robin.black_robin",
+    "norfolk robin": "all.bird.robin.norfolk_robin",
+    "north island robin": "all.bird.robin.north_island_robin",
+    "south island robin": "all.bird.robin.south_island_robin",
+    "rock wren": "all.bird.rock_wren",
+    "rook": "all.bird.rook",
+    "rooster": "all.bird.rooster",
+    "rosella": "all.bird.rosella",
+    "crimson rosella": "all.bird.rosella.crimson_rosella",
+    "eastern rosella": "all.bird.rosella.eastern_rosella",
+    "ruff": "all.bird.ruff",
+    "saddleback": "all.bird.saddleback",
+    "north island saddleback": "all.bird.saddleback.north_island_saddleback",
+    "south island saddleback": "all.bird.saddleback.south_island_saddleback",
+    "sanderling": "all.bird.sanderling",
+    "sandpiper": "all.bird.sandpiper",
+    "baird's sandpiper": "all.bird.sandpiper.bairds_sandpiper",
+    "broad-billed sandpiper": "all.bird.sandpiper.broadbilled_sandpiper",
+    "buff-breasted sandpiper": "all.bird.sandpiper.buffbreasted_sandpiper",
+    "common sandpiper": "all.bird.sandpiper.common_sandpiper",
+    "curlew sandpiper": "all.bird.sandpiper.curlew_sandpiper",
+    "least sandpiper": "all.bird.sandpiper.least_sandpiper",
+    "marsh sandpiper": "all.bird.sandpiper.marsh_sandpiper",
+    "pectoral sandpiper": "all.bird.sandpiper.pectoral_sandpiper",
+    "semipalmated sandpiper": "all.bird.sandpiper.semipalmated_sandpiper",
+    "sharp-tailed sandpiper": "all.bird.sandpiper.sharptailed_sandpiper",
+    "stilt sandpiper": "all.bird.sandpiper.stilt_sandpiper",
+    "terek sandpiper": "all.bird.sandpiper.terek_sandpiper",
+    "upland sandpiper": "all.bird.sandpiper.upland_sandpiper",
+    "western sandpiper": "all.bird.sandpiper.western_sandpiper",
+    "white-rumped sandpiper": "all.bird.sandpiper.whiterumped_sandpiper",
+    "scaup": "all.bird.scaup",
+    "new zealand scaup": "all.bird.scaup.new_zealand_scaup",
+    "shag": "all.bird.shag",
+    "auckland island shag": "all.bird.shag.auckland_island_shag",
+    "black shag": "all.bird.shag.black_shag",
+    "bounty island shag": "all.bird.shag.bounty_island_shag",
+    "campbell island shag": "all.bird.shag.campbell_island_shag",
+    "chatham island shag": "all.bird.shag.chatham_island_shag",
+    "little black shag": "all.bird.shag.little_black_shag",
+    "little shag": "all.bird.shag.little_shag",
+    "macquarie island shag": "all.bird.shag.macquarie_island_shag",
+    "new zealand king shag": "all.bird.shag.new_zealand_king_shag",
+    "pied shag": "all.bird.shag.pied_shag",
+    "pitt island shag": "all.bird.shag.pitt_island_shag",
+    "spotted shag": "all.bird.shag.spotted_shag",
+    "stewart island shag": "all.bird.shag.stewart_island_shag",
+    "shearwater": "all.bird.shearwater",
+    "buller's shearwater": "all.bird.shearwater.bullers_shearwater",
+    "christmas island shearwater": "all.bird.shearwater.christmas_island_shearwater",
+    "cory's shearwater": "all.bird.shearwater.corys_shearwater",
+    "flesh-footed shearwater": "all.bird.shearwater.fleshfooted_shearwater",
+    "fluttering shearwater": "all.bird.shearwater.fluttering_shearwater",
+    "great shearwater": "all.bird.shearwater.great_shearwater",
+    "hutton's shearwater": "all.bird.shearwater.huttons_shearwater",
+    "little shearwater": "all.bird.shearwater.little_shearwater",
+    "manx shearwater": "all.bird.shearwater.manx_shearwater",
+    "newell's shearwater": "all.bird.shearwater.newells_shearwater",
+    "pink-footed shearwater": "all.bird.shearwater.pinkfooted_shearwater",
+    "short-tailed shearwater": "all.bird.shearwater.shorttailed_shearwater",
+    "sooty shearwater": "all.bird.shearwater.sooty_shearwater",
+    "streaked shearwater": "all.bird.shearwater.streaked_shearwater",
+    "subantarctic little shearwater": "all.bird.shearwater.subantarctic_little_shearwater",
+    "wedge-tailed shearwater": "all.bird.shearwater.wedgetailed_shearwater",
+    "silvereye": "all.bird.silvereye",
+    "norfolk silvereye": "all.bird.silvereye.norfolk_silvereye",
+    "slender-billed white-eye": "all.bird.silvereye.slenderbilled_whiteeye",
+    "skua": "all.bird.skua",
+    "arctic skua": "all.bird.skua.arctic_skua",
+    "long-tailed skua": "all.bird.skua.longtailed_skua",
+    "pomarine skua": "all.bird.skua.pomarine_skua",
+    "south polar skua": "all.bird.skua.south_polar_skua",
+    "subantarctic skua": "all.bird.skua.subantarctic_skua",
+    "skylark": "all.bird.skylark",
+    "eurasian skylark": "all.bird.skylark.eurasian_skylark",
+    "snipe": "all.bird.snipe",
+    "chatham island snipe": "all.bird.snipe.chatham_island_snipe",
+    "japanese snipe": "all.bird.snipe.japanese_snipe",
+    "painted snipe": "all.bird.snipe.painted_snipe",
+    "snares island snipe": "all.bird.snipe.snares_island_snipe",
+    "subantarctic snipe": "all.bird.snipe.subantarctic_snipe",
+    "sparrow": "all.bird.sparrow",
+    "house sparrow": "all.bird.sparrow.house_sparrow",
+    "spoonbill": "all.bird.spoonbill",
+    "royal spoonbill": "all.bird.spoonbill.royal_spoonbill",
+    "yellow-billed spoonbill": "all.bird.spoonbill.yellowbilled_spoonbill",
+    "starling": "all.bird.starling",
+    "common starling": "all.bird.starling.common_starling",
+    "stilt": "all.bird.stilt",
+    "black stilt": "all.bird.stilt.black_stilt",
+    "pied stilt": "all.bird.stilt.pied_stilt",
+    "stint": "all.bird.stint",
+    "little stint": "all.bird.stint.little_stint",
+    "long-toed stint": "all.bird.stint.longtoed_stint",
+    "red-necked stint": "all.bird.stint.rednecked_stint",
+    "stitchbird": "all.bird.stitchbird",
+    "swallow": "all.bird.swallow",
+    "welcome swallow": "all.bird.swallow.welcome_swallow",
+    "swan": "all.bird.swan",
+    "black swan": "all.bird.swan.black_swan",
+    "mute swan": "all.bird.swan.mute_swan",
+    "swift": "all.bird.swift",
+    "fork-tailed swift": "all.bird.swift.forktailed_swift",
+    "takahe": "all.bird.takahe",
+    "south island takahe": "all.bird.takahe.south_island_takahe",
+    "tattler": "all.bird.tattler",
+    "grey-tailed tattler": "all.bird.tattler.greytailed_tattler",
+    "wandering tattler": "all.bird.tattler.wandering_tattler",
+    "teal": "all.bird.teal",
+    "auckland island teal": "all.bird.teal.auckland_island_teal",
+    "brown teal": "all.bird.teal.brown_teal",
+    "campbell island teal": "all.bird.teal.campbell_island_teal",
+    "chestnut teal": "all.bird.teal.chestnut_teal",
+    "grey teal": "all.bird.teal.grey_teal",
+    "tern": "all.bird.tern",
+    "antarctic tern": "all.bird.tern.antarctic_tern",
+    "arctic tern": "all.bird.tern.arctic_tern",
+    "black-fronted tern": "all.bird.tern.blackfronted_tern",
+    "bridled tern": "all.bird.tern.bridled_tern",
+    "caspian tern": "all.bird.tern.caspian_tern",
+    "common tern": "all.bird.tern.common_tern",
+    "crested tern": "all.bird.tern.crested_tern",
+    "fairy tern": "all.bird.tern.fairy_tern",
+    "grey-backed tern": "all.bird.tern.greybacked_tern",
+    "gull-billed tern": "all.bird.tern.gullbilled_tern",
+    "little tern": "all.bird.tern.little_tern",
+    "sooty tern": "all.bird.tern.sooty_tern",
+    "whiskered tern": "all.bird.tern.whiskered_tern",
+    "white tern": "all.bird.tern.white_tern",
+    "white-fronted tern": "all.bird.tern.whitefronted_tern",
+    "white-winged black tern": "all.bird.tern.whitewinged_black_tern",
+    "thrush": "all.bird.thrush",
+    "song thrush": "all.bird.thrush.song_thrush",
+    "tomtit": "all.bird.tomtit",
+    "triller": "all.bird.triller",
+    "white-winged triller": "all.bird.triller.whitewinged_triller",
+    "tropicbird": "all.bird.tropicbird",
+    "red-tailed tropicbird": "all.bird.tropicbird.redtailed_tropicbird",
+    "white-tailed tropicbird": "all.bird.tropicbird.whitetailed_tropicbird",
+    "tui": "all.bird.tui",
+    "turkey": "all.bird.turkey",
+    "turnstone": "all.bird.turnstone",
+    "ruddy turnstone": "all.bird.turnstone.ruddy_turnstone",
+    "wagtail": "all.bird.wagtail",
+    "willie wagtail": "all.bird.wagtail.willie_wagtail",
+    "warbler": "all.bird.warbler",
+    "australian reed warbler": "all.bird.warbler.australian_reed_warbler",
+    "chatham island warbler": "all.bird.warbler.chatham_island_warbler",
+    "grey warbler": "all.bird.warbler.grey_warbler",
+    "wattlebird": "all.bird.wattlebird",
+    "red wattlebird": "all.bird.wattlebird.red_wattlebird",
+    "weka": "all.bird.weka",
+    "whimbrel": "all.bird.whimbrel",
+    "little whimbrel": "all.bird.whimbrel.little_whimbrel",
+    "whitehead": "all.bird.whitehead",
+    "wild turkey": "all.bird.wild_turkey",
+    "woodswallow": "all.bird.woodswallow",
+    "masked woodswallow": "all.bird.woodswallow.masked_woodswallow",
+    "white-browed woodswallow": "all.bird.woodswallow.whitebrowed_woodswallow",
+    "wrybill": "all.bird.wrybill",
+    "yellowhammer": "all.bird.yellowhammer",
+    "yellowhead": "all.bird.yellowhead"
+}
diff --git a/load.py b/load.py
deleted file mode 100644
index e2eec0bb..00000000
--- a/load.py
+++ /dev/null
@@ -1,9 +0,0 @@
-#!/usr/bin/env python
-
-"""
-Script to Load Tagged tracks into a h5py file.
-"""
-
-from load.load import main
-
-main()
diff --git a/load/__init__.py b/load/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/load/cliploader.py b/load/cliploader.py
deleted file mode 100644
index ec7db318..00000000
--- a/load/cliploader.py
+++ /dev/null
@@ -1,357 +0,0 @@
-"" """
-classifier-pipeline - this is a server side component that manipulates cptv
-files and to create a classification model of animals present
-Copyright (C) 2018, The Cacophony Project
-
-This program is free software: you can redistribute it and/or modify
-it under the terms of the GNU General Public License as published by
-the Free Software Foundation, either version 3 of the License, or
-(at your option) any later version.
-
-This program is distributed in the hope that it will be useful,
-but WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-GNU General Public License for more details.
-
-You should have received a copy of the GNU General Public License
-along with this program. If not, see <http://www.gnu.org/licenses/>.
-"""
-
-import os
-import logging
-import time
-from multiprocessing import Process, Queue
-import traceback
-from ml_tools import tools
-
-from ml_tools.previewer import Previewer
-from ml_tools.trackdatabase import TrackDatabase
-from .clip import Clip
-from .irtrackextractor import IRTrackExtractor
-from .cliptrackextractor import ClipTrackExtractor
-from ml_tools.imageprocessing import clear_frame
-from track.track import Track
-import numpy as np
-import json
-
-
-def process_job(loader, queue):
-    i = 0
-    while True:
-        i += 1
-        filename, clip_id = queue.get()
-        try:
-            if filename == "DONE":
-                break
-            else:
-                loader.process_file(str(filename), clip_id)
-            if i % 50 == 0:
-                logging.info("%s jobs left", queue.qsize())
-        except Exception as e:
-            logging.error("Process_job error %s %s", filename, e)
-            traceback.print_exc()
-
-
-class ClipLoader:
-    def __init__(self, config, reprocess=False):
-        self.config = config
-        os.makedirs(self.config.tracks_folder, mode=0o775, exist_ok=True)
-        self.database = TrackDatabase(
-            os.path.join(self.config.tracks_folder, "dataset.hdf5")
-        )
-        self.reprocess = reprocess
-        self.compression = (
-            tools.gzip_compression if self.config.load.enable_compression else None
-        )
-        # number of threads to use when processing jobs.
-        self.workers_threads = config.worker_threads
-        self.previewer = Previewer.create_if_required(config, config.load.preview)
-
-    def process_all(self, root):
-        clip_id = self.database.get_unique_clip_id()
-        job_queue = Queue()
-        processes = []
-        for i in range(max(1, self.workers_threads)):
-            p = Process(
-                target=process_job,
-                args=(self, job_queue),
-            )
-            processes.append(p)
-            p.start()
-        if root is None:
-            root = self.config.source_folder
-        file_paths = []
-        for folder_path, _, files in os.walk(root):
-            for name in files:
-                if os.path.splitext(name)[1] in [".mp4", ".avi", ".cptv"]:
-                    full_path = os.path.join(folder_path, name)
-                    file_paths.append(full_path)
-        # allows us know the order of processing
-        file_paths.sort()
-        for file_path in file_paths:
-            job_queue.put((file_path, clip_id))
-            clip_id += 1
-
-        logging.info("Processing %d", job_queue.qsize())
-        for i in range(len(processes)):
-            job_queue.put(("DONE", 0))
-        for process in processes:
-            try:
-                process.join()
-            except KeyboardInterrupt:
-                logging.info("KeyboardInterrupt, terminating.")
-                for process in processes:
-                    process.terminate()
-                exit()
-
-    def _get_dest_folder(self, filename):
-        return os.path.join(self.config.tracks_folder, get_distributed_folder(filename))
-
-    def _export_tracks(self, full_path, clip):
-        """
-        Writes tracks to a track database.
-        :param database: database to write track to.
-        """
-        # overwrite any old clips.
-        # Note: we do this even if there are no tracks so there there will be a blank clip entry as a record
-        # that we have processed it.
-        self.database.create_clip(clip)
-        for track in clip.tracks:
-            start_time, end_time = clip.start_and_end_time_absolute(
-                track.start_s, track.end_s
-            )
-            original_thermal = []
-            cropped_data = []
-            prev_frame = 0
-            for i, region in enumerate(track.bounds_history):
-                frame = clip.frame_buffer.get_frame(region.frame_number)
-                if frame is None:
-                    prev_end = track.end_s
-                    track.end_s = (
-                        min(prev_frame, len(clip.frame_buffer.frames))
-                        / clip.frames_per_second
-                    )
-
-                    track.bounds_history = track.bounds_history[:i]
-                    logging.warn(
-                        "Could not get frame %s from %s track %s adjusting end_s from %s to %s",
-                        region.frame_number,
-                        clip.source_file,
-                        track.get_id(),
-                        prev_end,
-                        track.end_s,
-                    )
-                    break
-                original_thermal.append(frame.thermal)
-                cropped = frame.crop_by_region(region)
-                # zero out the filtered channel
-                if not self.config.load.include_filtered_channel:
-                    cropped.filtered = np.zeros(cropped.thermal.shape)
-                cropped_data.append(cropped)
-                prev_frame = region.frame_number
-            # sample_frames = get_sample_frames(
-            #     clip.ffc_frames,
-            #     [bounds.mass for bounds in track.bounds_history],
-            #     self.config.build.segment_min_avg_mass,
-            #     cropped_data,
-            # )
-            if track.start_s >= track.end_s:
-                logging.warn(
-                    "No data for %s track %s ", clip.source_file, track.get_id()
-                )
-                return
-            try:
-                self.database.add_track(
-                    clip.get_id(),
-                    track,
-                    cropped_data,
-                    opts=self.compression,
-                    original_thermal=original_thermal,
-                    start_time=start_time,
-                    end_time=end_time,
-                )
-            except:
-                self.database.remove_track(clip.get_id(), track.get_id())
-                logging.error(
-                    "Error adding track %s - %s",
-                    clip.get_id(),
-                    track.get_id(),
-                    exc_info=True,
-                )
-        self.database.finished_processing(clip.get_id())
-
-    def _filter_clip_tracks(self, clip_metadata, min_confidence):
-        """
-        Removes track metadata for tracks which are invalid. Tracks are invalid
-        if they aren't confident or they are in the excluded_tags list.
-        Returns valid tracks
-        """
-
-        tracks_meta = clip_metadata.get("Tracks")
-        if tracks_meta is None:
-            tracks_meta = clip_metadata.get("tracks", [])
-        valid_tracks = [
-            track
-            for track in tracks_meta
-            if self._track_meta_is_valid(track, min_confidence)
-        ]
-
-        clip_metadata["Tracks"] = valid_tracks
-        return valid_tracks
-
-    def _track_meta_is_valid(self, track_meta, min_confidence):
-        """
-        Tracks are valid if their confidence meets the threshold and they are
-        not in the excluded_tags list, defined in the config.
-        """
-        track_tags = []
-        if "tags" in track_meta:
-            track_tags = track_meta["tags"]
-        excluded_tags = [
-            tag
-            for tag in track_tags
-            if not tag.get("automatic", False)
-            and tag.get("what") in self.config.load.excluded_tags
-        ]
-        if len(excluded_tags) > 0:
-            return False
-
-        track_tag = Track.get_best_human_tag(
-            track_tags, self.config.load.tag_precedence, min_confidence
-        )
-
-        if track_tag is None:
-            return False
-        tag = track_tag.get("what")
-        confidence = track_tag.get("confidence", 0)
-        return tag and tag not in excluded_tags and confidence >= min_confidence
-
-    def get_track_extractor(self, filename):
-        ext = os.path.splitext(filename)[1]
-        if ext == ".cptv":
-            track_extractor = ClipTrackExtractor(
-                self.config.tracking,
-                self.config.use_opt_flow
-                or self.config.load.preview == Previewer.PREVIEW_TRACKING,
-                self.config.load.cache_to_disk,
-                high_quality_optical_flow=self.config.load.high_quality_optical_flow,
-                verbose=self.config.verbose,
-                do_tracking=False,
-            )
-        else:
-            track_extractor = IRTrackExtractor(
-                self.config.tracking,
-                self.config.load.cache_to_disk,
-                verbose=self.config.verbose,
-                do_tracking=True,
-                tracking_alg="subsense",
-            )
-        return track_extractor
-
-    def process_file(self, filename, clip_id=None, classifier=None):
-        start = time.time()
-        base_filename = os.path.splitext(os.path.basename(filename))[0]
-
-        logging.info(f"processing %s", filename)
-
-        track_extractor = self.get_track_extractor(filename)
-
-        destination_folder = self._get_dest_folder(base_filename)
-        # delete any previous files
-        tools.purge(destination_folder, base_filename + "*.mp4")
-
-        # read metadata
-        metadata_filename = os.path.join(
-            os.path.dirname(filename), base_filename + ".txt"
-        )
-
-        if not os.path.isfile(metadata_filename):
-            logging.error("No meta data found for %s", metadata_filename)
-            return
-
-        metadata = tools.load_clip_metadata(metadata_filename)
-        if "id" not in metadata and clip_id is not None:
-            metadata["id"] = clip_id
-            logging.info("Using clip id %s", clip_id)
-            with open(metadata_filename, "w") as f:
-                json.dump(metadata, f, indent=4, cls=tools.CustomJSONEncoder)
-        if not self.reprocess and self.database.has_clip(str(metadata["id"])):
-            logging.warning("Already loaded %s", filename)
-            return
-        valid_tracks = self._filter_clip_tracks(
-            metadata, track_extractor.config.min_tag_confidence
-        )
-
-        if not valid_tracks or len(valid_tracks) == 0:
-            logging.error("No valid track data found for %s", filename)
-            return
-        clip = Clip(track_extractor.config, filename)
-        clip.load_metadata(
-            metadata,
-            self.config.load.tag_precedence,
-        )
-        tracker_version = metadata.get("tracker_version", 0)
-        process_background = tracker_version < 10
-        logging.debug(
-            "Processing background? %s tracker_version %s",
-            process_background,
-            tracker_version,
-        )
-        if not track_extractor.parse_clip(clip, process_background=process_background):
-            logging.error("No valid clip found for %s", filename)
-            return
-
-        self._export_tracks(filename, clip)
-
-        # write a preview
-        if self.previewer:
-            os.makedirs(destination_folder, mode=0o775, exist_ok=True)
-
-            preview_filename = base_filename + "-preview" + ".mp4"
-            preview_filename = os.path.join(destination_folder, preview_filename)
-            self.previewer.create_individual_track_previews(preview_filename, clip)
-            self.previewer.export_clip_preview(preview_filename, clip)
-
-        if self.config.verbose:
-            num_frames = len(clip.frame_buffer.frames)
-            ms_per_frame = (time.time() - start) * 1000 / max(1, num_frames)
-            self._log_message(
-                "Tracks {}.  Frames: {}, Took {:.1f}ms per frame".format(
-                    len(clip.tracks), num_frames, ms_per_frame
-                )
-            )
-
-    def _log_message(self, message):
-        """Record message in stdout.  Will be printed if verbose is enabled."""
-        # note, python has really good logging... I should probably make use of this.
-        if self.config.verbose:
-            logging.info(message)
-
-
-def get_distributed_folder(name, num_folders=256, seed=31):
-    """Creates a hash of the name then returns the modulo num_folders"""
-    str_bytes = str.encode(name)
-    hash_code = 0
-    for byte in str_bytes:
-        hash_code = hash_code * seed + byte
-    return "{:02x}".format(hash_code % num_folders)
-
-
-def get_sample_frames(ffc_frames, mass_history, min_mass=None, frame_data=None):
-    clear_frames = []
-    lower_mass = np.percentile(mass_history, q=25)
-    upper_mass = np.percentile(mass_history, q=75)
-    for i, mass in enumerate(mass_history):
-        if i in ffc_frames:
-            continue
-        if (
-            min_mass is None
-            or mass >= min_mass
-            and mass >= lower_mass
-            and mass <= upper_mass
-        ):
-            if frame_data is not None:
-                if not clear_frame(frame_data[i]):
-                    continue
-            clear_frames.append(i)
-    return clear_frames
diff --git a/load/load.py b/load/load.py
deleted file mode 100644
index 56010c04..00000000
--- a/load/load.py
+++ /dev/null
@@ -1,106 +0,0 @@
-"""
-Processes a CPTV file identifying and tracking regions of interest, and saving them in the 'trk' format.
-"""
-
-import argparse
-import cv2
-import os
-
-from ml_tools.logs import init_logging
-from config.config import Config
-from .cliploader import ClipLoader
-
-
-def parse_params():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-target",
-        default=None,
-        help='Target to process, "all" processes all folders, "test" runs test cases, "clean" to remove banned clips from db, or a "cptv" file to run a single source.',
-    )
-
-    parser.add_argument(
-        "-p",
-        "--create-previews",
-        action="count",
-        help="Create MP4 previews for tracks (can be slow)",
-    )
-    parser.add_argument(
-        "-t",
-        "--test-file",
-        default="tests.txt",
-        help="File containing test cases to run",
-    )
-    parser.add_argument(
-        "-v", "--verbose", action="count", help="Display additional information."
-    )
-
-    parser.add_argument(
-        "-r",
-        "--reprocess",
-        action="count",
-        help="Re process clips that already exist in the database",
-    )
-
-    parser.add_argument(
-        "-i",
-        "--show-build-information",
-        action="count",
-        help="Show openCV build information and exit.",
-    )
-    parser.add_argument("-c", "--config-file", help="Path to config file to use")
-    args = parser.parse_args()
-
-    if args.show_build_information:
-        print(cv2.getBuildInformation())
-        return None, None
-
-    config = Config.load_from_file(args.config_file)
-    if args.create_previews:
-        config.loader.preview = "tracking"
-    if args.verbose:
-        config.tracking.verbose = True
-
-    return config, args
-
-
-def load_clips(config, args):
-    loader = ClipLoader(config, args.reprocess)
-    target = args.target
-    if target is None:
-        target = config.source_folder
-    if os.path.splitext(target)[1] in [".avi", ".cptv", ".mp4"]:
-        clip_id = loader.database.get_unique_clip_id()
-
-        loader.process_file(target, clip_id)
-    else:
-        loader.process_all(target)
-
-
-def print_opencl_info():
-    """Print information about opencv support for opencl."""
-    if cv2.ocl.haveOpenCL():
-        if cv2.ocl.useOpenCL():
-            print("OpenCL found and enabled, threads={}".format(cv2.getNumThreads()))
-        else:
-            print("OpenCL found but disabled")
-
-
-def main():
-    config, args = parse_params()
-    init_logging()
-
-    if config and args:
-        load_clips(config, args)
-
-
-if __name__ == "__main__":
-    # opencv sometimes uses too many threads which can reduce performance.  We are running a worker pool which makes
-    # better use of multiple cores, so best to leave thread count per process reasonably low.
-    # there is quite a big difference between 1 thread and 2, but after that gains are very minimal, and a lot of
-    # cpu time gets wasted, starving the other workers.
-    cv2.setNumThreads(2)
-
-    print_opencl_info()
-
-    main()
diff --git a/makedataset.py b/makedataset.py
new file mode 100644
index 00000000..7d91c590
--- /dev/null
+++ b/makedataset.py
@@ -0,0 +1,42 @@
+"""
+Processes a CPTV file identifying and tracking regions of interest, and saving them in the 'trk' format.
+"""
+
+import argparse
+import os
+
+from ml_tools.logs import init_logging
+from config.config import Config
+from mldataset.makedataset import ClipLoader
+
+from pathlib import Path
+
+
+def parse_params():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--out-folder",
+        default="clip_db",
+        help="Folder to put processed files",
+    )
+
+    parser.add_argument(
+        "target",
+        default=None,
+        help='Target to process, "all" processes all folders, "test" runs test cases, "clean" to remove banned clips from db, or a "cptv" file to run a single source.',
+    )
+    args = parser.parse_args()
+    config = Config.load_from_file(None)
+    args.out_folder = Path(args.out_folder)
+    return config, args
+
+
+def main():
+    config, args = parse_params()
+    init_logging()
+    loader = ClipLoader(config)
+    loader.process_all(args.target, args.out_folder)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/ml_tools/dataset.py b/ml_tools/dataset.py
index 65ba7d75..98cc4e02 100644
--- a/ml_tools/dataset.py
+++ b/ml_tools/dataset.py
@@ -15,9 +15,12 @@
 import gc
 from ml_tools.datasetstructures import TrackHeader, TrackingSample, SegmentType
 from ml_tools.trackdatabase import TrackDatabase
+from ml_tools.rawdb import RawDatabase
 from ml_tools import tools
 from track.region import Region
 import json
+from config.loadconfig import LoadConfig
+from pathlib import Path
 
 
 class Dataset:
@@ -28,42 +31,39 @@ class Dataset:
 
     def __init__(
         self,
-        db_file,
+        dataset_dir,
         name="Dataset",
         config=None,
         use_predictions=False,
         consecutive_segments=False,
         labels=[],
+        label_mapping=None,
+        raw=True,
+        ext=".cptv",
     ):
+        self.ext = ext
+        self.raw = raw
+        # if self.raw and self.ext == ".cptv":
+        # raise Exception("Raw CPTV not implemented yet")
+        self.dataset_dir = Path(dataset_dir)
         self.consecutive_segments = consecutive_segments
-        # self.camera_bins = {}
-        # database holding track data
-        self.db_file = db_file
-        self.db = None
-        self.load_db()
-        self.label_mapping = None
+        self.label_mapping = label_mapping
         # name of this dataset
         self.name = name
         # list of our tracks
-        self.samples_by_label = {}
         self.samples_by_bin = {}
         self.samples = []
-        self.sample_cdf = []
-        self.tracks = []
-        self.tracks_by_label = {}
-        self.tracks_by_bin = {}
-        self.tracks_by_id = {}
-        self.camera_names = set()
-        # self.cameras_by_id = {}
+        self.samples_by_id = {}
+        self.clips = []
 
         # list of label names
         self.labels = labels
-        self.label_mapping = None
 
         self.enable_augmentation = False
         self.label_caps = {}
         self.use_segments = True
         if config:
+            self.tag_precedence = config.load.tag_precedence
             self.type = config.train.type
             if config.train.type == "IR":
                 self.use_segments = False
@@ -82,14 +82,14 @@ def __init__(
             self.filter_by_lq = config.build.filter_by_lq
             self.segment_type = SegmentType.ALL_RANDOM
             self.max_segments = config.build.max_segments
-
         else:
+            self.tag_precedence = LoadConfig.DEFAULT_GROUPS
             self.filter_by_lq = True
             # number of seconds each segment should be
             self.segment_length = 25
             # number of seconds segments are spaced apart
             self.segment_spacing = 1
-            self.segment_min_avg_mass = None
+            self.segment_min_avg_mass = 10
             self.min_frame_mass = 16
             self.segment_type = SegmentType.ALL_RANDOM
         self.filtered_stats = {
@@ -108,28 +108,11 @@ def __init__(
         self.lbl_p = None
         self.numpy_data = None
 
-    def load_db(self):
-        self.db = TrackDatabase(self.db_file)
-
-    def set_read_only(self, read_only):
-        if self.db is not None:
-            self.db.set_read_only(read_only)
+        self.skip_ffc = True
 
     @property
     def sample_count(self):
-        return len(self.samples)
-
-    #
-    # def samples(self):
-    #     if self.use_segments:
-    #         return self.segments
-    #     return self.frame_samples
-
-    def set_samples(self, samples):
-        self.samples = samples
-
-    def set_samples_for(self, label, samples):
-        self.samples_by_label[label] = samples
+        return len(self.samples_by_id)
 
     def get_label_caps(self, labels, remapped=False):
         counts = []
@@ -169,31 +152,11 @@ def get_counts(self, label):
         bins = 0
         weight = 0
         samples_count = 0
-        if self.label_mapping:
-            for key, value in self.label_mapping.items():
-                if key == label or value == label:
-                    label_tracks = self.tracks_by_label.get(key, [])
-                    tracks += len(label_tracks)
-                    samples_count += len(self.samples_by_label.get(key, []))
-                    # segments += sum(len(track.segments) for track in label_tracks)
-                    # frames += sum(
-                    #     len(track.get_sample_frames())
-                    #     for track in label_tracks
-                    #     if track.sample_frames is not None
-                    # )
-
-                if key == label or value == label:
-                    label_tracks = []
-                    segments = self.segments_by_label.get(key, [])
-                    tracks += len(set([segment.track_id for segment in segments]))
-                    segments_count += len(segments)
-                    frames += sum([segment.frames for segment in segments])
-        else:
-            samples = self.samples_by_label.get(label, [])
-            tracks = len(set([sample.track_id for sample in samples]))
-            weight = self.get_label_weight(label)
-            bins = len(set([sample.bin_id for sample in samples]))
-            samples_count = len(samples)
+        samples = self.samples_by_label.get(label, [])
+        tracks = len(set([sample.track_id for sample in samples]))
+        weight = self.get_label_weight(label)
+        bins = len(set([sample.bin_id for sample in samples]))
+        samples_count = len(samples)
         return samples_count, tracks, bins, weight
 
     def load_clips(
@@ -202,36 +165,49 @@ def load_clips(
         before_date=None,
         after_date=None,
         label=None,
+        dont_filter_segment=False,
     ):
         """
         Loads track headers from track database with optional filter
         :return: [number of tracks added, total tracks].
         """
+
         counter = 0
         logging.info("Loading clips")
-        clip_ids = self.db.get_all_clip_ids(before_date, after_date, label)
-        if shuffle:
-            np.random.shuffle(clip_ids)
-        for clip_id in clip_ids:
-            self.load_clip(clip_id)
+        for db_clip in self.dataset_dir.glob(f"**/*{self.ext}"):
+            tracks_added = self.load_clip(db_clip, dont_filter_segment)
+            if tracks_added == 0:
+                logging.info("No tracks added for %s", db_clip)
             counter += 1
             if counter % 50 == 0:
-                logging.debug("Dataset loaded %s / %s", counter, len(clip_ids))
-        return [counter, len(clip_ids)]
+                logging.debug("Dataset loaded %s", counter)
+        return [counter, counter]
 
-    def load_clip(self, clip_id):
-        clip_meta = self.db.get_clip_meta(clip_id)
-        tracks = self.db.get_clip_tracks(clip_id)
+    def load_clip(self, db_clip, dont_filter_segment=False):
+        if self.raw:
+            db = RawDatabase(db_clip)
+        else:
+            db = TrackDatabase(db_clip)
+        try:
+            clip_header = db.get_clip_tracks(self.tag_precedence)
+        except:
+            logging.error("Could not load %s", db_clip, exc_info=True)
+            return 0
+        if clip_header is None or filter_clip(clip_header):
+            return 0
         filtered = 0
-        for track_meta in tracks:
-            if self.filter_track(clip_meta, track_meta):
-                filtered += 1
-                continue
-            track_header = TrackHeader.from_meta(clip_id, clip_meta, track_meta)
-            self.tracks.append(track_header)
+        added = 0
+        clip_header.tracks = [
+            track
+            for track in clip_header.tracks
+            if not filter_track(track, self.filtered_stats)
+        ]
+        self.clips.append(clip_header)
+        for track_header in clip_header.tracks:
+            added += 1
             if self.use_segments:
                 segment_frame_spacing = int(
-                    round(self.segment_spacing * track_header.frames_per_second)
+                    round(self.segment_spacing * clip_header.frames_per_second)
                 )
                 segment_width = self.segment_length
                 track_header.calculate_segments(
@@ -240,25 +216,30 @@ def load_clip(self, clip_id):
                     self.segment_type,
                     self.segment_min_avg_mass,
                     max_segments=self.max_segments,
+                    dont_filter=dont_filter_segment,
+                    skip_ffc=self.skip_ffc,
+                    ffc_frames=clip_header.ffc_frames,
                 )
                 self.filtered_stats["segment_mass"] += track_header.filtered_stats[
                     "segment_mass"
                 ]
-                for segment in track_header.segments:
-                    self.add_clip_sample_mappings(segment)
+
             else:
-                sample_frames = track_header.get_sample_frames()
-                skip_x = None
+                skip_last = None
                 if self.type == "IR":
-                    skip_last = int(len(sample_frames) * 0.1)
-                    sample_frames = sample_frames[:-skip_last]
-                for sample in sample_frames:
-                    if not self.filter_by_lq or (
-                        sample.mass >= track_header.lower_mass
-                        and sample.mass <= track_header.upper_mass
-                    ):
-                        self.add_clip_sample_mappings(sample)
-        return filtered
+                    # a lot of ir clips have bad tracking near end so just reduce track length
+                    skip_last = 0.1
+                track_header.calculate_sample_frames(
+                    min_mass=None if not self.filter_by_lq else track_header.lower_mass,
+                    max_mass=None if not self.filter_by_lq else track_header.upper_mass,
+                    ffc_frames=clip_header.ffc_frames,
+                    skip_last=skip_last,
+                )
+                if track_header.label not in self.labels:
+                    self.labels.append(track_header.label)
+            for sample in track_header.samples:
+                self.add_clip_sample_mappings(sample)
+        return added
 
     def add_samples(self, samples):
         """
@@ -271,96 +252,30 @@ def add_samples(self, samples):
                 result += 1
         return result
 
-    def filter_sample(self, sample):
-        if sample.label not in self.included_labels:
-            self.filtered_stats["tags"] += 1
-            return True
-
-        if self.min_frame_mass and sample.mass < self.min_frame_mass:
-            return True
-        return False
+    @property
+    def samples_by_label(self):
+        samples_by_label = {}
+        for sample in self.samples:
+            if sample.label not in samples_by_label:
+                samples_by_label[sample.label] = []
+            samples_by_label[sample.label].append(sample)
+        return samples_by_label
 
     def add_clip_sample_mappings(self, sample):
-        if self.filter_sample(sample):
-            return False
+        self.samples_by_id[sample.id] = sample
         self.samples.append(sample)
-        if self.label_mapping and sample.label in self.label_mapping:
-            sample.label = self.mapped_label(sample.label)
+
+        if self.label_mapping:
+            sample.remapped_label = self.label_mapping.get(
+                sample.original_label, sample.original_label
+            )
 
         if sample.label not in self.labels:
             self.labels.append(sample.label)
-
-        if sample.label not in self.samples_by_label:
-            self.samples_by_label[sample.label] = []
-        self.samples_by_label[sample.label].append(sample)
-
         bins = self.samples_by_bin.setdefault(sample.bin_id, [])
         bins.append(sample)
-        self.camera_names.add(sample.camera)
         return True
 
-    def filter_track(self, clip_meta, track_meta):
-        # some clips are banned for various reasons4
-        source = os.path.basename(clip_meta["filename"])
-        if self.banned_clips and source in self.banned_clips:
-            self.filtered_stats["banned"] += 1
-            return True
-        if "tag" not in track_meta:
-            self.filtered_stats["notags"] += 1
-            return True
-        if track_meta["tag"] not in self.included_labels:
-            self.filtered_stats["tags"] += 1
-            self.filtered_stats["tag_names"].add(track_meta["tag"])
-            return True
-        track_tags = track_meta.get("track_tags")
-        if track_tags is not None:
-            try:
-                track_tags = json.loads(track_tags)
-            except:
-                logging.error(
-                    "Error loading track tags json for %s clip %s track %s",
-                    track_tags,
-                    clip_meta.get("id"),
-                    track_meta.get("id"),
-                )
-                self.filtered_stats["bad_track_json"] += 1
-
-                return True
-            excluded_tags = [
-                tag["what"]
-                for tag in track_tags
-                if not tag.get("automatic", False)
-                and tag.get("what") in self.excluded_tags
-            ]
-            if len(excluded_tags) > 0:
-                self.filtered_stats["tag_names"] |= set(excluded_tags)
-
-                self.filtered_stats["tags"] += 1
-                return True
-        # always let the false-positives through as we need them even though they would normally
-        # be filtered out.
-        if "bounds_history" not in track_meta or len(track_meta["bounds_history"]) == 0:
-            self.filtered_stats["no_data"] += 1
-            return True
-
-        if track_meta["tag"] == "false-positive":
-            return False
-
-        # for some reason we get some records with a None confidence?
-        if track_meta.get("confidence", 1.0) <= 0.6:
-            self.filtered_stats["confidence"] += 1
-            return True
-
-        # remove tracks of trapped animals
-        if (
-            "trap" in clip_meta.get("event", "").lower()
-            or "trap" in clip_meta.get("trap", "").lower()
-        ):
-            self.filtered_stats["trap"] += 1
-            return True
-
-        return False
-
     def epoch_samples(
         self, cap_samples=None, replace=True, random=True, cap_at=None, label_cap=None
     ):
@@ -396,44 +311,51 @@ def cdf(self):
     def label_cdf(self, label):
         return self.sample_label_cdf.get(label, [])
 
-    def get_sample(self, cap=None, replace=True, label=None, random=True):
-        """Returns a random frames from weighted list."""
-        if label:
-            samples = self.samples_for(label, remapped=True)
-            cdf = self.label_cdf(label)
-        else:
-            samples = self.samples()
-            cdf = self.cdf()
-        if not samples:
-            return None
-        if cap is None:
-            return samples
-        if random:
-            return np.random.choice(samples, cap, replace=replace, p=cdf)
-        else:
-            cap = min(cap, len(samples))
-            return samples[:cap]
+    def get_samples_by_source(self):
+        samples_by_source = {}
+        for s in self.samples_by_id.values():
+            samples_by_source.setdefault(s.source_file, []).append(s)
+        return samples_by_source
 
-    def balance_bins(self, max_bin_weight=None):
-        """
-        Adjusts weights so that bins with a number number of segments aren't sampled so frequently.
-        :param max_bin_weight: bins with more weight than this number will be scaled back to this weight.
-        """
-
-        for bin_name, samples in self.samples_by_bin.items():
-            bin_weight = sum(sample.sample_weight for sample in samples)
-            if bin_weight == 0:
-                continue
-            if max_bin_weight is None:
-                scale_factor = 1 / bin_weight
-                # means each bin has equal possiblity
-            elif bin_weight > max_bin_weight:
-                scale_factor = max_bin_weight / bin_weight
-            else:
-                scale_factor = 1
-            for sample in samples:
-                sample.weight = np.float16(sample.sample_weight * scale_factor)
-        self.rebuild_cdf()
+    #
+    # def get_sample(self, cap=None, replace=True, label=None, random=True):
+    #     """Returns a random frames from weighted list."""
+    #     if label:
+    #         samples = self.samples_for(label, remapped=True)
+    #         cdf = self.label_cdf(label)
+    #     else:
+    #         samples = self.samples()
+    #         cdf = self.cdf()
+    #     if not samples:
+    #         return None
+    #     if cap is None:
+    #         return samples
+    #     if random:
+    #         return np.random.choice(samples, cap, replace=replace, p=cdf)
+    #     else:
+    #         cap = min(cap, len(samples))
+    #         return samples[:cap]
+    #
+    # def balance_bins(self, max_bin_weight=None):
+    #     """
+    #     Adjusts weights so that bins with a number number of segments aren't sampled so frequently.
+    #     :param max_bin_weight: bins with more weight than this number will be scaled back to this weight.
+    #     """
+    #
+    #     for bin_name, samples in self.samples_by_bin.items():
+    #         bin_weight = sum(sample.sample_weight for sample in samples)
+    #         if bin_weight == 0:
+    #             continue
+    #         if max_bin_weight is None:
+    #             scale_factor = 1 / bin_weight
+    #             # means each bin has equal possiblity
+    #         elif bin_weight > max_bin_weight:
+    #             scale_factor = max_bin_weight / bin_weight
+    #         else:
+    #             scale_factor = 1
+    #         for sample in samples:
+    #             sample.weight = np.float16(sample.sample_weight * scale_factor)
+    #     self.rebuild_cdf()
 
     def remove_label(self, label_to_remove):
         """
@@ -453,55 +375,56 @@ def mapped_label(self, label):
             return self.label_mapping.get(label, label)
         return label
 
-    def rebuild_cdf(self, lbl_p=None):
-        """Calculates the CDF used for fast random sampling for frames and
-        segments, if balance labels is set each label has an equal chance of
-        being chosen
-        """
-        if lbl_p is None:
-            lbl_p = self.lbl_p
-
-        self.sample_cdf = []
-        total = 0
-        self.sample_label_cdf = {}
-
-        for sample in self.samples:
-            sample_weight = sample.sample_weight
-            if lbl_p and sample.label in lbl_p:
-                sample_weight *= lbl_p[sample.label]
-            total += sample_weight
-
-            self.sample_cdf.append(sample_weight)
-
-            cdf = self.sample_label_cdf.setdefault(sample.label, [])
-            cdf.append(sample.sample_weight)
-
-        if len(self.sample_cdf) > 0:
-            self.sample_cdf = [x / total for x in self.sample_cdf]
-
-        for key, cdf in self.sample_label_cdf.items():
-            total = sum(cdf)
-            self.sample_label_cdf[key] = [x / total for x in cdf]
-
-        if self.label_mapping:
-            mapped_cdf = {}
-            labels = list(self.label_mapping.keys())
-            labels.sort()
-            for label in labels:
-                if label not in self.sample_label_cdf:
-                    continue
-                label_cdf = self.sample_label_cdf[label]
-                new_label = self.label_mapping[label]
-                if lbl_p and label in lbl_p:
-                    label_cdf = np.float64(label_cdf)
-                    label_cdf *= lbl_p[label]
-                cdf = mapped_cdf.setdefault(new_label, [])
-                cdf.extend(label_cdf)
-
-            for key, cdf in mapped_cdf.items():
-                total = sum(cdf)
-                mapped_cdf[key] = [x / total for x in cdf]
-            self.sample_label_cdf = mapped_cdf
+    #
+    # def rebuild_cdf(self, lbl_p=None):
+    #     """Calculates the CDF used for fast random sampling for frames and
+    #     segments, if balance labels is set each label has an equal chance of
+    #     being chosen
+    #     """
+    #     if lbl_p is None:
+    #         lbl_p = self.lbl_p
+    #
+    #     self.sample_cdf = []
+    #     total = 0
+    #     self.sample_label_cdf = {}
+    #
+    #     for sample in self.samples:
+    #         sample_weight = sample.sample_weight
+    #         if lbl_p and sample.label in lbl_p:
+    #             sample_weight *= lbl_p[sample.label]
+    #         total += sample_weight
+    #
+    #         self.sample_cdf.append(sample_weight)
+    #
+    #         cdf = self.sample_label_cdf.setdefault(sample.label, [])
+    #         cdf.append(sample.sample_weight)
+    #
+    #     if len(self.sample_cdf) > 0:
+    #         self.sample_cdf = [x / total for x in self.sample_cdf]
+    #
+    #     for key, cdf in self.sample_label_cdf.items():
+    #         total = sum(cdf)
+    #         self.sample_label_cdf[key] = [x / total for x in cdf]
+    #
+    #     if self.label_mapping:
+    #         mapped_cdf = {}
+    #         labels = list(self.label_mapping.keys())
+    #         labels.sort()
+    #         for label in labels:
+    #             if label not in self.sample_label_cdf:
+    #                 continue
+    #             label_cdf = self.sample_label_cdf[label]
+    #             new_label = self.label_mapping[label]
+    #             if lbl_p and label in lbl_p:
+    #                 label_cdf = np.float64(label_cdf)
+    #                 label_cdf *= lbl_p[label]
+    #             cdf = mapped_cdf.setdefault(new_label, [])
+    #             cdf.extend(label_cdf)
+    #
+    #         for key, cdf in mapped_cdf.items():
+    #             total = sum(cdf)
+    #             mapped_cdf[key] = [x / total for x in cdf]
+    #         self.sample_label_cdf = mapped_cdf
 
     def get_label_weight(self, label):
         """Returns the total weight for all segments of given label."""
@@ -516,77 +439,78 @@ def regroup(
         """
         regroups the dataset so multiple animals can be under a single label
         """
-        self.label_mapping = {}
-        counts = []
+        self.label_mapping = groups
         samples_by_bin = {}
-        samples = []
-        for mapped_label, labels in groups.items():
+        samples_by_label = {}
+
+        for lbl, remapped_lbl in groups.items():
             count = 0
-            for label in labels:
-                lbl_samples = self.samples_for(label)
-                count += len(lbl_samples)
-                samples.extend(lbl_samples)
-                self.label_mapping[label] = mapped_label
-                for sample in lbl_samples:
-                    samples_by_bin[sample.bin_id] = sample
-            counts.append(count)
+            samples = self.samples_by_label[lbl]
+            for s in samples:
+                s.remapped_label = remapped_lbl
+                samples_by_bin.setdefault(s.bin_id, []).append(s)
+            samples_by_label.setdefault(remapped_lbl, []).extend(samples)
 
         self.labels = list(groups.keys())
         self.labels.sort()
         self.samples_by_bin = samples_by_bin
-        self.set_samples(samples)
-        # self.samples_by_id == {}
-
-        if shuffle:
-            np.random.shuffle(self.samples)
-        self.rebuild_cdf()
+        self.samples_by_label = samples_by_label
+        #
+        # if shuffle:
+        #     np.random.shuffle(self.samples)
+        # self.rebuild_cdf()
 
     def has_data(self):
-        return len(self.samples) > 0
+        return len(self.samples_by_id) > 0
 
-    def recalculate_segments(self, segment_type=SegmentType.ALL_RANDOM):
-        self.samples_by_bin.clear()
-        self.samples_by_label.clear()
-        del self.samples[:]
-        del self.samples
-        self.samples = []
-        self.samples_by_label = {}
-        self.samples_by_bin = {}
-        logging.info("%s generating segments  type %s", self.name, segment_type)
-        start = time.time()
-        empty_tracks = []
-        filtered_stats = 0
-
-        for track in self.tracks:
-            segment_frame_spacing = int(
-                round(self.segment_spacing * track.frames_per_second)
-            )
-            segment_width = self.segment_length
-            track.calculate_segments(
-                segment_frame_spacing,
-                segment_width,
-                segment_type,
-                segment_min_mass=segment_min_avg_mass,
-            )
-            filtered_stats = filtered_stats + track.filtered_stats["segment_mass"]
-            if len(track.segments) == 0:
-                empty_tracks.append(track)
-                continue
-            for sample in track.segments:
-                self.add_clip_sample_mappings(sample)
-
-        self.rebuild_cdf()
-        logging.info(
-            "%s #segments %s filtered stats are %s took  %s",
-            self.name,
-            len(self.samples),
-            filtered_stats,
-            time.time() - start,
-        )
+    #
+    # def recalculate_segments(self, segment_type=SegmentType.ALL_RANDOM):
+    #     self.samples_by_bin.clear()
+    #     self.samples_by_label.clear()
+    #     del self.samples[:]
+    #     del self.samples
+    #     self.samples = []
+    #     self.samples_by_label = {}
+    #     self.samples_by_bin = {}
+    #     logging.info("%s generating segments  type %s", self.name, segment_type)
+    #     start = time.time()
+    #     empty_tracks = []
+    #     filtered_stats = 0
+    #
+    #     for track in self.tracks:
+    #         segment_frame_spacing = int(
+    #             round(self.segment_spacing * track.frames_per_second)
+    #         )
+    #         segment_width = self.segment_length
+    #         track.calculate_segments(
+    #             segment_frame_spacing,
+    #             segment_width,
+    #             segment_type,
+    #             segment_min_mass=segment_min_avg_mass,
+    #         )
+    #         filtered_stats = filtered_stats + track.filtered_stats["segment_mass"]
+    #         if len(track.segments) == 0:
+    #             empty_tracks.append(track)
+    #             continue
+    #         for sample in track.segments:
+    #             self.add_clip_sample_mappings(sample)
+    #
+    #     self.rebuild_cdf()
+    #     logging.info(
+    #         "%s #segments %s filtered stats are %s took  %s",
+    #         self.name,
+    #         len(self.samples),
+    #         filtered_stats,
+    #         time.time() - start,
+    #     )
 
     def remove_sample(self, sample):
-        self.samples.remove(sample)
-        self.samples_by_label[sample.label].remove(sample)
+        del self.samples_by_id[sample.id]
+        try:
+            # not nessesarily there if splitting by clip hack
+            self.samples_by_bin[sample.bin_id].remove(sample)
+        except:
+            pass
 
     def fetch_track(
         self,
@@ -610,3 +534,59 @@ def fetch_track(
             frame_numbers=frame_numbers,
         )
         return frames
+
+
+def filter_track(track_header, excluded_tags, filtered_stats={}):
+    # some clips are banned for various reasons4
+    if track_header.label is None:
+        filtered_stats.setdefault("notags", 0)
+        filtered_stats["notags"] += 1
+        return True
+    if track_header.label in excluded_tags:
+        filtered_stats.setdefault("tags", 0)
+
+        filtered_stats["tags"] += 1
+        filter_tags = filtered_stats.setdefault("tag_names", set())
+        filter_tags.add(track_header.label)
+        return True
+
+    if track_header.human_tags is not None:
+        found_tags = [tag for tag in track_header.human_tags if tag in excluded_tags]
+        if len(found_tags) > 0:
+            filter_tags = filtered_stats.setdefault("tag_names", set())
+            filter_tags |= set(found_tags)
+            filtered_stats.setdefault("tags", 0)
+            filtered_stats["tags"] += 1
+            return True
+    # always let the false-positives through as we need them even though they would normally
+    # be filtered out.
+
+    if track_header.regions_by_frame is None or len(track_header.regions_by_frame) == 0:
+        filtered_stats.setdefault("no_data", 0)
+        filtered_stats["no_data"] += 1
+        logging.info("No region data")
+        return True
+
+    # dont think we need this gp 28/08/2023
+    # if track_meta["human_tag"] == "false-positive":
+    # return False
+
+    # for some reason we get some records with a None confidence?
+    if track_header.confidence <= 0.6:
+        filtered_stats.setdefault("confidence", 0)
+        filtered_stats["confidence"] += 1
+        logging.info("Low confidence")
+        return True
+
+    return False
+
+
+def filter_clip(clip, filtered_stats={}):
+    # remove tracks of trapped animals
+    if (clip.events is not None and "trap" in clip.events.lower()) or (
+        clip.trap is not None and "trap" in clip.trap.lower()
+    ):
+        self.filtered_stats["trap"] += 1
+        logging.info("Filtered because in trap")
+        return True
+    return False
diff --git a/ml_tools/datasetstructures.py b/ml_tools/datasetstructures.py
index 990ad8ed..00b3ddb5 100644
--- a/ml_tools/datasetstructures.py
+++ b/ml_tools/datasetstructures.py
@@ -12,6 +12,7 @@
 from ml_tools import imageprocessing
 import datetime
 from enum import Enum
+import attr
 
 FRAMES_PER_SECOND = 9
 
@@ -32,13 +33,25 @@ class SegmentType(Enum):
     ALL_RANDOM_NOMIN = 7
 
 
-class Sample(ABC):
+class BaseSample(ABC):
     @property
     @abstractmethod
     def track_bounds(cls):
         """Get all regions for this sample"""
         ...
 
+    @property
+    @abstractmethod
+    def source_file(self):
+        """Source file."""
+        ...
+
+    @property
+    @abstractmethod
+    def track_median_mass(self):
+        """The function gets median mass of the track."""
+        ...
+
     @property
     @abstractmethod
     def frame_indices(self):
@@ -76,12 +89,44 @@ def mass(cls):
         ...
 
 
+class Sample(BaseSample):
+    def __init__(self, label):
+        self.original_label = label
+        self.remapped_label = label
+
+    @property
+    def label(self):
+        return self.remapped_label
+
+
 EDGE = 1
 
 res_x = 120
 res_y = 160
 
 
+@attr.s
+class ClipHeader:
+    clip_id = attr.ib()
+    location = attr.ib()
+    station_id = attr.ib()
+    rec_time = attr.ib()
+    source_file = attr.ib()
+    frames_per_second = attr.ib()
+    camera = attr.ib()
+    events = attr.ib()
+    trap = attr.ib()
+    tracks = attr.ib()
+    ffc_frames = attr.ib()
+    frame_temp_median = attr.ib(default=None)
+
+    def get_samples(self):
+        samples = []
+        for t in self.tracks:
+            samples.extend(t.samples)
+        return samples
+
+
 class TrackHeader:
     """Header for track."""
 
@@ -90,50 +135,34 @@ def __init__(
         clip_id,
         track_id,
         label,
-        start_time,
         num_frames,
-        duration,
-        camera,
-        location,
-        score,
         regions,
-        frame_temp_median,
-        frames_per_second,
-        predictions,
-        correct_prediction,
         start_frame,
-        res_x=CPTV_FILE_WIDTH,
-        res_y=CPTV_FILE_HEIGHT,
         ffc_frames=None,
         sample_frames_indices=None,
-        skipped_frames=None,
         station_id=None,
+        rec_time=None,
+        source_file=None,
+        camera=None,
+        confidence=None,
+        human_tags=None,
     ):
         self.station_id = station_id
-        self.res_x = np.uint8(res_x)
-        self.res_y = np.uint8(res_y)
+        self.clip_id = clip_id
+        self.source_file = source_file
+        self.human_tags = human_tags
+        self.confidence = confidence
         self.filtered_stats = {"segment_mass": 0}
         # reference to clip this segment came from
-        self.clip_id = clip_id
         # reference to track this segment came from
         self.track_id = track_id
-        # list of segments that belong to this track
-        self.segments = []
         # label for this track
         self.label = label
         # date and time of the start of the track
-        self.start_time = start_time
         self.start_frame = np.uint16(start_frame)
-        # duration in seconds
-        self.duration = duration
-        # camera this track came fromsegment
-        self.camera = camera
 
-        self.location = np.float16(location)
-        # score of track
-        self.score = score
         # thermal reference point for each frame.
-        self.frame_temp_median = np.uint16(frame_temp_median)
+        # self.frame_temp_median = np.uint16(frame_temp_median)
         # tracking frame movements for each frame, array of tuples (x-vel, y-vel)
         self.frame_velocity = None
         # original tracking bounds
@@ -141,66 +170,36 @@ def __init__(
         # what fraction of pixels are from out of bounds
         self.frame_crop = None
         self.num_frames = num_frames
-        self.frames_per_second = frames_per_second
         self.important_predicted = 0
-        mass_history = [region.mass for region in self.regions_by_frame.values()]
-        self.lower_mass = np.uint16(np.percentile(mass_history, q=25))
-        self.upper_mass = np.uint16(np.percentile(mass_history, q=75))
-        self.median_mass = np.uint16(np.median(mass_history))
-        self.mean_mass = np.uint16(np.mean(mass_history))
-        self.ffc_frames = np.uint16(ffc_frames)
-        self.skipped_frames = skipped_frames
-        self.sample_frames = []
-        sample_frames_indices is None
-        if sample_frames_indices is not None:
-            # shouldnt need to use this ever
-            for frame_num, frame_temp in zip(
-                sample_frames_indices, self.frame_temp_median
-            ):
-                region = self.regions_by_frame[frame_num]
-                assert region.frame_number == frame_num
-                f = FrameSample(
-                    self.clip_id,
-                    self.track_id,
-                    frame_num,
-                    self.label,
-                    frame_temp,
-                    None,
-                    region,
-                    weight=1,
-                    camera=self.camera,
-                    start_time=self.start_time
-                    + datetime.timedelta(seconds=frame_num / self.frames_per_second),
-                )
-                self.sample_frames.append(f)
-        else:
-            frame_numbers = list(self.regions_by_frame.keys())
-            frame_numbers = [
-                frame
-                for frame in frame_numbers
-                if (skipped_frames is None or frame not in skipped_frames)
-                and (ffc_frames is None or frame not in ffc_frames)
-            ]
-            frame_numbers.sort()
 
-            for frame_num, frame_temp in zip(frame_numbers, self.frame_temp_median):
-                region = self.regions_by_frame[frame_num]
-                if region.mass == 0 or region.blank:
-                    continue
-                f = FrameSample(
-                    self.clip_id,
-                    self.track_id,
-                    region.frame_number,
-                    self.label,
-                    frame_temp,
-                    None,
-                    region,
-                    weight=1,
-                    camera=self.camera,
-                    start_time=self.start_time
-                    + datetime.timedelta(seconds=frame_num / self.frames_per_second),
-                )
-                self.sample_frames.append(f)
+        mass_history = np.uint16(
+            [region.mass for region in self.regions_by_frame.values()]
+        )
+        mass_history = [
+            region.frame_number
+            for region in self.regions_by_frame.values()
+            if region.mass > 0
+            and (
+                ffc_frames is None
+                or skip_ffc is False
+                or region.frame_number not in ffc_frames
+            )
+            and not region.blank
+            and region.width > 0
+            and region.height > 0
+        ]
+        self.has_no_mass = np.sum(mass_history) == 0
+        if len(mass_history) == 0:
+            self.lower_mass = 0
+            self.upper_mass = 0
+            self.median_mass = 0
+            self.mean_mass = 0
+        else:
+            self.lower_mass = np.uint16(np.percentile(mass_history, q=25))
+            self.upper_mass = np.uint16(np.percentile(mass_history, q=75))
+            self.median_mass = np.uint16(np.median(mass_history))
+            self.mean_mass = np.uint16(np.mean(mass_history))
+        self.samples = []
 
     @property
     def bounds_history(self):
@@ -230,34 +229,64 @@ def toJSON(self, clip_meta):
 
         return json.dumps(meta_dict, indent=3, cls=tools.CustomJSONEncoder)
 
-    def add_sample(self, sample, use_segments):
-        if use_segments:
-            self.segments.append(sample)
-        else:
-            self.sample_frames.append(sample)
+    def add_sample(self, sample):
+        self.samples.append(sample)
 
-    def get_sample_frames(self):
-        return self.sample_frames
+    def calculate_sample_frames(
+        self, min_mass=None, max_mass=None, ffc_frames=None, skip_last=None
+    ):
+        frame_numbers = list(self.regions_by_frame.keys())
+        if skip_last is not None:
+            skip_x = int(len(frame_numbers) * skip_last)
+            frame_numbers = frame_numbers[:-skip_x]
+        frame_numbers = [
+            frame
+            for frame in frame_numbers
+            if (ffc_frames is None or frame not in ffc_frames)
+        ]
+        frame_numbers.sort()
 
-    def remove_sample_frame(self, f):
-        self.sample_frames.remove(f)
+        for frame_num in frame_numbers:
+            region = self.regions_by_frame[frame_num]
+            if region.mass == 0 or region.blank:
+                continue
+            if min_mass is not None and region.mass < min_mass:
+                continue
+            if max_mass is not None and region.mass > max_mass:
+                continue
+            f = FrameSample(
+                self.clip_id,
+                self.track_id,
+                region.frame_number,
+                self.label,
+                None,
+                region,
+                source_file=self.source_file,
+                weight=1,
+            )
+            self.samples.append(f)
 
-    def get_sample_frame(self, i=0, remove=False):
-        if len(self.sample_frames) == 0:
-            return None
-        f = self.sample_frames[i]
+    def remove_sample(self, f):
+        self.samples.remove(f)
 
-        if remove:
-            del self.sample_frames[i]
-        return f
+    #
+    # def get_sample_frame(self, i=0, remove=False):
+    #     if len(self.sample_frames) == 0:
+    #         return None
+    #     f = self.sample_frames[i]
+    #
+    #     if remove:
+    #         del self.sample_frames[i]
+    #     return f
 
     @property
     def frame_weight(self):
         return 1 / self.num_sample_frames
 
-    @property
-    def num_sample_frames(self):
-        return len(self.sample_frames)
+    #
+    # @property
+    # def num_samples(self):
+    #     return len(self.sample_frames)
 
     def calculate_frame_crop(self):
         # frames are always square, but bounding rect may not be, so to see how much we clipped I need to create a square
@@ -286,11 +315,14 @@ def calculate_segments(
         self,
         segment_frame_spacing,
         segment_width,
-        segment_type,
+        segment_type=SegmentType.ALL_RANDOM,
         segment_min_mass=None,
-        use_important=False,
         repeats=1,
         max_segments=None,
+        dont_filter=False,
+        skip_ffc=True,
+        ffc_frames=None,
+        location=None,
     ):
         min_frames = segment_width
         if self.label == "vehicle" or self.label == "human":
@@ -298,29 +330,30 @@ def calculate_segments(
 
         # in python3.7+ can just take the values and it guarantees order it was added to dict
         regions = self.bounds_history
-        self.segments, self.filtered_stats = get_segments(
+        self.samples, self.filtered_stats = get_segments(
             self.clip_id,
             self.track_id,
             self.start_frame,
-            segment_frame_spacing,
-            segment_width,
+            segment_frame_spacing=segment_frame_spacing,
+            segment_width=segment_width,
             label=self.label,
             regions=np.array(regions),
-            frame_temp_median=self.frame_temp_median,
             segment_min_mass=segment_min_mass,
-            sample_frames=self.sample_frames if use_important else None,
-            ffc_frames=self.ffc_frames,
+            ffc_frames=ffc_frames,
             lower_mass=self.lower_mass,
             repeats=repeats,
             min_frames=min_frames,
-            camera=self.camera,
-            skipped_frames=self.skipped_frames,
-            start_time=self.start_time,
             segment_type=segment_type,
             max_segments=max_segments,
-            location=self.location,
             station_id=self.station_id,
+            source_file=self.source_file,
+            dont_filter=dont_filter,
+            skip_ffc=skip_ffc,
         )
+        # GP could get this from the tracks when writing
+        # but might be best to keep samples independent for ease
+        for s in self.samples:
+            s._track_median_mass = self.median_mass
 
     @property
     def camera_id(self):
@@ -342,19 +375,20 @@ def unique_id(self):
         return "{}-{}".format(self.clip_id, self.track_id)
 
     @staticmethod
-    def from_meta(clip_id, clip_meta, track_meta, predictions=None):
+    def from_meta(source_file, clip_id, clip_meta, track_meta, predictions=None):
         """Creates a track header from given metadata."""
         correct_prediction = track_meta.get("correct_prediction", None)
-        start_time = dateutil.parser.parse(track_meta["start_time"])
-        end_time = dateutil.parser.parse(track_meta["end_time"])
-        duration = (end_time - start_time).total_seconds()
+        rec_time = dateutil.parser.parse(clip_meta["start_time"])
+        # end_time = dateutil.parser.parse(track_meta["end_time"])
         location = clip_meta.get("location")
-        num_frames = track_meta["frames"]
-        camera = clip_meta["device"]
+        camera = clip_meta.get("device_id")
+        # clip_meta["device"]
         station_id = clip_meta.get("station_id")
         frames_per_second = clip_meta.get("frames_per_second", FRAMES_PER_SECOND)
         # get the reference levels from clip_meta and load them into the track.
-        track_start_frame = track_meta["start_frame"]
+        track_start_frame = int(round(track_meta["start_frame"]))
+        track_end_frame = int(round(track_meta["end_frame"]))
+        num_frames = track_end_frame - track_start_frame + 1
         frame_temp_median = np.float32(
             clip_meta["frame_temp_median"][
                 track_start_frame : num_frames + track_start_frame
@@ -366,32 +400,39 @@ def from_meta(clip_id, clip_meta, track_meta, predictions=None):
         # sample_frames = track_meta.get("sample_frames")
         # if sample_frames is not None:
         #     sample_frames = sample_frames + track_start_frame
-        skipped_frames = track_meta.get("skipped_frames")
         regions = {}
         f_i = 0
-        for bounds, mass in zip(
-            track_meta["bounds_history"], track_meta["mass_history"]
-        ):
+        first_region = None
+        for bounds in track_meta["regions"]:
             r = Region.region_from_array(bounds)
-            if r.frame_number is None:
-                r.frame_number = np.uint16(f_i + track_start_frame)
-            r.mass = np.uint16(mass)
-
-            if r.mass == 0:
-                r.blank = True
+            if first_region is None:
+                first_region = r
             regions[r.frame_number] = r
             f_i += 1
+        if track_start_frame != first_region.frame_number:
+            logging.warn(
+                "Have wrong region or start frame for %s clip %s track %s adjusting regions first region %s becomes %s",
+                source_file,
+                clip_id,
+                track_meta["id"],
+                first_region.frame_number,
+                track_start_frame,
+            )
 
+            region_list = list(regions.values())
+            regions = {}
+            frame_number = track_start_frame
+            for r in region_list:
+                r = Region.region_from_array(bounds)
+                r.frame_number = frame_number
+                frame_number += 1
+                regions[r.frame_number] = r
         header = TrackHeader(
             clip_id=int(clip_id),
             track_id=int(track_meta["id"]),
-            label=track_meta["tag"],
-            start_time=start_time,
+            label=track_meta["human_tag"],
             num_frames=num_frames,
-            duration=duration,
-            camera=camera,
             location=location,
-            score=float(track_meta["score"]),
             regions=regions,
             frame_temp_median=frame_temp_median,
             frames_per_second=frames_per_second,
@@ -402,8 +443,10 @@ def from_meta(clip_id, clip_meta, track_meta, predictions=None):
             res_y=clip_meta.get("res_y", CPTV_FILE_HEIGHT),
             ffc_frames=ffc_frames,
             sample_frames_indices=sample_frames,
-            skipped_frames=skipped_frames,
             station_id=station_id,
+            rec_time=rec_time,
+            source_file=source_file,
+            camera=camera,
         )
         return header
 
@@ -557,27 +600,28 @@ def __init__(
         track_id,
         frame_num,
         label,
-        temp_median,
         velocity,
         region,
         weight,
-        camera,
-        start_time,
+        source_file,
         augment=False,
     ):
+        super().__init__(label)
+        self.clip_id = clip_id
         self.id = FrameSample._frame_id
         FrameSample._frame_id += 1
-        self.clip_id = clip_id
+
         self.track_id = track_id
         self.frame_number = frame_num
-        self.label = label
-        self.temp_median = temp_median
         self.velocity = velocity
         self.region = region
         self.weight = weight
-        self.camera = camera
-        self.start_time = start_time
         self.augment = augment
+        self._source_file = source_file
+
+    @property
+    def source_file(self):
+        return self._source_file
 
     def copy(self):
         f = FrameSample(
@@ -592,6 +636,9 @@ def copy(self):
             camera=self.camera,
             start_time=self.start_time,
             augment=self.augment,
+            source_file=self.source_file,
+            station_id=self.station_id,
+            rec_time=self.rec_time,
         )
         FrameSample._frame_id += 1
 
@@ -644,7 +691,6 @@ def __init__(
         mass,
         label,
         regions,
-        frame_temp_median,
         frame_indices=None,
         movement_data=None,
         best_mass=False,
@@ -653,7 +699,14 @@ def __init__(
         camera=None,
         location=None,
         station_id=None,
+        rec_time=None,
+        source_file=None,
+        filtered=False,
+        track_median_mass=None,
     ):
+        super().__init__(label)
+        self.filtered = filtered
+        self.rec_time = rec_time
         self.location = location
         self.station_id = station_id
         self.movement_data = movement_data
@@ -667,12 +720,7 @@ def __init__(
         self.frame_numbers = np.uint16(frame_indices)
         self.start_time = start_time
         self.regions = regions
-        self.frame_temp_median = np.uint16(frame_temp_median)
-        # for i, frame in enumerate(frame_indices):
-        #     self.track_bounds[frame] = regions[i]
-        #     self.frame_temp_median[frame] = frame_temp_median[i]
-        self.label = label
-        # first frame of this segment referenced by start of track
+
         self.start_frame = start_frame
         # length of segment in frames
         self.frames = np.uint16(frames)
@@ -681,6 +729,16 @@ def __init__(
 
         self._mass = np.uint16(mass)
         self.camera = camera
+        self._source_file = source_file
+        self._track_median_mass = track_median_mass
+
+    @property
+    def track_median_mass(self):
+        return self._track_median_mass
+
+    @property
+    def source_file(self):
+        return self._source_file
 
     @property
     def mass(self):
@@ -732,7 +790,7 @@ def end_frame(self):
     @property
     def bin_id(self):
         """Unique name of this segments track."""
-        return f"{self.clip_id}-{self.location}"
+        return f"{self.station_id}"
 
     def __str__(self):
         return "{0} label {1} offset:{2} weight:{3:.1f}".format(
@@ -826,26 +884,25 @@ def get_segments(
     clip_id,
     track_id,
     start_frame,
-    segment_frame_spacing,
-    segment_width,
     regions,
-    frame_temp_median,
+    segment_frame_spacing=9,
+    segment_width=25,
     label=None,
     segment_min_mass=None,
-    sample_frames=None,
     ffc_frames=[],
     lower_mass=0,
     repeats=1,
     min_frames=None,
-    skipped_frames=None,
     segment_frames=None,
-    ignore_mass=False,
-    camera=None,
-    start_time=None,
     segment_type=SegmentType.ALL_RANDOM,
     max_segments=None,
     location=None,
     station_id=None,
+    camera=None,
+    rec_time=None,
+    source_file=None,
+    dont_filter=False,
+    skip_ffc=True,
 ):
     if segment_type == SegmentType.ALL_RANDOM_NOMIN:
         segment_min_mass = None
@@ -854,43 +911,47 @@ def get_segments(
     segments = []
     mass_history = np.uint16([region.mass for region in regions])
     filtered_stats = {"segment_mass": 0, "too short": 0}
-    if sample_frames is not None:
-        frame_indices = [frame.frame_number for frame in sample_frames]
+
+    has_no_mass = np.sum(mass_history) == 0
+    frame_indices = [
+        region.frame_number
+        for region in regions
+        if (has_no_mass or region.mass > 0)
+        and (
+            ffc_frames is None
+            or skip_ffc is False
+            or region.frame_number not in ffc_frames
+        )
+        and not region.blank
+        and region.width > 0
+        and region.height > 0
+    ]
+    if len(frame_indices) == 0:
+        logging.warn("Nothing to load for %s - %s", clip_id, track_id)
+        return [], filtered_stats
+    if segment_min_mass is not None:
+        segment_min_mass = min(
+            segment_min_mass,
+            np.median(mass_history[frame_indices - start_frame]),
+        )
     else:
-        frame_indices = [
-            region.frame_number
-            for region in regions
-            if (ignore_mass or region.mass > 0)
-            and region.frame_number not in ffc_frames
-            and (skipped_frames is None or region.frame_number not in skipped_frames)
-            and not region.blank
-        ]
-        if segment_min_mass is not None:
-            if len(frame_indices) > 0:
-                segment_min_mass = min(
-                    segment_min_mass,
-                    np.median(mass_history[frame_indices - start_frame]),
-                )
-        else:
-            segment_min_mass = 1
-            # remove blank frames
-
-        if segment_type == SegmentType.TOP_RANDOM:
-            # take top 50 mass frames
-            frame_indices = sorted(
-                frame_indices,
-                key=lambda f_i: mass_history[f_i - start_frame],
-                reverse=True,
-            )
-            frame_indices = frame_indices[:50]
-            frame_indices.sort()
-    # 1 / 0
+        segment_min_mass = 1
+        # remove blank frames
+
+    if segment_type == SegmentType.TOP_RANDOM:
+        # take top 50 mass frames
+        frame_indices = sorted(
+            frame_indices,
+            key=lambda f_i: mass_history[f_i - start_frame],
+            reverse=True,
+        )
+        frame_indices = frame_indices[:50]
+        frame_indices.sort()
     if segment_type == SegmentType.TOP_SEQUENTIAL:
         return get_top_mass_segments(
             clip_id,
             track_id,
             label,
-            frame_temp_median,
             camera,
             segment_width,
             segment_frame_spacing,
@@ -898,14 +959,13 @@ def get_segments(
             ffc_frames,
             regions,
             start_frame,
-            start_time,
             lower_mass,
             segment_min_mass,
-            ignore_mass,
+            source_file=source_file,
         )
-    if len(frame_indices) < min_frames:
-        filtered_stats["too short"] += 1
-        return segments, filtered_stats
+    # if len(frame_indices) < min_frames:
+    # filtered_stats["too short"] += 1
+    # return segments, filtered_stats
     frame_indices = np.array(frame_indices)
     segment_count = max(1, len(frame_indices) // segment_frame_spacing)
     segment_count = int(segment_count)
@@ -926,10 +986,12 @@ def get_segments(
             # random_frames and not random_sections:
             np.random.shuffle(frame_indices)
         for i in range(segment_count):
-            if (len(frame_indices) < segment_width and len(segments) > 1) or len(
-                frame_indices
-            ) < (segment_width / 4.0):
-                break
+            # always get atleast one segmnet
+            if i > 0:
+                if (len(frame_indices) < segment_width and len(segments) > 1) or len(
+                    frame_indices
+                ) < (segment_width / 4.0):
+                    break
 
             if segment_type == SegmentType.ALL_SECTIONS:
                 # random frames from section 2.2 * segment_width
@@ -942,6 +1004,7 @@ def get_segments(
                 frames = section[indices]
                 frame_indices = frame_indices[segment_frame_spacing:]
             elif random_frames:
+                # frame indices already randomized so just need to grab some
                 frames = frame_indices[:segment_width]
                 frame_indices = frame_indices[segment_width:]
             else:
@@ -961,32 +1024,27 @@ def get_segments(
                 frames = np.concatenate([frames, extra_frames])
             frames.sort()
             relative_frames = frames - start_frame
-
             mass_slice = mass_history[relative_frames]
             segment_mass = np.sum(mass_slice)
-
             segment_avg_mass = segment_mass / len(mass_slice)
-            if (
-                not ignore_mass
-                and segment_min_mass
-                and segment_avg_mass < segment_min_mass
-            ):
-                filtered_stats["segment_mass"] += 1
-                continue
+            filtered = False
+            if segment_min_mass and segment_avg_mass < segment_min_mass:
+                if dont_filter:
+                    filtered = True
+                else:
+                    filtered_stats["segment_mass"] += 1
+                    continue
+
+            # temp_slice = frame_temp_median[relative_frames]
+            region_slice = regions[relative_frames]
+            movement_data = None
             if segment_avg_mass < 50:
                 segment_weight_factor = 0.75
             elif segment_avg_mass < 100:
                 segment_weight_factor = 1
             else:
                 segment_weight_factor = 1.2
-            # if we want to use movement_data
-            # movement_data = get_movement_data(
-            #     self.track_bounds[frames],
-            #     mass_history[frames],
-            # )
-            temp_slice = frame_temp_median[relative_frames]
-            region_slice = regions[relative_frames]
-            movement_data = None
+
             for z, f in enumerate(frames):
                 assert region_slice[z].frame_number == f
             segment = SegmentHeader(
@@ -998,13 +1056,14 @@ def get_segments(
                 mass=segment_mass,
                 label=label,
                 regions=region_slice,
-                frame_temp_median=temp_slice,
                 frame_indices=frames,
                 movement_data=movement_data,
                 camera=camera,
-                start_time=start_time,
                 location=location,
                 station_id=station_id,
+                rec_time=rec_time,
+                source_file=source_file,
+                filtered=filtered,
             )
             segments.append(segment)
     return segments, filtered_stats
@@ -1014,7 +1073,6 @@ def get_top_mass_segments(
     clip_id,
     track_id,
     label,
-    frame_temp_median,
     camera,
     segment_width,
     segment_frame_spacing,
@@ -1022,11 +1080,12 @@ def get_top_mass_segments(
     ffc_frames,
     regions,
     start_frame,
-    start_time,
     lower_mass,
     segment_min_mass,
-    ignore_mass=False,
+    source_file=None,
 ):
+    # Needs re testing if want to use
+    logging.warning("Top mass segments hasn't been tested may not work")
     filtered_stats = {"segment_mass": 0, "too short": 0}
 
     segments = []
@@ -1054,9 +1113,7 @@ def get_top_mass_segments(
         segment_info = sorted_mass[0]
         index = segment_info[0]
         avg_mass = segment_info[1] / segment_width
-        if not best_mass and (
-            ignore_mass or (avg_mass < lower_mass or avg_mass < segment_min_mass)
-        ):
+        if not best_mass and (avg_mass < lower_mass or avg_mass < segment_min_mass):
             break
         movement_data = get_movement_data(regions[index : index + segment_width])
         width = min(segment_width, len(regions))
@@ -1070,13 +1127,12 @@ def get_top_mass_segments(
             mass=segment_info[1],
             label=label,
             regions=regions[frames],
-            frame_temp_median=frame_temp_median[frames],
             frame_indices=frames + start_frame,
             movement_data=movement_data,
             best_mass=best_mass,
             top_mass=True,
             camera=camera,
-            start_time=start_time,
+            source_file=source_file,
         )
         best_mass = False
         segments.append(segment)
@@ -1106,7 +1162,6 @@ def __init__(
         labels,
         temp_median,
         region,
-        start_time,
         camera,
         filename,
     ):
@@ -1115,7 +1170,6 @@ def __init__(
         self.clip_id = clip_id
         self.filename = filename
         self.track_id = track_id
-        self.start_time = dateutil.parser.parse(start_time)
         # needed for original frame will change this GP
         self.frame_number = frame_num
         self.camera = camera
diff --git a/ml_tools/forestmodel.py b/ml_tools/forestmodel.py
index 69273560..6590417f 100644
--- a/ml_tools/forestmodel.py
+++ b/ml_tools/forestmodel.py
@@ -132,7 +132,7 @@ def classify_track(self, clip, track, last_x_frames=None, segment_frames=None):
         return track_prediction
 
     def shape(self):
-        return (1, 188)
+        return 1, (1, 188)
 
     def preprocess(self):
         return
@@ -161,15 +161,15 @@ def predict_track(self, clip, track, **args):
 
 
 def process_track(
-    clip, track, last_x_frames=None, buf_len=5, scale=None, normalize=False
+    clip, track, predict_from_last=None, buf_len=5, scale=None, normalize=False
 ):
     background = clip.background
     all_frames = None
     frame_temp_median = {}
-    if last_x_frames is None:
+    if predict_from_last is None:
         bounds = track.bounds_history
     else:
-        bounds = track.bounds_history[-last_x_frames:]
+        bounds = track.bounds_history[-predict_from_last:]
         all_frames = clip.frame_buffer.get_last_x(len(bounds))
 
     logging.debug(
diff --git a/ml_tools/frame.py b/ml_tools/frame.py
index fca03016..2781b89a 100644
--- a/ml_tools/frame.py
+++ b/ml_tools/frame.py
@@ -198,7 +198,7 @@ def contrast_adjust(self, adjust):
         if self.filtered is not None:
             self.filtered *= adjust
 
-    def crop_by_region(self, region, out=None):
+    def crop_by_region(self, region, only_thermal=False, out=None):
         # make a new frame cropped by region
         thermal = None
         filtered = None
@@ -206,12 +206,13 @@ def crop_by_region(self, region, out=None):
         flow = None
         if self.thermal is not None:
             thermal = region.subimage(self.thermal)
-        if self.filtered is not None:
-            filtered = region.subimage(self.filtered)
-        if self.mask is not None:
-            mask = region.subimage(self.mask)
-        if self.flow is not None:
-            flow = region.subimage(self.flow)
+        if not only_thermal:
+            if self.filtered is not None:
+                filtered = region.subimage(self.filtered)
+            if self.mask is not None:
+                mask = region.subimage(self.mask)
+            if self.flow is not None:
+                flow = region.subimage(self.flow)
         if out:
             out.thermal = thermal
             out.filtered = filtered
diff --git a/ml_tools/hyperparams.py b/ml_tools/hyperparams.py
index f9f13787..5e9827f7 100644
--- a/ml_tools/hyperparams.py
+++ b/ml_tools/hyperparams.py
@@ -23,11 +23,13 @@ def insert_defaults(self):
         self["use_segments"] = self.use_segments
         self["square_width"] = self.square_width
         self["frame_size"] = self.frame_size
+        self["segment_width"] = self.segment_width
 
         self["shuffle"] = self.shuffle
         self["channel"] = self.channel
         self["type"] = self.type
         self["segment_type"] = self.segment_type
+        self["multi_label"] = False
 
     @property
     def output_dim(self):
@@ -39,6 +41,10 @@ def output_dim(self):
             )
         return (self.frame_size, self.frame_size, 3)
 
+    @property
+    def multi_label(self):
+        return self.get("multi_label", False)
+
     @property
     def keep_aspect(self):
         return self.get("keep_aspect", False)
@@ -51,6 +57,10 @@ def use_background_filtered(self):
     def keep_edge(self):
         return self.get("keep_edge", False)
 
+    @property
+    def segment_width(self):
+        return self.get("segment_width", 25 if self.use_segments else 1)
+
     @property
     def segment_type(self):
         segment_type = self.get("segment_type", SegmentType.ALL_RANDOM.name)
diff --git a/ml_tools/imageprocessing.py b/ml_tools/imageprocessing.py
index 1b124fef..6b2e4fbf 100644
--- a/ml_tools/imageprocessing.py
+++ b/ml_tools/imageprocessing.py
@@ -76,7 +76,7 @@ def resize_cv(image, dim, interpolation=cv2.INTER_LINEAR, extra_h=0, extra_v=0):
     )
 
 
-def square_clip(data, frames_per_row, tile_dim):
+def square_clip(data, frames_per_row, tile_dim, normalize=True):
     # lay each frame out side by side in rows
     new_frame = np.zeros((frames_per_row * tile_dim[0], frames_per_row * tile_dim[1]))
     i = 0
@@ -88,9 +88,12 @@ def square_clip(data, frames_per_row, tile_dim):
             else:
                 frame = data[i]
 
-            frame, stats = normalize(frame)
-            if not stats[0]:
-                continue
+            # cv2.imshow("frame", np.uint8(frame))
+            # cv2.waitKey(0)
+            if normalize:
+                frame, stats = normalize(frame, new_max=255)
+                if not stats[0]:
+                    continue
             success = True
             new_frame[
                 x * tile_dim[0] : (x + 1) * tile_dim[0],
diff --git a/ml_tools/interpreter.py b/ml_tools/interpreter.py
index 8f194ea8..a2376634 100644
--- a/ml_tools/interpreter.py
+++ b/ml_tools/interpreter.py
@@ -25,7 +25,7 @@ def load_json(self, filename):
 
     @abstractmethod
     def shape(self):
-        """Prediction shape"""
+        """Num Inputs, Prediction shape"""
         ...
 
     @abstractmethod
@@ -33,126 +33,215 @@ def predict(self, frames):
         """predict"""
         ...
 
-    def predict_track(self, clip, track, **args):
-        frames, preprocessed, mass = self.preprocess(clip, track, args)
-        # print("preprocess is %s", preprocessed)
-        if preprocessed is None or len(preprocessed) == 0:
-            return None, None, None
-        pred = self.predict(np.array(preprocessed))
-        return frames, pred, mass
-
-    def preprocess(self, clip, track, args):
-        if self.TYPE == "RandomForest":
-            return
-        last_x_frames = args.get("last_x_frames", 1)
+    def preprocess(self, clip, track, **args):
         scale = args.get("scale", None)
+        num_predictions = args.get("num_predictions", None)
+        predict_from_last = args.get("predict_from_last", None)
+        segment_frames = args.get("segment_frames")
+        frames_per_classify = args.get("frames_per_classify", 25)
+        available_frames = (
+            min(len(track.bounds_history), clip.frame_buffer.max_frames)
+            if clip.frame_buffer.max_frames is not None
+            else len(track.bounds_history)
+        )
+        if predict_from_last is not None:
+            predict_from_last = min(predict_from_last, available_frames)
+        # this might be a little slower as it checks some massess etc
+        # but keeps it the same for all ways of classifying
+        if frames_per_classify > 1:
+            if predict_from_last is not None and segment_frames is None:
+                logging.debug(
+                    "Prediction from last available frames %s track is of length %s",
+                    available_frames,
+                    len(track.bounds_history),
+                )
+                regions = track.bounds_history[-available_frames:]
+                valid_regions = 0
+                if available_frames > predict_from_last:
+                    # want to get rid of any blank frames
+                    predict_from_last = 0
+                    for i, r in enumerate(
+                        reversed(track.bounds_history[-available_frames:])
+                    ):
+                        if r.blank:
+                            continue
+                        valid_regions += 1
+                        predict_from_last = i + 1
+                        if valid_regions >= predict_from_last:
+                            break
+                logging.debug(
+                    "After checking blanks have predict from last %s from last available frames %s track is of length %s",
+                    predict_from_last,
+                    available_frames,
+                    len(track.bounds_history),
+                )
+            frames, preprocessed, masses = self.preprocess_segments(
+                clip,
+                track,
+                num_predictions,
+                predict_from_last,
+                segment_frames=segment_frames,
+            )
+        else:
+            frames, preprocessed, masses = self.preprocess_frames(
+                clip, track, num_predictions, segment_frames=segment_frames
+            )
+        return frames, preprocessed, masses
 
-        frame_ago = 0
-        # get non blank frames
-        regions = []
-        frames = []
-        for r in reversed(track.bounds_history):
-            if not r.blank:
-                frame = clip.frame_buffer.get_frame_ago(frame_ago)
-                if frame is None:
-                    break
-                frame = frame
-                regions.append(r)
-                frames.append(frame)
-                assert frame.frame_number == r.frame_number
-                if len(regions) == last_x_frames:
-                    break
-            frame_ago += 1
-        if len(frames) == 0:
+    def predict_track(self, clip, track, **args):
+        frames, preprocessed, masses = self.preprocess(clip, track, **args)
+        if preprocessed is None or len(preprocessed) == 0:
             return None, None, None
-
-        if self.data_type == "IR":
-            logging.info("Preprocess IR scale %s last_x %s", scale, last_x_frames)
-            from ml_tools.preprocess import (
-                preprocess_ir,
+        pred = self.predict(preprocessed)
+        return frames, pred, masses
+
+    def preprocess_frames(
+        self,
+        clip,
+        track,
+        max_frames=None,
+        segment_frames=None,
+    ):
+        from ml_tools.preprocess import preprocess_single_frame
+
+        data = []
+        frames_used = []
+
+        for i, region in enumerate(reversed(track.bounds_history)):
+            if region.blank:
+                continue
+            if region.width == 0 or region.height == 0:
+                logging.warn(
+                    "No width or height for frame %s regoin %s",
+                    region.frame_number,
+                    region,
+                )
+                continue
+            frame = clip.frame_buffer.get_frame(region.frame_number)
+            if frame is None:
+                logging.error(
+                    "Clasifying clip %s track %s can't get frame %s",
+                    clip.get_id(),
+                    track.get_id(),
+                    region.frame_number,
+                )
+                raise Exception(
+                    "Clasifying clip {} track {} can't get frame {}".format(
+                        clip.get_id(), track.get_id(), region.frame_number
+                    )
+                )
+            logging.debug(
+                "classifying single frame with preprocess %s size %s crop? %s f shape %s region %s",
+                self.preprocess_fn.__module__,
+                self.params.frame_size,
+                crop,
+                frame.thermal.shape,
+                region,
+            )
+            preprocessed = preprocess_single_frame(
+                frame,
+                (
+                    self.params.frame_size,
+                    self.params.frame_size,
+                ),
+                region,
+                self.preprocess_fn,
+                save_info=f"{region.frame_number} - {region}",
             )
 
-            preprocessed = []
-            masses = []
-            for region, frame in zip(regions, frames):
-                if (
-                    frame is None
-                    or region.width == 0
-                    or region.height == 0
-                    or region.blank
-                ):
-                    continue
-                params = self.params
-
-                pre_f = preprocess_ir(
-                    frame.copy(),
-                    (
-                        params.frame_size,
-                        params.frame_size,
-                    ),
-                    region=region,
-                    preprocess_fn=self.preprocess_fn,
+            frames_used.append(region.frame_number)
+            data.append(preprocessed)
+            if max_frames is not None and len(data) >= max_frames:
+                break
+        return frames_used, np.array(data), mass
+
+    def preprocess_segments(
+        self,
+        clip,
+        track,
+        max_segments=None,
+        predict_from_last=None,
+        segment_frames=None,
+    ):
+        from ml_tools.preprocess import preprocess_frame, preprocess_movement
+
+        track_data = {}
+        segments = track.get_segments(
+            clip.ffc_frames,
+            self.params.square_width**2,
+            repeats=1,
+            segment_frames=segment_frames,
+            segment_type=self.params.segment_type,
+            from_last=predict_from_last,
+            max_segments=max_segments,
+        )
+        frame_indices = set()
+        for segment in segments:
+            frame_indices.update(set(segment.frame_indices))
+        frame_indices = list(frame_indices)
+        frame_indices.sort()
+        for frame_index in frame_indices:
+            region = track.bounds_history[frame_index - track.start_frame]
+
+            frame = clip.frame_buffer.get_frame(region.frame_number)
+            # filtered is calculated slightly different for tracking, set to null so preprocess can recalc it
+            if frame is None:
+                logging.error(
+                    "Clasifying clip %s track %s can't get frame %s",
+                    clip.get_id(),
+                    track.get_id(),
+                    region.frame_number,
                 )
-                if pre_f is None:
-                    continue
-                preprocessed.append(pre_f)
-                masses.append(1)
-            return [frame.frame_number for f in frames], preprocessed, masses
-        elif self.data_type == "thermal":
-            from ml_tools.preprocess import (
-                preprocess_movement,
+                raise Exception(
+                    "Clasifying clip {} track {} can't get frame {}".format(
+                        clip.get_id(), track.get_id(), region.frame_number
+                    )
+                )
+            cropped_frame = preprocess_frame(
+                frame,
+                (self.params.frame_size, self.params.frame_size),
+                region,
+                clip.background,
+                clip.crop_rectangle,
             )
+            track_data[frame.frame_number] = cropped_frame
+        features = None
+        if self.params.mvm:
+            from ml_tools.forestmodel import process_track as forest_process_track
 
-            frames_per_classify = args.get("frames_per_classify", 25)
-            logging.info(
-                "Preprocess thermal scale %s frames_per_classify %s last_x %s",
-                scale,
-                frames_per_classify,
-                last_x_frames,
+            features = forest_process_track(
+                clip, track, normalize=True, predict_from_last=predict_from_last
             )
 
-            indices = np.random.choice(
-                len(regions),
-                min(frames_per_classify, len(regions)),
-                replace=False,
-            )
-            indices.sort()
-            frames = np.array(frames)[indices]
-            regions = np.array(regions)[indices]
-
-            refs = []
-            segment_data = []
-            mass = 0
-            params = self.params
-
-            for frame, region in zip(frames, regions):
-                if region.blank:
-                    continue
-                refs.append(np.median(frame.thermal))
-                thermal_reference = np.median(frame.thermal)
-                f = frame.crop_by_region(region)
-                mass += region.mass
-                f.resize_with_aspect(
-                    (params.frame_size, params.frame_size),
-                    clip.crop_rectangle,
-                    True,
-                )
-                segment_data.append(f)
-
-            preprocessed = preprocess_movement(
-                segment_data,
-                params.square_width,
-                params.frame_size,
-                red_type=params.red_type,
-                green_type=params.green_type,
-                blue_type=params.blue_type,
-                preprocess_fn=self.preprocess_fn,
-                reference_level=refs,
-                keep_edge=params.keep_edge,
+        preprocessed = []
+        masses = []
+        for segment in segments:
+            segment_frames = []
+            for frame_i in segment.frame_indices:
+                f = track_data[frame_i]
+                # probably no need to copy
+                segment_frames.append(f.copy())
+            frames = preprocess_movement(
+                segment_frames,
+                self.params.square_width,
+                self.params.frame_size,
+                self.params.red_type,
+                self.params.green_type,
+                self.params.blue_type,
+                self.preprocess_fn,
             )
-            if preprocessed is None:
-                return None, None, mass
-            return [f.frame_number for f in frames], [preprocessed], [mass]
+            if frames is None:
+                logging.warn("No frames to predict on")
+                continue
+            preprocessed.append(frames)
+            masses.append(segment.mass)
+        preprocessed = np.array(preprocessed)
+        if self.params.mvm:
+            features = features[np.newaxis, :]
+            features = np.repeat(features, len(preprocessed), axis=0)
+            preprocessed = [preprocessed, features]
+
+        return [s.frame_indices for s in segments], preprocessed, masses
 
 
 class NeuralInterpreter(Interpreter):
@@ -193,7 +282,7 @@ def predict(self, input_x):
         return res
 
     def shape(self):
-        return self.input_shape
+        return 1, self.input_shape
 
 
 class LiteInterpreter(Interpreter):
@@ -225,7 +314,7 @@ def predict(self, input_x):
         return pred
 
     def shape(self):
-        return self.input["shape"]
+        return 1, self.input["shape"]
 
 
 def inc3_preprocess(x):
diff --git a/ml_tools/irdataset.py b/ml_tools/irdataset.py
index db225c1a..f8d5fd7c 100644
--- a/ml_tools/irdataset.py
+++ b/ml_tools/irdataset.py
@@ -33,7 +33,9 @@ def load_dataset(filenames, remap_lookup, num_labels, args):
     ignore_order.experimental_deterministic = (
         deterministic  # disable order, increase speed
     )
-    dataset = tf.data.TFRecordDataset(filenames)
+    dataset = tf.data.TFRecordDataset(
+        filenames, compression_type="GZIP", num_parallel_reads=4
+    )
     dataset = dataset.with_options(
         ignore_order
     )  # uses data as soon as it streams in, rather than in its original order
@@ -48,7 +50,7 @@ def load_dataset(filenames, remap_lookup, num_labels, args):
         partial(
             read_irrecord,
             remap_lookup=remap_lookup,
-            num_labels=num_labels,
+            num_labels=len(labels),
             image_size=image_size,
             labeled=labeled,
             augment=augment,
diff --git a/ml_tools/irwriter.py b/ml_tools/irwriter.py
index f8358000..8100b28e 100644
--- a/ml_tools/irwriter.py
+++ b/ml_tools/irwriter.py
@@ -12,19 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-r"""Convert raw COCO 2017 dataset to TFRecord.
-
-Example usage:
-    python create_coco_tf_record.py --logtostderr \
-      --image_dir="${TRAIN_IMAGE_DIR}" \
-      --image_info_file="${TRAIN_IMAGE_INFO_FILE}" \
-      --object_annotations_file="${TRAIN_ANNOTATIONS_FILE}" \
-      --caption_annotations_file="${CAPTION_ANNOTATIONS_FILE}" \
-      --output_file_prefix="${OUTPUT_DIR/FILE_PREFIX}" \
-      --num_shards=100
-"""
 from PIL import Image
 from pathlib import Path
+from multiprocessing import Process, Queue
 
 import collections
 import hashlib
@@ -32,7 +22,7 @@
 import json
 import multiprocessing
 import os
-
+import time
 from absl import app
 from absl import flags
 from absl import logging
@@ -48,10 +38,8 @@
 import random
 import math
 
-crop_rectangle = tools.Rectangle(0, 0, 640 - 1, 480 - 1)
-
 
-def create_tf_example(frame, image_dir, sample, labels, filename):
+def create_tf_example(sample, thermal, filtered, labels):
     """Converts image and annotations to a tf.Example proto.
 
     Args:
@@ -82,9 +70,9 @@ def create_tf_example(frame, image_dir, sample, labels, filename):
     Raises:
       ValueError: if the image pointed to by data['filename'] is not a valid JPEG
     """
-    image_height, image_width = frame.thermal.shape
+    image_height, image_width = thermal.shape
 
-    image = Image.fromarray(np.uint8(frame.thermal))
+    image = Image.fromarray(np.uint8(thermal))
 
     image_id = sample.unique_id
 
@@ -94,26 +82,20 @@ def create_tf_example(frame, image_dir, sample, labels, filename):
     encoded_thermal = encoded_jpg_io.getvalue()
     thermal_key = hashlib.sha256(encoded_thermal).hexdigest()
 
-    image = Image.fromarray(np.uint8(frame.filtered))
+    image = Image.fromarray(np.uint8(filtered))
 
     encoded_jpg_io = io.BytesIO()
     image.save(encoded_jpg_io, format="PNG", quality=100, subsampling=0)
     encoded_filtered = encoded_jpg_io.getvalue()
     filtered_key = hashlib.sha256(encoded_filtered).hexdigest()
-    # if frame.mask:
-    # image = Image.fromarray(frame.mask)
-    # image = ImageOps.grayscale(image)
-    # image_id = sample.id
-    #
-    # encoded_jpg_io = io.BytesIO()
-    # image.save(encoded_jpg_io, format="JPEG")
-    # encoded_mask = encoded_jpg_io.getvalue()
-    # mask_key = hashlib.sha256(encoded_mask).hexdigest()
+
     feature_dict = {
         "image/augmented": tfrecord_util.int64_feature(sample.augment),
         "image/height": tfrecord_util.int64_feature(image_height),
         "image/width": tfrecord_util.int64_feature(image_width),
-        "image/filename": tfrecord_util.bytes_feature(filename.encode("utf8")),
+        "image/filename": tfrecord_util.bytes_feature(
+            str(sample.source_file).encode("utf8")
+        ),
         "image/source_id": tfrecord_util.bytes_feature(str(image_id).encode("utf8")),
         "image/thermalkey/sha256": tfrecord_util.bytes_feature(
             thermal_key.encode("utf8")
@@ -125,205 +107,70 @@ def create_tf_example(frame, image_dir, sample, labels, filename):
         "image/clip_id": tfrecord_util.int64_feature(sample.clip_id),
         "image/track_id": tfrecord_util.int64_feature(sample.track_id),
         "image/filteredencoded": tfrecord_util.bytes_feature(encoded_filtered),
-        # "image/maskkey/sha256": tfrecord_util.bytes_feature(mask_key.encode("utf8")),
-        # "image/maskencoded": tfrecord_util.bytes_feature(encoded_mask),
         "image/format": tfrecord_util.bytes_feature("jpeg".encode("utf8")),
         "image/class/text": tfrecord_util.bytes_feature(sample.label.encode("utf8")),
         "image/class/label": tfrecord_util.int64_feature(labels.index(sample.label)),
-        # "image/object/bbox/xmin": tfrecord_util.float_list_feature([0]),
-        # "image/object/bbox/xmax": tfrecord_util.float_list_feature([1]),
-        # "image/object/bbox/ymin": tfrecord_util.float_list_feature([0]),
-        # "image/object/bbox/ymax": tfrecord_util.float_list_feature([1]),
     }
-    #
-    # num_annotations_skipped = 0
-    # if len(sample.regions) > 0:
-    #     xmin = []
-    #     xmax = []
-    #     ymin = []
-    #     ymax = []
-    #     is_crowd = []
-    #     category_names = []
-    #     category_ids = []
-    #     area = []
-    #     encoded_mask_png = []
-    #     for r, label in zip(sample.regions, sample.labels):
-    #         (x, y, width, height) = r.x, r.y, r.width, r.height
-    #         if width <= 0 or height <= 0:
-    #             num_annotations_skipped += 1
-    #             continue
-    #         if x + width > image_width or y + height > image_height:
-    #             num_annotations_skipped += 1
-    #             continue
-    #         xmin.append(float(x) / image_width)
-    #         xmax.append(float(x + width) / image_width)
-    #         ymin.append(float(y) / image_height)
-    #         ymax.append(float(y + height) / image_height)
-    #         category_id = int(labels.index(label) + 1)
-    #         category_ids.append(category_id)
-    #         category_names.append(label.encode("utf8"))
-    #         area.append(r.area)
-    #     feature_dict.update(
-    #         {
-    #             "image/object/bbox/xmin": tfrecord_util.float_list_feature(xmin),
-    #             "image/object/bbox/xmax": tfrecord_util.float_list_feature(xmax),
-    #             "image/object/bbox/ymin": tfrecord_util.float_list_feature(ymin),
-    #             "image/object/bbox/ymax": tfrecord_util.float_list_feature(ymax),
-    #             "image/object/class/text": tfrecord_util.bytes_list_feature(
-    #                 category_names
-    #             ),
-    #             "image/object/class/label": tfrecord_util.int64_list_feature(
-    #                 category_ids
-    #             ),
-    #             "image/object/area": tfrecord_util.float_list_feature(area),
-    #         }
-    #     )
     example = tf.train.Example(features=tf.train.Features(feature=feature_dict))
-    return example, 0
-
-
-def create_tf_records(
-    dataset, output_path, labels, back_thresh, num_shards=1, cropped=True, augment=False
-):
-    output_path = Path(output_path)
-    if output_path.is_dir():
-        logging.info("Clearing dir %s", output_path)
-        for child in output_path.glob("*"):
-            if child.is_file():
-                child.unlink()
-    output_path.mkdir(parents=True, exist_ok=True)
-    samples = dataset.samples
-    np.random.shuffle(samples)
-
-    dataset.load_db()
-    db = dataset.db
-    total_num_annotations_skipped = 0
-    num_labels = len(dataset.labels)
-    logging.info("writing to output path: %s for %s samples", output_path, len(samples))
-    lbl_counts = [0] * num_labels
-
-    writers = []
-    for i in range(num_shards):
-        writers.append(
-            tf.io.TFRecordWriter(
-                str(output_path / (f"%05d-of-%05d.tfrecord" % (i, num_shards)))
+    return example
+
+
+def get_data(samples, back_thresh):
+    vidcap = cv2.VideoCapture(str(samples[0].source_file))
+    frames = {}
+    backgorund = None
+    frame_num = 0
+    print("Loading ", str(samples[0].source_file))
+    while True:
+        success, image = vidcap.read()
+        is_background_frame = False
+        if not success:
+            break
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        if backgorund is None:
+            is_background_frame = np.all(image[:, :, 0] == image[:, :, 1]) and np.all(
+                image[:, :, 1] == image[:, :, 2]
             )
-        )
-    load_first = 200
+            background = np.uint8(gray)
+        if not is_background_frame:
+            frames[frame_num] = gray
+            frame_num += 1
+            # append(gray)
+    data = []
+    for sample in samples:
+        frame = frames[sample.region.frame_number]
+        gray_sub = sample.region.subimage(frame)
+        back_sub = sample.region.subimage(background)
+        filtered = get_diff_back_filtered(back_sub, gray_sub, back_thresh)
+        gray_sub, stats = normalize(gray_sub, new_max=255)
+        if not stats[0]:
+            continue
+        filtered, stats = normalize(filtered, new_max=255)
+        cv2.imwrite(f"{sample.id}-{sample.track_id}-{sample.label}.png", gray_sub)
+        if not stats[0]:
+            continue
+        data.append((sample, gray_sub, filtered))
+
+    return data
+
+
+def save_data(samples, writer, labels, extra_args):
+    sample_data = get_data(samples, extra_args["back_thresh"])
+    if sample_data is None:
+        return 0
+    saved = 0
     try:
-        count = 0
-        while len(samples) > 0:
-            local_set = samples[:load_first]
-            samples = samples[load_first:]
-            loaded = []
-
-            for sample in local_set:
-                try:
-                    f = get_data(db, sample, cropped, back_thresh)
-                    if f is None:
-                        continue
-                    loaded.append((f, sample))
-                except Exception as e:
-                    logging.error("Got exception", exc_info=True)
-
-            loaded = np.array(loaded)
-            np.random.shuffle(loaded)
-            for data, sample in loaded:
-                try:
-                    tf_example, num_annotations_skipped = create_tf_example(
-                        data,
-                        output_path,
-                        sample,
-                        dataset.labels,
-                        sample.unique_track_id,
-                    )
-                    l_i = labels.index(sample.label)
-
-                    total_num_annotations_skipped += num_annotations_skipped
-                    writer = writers[count % num_shards]
-                    writer.write(tf_example.SerializeToString())
-                    count += 1
-                    lbl_counts[l_i] += 1
-
-                    if count % 100 == 0:
-                        logging.info("saved %s", count)
-                except Exception as e:
-                    logging.error("Error saving ", exc_info=True)
-            # break
+        for sample, thermal, filtered in sample_data:
+            tf_example = create_tf_example(
+                sample,
+                thermal,
+                filtered,
+                labels,
+            )
+            writer.write(tf_example.SerializeToString())
+            saved += 1
     except:
-        logging.error("Error saving track info", exc_info=True)
-    for writer in writers:
-        writer.close()
-
-    logging.info(
-        "Finished writing, skipped %d annotations.", total_num_annotations_skipped
-    )
-
-
-def get_data(db, sample, cropped, back_thresh):
-    background = db.get_clip_background(sample.clip_id)
-    try:
-        f = db.get_clip(
-            sample.clip_id,
-            [sample.frame_number],
-        )[0]
-    except Exception as e:
         logging.error(
-            "Error gettin clip %s %s",
-            sample.clip_id,
-            sample,
-            exc_info=True,
+            "Could not save data for %s", samples[0].source_file, exc_info=True
         )
-        return None
-    prev = f.thermal.copy()
-    prev = sample.region.subimage(prev)
-    if cropped:
-        if sample.augment:
-            frame_height, frame_width = f.thermal.shape
-            percent_offset = 0.02
-            height_range = int(math.ceil(percent_offset * sample.region.height) / 2.0)
-            width_range = int(math.ceil(percent_offset * sample.region.width) / 2.0)
-            top_offset = random.randint(-height_range, height_range)
-            height_offset = random.randint(-height_range, height_range)
-            left_offset = random.randint(-width_range, width_range)
-            width_offset = random.randint(-width_range, width_range)
-            sample.region.left += left_offset
-            sample.region.top += top_offset
-            sample.region.width += width_offset
-            sample.region.height += height_offset
-            sample.region.crop(crop_rectangle)
-        region = sample.region.copy()
-        f.crop_by_region(region, out=f)
-        background = region.subimage(background)
-    f.mask = f.filtered
-    f.filtered = get_diff_back_filtered(background, f.thermal, back_thresh)
-
-    assert f.thermal.shape == f.filtered.shape
-    f.normalize()
-
-    if sample.augment:
-        degrees = random.randint(-45, 45)
-
-        # degrees = int(chance * 40) - 20
-        brightness_adjust = random.randint(-20, 20)
-
-        contrast_adjust = random.random() * 0.4 + 1 - 0.4 / 2.0
-        # contrast_adjust *=
-        f.float_arrays()
-        f.brightness_adjust(brightness_adjust)
-        f.contrast_adjust(contrast_adjust)
-        f.filtered = np.fliplr(f.filtered)
-        f.thermal = np.fliplr(f.thermal)
-        f.filtered = rotate(f.filtered, degrees)
-        f.thermal = rotate(f.thermal, degrees)
-        np.clip(f.filtered, a_min=0, a_max=255, out=f.filtered)
-        np.clip(f.thermal, a_min=0, a_max=255, out=f.thermal)
-        #
-        # cv2.imshow("prev", np.uint8(prev))
-        #
-        # cv2.imshow("aug", np.uint8(f.thermal))
-        # cv2.moveWindow(f"prev", 0, 0)
-        # cv2.moveWindow(f"aug", 0, 500)
-        # cv2.waitKey(3000)
-    # f.no
-    return f
+    return saved
diff --git a/ml_tools/kerasmodel.py b/ml_tools/kerasmodel.py
index c94e5f65..5956a381 100644
--- a/ml_tools/kerasmodel.py
+++ b/ml_tools/kerasmodel.py
@@ -19,17 +19,16 @@
 from ml_tools import tools
 from ml_tools.datasetstructures import SegmentType
 
-from ml_tools.preprocess import preprocess_movement, preprocess_frame, preprocess_ir
+from ml_tools.preprocess import (
+    preprocess_movement,
+    preprocess_frame,
+)
 from ml_tools.interpreter import Interpreter
 from classify.trackprediction import TrackPrediction
 from ml_tools.hyperparams import HyperParams
 from ml_tools import thermaldataset
+from ml_tools.resnet.wr_resnet import WRResNet
 
-# import (
-#    get_resampled_by_label as get_thermal_dataset_by_label,
-#    get_dataset as get_thermal_dataset,
-#    get_weighting,
-# )
 from ml_tools import irdataset
 from ml_tools.tfdataset import get_weighting, get_distribution, get_dataset as get_tf
 import tensorflow_decision_forests as tfdf
@@ -85,30 +84,34 @@ def load_training_meta(self, base_dir):
     def shape(self):
         if self.model is None:
             return None
-        return self.model.get_config()["layers"][0]["config"]["batch_input_shape"]
+        inputs = self.model.inputs
+        shape = []
+        for input in inputs:
+            shape.append(tuple(input.shape.as_list()))
+        return len(shape), shape
 
-    def get_base_model(self, input_shape, weights="imagenet"):
+    def get_base_model(self, input, weights="imagenet"):
         pretrained_model = self.params.model_name
         if pretrained_model == "resnet":
             return (
                 tf.keras.applications.ResNet50(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=inputs,
                 ),
                 tf.keras.applications.resnet.preprocess_input,
             )
         elif pretrained_model == "resnetv2":
             return (
                 tf.keras.applications.ResNet50V2(
-                    weights=weights, include_top=False, input_shape=input_shape
+                    weights=weights, include_top=False, input_tensor=input
                 ),
                 tf.keras.applications.resnet_v2.preprocess_input,
             )
         elif pretrained_model == "resnet152":
             return (
                 tf.keras.applications.ResNet152(
-                    weights=weights, include_top=False, input_shape=input_shape
+                    weights=weights, include_top=False, input_tensor=input
                 ),
                 tf.keras.applications.resnet.preprocess_input,
             )
@@ -117,7 +120,7 @@ def get_base_model(self, input_shape, weights="imagenet"):
                 tf.keras.applications.VGG16(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=input,
                 ),
                 tf.keras.applications.vgg16.preprocess_input,
             )
@@ -126,7 +129,7 @@ def get_base_model(self, input_shape, weights="imagenet"):
                 tf.keras.applications.VGG19(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=input,
                 ),
                 tf.keras.applications.vgg19.preprocess_input,
             )
@@ -135,7 +138,7 @@ def get_base_model(self, input_shape, weights="imagenet"):
                 tf.keras.applications.MobileNetV2(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=input,
                 ),
                 tf.keras.applications.mobilenet_v2.preprocess_input,
             )
@@ -144,7 +147,7 @@ def get_base_model(self, input_shape, weights="imagenet"):
                 tf.keras.applications.DenseNet121(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=input,
                 ),
                 tf.keras.applications.densenet.preprocess_input,
             )
@@ -153,7 +156,7 @@ def get_base_model(self, input_shape, weights="imagenet"):
                 tf.keras.applications.InceptionResNetV2(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=input,
                 ),
                 tf.keras.applications.inception_resnet_v2.preprocess_input,
             )
@@ -162,7 +165,7 @@ def get_base_model(self, input_shape, weights="imagenet"):
                 tf.keras.applications.InceptionV3(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=input,
                 ),
                 tf.keras.applications.inception_v3.preprocess_input,
             )
@@ -171,7 +174,7 @@ def get_base_model(self, input_shape, weights="imagenet"):
                 tf.keras.applications.EfficientNetB5(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=input,
                 ),
                 None,
             )
@@ -180,7 +183,7 @@ def get_base_model(self, input_shape, weights="imagenet"):
                 tf.keras.applications.EfficientNetB0(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=input,
                 ),
                 None,
             )
@@ -189,17 +192,27 @@ def get_base_model(self, input_shape, weights="imagenet"):
                 tf.keras.applications.EfficientNetB1(
                     weights=weights,
                     include_top=False,
-                    input_shape=input_shape,
+                    input_tensor=input,
                 ),
                 None,
             )
-        raise Exception("Could not find model" + pretrained_model)
+        elif pretrained_model == "nasnet":
+            return (
+                tf.keras.applications.nasnet.NASNetLarge(
+                    weights=weights, include_top=False, input_tensor=input
+                ),
+                tf.keras.applications.nasnet.preprocess_input,
+            )
+        elif pretrained_model == "wr-resnet":
+            return WRResNet(input), None
+        raise Exception("Could not find model " + pretrained_model)
 
     def get_preprocess_fn(self):
         pretrained_model = self.params.model_name
         if pretrained_model == "resnet":
             return tf.keras.applications.resnet.preprocess_input
-
+        elif pretrained_model == "nasnet":
+            return tf.keras.applications.nasnet.preprocess_input
         elif pretrained_model == "resnetv2":
             return tf.keras.applications.resnet_v2.preprocess_input
 
@@ -256,9 +269,12 @@ def build_model(
         width = self.params.output_dim[0]
         inputs = tf.keras.Input(shape=(width, width, 3), name="input")
         weights = None if self.params.base_training else "imagenet"
-        base_model, preprocess = self.get_base_model((width, width, 3), weights=weights)
+        base_model, preprocess = self.get_base_model(inputs, weights=weights)
         self.preprocess_fn = preprocess
-        x = base_model(inputs, training=self.params.base_training)
+
+        # inputs = base_model.input
+        x = base_model.output
+        # x = base_model(inputs, training=self.params.base_training)
 
         if self.params.lstm:
             x = tf.keras.layers.GlobalAveragePooling2D()(x)
@@ -303,8 +319,13 @@ def build_model(
                     x = tf.keras.layers.Dense(i, activation="relu")(x)
             if dropout:
                 x = tf.keras.layers.Dropout(dropout)(x)
+
+            activation = "softmax"
+            if self.params.multi_label:
+                activation = "sigmoid"
+            logging.info("Using %s activation", activation)
             preds = tf.keras.layers.Dense(
-                len(self.labels), activation="softmax", name="prediction"
+                len(self.labels), activation=activation, name="prediction"
             )(x)
             self.model = tf.keras.models.Model(inputs, outputs=preds)
         if retrain_from is None:
@@ -319,11 +340,17 @@ def build_model(
                     layer.trainable = i >= retrain_from
         else:
             base_model.trainable = self.params.base_training
+
+        if self.params.multi_label:
+            acc = tf.metrics.binary_accuracy
+        else:
+            acc = tf.metrics.categorical_accuracy
+
         self.model.compile(
             optimizer=optimizer(self.params),
             loss=loss(self.params),
             metrics=[
-                "accuracy",
+                acc,
                 tf.keras.metrics.AUC(),
                 tf.keras.metrics.Recall(),
                 tf.keras.metrics.Precision(),
@@ -362,8 +389,9 @@ def load_model(self, model_path, training=False, weights=None):
         self.model.trainable = training
         self.load_meta(dir_name)
         if weights is not None:
-            self.model.load_weights(weights).expect_partial()
-        logging.info("Loaded weight %s", weights)
+            self.model.load_weights(weights)
+            logging.info("Loaded weight %s", weights)
+        print(self.model.summary())
 
     def load_meta(self, dir_name):
         meta = json.load((dir_name / "saved_model.json").open("r"))
@@ -374,7 +402,6 @@ def load_meta(self, dir_name):
         self.label_probabilities = meta.get("label_probabilities")
         self.preprocess_fn = self.get_preprocess_fn()
         self.type = meta.get("type", "thermal")
-
         logging.debug(
             "using types r %s g %s b %s type %s",
             self.params.red_type,
@@ -460,13 +487,17 @@ def train_model_tfrecords(
         logging.info(
             "%s Training model for %s epochs with weights %s", run_name, epochs, weights
         )
-        self.excluded_labels, self.remapped_labels = get_excluded(self.type)
+        self.excluded_labels, self.remapped_labels = get_excluded(
+            self.type, self.params.multi_label
+        )
         train_files = os.path.join(self.data_dir, "train")
         validate_files = os.path.join(self.data_dir, "validation")
         logging.info(
             "Excluding %s remapping %s", self.excluded_labels, self.remapped_labels
         )
 
+        if self.params.multi_label:
+            self.labels.append("land-bird")
         orig_labels = self.labels.copy()
         for l in self.excluded_labels:
             if l in self.labels:
@@ -486,7 +517,7 @@ def train_model_tfrecords(
                 run_name=run_name,
             )
         self.model.summary()
-        self.train, remapped, new_labels = get_dataset(
+        self.train, remapped, new_labels, epoch_size = get_dataset(
             train_files,
             self.type,
             orig_labels,
@@ -500,9 +531,10 @@ def train_model_tfrecords(
             excluded_labels=self.excluded_labels,
             remapped_labels=self.remapped_labels,
             # dist=self.dataset_counts["train"],
+            multi_label=self.params.multi_label,
         )
         self.remapped = remapped
-        self.validate, remapped, _ = get_dataset(
+        self.validate, remapped, _, _ = get_dataset(
             validate_files,
             self.type,
             orig_labels,
@@ -514,6 +546,7 @@ def train_model_tfrecords(
             include_features=self.params.mvm,
             excluded_labels=self.excluded_labels,
             remapped_labels=self.remapped_labels,
+            multi_label=self.params.multi_label
             # dist=self.dataset_counts["validation"],
         )
 
@@ -549,7 +582,7 @@ def train_model_tfrecords(
         test_files = os.path.join(self.data_dir, "test")
 
         if len(test_files) > 0:
-            self.test, _, _ = get_dataset(
+            self.test, _, _, _ = get_dataset(
                 test_files,
                 self.type,
                 orig_labels,
@@ -562,6 +595,7 @@ def train_model_tfrecords(
                 resample=False,
                 excluded_labels=self.excluded_labels,
                 remapped_labels=self.remapped_labels,
+                multi_label=self.params.multi_label,
             )
             if self.test:
                 test_accuracy = self.model.evaluate(self.test)
@@ -569,6 +603,11 @@ def train_model_tfrecords(
         self.save(run_name, history=history, test_results=test_accuracy)
 
     def checkpoints(self, run_name):
+        checkpoint_file = os.path.join(self.checkpoint_folder, run_name, "cp.ckpt")
+
+        cp_callback = tf.keras.callbacks.ModelCheckpoint(
+            filepath=checkpoint_file, save_weights_only=True, verbose=1
+        )
         val_loss = os.path.join(self.checkpoint_folder, run_name, "val_loss")
 
         checkpoint_loss = tf.keras.callbacks.ModelCheckpoint(
@@ -579,11 +618,15 @@ def checkpoints(self, run_name):
             save_weights_only=True,
             mode="auto",
         )
-        val_acc = os.path.join(self.checkpoint_folder, run_name, "val_acc")
+        val_acc = os.path.join(
+            self.checkpoint_folder,
+            run_name,
+            "val_acc",
+        )
 
         checkpoint_acc = tf.keras.callbacks.ModelCheckpoint(
             val_acc,
-            monitor="val_accuracy",
+            monitor="val_binary_accuracy" if self.params.multi_label else "val_acc",
             verbose=1,
             save_best_only=True,
             save_weights_only=True,
@@ -601,7 +644,8 @@ def checkpoints(self, run_name):
             mode="max",
         )
         earlyStopping = tf.keras.callbacks.EarlyStopping(
-            patience=22, monitor="val_accuracy"
+            patience=22,
+            monitor="val_binary_accuracy" if self.params.multi_label else "val_acc",
         )
         # havent found much use in this just takes training time
         # file_writer_cm = tf.summary.create_file_writer(
@@ -621,9 +665,16 @@ def checkpoints(self, run_name):
         #   "verbose": 1
         # },
         reduce_lr_callback = tf.keras.callbacks.ReduceLROnPlateau(
-            monitor="val_accuracy", verbose=1
+            monitor="val_binary_accuracy" if self.params.multi_label else "val_acc",
+            verbose=1,
         )
-        return [earlyStopping, checkpoint_acc, checkpoint_loss, reduce_lr_callback]
+        return [
+            earlyStopping,
+            checkpoint_acc,
+            checkpoint_loss,
+            reduce_lr_callback,
+            cp_callback,
+        ]
 
     def regroup(self, shuffle=True):
         # can use this to put animals into groups i.e. wallaby vs not
@@ -702,64 +753,33 @@ def add_lstm(self, cnn):
         model = tf.keras.models.Model(input_layer, preds)
         return model
 
-    def classify_thermal_track(self, clip, track, keep_all=True, segment_frames=None):
-        track_data = {}
-        thermal_median = np.empty(len(track.bounds_history), dtype=np.uint16)
-        for i, region in enumerate(track.bounds_history):
-            if region.width == 0 or region.height == 0:
-                continue
-
-            frame = clip.frame_buffer.get_frame(region.frame_number)
-            if frame is None:
-                logging.error(
-                    "Clasifying clip %s track %s can't get frame %s",
-                    clip.get_id(),
-                    track.get_id(),
-                    region.frame_number,
-                )
-                raise Exception(
-                    "Clasifying clip {} track {} can't get frame {}".format(
-                        clip.get_id(), track.get_id(), region.frame_number
-                    )
-                )
-            # AI just uses a background subtraction
-            frame.filtered = frame.thermal - clip.background
-            cropped_frame = frame.crop_by_region(region)
-            thermal_median[i] = np.median(frame.thermal)
-            cropped_frame.resize_with_aspect(
-                (self.params.frame_size, self.params.frame_size),
-                clip.crop_rectangle,
-                True,
-            )
-            track_data[frame.frame_number] = cropped_frame
-
-        segments = track.get_segments(
-            clip.ffc_frames,
-            thermal_median,
-            self.params.square_width**2,
-            repeats=1,
+    def classify_track(self, clip, track, segment_frames=None):
+        # logging.debug("Classifying track %s", self.type)
+        # if self.params.square_width == 1:
+        #     frames, preprocessed, masses = self.preprocess_frames(clip, track)
+        # else:
+        #     frames, preprocessed, masses = self.preprocess_segments(
+        #         clip, track, segment_frames=segment_frames
+        #     )
+        start = time.time()
+        prediction_frames, output, masses = self.predict_track(
+            clip,
+            track,
             segment_frames=segment_frames,
-            segment_type=self.params.segment_type,
+            frames_per_classify=self.params.square_width**2,
         )
-        features = None
-        if self.params.mvm:
-            features = forestmodel.process_track(clip, track, normalize=True)
-        return self.classify_track_data(
-            track.get_id(),
-            track_data,
-            segments,
-            features,
+        track_prediction = TrackPrediction(track.get_id(), self.labels)
+        # self.model.predict(preprocessed)
+        masses = np.array(masses)
+        masses = masses[:, None]
+        track_prediction.classified_clip(
+            output, output * output * masses, prediction_frames
         )
+        track_prediction.classify_time = time.time() - start
+        return track_prediction
 
-    def classify_track(self, clip, track, keep_all=True, segment_frames=None):
-        logging.debug("Classifying track %s", self.type)
-        if self.type == "IR":
-            return self.classify_ir(clip, track)
-        return self.classify_thermal_track(clip, track, keep_all, segment_frames)
-
-    def classify_ir(self, clip, track, keep_all=True, segment_frames=None):
+    def classify_ir(self, clip, track, segment_frames=None):
         data = []
-        crop = True
         thermal_median = np.empty(len(track.bounds_history), dtype=np.uint16)
         frames_used = []
         for i, region in enumerate(track.bounds_history):
@@ -794,12 +814,11 @@ def classify_ir(self, clip, track, keep_all=True, segment_frames=None):
                 region,
             )
             preprocessed = preprocess_ir(
-                frame.copy(),
+                frame,
                 (
                     self.params.frame_size,
                     self.params.frame_size,
                 ),
-                crop,
                 region,
                 self.preprocess_fn,
                 save_info=f"{region.frame_number} - {region}",
@@ -822,87 +841,106 @@ def classify_ir(self, clip, track, keep_all=True, segment_frames=None):
         track_prediction.normalize_score()
         return track_prediction
 
-    def classify_track_data(
-        self,
-        track_id,
-        data,
-        segments,
-        features=None,
-    ):
-        track_prediction = TrackPrediction(track_id, self.labels)
-        start = time.time()
-        predictions = []
-        smoothed_predictions = []
-        predict_me = []
-        prediction_frames = []
-        mass = []
-        for segment in segments:
-            segment_frames = []
-            median = np.zeros((len(segment.frame_indices)))
-            for frame_i in segment.frame_indices:
-                f = data[frame_i]
-                segment_frames.append(f.copy())
-            frames = preprocess_movement(
-                segment_frames,
-                self.params.square_width,
-                self.params.frame_size,
-                self.params.red_type,
-                self.params.green_type,
-                self.params.blue_type,
-                self.preprocess_fn,
-                reference_level=segment.frame_temp_median,
-                keep_edge=self.params.keep_edge,
-            )
-            if frames is None:
-                logging.warn("No frames to predict on")
-                continue
-
-            predict_me.append(frames)
-            prediction_frames.append(segment.frame_indices)
-            mass.append(segment.mass)
-        if len(predict_me) > 0:
-            mass = np.array(mass)
-            mass = mass[:, None]
-            predict_me = np.array(predict_me)
-            if self.params.mvm:
-                features = features[np.newaxis, :]
-                features = np.repeat(features, len(predict_me), axis=0)
-                output = self.model.predict([predict_me, features])
-            else:
-                output = self.model.predict(predict_me)
-
-            track_prediction.classified_clip(
-                output, output * output * mass, prediction_frames
-            )
-
-        track_prediction.classify_time = time.time() - start
-        return track_prediction
+    #
+    # def classify_track_data(
+    #     self,
+    #     track_id,
+    #     data,
+    #     segments,
+    #     features=None,
+    # ):
+    #     track_prediction = TrackPrediction(track_id, self.labels)
+    #     start = time.time()
+    #     predictions = []
+    #     smoothed_predictions = []
+    #     predict_me = []
+    #     prediction_frames = []
+    #     mass = []
+    #     for segment in segments:
+    #         segment_frames = []
+    #         for frame_i in segment.frame_indices:
+    #             f = data[frame_i]
+    #             segment_frames.append(f.copy())
+    #         frames = preprocess_movement(
+    #             segment_frames,
+    #             self.params.square_width,
+    #             self.params.frame_size,
+    #             self.params.red_type,
+    #             self.params.green_type,
+    #             self.params.blue_type,
+    #             self.preprocess_fn,
+    #         )
+    #         if frames is None:
+    #             logging.warn("No frames to predict on")
+    #             continue
+    #
+    #         predict_me.append(frames)
+    #         prediction_frames.append(segment.frame_indices)
+    #         mass.append(segment.mass)
+    #     if len(predict_me) > 0:
+    #         mass = np.array(mass)
+    #         mass = mass[:, None]
+    #         predict_me = np.float32(predict_me)
+    #         if self.params.mvm:
+    #             features = features[np.newaxis, :]
+    #             features = np.repeat(features, len(predict_me), axis=0)
+    #             output = self.model.predict([predict_me, features])
+    #         else:
+    #             output = self.model.predict(predict_me)
+    #
+    #         track_prediction.classified_clip(
+    #             output, output * output * mass, prediction_frames
+    #         )
+    #
+    #     track_prediction.classify_time = time.time() - start
+    #     return track_prediction
 
     def predict(self, frames):
         return self.model.predict(frames)
 
-    def classify_frame(self, frame, thermal_median, preprocess=True):
-        if preprocess:
-            frame = preprocess_frame(
-                frame,
-                False,
-                thermal_median,
-                0,
-                self.params.output_dim,
-                preprocess_fn=self.preprocess_fn,
-            )
-            if frame is None:
-                return None
-        output = self.model.predict(frame[np.newaxis, :])
-        return output[0]
-
     def confusion_tfrecords(self, dataset, filename):
         true_categories = tf.concat([y for x, y in dataset], axis=0)
         if len(true_categories) > 1:
-            true_categories = np.int64(tf.argmax(true_categories, axis=1))
+            if self.params.multi_label:
+                # multi = []
+                # for y in true_categories:
+                # multi.append(tf.where(y).numpy().ravel())
+                # print(y, tf.where(y))
+                # true_categories = np.int64(true_categories)
+                pass
+            else:
+                true_categories = np.int64(tf.argmax(true_categories, axis=1))
         y_pred = self.model.predict(dataset)
-
-        predicted_categories = np.int64(tf.argmax(y_pred, axis=1))
+        if self.params.multi_label:
+            self.labels.append("nothing")
+            # predicted_categori/es = []
+            # for p in y_pred:
+            # predicted_categories.append(tf.where(p >= 0.8).numpy().ravel())
+            # predicted_categories = np.int64(predicted_categories)
+
+            flat_p = []
+            flat_y = []
+            for y, p in zip(true_categories, y_pred):
+                index = 0
+                for y_l, p_l in zip(y, p):
+                    predicted = p_l >= 0.8
+                    if y_l == 0 and predicted:
+                        flat_y.append(len(self.labels) - 1)
+                        flat_p.append(index)
+                    elif y_l == 1 and predicted:
+                        flat_y.append(index)
+                        flat_p.append(index)
+                    elif y_l == 1 and not predicted:
+                        flat_y.append(index)
+                        flat_p.append(len(self.labels) - 1)
+                    # elif y_l == 0 and not predicted:
+                    # all good
+                    # continue
+                    index += 1
+            true_categories = np.int64(flat_p)
+            predicted_categories = np.int64(flat_y)
+        else:
+            predicted_categories = np.int64(tf.argmax(y_pred, axis=1))
 
         cm = confusion_matrix(
             true_categories, predicted_categories, labels=np.arange(len(self.labels))
@@ -932,88 +970,89 @@ def evaluate(self, dataset):
         test_accuracy = self.model.evaluate(dataset)
         logging.info("Test accuracy is %s", test_accuracy)
 
-    # needs to be updated to work with tfrecord datagen
-    def track_accuracy(self, dataset, confusion="confusion.png"):
-        dataset.set_read_only(True)
-        dataset.use_segments = self.params.use_segments
-        predictions = []
-        actual = []
-        raw_predictions = []
-        total = 0
-        correct = 0
-        samples_by_label = {}
-        incorrect_labels = {}
-        for sample in dataset.segments:
-            label_samples = samples_by_label.setdefault(sample.label, {})
-            if sample.track_id in label_samples:
-                label_samples[sample.track_id].append(sample)
-            else:
-                label_samples[sample.track_id] = [sample]
-        bird_tracks = len(label_samples.get("bird", []))
-
-        for label in dataset.label_mapping.keys():
-            incorrect = {}
-            incorrect_labels[label] = incorrect
-            track_samples = samples_by_label.get(label)
-            if not track_samples:
-                logging.warn("No samples for %s", label)
-                continue
-            track_samples = track_samples.values()
-            if label == "insect" or label == "false-positive":
-                track_samples = np.random.choice(
-                    list(track_samples),
-                    min(len(track_samples), bird_tracks),
-                    replace=False,
-                )
-            logging.info("taking %s tracks for %s", len(track_samples), label)
-            mapped_label = dataset.mapped_label(label)
-            for track_segments in track_samples:
-                segment_db = dataset.numpy_data.load_segments(track_segments)
-                frame_db = {}
-                for frames in segment_db.values():
-                    for f in frames:
-                        frame_db[f.frame_number] = f
-                track_prediction = self.classify_track_data(
-                    track_segments[0].track_id,
-                    frame_db,
-                    segments=track_segments,
-                )
-
-                total += 1
-                if track_prediction is None or len(track_prediction.predictions) == 0:
-                    logging.warn("No predictions for %s", track_segments[0].track_id)
-                    continue
-                avg = np.mean(track_prediction.predictions, axis=0)
-                actual.append(self.labels.index(mapped_label))
-                predictions.append(track_prediction.best_label_index)
-
-                raw_predictions.append(avg)
-                if actual[-1] == predictions[-1]:
-                    correct += 1
-                else:
-                    if track_prediction.predicted_tag() in incorrect:
-                        incorrect[track_prediction.predicted_tag()].append(
-                            track_segments[0].unique_track_id
-                        )
-                    else:
-                        incorrect[track_prediction.predicted_tag()] = [
-                            track_segments[0].unique_track_id
-                        ]
-
-                if total % 50 == 0:
-                    logging.info("Processed %s", total)
-        for label, incorrect in incorrect_labels.items():
-            logging.info("Incorrect ************ %s", label)
-            logging.info(incorrect.get("false-positive"))
-        logging.info("Predicted correctly %s", round(100 * correct / total))
-        self.f1(actual, raw_predictions)
-
-        if confusion is not None:
-            cm = confusion_matrix(
-                actual, predictions, labels=np.arange(len(self.labels))
-            )
-            figure = plot_confusion_matrix(cm, class_names=self.labels)
-            plt.savefig(confusion, format="png")
+    #
+    # # needs to be updated to work with tfrecord datagen
+    # def track_accuracy(self, dataset, confusion="confusion.png"):
+    #     dataset.set_read_only(True)
+    #     dataset.use_segments = self.params.use_segments
+    #     predictions = []
+    #     actual = []
+    #     raw_predictions = []
+    #     total = 0
+    #     correct = 0
+    #     samples_by_label = {}
+    #     incorrect_labels = {}
+    #     for sample in dataset.segments:
+    #         label_samples = samples_by_label.setdefault(sample.label, {})
+    #         if sample.track_id in label_samples:
+    #             label_samples[sample.track_id].append(sample)
+    #         else:
+    #             label_samples[sample.track_id] = [sample]
+    #     bird_tracks = len(label_samples.get("bird", []))
+    #
+    #     for label in dataset.label_mapping.keys():
+    #         incorrect = {}
+    #         incorrect_labels[label] = incorrect
+    #         track_samples = samples_by_label.get(label)
+    #         if not track_samples:
+    #             logging.warn("No samples for %s", label)
+    #             continue
+    #         track_samples = track_samples.values()
+    #         if label == "insect" or label == "false-positive":
+    #             track_samples = np.random.choice(
+    #                 list(track_samples),
+    #                 min(len(track_samples), bird_tracks),
+    #                 replace=False,
+    #             )
+    #         logging.info("taking %s tracks for %s", len(track_samples), label)
+    #         mapped_label = dataset.mapped_label(label)
+    #         for track_segments in track_samples:
+    #             segment_db = dataset.numpy_data.load_segments(track_segments)
+    #             frame_db = {}
+    #             for frames in segment_db.values():
+    #                 for f in frames:
+    #                     frame_db[f.frame_number] = f
+    #             track_prediction = self.classify_track_data(
+    #                 track_segments[0].track_id,
+    #                 frame_db,
+    #                 segments=track_segments,
+    #             )
+    #
+    #             total += 1
+    #             if track_prediction is None or len(track_prediction.predictions) == 0:
+    #                 logging.warn("No predictions for %s", track_segments[0].track_id)
+    #                 continue
+    #             avg = np.mean(track_prediction.predictions, axis=0)
+    #             actual.append(self.labels.index(mapped_label))
+    #             predictions.append(track_prediction.best_label_index)
+    #
+    #             raw_predictions.append(avg)
+    #             if actual[-1] == predictions[-1]:
+    #                 correct += 1
+    #             else:
+    #                 if track_prediction.predicted_tag() in incorrect:
+    #                     incorrect[track_prediction.predicted_tag()].append(
+    #                         track_segments[0].unique_track_id
+    #                     )
+    #                 else:
+    #                     incorrect[track_prediction.predicted_tag()] = [
+    #                         track_segments[0].unique_track_id
+    #                     ]
+    #
+    #             if total % 50 == 0:
+    #                 logging.info("Processed %s", total)
+    #     for label, incorrect in incorrect_labels.items():
+    #         logging.info("Incorrect ************ %s", label)
+    #         logging.info(incorrect.get("false-positive"))
+    #     logging.info("Predicted correctly %s", round(100 * correct / total))
+    #     self.f1(actual, raw_predictions)
+    #
+    #     if confusion is not None:
+    #         cm = confusion_matrix(
+    #             actual, predictions, labels=np.arange(len(self.labels))
+    #         )
+    #         figure = plot_confusion_matrix(cm, class_names=self.labels)
+    #         plt.savefig(confusion, format="png")
 
 
 # from tensorflow examples
@@ -1089,10 +1128,13 @@ def plot_to_image(figure):
 
 
 def loss(params):
-    softmax = tf.keras.losses.CategoricalCrossentropy(
+    if params.multi_label:
+        return tf.keras.losses.BinaryCrossentropy(
+            label_smoothing=params.label_smoothing,
+        )
+    return tf.keras.losses.CategoricalCrossentropy(
         label_smoothing=params.label_smoothing,
     )
-    return softmax
 
 
 def optimizer(params):
@@ -1326,9 +1368,9 @@ def on_epoch_end(self, epoch, logs=None):
         tf.keras.backend.clear_session()
 
 
-def get_excluded(type):
+def get_excluded(type, multi_label=False):
     if type == "thermal":
-        return thermaldataset.get_excluded(), thermaldataset.get_remapped()
+        return thermaldataset.get_excluded(), thermaldataset.get_remapped(multi_label)
     else:
         return irdataset.get_excluded(), irdataset.get_remapped()
 
diff --git a/ml_tools/preprocess.py b/ml_tools/preprocess.py
index 842e7405..441bc49b 100644
--- a/ml_tools/preprocess.py
+++ b/ml_tools/preprocess.py
@@ -53,237 +53,145 @@ def augement_frame(frame, frame_size, dim):
     return image.numpy()
 
 
-def preprocess_segment(
-    frames,
-    frame_size,
-    reference_level=None,
-    frame_velocity=None,
-    augment=False,
-    default_inset=0,
-    keep_edge=False,
-):
-    """
-    Preprocesses the raw track data, scaling it to correct size, and adjusting to standard levels
-    :param frames: a list of Frames
-    :param reference_level: thermal reference level for each frame in data
-    :param frame_velocity: velocity (x,y) for each frame.
-    :param augment: if true applies a slightly random crop / scale
-    :param default_inset: the default number of pixels to inset when no augmentation is applied.
-    """
-    if reference_level is not None:
-        # -------------------------------------------
-        # next adjust temperature and flow levels
-        # get reference level for thermal channel
-        assert len(frames) == len(
-            reference_level
-        ), "Reference level shape and data shape not match."
-
-    crop_rectangle = tools.Rectangle(EDGE, EDGE, res_x - 2 * EDGE, res_y - 2 * EDGE)
-
-    # -------------------------------------------
-    # first we scale to the standard size
-    top_offset = 0
-    bottom_offset = 0
-    left_offset = 0
-    right_offset = 0
-    data = []
-    flip = False
-    chance = random.random()
-    if augment:
-        contrast_adjust = None
-        level_adjust = None
-        if chance <= 0.75:
-            # we will adjust contrast and levels, but only within these bounds.
-            # that is a bright input may have brightness reduced, but not increased.
-            LEVEL_OFFSET = 4
-
-            # apply level and contrast shift
-            level_adjust = float(random.normalvariate(0, LEVEL_OFFSET))
-            contrast_adjust = float(tools.random_log(0.9, (1 / 0.9)))
-        if chance <= 0.50:
-            flip = True
-
-    for i, frame in enumerate(frames):
-        frame.float_arrays()
-        frame_height, frame_width = frame.thermal.shape
-        # adjusting the corners makes the algorithm robust to tracking differences.
-        # gp changed to 0,1 maybe should be a percent of the frame size
-        if augment or frame_height > frame_size or frame_width > frame_size:
-            max_height_offset = frame_height - frame_size
-            max_width_offset = frame_width - frame_size
-            top_offset = int(random.random() * max_height_offset)
-            # bottom_offset = int(random.random() * max_height_offset)
-            left_offset = int(random.random() * max_width_offset)
-            # right_offset = int(random.random() * max_width_offset)
-            crop_region = tools.Rectangle(
-                left_offset, top_offset, frame_size, frame_size
-            )
-            frame.crop_by_region(crop_region, out=frame)
-        if augment and chance <= 0.75:
-            # degress = 0
-
-            degrees = int(chance * 40) - 20
-            frame.rotate(degrees)
-        if frame_height < MIN_SIZE or frame_width < MIN_SIZE:
-            continue
-        if reference_level is not None:
-            frame.thermal -= reference_level[i]
-            np.clip(frame.thermal, a_min=0, a_max=None, out=frame.thermal)
-
-        frame.normalize()
-
-        if augment:
-            if level_adjust is not None:
-                frame.brightness_adjust(level_adjust)
-            if contrast_adjust is not None:
-                frame.contrast_adjust(contrast_adjust)
-            if flip:
-                frame.flip()
-        data.append(frame)
-
-    return data, flip
-
-
 def preprocess_frame(
     frame,
-    frame_size,
-    augment,
-    thermal_median,
-    velocity,
-    output_dim,
-    preprocess_fn=None,
-    sample=None,
+    out_dim,
+    region,
+    background=None,
+    crop_rectangle=None,
+    calculate_filtered=True,
 ):
-    processed_frame, flipped = preprocess_segment(
-        [frame],
-        frame_size,
-        reference_level=[thermal_median],
-        augment=augment,
-        default_inset=0,
+    median = np.median(frame.thermal)
+    cropped_frame = frame.crop_by_region(region, only_thermal=True)
+    cropped_frame.thermal = np.float32(cropped_frame.thermal)
+    if calculate_filtered:
+        if background is None:
+            logging.warning(
+                "Not calculating filtered frame as no background was supplied"
+            )
+        else:
+            cropped_frame.filtered = cropped_frame.thermal - region.subimage(background)
+    cropped_frame.resize_with_aspect(
+        out_dim,
+        crop_rectangle,
+        True,
     )
-    if len(processed_frame) == 0:
-        return
-    processed_frame = processed_frame[0]
-    thermal = processed_frame.get_channel(TrackChannels.thermal)
-    filtered = processed_frame.get_channel(TrackChannels.filtered)
-    thermal, stats = imageprocessing.normalize(thermal, min=0)
-    if not stats[0]:
-        return None
-    filtered, stats = imageprocessing.normalize(filtered, min=0)
-    if not stats[0]:
-        return None
-
-    data = np.empty((*thermal.shape, 3))
-    data[:, :, 0] = thermal
-    data[:, :, 1] = filtered
-    data[:, :, 2] = filtered
-    # for testing
-    # tools.saveclassify_image(
-    #     data,
-    #     f"samples/{sample.label}-{sample.clip_id}-{sample.track_id}",
-    # )
-
-    # preprocess expects values in range 0-255
-    if preprocess_fn:
-        data = data * 255
-        data = preprocess_fn(data)
-    return data
+    cropped_frame.thermal -= median
+    np.clip(cropped_frame.thermal, 0, None, out=cropped_frame.thermal)
+    cropped_frame.normalize()
+    return cropped_frame
+
+
+#
+#
+# def preprocess_frame(
+#     frame,
+#     frame_size,
+#     thermal_median,
+#     velocity,
+#     output_dim,
+#     preprocess_fn=None,
+#     sample=None,
+# ):
+#     processed_frame, flipped = preprocess_segment(
+#         [frame],
+#         frame_size,
+#         reference_level=[thermal_median],
+#         augment=augment,
+#         default_inset=0,
+#     )
+#     if len(processed_frame) == 0:
+#         return
+#     processed_frame = processed_frame[0]
+#     thermal = processed_frame.get_channel(TrackChannels.thermal)
+#     filtered = processed_frame.get_channel(TrackChannels.filtered)
+#     thermal, stats = imageprocessing.normalize(thermal, min=0)
+#     if not stats[0]:
+#         return None
+#     filtered, stats = imageprocessing.normalize(filtered, min=0)
+#     if not stats[0]:
+#         return None
+#
+#     data = np.empty((*thermal.shape, 3))
+#     data[:, :, 0] = thermal
+#     data[:, :, 1] = filtered
+#     data[:, :, 2] = filtered
+#     # for testing
+#     # tools.saveclassify_image(
+#     #     data,
+#     #     f"samples/{sample.label}-{sample.clip_id}-{sample.track_id}",
+#     # )
+#
+#     # preprocess expects values in range 0-255
+#     if preprocess_fn:
+#         data = data * 255
+#         data = preprocess_fn(data)
+#     return data
 
 
 index = 0
 
 
-def preprocess_ir(
+def preprocess_single_frame(
     frame,
     frame_size,
-    crop=True,
     region=None,
     preprocess_fn=None,
     save_info="",
 ):
-    if crop:
-        crop_rectangle = tools.Rectangle(
-            0, 0, frame.thermal.shape[1], frame.thermal.shape[0]
-        )
-        frame.crop_by_region(region, out=frame)
-    frame.normalize()
-    image = np.stack((frame.thermal, frame.thermal, frame.thermal), axis=2)
-    image = np.float32(image)
-    image = imageprocessing.resize_and_pad(
-        image,
-        (frame_size[0], frame_size[1], 3),
-        frame.region,
-        None,
-        True,
+    preprocessed_frame = preprocess_frame(
+        frame,
+        frame_size,
+        region,
     )
 
+    image = np.stack(
+        (
+            preprocessed_frame.thermal,
+            preprocessed_frame.thermal,
+            preprocessed_frame.thermal,
+        ),
+        axis=2,
+    )
     if preprocess_fn:
         image = preprocess_fn(image)
     return image
 
 
+# index = 0
+
+
+#
 def preprocess_movement(
-    segment,
+    preprocess_frames,
     frames_per_row,
     frame_size,
     red_type,
     green_type,
     blue_type,
     preprocess_fn=None,
-    augment=False,
-    reference_level=None,
     sample=None,
-    keep_edge=False,
 ):
-    segment, flipped = preprocess_segment(
-        segment,
-        frame_size,
-        reference_level=reference_level,
-        augment=augment,
-        default_inset=0,
-        keep_edge=keep_edge,
-    )
     frame_types = {}
     channel_types = set([green_type, blue_type, red_type])
     for type in channel_types:
-        if type == FrameTypes.overlay:
-            if overlay is None:
-                overlay = imageprocessing.overlay_image(
-                    data,
-                    regions,
-                    dim=(frames_per_row * frame_size, frames_per_row * frame_size),
-                    require_movement=True,
-                )
-                channel_data, stats = imageprocessing.normalize(overlay, min=0)
-                if not stats[0]:
-                    return None
-            else:
-                channel_data = np.zeros((square.shape[0], square.shape[1]))
-                channel_data[: overlay.shape[0], : overlay.shape[1]] = overlay
-
-            if flipped:
-                channel_data = np.flip(channel_data, axis=1)
-        elif type == FrameTypes.flow_tiled:
-            channel_segment = [
-                frame.get_channel(TrackChannels.flow) for frame in segment
-            ]
-            channel_data, success = imageprocessing.square_clip_flow(
-                channel_segment, frames_per_row, (frame_size, frame_size)
-            )
-            if not success:
-                return None
+        if type == FrameTypes.thermal_tiled:
+            channel = TrackChannels.thermal
+        elif type == FrameTypes.filtered_tiled:
+            channel = TrackChannels.filtered
         else:
-            if type == FrameTypes.thermal_tiled:
-                channel = TrackChannels.thermal
-            else:
-                channel = TrackChannels.filtered
-            channel_segment = [frame.get_channel(channel) for frame in segment]
-            channel_data, success = imageprocessing.square_clip(
-                channel_segment, frames_per_row, (frame_size, frame_size)
-            )
+            raise Exception("Cannot process type %s for movement", type)
+        channel_segment = [frame.get_channel(channel) for frame in preprocess_frames]
+        channel_data, success = imageprocessing.square_clip(
+            channel_segment,
+            frames_per_row,
+            (frame_size, frame_size),
+            normalize=False,
+        )
+        # already done normalization
 
-            if not success:
-                return None
+        if not success:
+            return None
 
         frame_types[type] = channel_data
 
@@ -291,15 +199,14 @@ def preprocess_movement(
         (frame_types[red_type], frame_types[green_type], frame_types[blue_type]), axis=2
     )
     #
+    # # # # # for testing
     # global index
     # index += 1
-    # # # # # for testing
     # tools.saveclassify_image(
-    #     data * 255,
+    #     data,
     #     f"samples/{index}",
     # )
 
     if preprocess_fn:
-        data = data * 255
         data = preprocess_fn(data)
-    return data
+    return np.float32(data)
diff --git a/ml_tools/previewer.py b/ml_tools/previewer.py
index c26204c3..cacebe25 100644
--- a/ml_tools/previewer.py
+++ b/ml_tools/previewer.py
@@ -24,7 +24,7 @@
 import cv2
 from PIL import Image, ImageDraw, ImageFont
 
-from load.clip import Clip
+from track.clip import Clip
 from ml_tools import tools
 from ml_tools.tools import Rectangle
 import ml_tools.globals as globs
diff --git a/ml_tools/rawdb.py b/ml_tools/rawdb.py
new file mode 100644
index 00000000..891c505a
--- /dev/null
+++ b/ml_tools/rawdb.py
@@ -0,0 +1,185 @@
+"""
+
+Author GP
+
+Date 2023
+
+
+"""
+import os
+import logging
+import filelock
+import numpy as np
+from dateutil.parser import parse as parse_date
+from .frame import Frame, TrackChannels
+import json
+from pathlib import Path
+import numpy as np
+from track.region import Region
+from ml_tools.datasetstructures import TrackHeader, ClipHeader
+from track.track import Track
+from track.cliptrackextractor import is_affected_by_ffc
+from cptv import CPTVReader
+
+special_datasets = [
+    "tag_frames",
+    "original_frames",
+    "background_frame",
+    "predictions",
+    "overlay",
+]
+
+
+class RawDatabase:
+    def __init__(self, database_filename):
+        self.file = Path(database_filename)
+        self.meta_data_file = self.file.with_suffix(".txt")
+        self._meta_data = None
+        self.background = None
+        self.ffc_frames = None
+        self.frames = None
+
+    def get_frames(self):
+        return self.frames
+
+    def get_clip_background(self):
+        return self.background
+
+    def load_frames(self):
+        ffc_frames = []
+        cptv_frames = []
+        background = None
+        tracker_version = self.meta_data.get("tracker_version")
+        frame_i = 0
+        with open(self.file, "rb") as f:
+            reader = CPTVReader(f)
+            for frame in reader:
+                if frame.background_frame:
+                    background = frame.pix
+                    # bug in previous tracker version where background was first frame
+                    if tracker_version >= 10:
+                        continue
+                ffc = is_affected_by_ffc(frame)
+                if ffc:
+                    ffc_frames.append(frame_i)
+                cptv_frames.append(frame.pix)
+                frame_i += 1
+        frames = np.uint16(cptv_frames)
+        if background is None:
+            background = np.mean(frames, axis=0)
+
+        self.frames = []
+        for i, f in enumerate(frames):
+            self.frames.append(Frame(f, None, None, i))
+        self.ffc_frames = ffc_frames
+        self.background = background
+
+    @property
+    def meta_data(self):
+        if self._meta_data is not None:
+            return self._meta_data
+        if not self.meta_data_file.is_file():
+            logging.warn("Could not load meta data for %s", self.meta_data_file)
+            return None
+        with open(self.meta_data_file, "r") as t:
+            # add in some metadata stats
+            self._meta_data = json.load(t)
+        return self._meta_data
+
+    def get_clip_tracks(self, tag_precedence):
+        metadata = self.meta_data
+        if metadata is None:
+            return None
+        clip_header = ClipHeader(
+            clip_id=int(metadata["id"]),
+            station_id=metadata.get("stationId"),
+            source_file=self.file,
+            location=metadata.get("location"),
+            camera=metadata.get("deviceId"),
+            rec_time=parse_date(metadata["recordingDateTime"]),
+            frames_per_second=10 if self.file.suffix == "mp4" else 9,
+            events=metadata.get("event", ""),
+            trap=metadata.get("trap", ""),
+            tracks=[],
+            ffc_frames=self.ffc_frames,
+        )
+        tracks = metadata.get("Tracks", [])
+        meta = []
+        for track_meta in tracks:
+            tags = track_meta.get("tags", [])
+            tag = Track.get_best_human_tag(tags, tag_precedence, 0)
+            human_tag = None
+            human_tag_confidence = None
+            master_tag = [
+                t
+                for t in tags
+                if t.get("automatic")
+                and not isinstance(t.get("data", ""), str)
+                and t.get("data", {}).get("name") == "Master"
+            ]
+            if len(master_tag) > 0:
+                master_tag = master_tag[0]
+                track_meta["ai_tag"] = master_tag["what"]
+                track_meta["ai_tag_confidence"] = master_tag["confidence"]
+
+            if tag is not None:
+                human_tag = tag["what"]
+                human_tag_confidence = tag["confidence"]
+            human_tags = [
+                (t.get("what"), t["confidence"])
+                for t in tags
+                if t.get("automatic", False) != True
+            ]
+
+            start = None
+            end = None
+
+            prev_frame = None
+            regions = {}
+            for i, r in enumerate(track_meta.get("positions")):
+                if isinstance(r, list):
+                    region = Region.region_from_array(r[1])
+                    if region.frame_number is None:
+                        if i == 0:
+                            frame_number = round(r[0] * FPS)
+                            region.frame_number = frame_number
+                        else:
+                            region.frame_number = prev_frame + 1
+                else:
+                    region = Region.region_from_json(r)
+                if region.frame_number is None:
+                    if "frameTime" in r:
+                        if i == 0:
+                            region.frame_number = round(r["frameTime"] * 9)
+                        else:
+                            region.frame_number = prev_frame + 1
+                prev_frame = region.frame_number
+                region.frame_number = region.frame_number
+                assert region.frame_number >= 0
+                regions[region.frame_number] = region
+                if start is None:
+                    start = region.frame_number
+                end = region.frame_number
+            header = TrackHeader(
+                clip_id=clip_header.clip_id,
+                track_id=int(track_meta["id"]),
+                label=human_tag,
+                num_frames=len(regions),
+                regions=regions,
+                start_frame=start,
+                confidence=human_tag_confidence,
+                human_tags=human_tags,
+                source_file=self.file,
+                # frame_temp_median=frame_temp_median,
+            )
+            clip_header.tracks.append(header)
+        return clip_header
+
+    def get_clip_meta(self, tag_precedence):
+        return self.get_clip_tracks(tag_precedence)
+        #
+        # clip_meta = {}
+        # clip_meta["clip_id"] = r_id
+        # stationId = metadata.get("stationId", 0)
+        # clip_meta["station_id"] = stationId
+        # clip_meta["crop_rectangle"] = np.uint8(clip.crop_rectangle.to_ltrb())
diff --git a/ml_tools/resnet/resnet.py b/ml_tools/resnet/resnet.py
new file mode 100644
index 00000000..1d172720
--- /dev/null
+++ b/ml_tools/resnet/resnet.py
@@ -0,0 +1,224 @@
+# https://medium.com/analytics-vidhya/understanding-and-implementation-of-residual-networks-resnets-b80f9a507b9c
+import tensorflow as tf
+
+
+def ResNet50(input_shape=(64, 64, 3), classes=6):
+    """
+    Implementation of the popular ResNet50 the following architecture:
+    CONV2D -> BATCHNORM -> RELU -> MAXPOOL -> CONVBLOCK -> IDBLOCK*2 -> CONVBLOCK -> IDBLOCK*3
+    -> CONVBLOCK -> IDBLOCK*5 -> CONVBLOCK -> IDBLOCK*2 -> AVGPOOL -> TOPLAYER
+
+    Arguments:
+    input_shape -- shape of the images of the dataset
+    classes -- integer, number of classes
+
+    Returns:
+    model -- a Model() instance in Keras
+    """
+
+    # Define the input as a tensor with shape input_shape
+    X_input = tf.keras.Input(input_shape)
+
+    # Zero-Padding
+    X = tf.keras.layers.ZeroPadding2D((3, 3))(X_input)
+
+    # Stage 1
+    X = tf.keras.layers.Conv2D(
+        64,
+        (7, 7),
+        strides=(2, 2),
+        name="conv1",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X)
+    X = tf.keras.layers.BatchNormalization(axis=3, name="bn_conv1")(X)
+    X = tf.keras.layers.Activation("relu")(X)
+    X = tf.keras.layers.MaxPooling2D((3, 3), strides=(2, 2))(X)
+
+    # Stage 2
+    X = convolutional_block(X, f=3, filters=[64, 64, 256], stage=2, block="a", s=1)
+    X = identity_block(X, 3, [64, 64, 256], stage=2, block="b")
+    X = identity_block(X, 3, [64, 64, 256], stage=2, block="c")
+
+    # Stage 3
+    X = convolutional_block(X, f=3, filters=[128, 128, 512], stage=3, block="a", s=2)
+    X = identity_block(X, 3, [128, 128, 512], stage=3, block="b")
+    X = identity_block(X, 3, [128, 128, 512], stage=3, block="c")
+    X = identity_block(X, 3, [128, 128, 512], stage=3, block="d")
+
+    # Stage 4
+    X = convolutional_block(X, f=3, filters=[256, 256, 1024], stage=4, block="a", s=2)
+    X = identity_block(X, 3, [256, 256, 1024], stage=4, block="b")
+    X = identity_block(X, 3, [256, 256, 1024], stage=4, block="c")
+    X = identity_block(X, 3, [256, 256, 1024], stage=4, block="d")
+    X = identity_block(X, 3, [256, 256, 1024], stage=4, block="e")
+    X = identity_block(X, 3, [256, 256, 1024], stage=4, block="f")
+
+    # Stage 5
+    X = convolutional_block(X, f=3, filters=[512, 512, 2048], stage=5, block="a", s=2)
+    X = identity_block(X, 3, [512, 512, 2048], stage=5, block="b")
+    X = identity_block(X, 3, [512, 512, 2048], stage=5, block="c")
+
+    # AVGPOOL . Use "X = AveragePooling2D(...)(X)"
+    X = tf.keras.layers.AveragePooling2D()(X)
+
+    # output layer
+    X = tf.keras.layers.Flatten()(X)
+    # X = tf.keras.layers.Dense(
+    #     classes,
+    #     activation="softmax",
+    #     name="fc" + str(classes),
+    #     kernel_initializer=glorot_uniform(seed=0),
+    # )(X)
+
+    # Create model
+    model = tf.keras.Model(inputs=X_input, outputs=X, name="ResNet50")
+
+    return model
+
+
+def identity_block(X, f, filters, stage, block):
+    """
+    Implementation of the identity block
+
+    Arguments:
+    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
+    f -- integer, specifying the shape of the middle CONV's window for the main path
+    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
+    stage -- integer, used to name the layers, depending on their position in the network
+    block -- string/character, used to name the layers, depending on their position in the network
+
+    Returns:
+    X -- output of the identity block, tensor of shape (n_H, n_W, n_C)
+    """
+
+    # defining name basis
+    conv_name_base = "res" + str(stage) + block + "_branch"
+    bn_name_base = "bn" + str(stage) + block + "_branch"
+
+    # Retrieve Filters
+    F1, F2, F3 = filters
+
+    # Save the input value. You'll need this later to add back to the main path.
+    X_shortcut = X
+
+    # First component of main path
+    X = tf.keras.layers.Conv2D(
+        filters=F1,
+        kernel_size=(1, 1),
+        strides=(1, 1),
+        padding="valid",
+        name=conv_name_base + "2a",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X)
+    X = tf.keras.layers.BatchNormalization(axis=3, name=bn_name_base + "2a")(X)
+    X = tf.keras.layers.Activation("relu")(X)
+
+    # Second component of main path
+    X = tf.keras.layers.Conv2D(
+        filters=F2,
+        kernel_size=(f, f),
+        strides=(1, 1),
+        padding="same",
+        name=conv_name_base + "2b",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X)
+    X = tf.keras.layers.BatchNormalization(axis=3, name=bn_name_base + "2b")(X)
+    X = tf.keras.layers.Activation("relu")(X)
+
+    # Third component of main path
+    X = tf.keras.layers.Conv2D(
+        filters=F3,
+        kernel_size=(1, 1),
+        strides=(1, 1),
+        padding="valid",
+        name=conv_name_base + "2c",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X)
+    X = tf.keras.layers.BatchNormalization(axis=3, name=bn_name_base + "2c")(X)
+
+    # Final step: Add shortcut value to main path, and pass it through a RELU activation
+    X = tf.keras.layers.Add()([X, X_shortcut])
+    X = tf.keras.layers.Activation("relu")(X)
+
+    return X
+
+
+def convolutional_block(X, f, filters, stage, block, s=2):
+    """
+    Implementation of the convolutional block
+
+    Arguments:
+    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
+    f -- integer, specifying the shape of the middle CONV's window for the main path
+    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
+    stage -- integer, used to name the layers, depending on their position in the network
+    block -- string/character, used to name the layers, depending on their position in the network
+    s -- Integer, specifying the stride to be used
+
+    Returns:
+    X -- output of the convolutional block, tensor of shape (n_H, n_W, n_C)
+    """
+
+    # defining name basis
+    conv_name_base = "res" + str(stage) + block + "_branch"
+    bn_name_base = "bn" + str(stage) + block + "_branch"
+
+    # Retrieve Filters
+    F1, F2, F3 = filters
+
+    # Save the input value
+    X_shortcut = X
+
+    ##### MAIN PATH #####
+    # First component of main path
+    X = tf.keras.layers.Conv2D(
+        F1,
+        (1, 1),
+        strides=(s, s),
+        name=conv_name_base + "2a",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X)
+    X = tf.keras.layers.BatchNormalization(axis=3, name=bn_name_base + "2a")(X)
+    X = tf.keras.layers.Activation("relu")(X)
+
+    # Second component of main path
+    X = tf.keras.layers.Conv2D(
+        F2,
+        (f, f),
+        strides=(1, 1),
+        padding="same",
+        name=conv_name_base + "2b",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X)
+    X = tf.keras.layers.BatchNormalization(axis=3, name=bn_name_base + "2b")(X)
+    X = tf.keras.layers.Activation("relu")(X)
+
+    # Third component of main path
+    X = tf.keras.layers.Conv2D(
+        F3,
+        (1, 1),
+        strides=(1, 1),
+        padding="valid",
+        name=conv_name_base + "2c",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X)
+    X = tf.keras.layers.BatchNormalization(axis=3, name=bn_name_base + "2c")(X)
+
+    ##### SHORTCUT PATH ####
+    X_shortcut = tf.keras.layers.Conv2D(
+        F3,
+        (1, 1),
+        strides=(s, s),
+        padding="valid",
+        name=conv_name_base + "1",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X_shortcut)
+    X_shortcut = tf.keras.layers.BatchNormalization(axis=3, name=bn_name_base + "1")(
+        X_shortcut
+    )
+
+    # Final step: Add shortcut value to main path, and pass it through a RELU activation
+    X = tf.keras.layers.Add()([X, X_shortcut])
+    X = tf.keras.layers.Activation("relu")(X)
+
+    return X
diff --git a/ml_tools/resnet/wr_resnet.py b/ml_tools/resnet/wr_resnet.py
new file mode 100644
index 00000000..f3db71dc
--- /dev/null
+++ b/ml_tools/resnet/wr_resnet.py
@@ -0,0 +1,98 @@
+import tensorflow as tf
+
+
+# https://arxiv.org/pdf/1605.07146.pdf
+def WRResNet(X_input, depth=22, k=4):
+    filters = [16, 16 * k, 32 * k, 64 * k]
+    # Define the input as a tensor with shape input_shape
+    # X_input = tf.keras.Input(input_shape, name="wr-input")
+    n = int((depth - 4) / 6)
+    for stage, f in enumerate(filters):
+        if stage == 0:
+            X = tf.keras.layers.Conv2D(
+                f,
+                (3, 3),
+                strides=1,
+                padding="same",
+                name=f"conv1_{stage+1}",
+                kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+            )(X_input)
+        else:
+            X = wr_block(
+                X, 3, (f, f), stage=stage + 1, block="b", stride=stage, depth=n
+            )
+    #
+    X = tf.keras.layers.BatchNormalization(axis=3, name="final_bn")(X)
+    X = tf.keras.layers.Activation("relu")(X)
+    # LME?
+    # X = tf.keras.layers.GlobalAveragePooling2D()(X)
+
+    # X = tf.keras.layers.Flatten()(X)
+    # X = tf.keras.layers.Dense(classes, activation="sigmoid", name="prediction")(X)
+    model = tf.keras.Model(inputs=X_input, outputs=X, name="WRResNet")
+    return model
+
+
+def wr_block(X, f, filters, stage, block, stride=1, depth=1):
+    s_block = f"{block}0"
+
+    X = basic_block(X, f, filters, stage, s_block, stride)
+    for d in range(depth - 1):
+        s_block = f"{block}{d+1}"
+        X = basic_block(X, f, filters, stage, s_block, 1)
+    return X
+
+
+def basic_block(X, f, filters, stage, block, stride=1):
+    # defining name basis
+    conv_name_base = "res" + str(stage) + block + "_branch"
+    bn_name_base = "bn" + str(stage) + block + "_branch"
+
+    # Retrieve Filters
+    F1, F2 = filters
+    # Save the input value. You'll need this later to add back to the main path.
+    X_shortcut = X
+    # First component of main path
+    X = tf.keras.layers.BatchNormalization(axis=3, name=bn_name_base + "2a")(X)
+    X = tf.keras.layers.Activation("relu")(X)
+    X = tf.keras.layers.Conv2D(
+        filters=F1,
+        kernel_size=(3, 3),
+        strides=(stride, stride),
+        padding="same",
+        name=conv_name_base + "2a",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X)
+
+    X = tf.keras.layers.Dropout(rate=0.1)(X)
+    # , training=istraining_ph)
+
+    # Second component of main path
+    X = tf.keras.layers.BatchNormalization(axis=3, name=bn_name_base + "2b")(X)
+    X = tf.keras.layers.Activation("relu")(X)
+    X = tf.keras.layers.Conv2D(
+        filters=F2,
+        kernel_size=(3, 3),
+        strides=1,
+        padding="same",
+        name=conv_name_base + "2b",
+        kernel_initializer=tf.keras.initializers.GlorotUniform(seed=0),
+    )(X)
+    if X.shape[-1] == X_shortcut.shape[-1]:
+        X_shortcut = tf.keras.layers.Identity()(X_shortcut)
+    else:
+        X_shortcut = tf.keras.layers.Conv2D(
+            X.shape[-1], strides=(stride, stride), kernel_size=1
+        )(X_shortcut)
+    # Final step: Add shortcut value to main path, and pass it through a RELU activation
+    X = tf.keras.layers.Add()([X, X_shortcut])
+    X = tf.keras.layers.Activation("relu")(X)
+    return X
+
+
+def main():
+    WRResNet()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/ml_tools/tfdataset.py b/ml_tools/tfdataset.py
index cdba7d97..c42a2e47 100644
--- a/ml_tools/tfdataset.py
+++ b/ml_tools/tfdataset.py
@@ -3,6 +3,7 @@
 from functools import partial
 import numpy as np
 import logging
+import random
 
 AUTOTUNE = tf.data.AUTOTUNE
 
@@ -53,9 +54,21 @@ def get_distribution(dataset, num_labels, batched=True):
 
 
 def get_dataset(load_function, base_dir, labels, **args):
+    land_birds = [
+        "pukeko",
+        "california quail",
+        "brown quail",
+        "black swan",
+        "quail",
+        "pheasant",
+        "penguin",
+        "duck",
+        "chicken",
+        "rooster",
+    ]
     excluded_labels = args.get("excluded_labels", [])
     to_remap = args.get("remapped_labels", {})
-
+    logging.info("Excluding %s", excluded_labels)
     remapped = {}
     keys = []
     values = []
@@ -65,9 +78,9 @@ def get_dataset(load_function, base_dir, labels, **args):
     for excluded in excluded_labels:
         if excluded in labels:
             new_labels.remove(excluded)
-    for excluded in to_remap.keys():
-        if excluded in labels:
-            new_labels.remove(excluded)
+    for remapped_lbl in to_remap.keys():
+        if remapped_lbl in labels:
+            new_labels.remove(remapped_lbl)
     for l in labels:
         keys.append(labels.index(l))
         if l not in new_labels:
@@ -99,12 +112,18 @@ def get_dataset(load_function, base_dir, labels, **args):
 
     # 1 / 0
     filenames = tf.io.gfile.glob(f"{base_dir}/*.tfrecord")
-    dataset = load_function(filenames, remap_lookup, num_labels, args)
+    if not args.get("deterministic"):
+        random.shuffle(filenames)
+
+    dataset = load_function(filenames, remap_lookup, new_labels, args)
+
+    filter_excluded = lambda x, y: not tf.math.equal(tf.math.count_nonzero(y), 0)
+    dataset = dataset.filter(filter_excluded)
     if dataset is None:
         logging.warn("No dataset for %s", filenames)
         return None, None
 
-    dataset = dataset.cache()
+    # dataset = dataset.cache()
     if not args.get("only_features") and args.get("shuffle", True):
         logging.info("shuffling data")
         dataset = dataset.shuffle(
@@ -129,7 +148,7 @@ def get_dataset(load_function, base_dir, labels, **args):
         dataset = dataset.batch(batch_size)
     dataset = dataset.prefetch(buffer_size=AUTOTUNE)
 
-    return dataset, remapped, new_labels
+    return dataset, remapped, new_labels, epoch_size
 
 
 def resample(dataset, labels):
diff --git a/ml_tools/tfwriter.py b/ml_tools/tfwriter.py
new file mode 100644
index 00000000..d1bf690f
--- /dev/null
+++ b/ml_tools/tfwriter.py
@@ -0,0 +1,142 @@
+# Copyright 2020 Google Research. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from PIL import Image
+from pathlib import Path
+from multiprocessing import Process, Queue
+
+import collections
+import hashlib
+import io
+import json
+import multiprocessing
+import os
+import time
+from absl import app
+from absl import flags
+from absl import logging
+import numpy as np
+from PIL import Image, ImageOps
+
+import tensorflow as tf
+from . import tfrecord_util
+from ml_tools import tools
+from ml_tools.imageprocessing import normalize, rotate
+from track.cliptracker import get_diff_back_filtered
+import cv2
+import random
+import math
+
+
+def process_job(queue, labels, base_dir, save_data, extra_args):
+    import gc
+
+    pid = os.getpid()
+
+    writer_i = 1
+    name = f"{writer_i}-{pid}.tfrecord"
+
+    options = tf.io.TFRecordOptions(compression_type="GZIP")
+    writer = tf.io.TFRecordWriter(str(base_dir / name), options=options)
+    i = 0
+    saved = 0
+    files = 0
+    while True:
+        i += 1
+        samples = queue.get()
+        try:
+            if samples == "DONE":
+                writer.close()
+                break
+            else:
+                if len(samples) == 0:
+                    continue
+                saved += save_data(samples, writer, labels, extra_args)
+                files += 1
+                del samples
+                if saved > 10000:
+                    logging.info("Closing old writer")
+                    writer.close()
+                    writer_i += 1
+                    name = f"{writer_i}-{pid}.tfrecord"
+                    logging.info("Opening %s", name)
+                    saved = 0
+                    writer = tf.io.TFRecordWriter(str(base_dir / name), options=options)
+                if i % 100 == 0:
+                    logging.info("Saved %s ", files)
+                    gc.collect()
+                    writer.flush()
+        except:
+            logging.error("Process_job error %s", samples[0].source_file, exc_info=True)
+
+
+def create_tf_records(
+    dataset,
+    output_path,
+    labels,
+    save_data,
+    num_shards=1,
+    augment=False,
+    **extra_args,
+):
+    output_path = Path(output_path)
+    if output_path.is_dir():
+        logging.info("Clearing dir %s", output_path)
+        for child in output_path.glob("*"):
+            if child.is_file():
+                child.unlink()
+    output_path.mkdir(parents=True, exist_ok=True)
+    samples_by_source = dataset.get_samples_by_source()
+    source_files = list(samples_by_source.keys())
+    np.random.shuffle(source_files)
+
+    num_labels = len(dataset.labels)
+    logging.info(
+        "writing to output path: %s for %s samples", output_path, len(samples_by_source)
+    )
+    num_processes = 1
+    try:
+        job_queue = Queue()
+        processes = []
+        for i in range(num_processes):
+            p = Process(
+                target=process_job,
+                args=(job_queue, labels, output_path, save_data, extra_args),
+            )
+            processes.append(p)
+            p.start()
+            added = 0
+        for source_file in source_files:
+            job_queue.put((samples_by_source[source_file]))
+            added += 1
+            while job_queue.qsize() > num_processes * 3:
+                logging.info("Sleeping for %s", 10)
+                # give it a change to catch up
+                time.sleep(10)
+
+        logging.info("Processing %d", job_queue.qsize())
+        for i in range(len(processes)):
+            job_queue.put(("DONE"))
+        for process in processes:
+            try:
+                process.join()
+            except KeyboardInterrupt:
+                logging.info("KeyboardInterrupt, terminating.")
+                for process in processes:
+                    process.terminate()
+                exit()
+        logging.info("Saved %s", len(dataset.samples_by_id))
+
+    except:
+        logging.error("Error saving track info", exc_info=True)
diff --git a/ml_tools/thermaldataset.py b/ml_tools/thermaldataset.py
index d8f6c63c..d7ddfbd0 100644
--- a/ml_tools/thermaldataset.py
+++ b/ml_tools/thermaldataset.py
@@ -24,22 +24,65 @@
 
 
 def get_excluded():
-    return []
+    return [
+        "goat",
+        "lizard",
+        "not identifiable",
+        "other",
+        "pest",
+        "pig",
+        "sealion",
+    ]
+
+
+def get_remapped(multi_label=False):
+    land_bird = "land-bird" if multi_label else "bird"
+    return {
+        "insect": "false-positive",
+        "allbirds": "bird",
+        "black swan": land_bird,
+        "brown quail": land_bird,
+        "california quail": land_bird,
+        "duck": land_bird,
+        "pheasant": land_bird,
+        "pukeko": land_bird,
+        "quail": land_bird,
+    }
 
 
-def get_remapped():
-    return {"insect": "false-positive"}
+def get_extra_mappings(labels):
+    land_birds = ["land-bird"]
+    if "bird" not in labels:
+        return None
+    bird_index = labels.index("bird")
+    values = []
+    keys = []
+    for l in land_birds:
+        if l in labels:
+            l_i = labels.index(l)
+            keys.append(l_i)
+            values.append(bird_index)
+    extra_label_map = tf.lookup.StaticHashTable(
+        initializer=tf.lookup.KeyValueTensorInitializer(
+            keys=tf.constant(keys),
+            values=tf.constant(values),
+        ),
+        default_value=tf.constant(-1),
+        name="extra_label_map",
+    )
+    logging.info("Extra label mapping is %s to %s ", keys, values)
+    return extra_label_map
 
 
 def load_dataset(filenames, remap_lookup, labels, args):
     deterministic = args.get("deterministic", False)
-
     ignore_order = tf.data.Options()
     ignore_order.experimental_deterministic = (
         deterministic  # disable order, increase speed
     )
-    dataset = tf.data.TFRecordDataset(filenames)
-
+    dataset = tf.data.TFRecordDataset(
+        filenames, compression_type="GZIP", num_parallel_reads=4
+    )
     dataset = dataset.with_options(
         ignore_order
     )  # uses data as soon as it streams in, rather than in its original order
@@ -52,19 +95,23 @@ def load_dataset(filenames, remap_lookup, labels, args):
     only_features = args.get("only_features", False)
     one_hot = args.get("one_hot", True)
     dataset = dataset.apply(tf.data.experimental.ignore_errors())
-
+    extra_label_map = None
+    if args.get("multi_label"):
+        extra_label_map = get_extra_mappings(labels)
+        logging.info("Using multi label")
     dataset = dataset.map(
         partial(
             read_tfrecord,
             image_size=image_size,
             remap_lookup=remap_lookup,
-            num_labels=len(new_labels),
+            num_labels=len(labels),
             labeled=labeled,
             augment=augment,
             preprocess_fn=preprocess_fn,
             include_features=include_features,
             only_features=only_features,
             one_hot=one_hot,
+            extra_label_map=extra_label_map,
         ),
         num_parallel_calls=AUTOTUNE,
         deterministic=deterministic,
@@ -76,10 +123,21 @@ def load_dataset(filenames, remap_lookup, labels, args):
 
     dataset = dataset.filter(filter_nan)
 
-    filter_excluded = lambda x, y: not tf.math.equal(tf.math.count_nonzero(y), 0)
-    dataset = dataset.filter(filter_excluded)
+    return dataset
+
 
-    return dataset, remapped, new_labels
+rotation_augmentation = tf.keras.Sequential(
+    [
+        tf.keras.layers.RandomRotation(0.1, fill_mode="nearest", fill_value=0),
+    ]
+)
+data_augmentation = tf.keras.Sequential(
+    [
+        tf.keras.layers.RandomFlip("horizontal"),
+        tf.keras.layers.RandomBrightness(0.2),  # better per frame or per sequence??
+        tf.keras.layers.RandomContrast(0.5),
+    ]
+)
 
 
 def read_tfrecord(
@@ -93,6 +151,7 @@ def read_tfrecord(
     only_features=False,
     one_hot=True,
     include_features=False,
+    extra_label_map=None,
 ):
     logging.info(
         "Read tf record with image %s lbls %s labeld %s aug  %s  prepr %s only features %s one hot %s include fetures %s",
@@ -130,11 +189,7 @@ def read_tfrecord(
         thermals = tf.reshape(thermalencoded, [25, 32, 32, 1])
         filtered = tf.reshape(filteredencoded, [25, 32, 32, 1])
         rgb_images = tf.concat((thermals, thermals, filtered), axis=3)
-        rotation_augmentation = tf.keras.Sequential(
-            [
-                tf.keras.layers.RandomRotation(0.1, fill_mode="nearest", fill_value=0),
-            ]
-        )
+
         # rotation augmentation before tiling
         if augment:
             logging.info("Augmenting")
@@ -143,15 +198,6 @@ def read_tfrecord(
         image = tile_images(rgb_images)
 
         if augment:
-            data_augmentation = tf.keras.Sequential(
-                [
-                    tf.keras.layers.RandomFlip("horizontal"),
-                    tf.keras.layers.RandomBrightness(
-                        0.2
-                    ),  # better per frame or per sequence??
-                    tf.keras.layers.RandomContrast(0.5),
-                ]
-            )
             image = data_augmentation(image)
         if preprocess_fn is not None:
             logging.info(
@@ -163,12 +209,13 @@ def read_tfrecord(
     if labeled:
         label = tf.cast(example["image/class/label"], tf.int32)
         label = remap_lookup.lookup(label)
-
+        if extra_label_map is not None:
+            extra = extra_label_map.lookup(label)
+            label = tf.stack([label, extra], axis=0)
         if one_hot:
-            if tf.math.equal(label, -1):
-                label = tf.zeros(num_labels)
-            else:
-                label = tf.one_hot(label, num_labels)
+            label = tf.one_hot(label, num_labels)
+            if extra_label_map is not None:
+                label = tf.reduce_max(label, axis=0)
         if include_features or only_features:
             features = example["image/features"]
             if only_features:
@@ -214,6 +261,8 @@ def tile_images(images):
 def main():
     init_logging()
     config = Config.load_from_file()
+    from .tfdataset import get_dataset, get_distribution
+
     # file = "/home/gp/cacophony/classifier-data/thermal-training/cp-training/training-meta.json"
     file = f"{config.tracks_folder}/training-meta.json"
     with open(file, "r") as f:
@@ -221,29 +270,26 @@ def main():
     labels = meta.get("labels", [])
     datasets = []
 
-    resampled_ds, remapped, labels = get_dataset(
+    resampled_ds, remapped, labels, epoch_size = get_dataset(
         # dir,
-        f"{config.tracks_folder}/training-data/validation",
+        load_dataset,
+        f"{config.tracks_folder}/training-data/train",
         labels,
         batch_size=32,
         image_size=(160, 160),
         # augment=True,
         # preprocess_fn=tf.keras.applications.inception_v3.preprocess_input,
         resample=False,
-        include_features=True,
+        include_features=False,
+        remapped_labels=get_remapped(),
+        excluded_labels=get_excluded(),
     )
-    # print(get_distribution(resampled_ds))
-    #
+    print("Ecpoh size is", epoch_size)
+    print(get_distribution(resampled_ds, len(labels)))
+    # return
     #
     for e in range(2):
         print("epoch", e)
-        true_categories = [y for x, y in resampled_ds]
-        true_categories = tf.concat(true_categories, axis=0)
-        true_categories = np.int64(tf.argmax(true_categories, axis=1))
-        c = Counter(list(true_categories))
-        print("epoch is size", len(true_categories))
-        for i in range(len(labels)):
-            print("after have", labels[i], c[i])
         for x, y in resampled_ds:
             show_batch(x, y, labels)
 
@@ -251,7 +297,6 @@ def main():
 
 
 def show_batch(image_batch, label_batch, labels):
-    image_batch = image_batch[0]
     plt.figure(figsize=(10, 10))
     print("images in batch", len(image_batch), len(label_batch))
     num_images = min(len(image_batch), 25)
diff --git a/ml_tools/thermalwriter.py b/ml_tools/thermalwriter.py
index 531412bd..d25cb024 100644
--- a/ml_tools/thermalwriter.py
+++ b/ml_tools/thermalwriter.py
@@ -25,13 +25,13 @@
 """
 from PIL import Image
 from pathlib import Path
-
+import time
 import collections
 import hashlib
 import io
 import json
-import multiprocessing
 import os
+from multiprocessing import Process, Queue
 
 from absl import app
 from absl import flags
@@ -46,12 +46,14 @@
 from ml_tools.forestmodel import forest_features
 from ml_tools import imageprocessing
 from ml_tools.frame import TrackChannels
+from ml_tools.trackdatabase import TrackDatabase
+from ml_tools.rawdb import RawDatabase
 
 crop_rectangle = tools.Rectangle(0, 0, 640, 480)
 from functools import lru_cache
 
 
-def create_tf_example(data, image_dir, sample, labels, filename, num_frames):
+def create_tf_example(sample, data, features, labels, num_frames):
     """Converts image and annotations to a tf.Example proto.
 
     Args:
@@ -84,8 +86,6 @@ def create_tf_example(data, image_dir, sample, labels, filename, num_frames):
     """
     average_dim = [r.area for r in sample.track_bounds]
     average_dim = int(round(np.mean(average_dim) ** 0.5))
-    features = data[1]
-    data = data[0]
     thermals = list(data[0])
     filtereds = list(data[1])
     image_id = sample.unique_track_id
@@ -98,14 +98,21 @@ def create_tf_example(data, image_dir, sample, labels, filename, num_frames):
     filtereds = np.array(filtereds)
     thermal_key = hashlib.sha256(thermals).hexdigest()
     filtered_key = hashlib.sha256(filtereds).hexdigest()
-
+    avg_mass = int(round(sample.mass / len(sample.frame_numbers)))
     feature_dict = {
+        "image/filtered": tfrecord_util.int64_feature(1 if sample.filtered else 0),
+        "image/avg_mass": tfrecord_util.int64_feature(avg_mass),
+        "image/track_median_mass": tfrecord_util.int64_feature(
+            int(sample.track_median_mass)
+        ),
         "image/avg_dim": tfrecord_util.int64_feature(average_dim),
         "image/height": tfrecord_util.int64_feature(image_height),
         "image/width": tfrecord_util.int64_feature(image_width),
         "image/clip_id": tfrecord_util.int64_feature(sample.clip_id),
         "image/track_id": tfrecord_util.int64_feature(sample.track_id),
-        "image/filename": tfrecord_util.bytes_feature(filename.encode("utf8")),
+        "image/filename": tfrecord_util.bytes_feature(
+            str(sample.source_file).encode("utf8")
+        ),
         "image/source_id": tfrecord_util.bytes_feature(str(image_id).encode("utf8")),
         "image/thermalencoded": tfrecord_util.float_list_feature(thermals.ravel()),
         "image/filteredencoded": tfrecord_util.float_list_feature(filtereds.ravel()),
@@ -122,111 +129,16 @@ def create_tf_example(data, image_dir, sample, labels, filename, num_frames):
     }
 
     example = tf.train.Example(features=tf.train.Features(feature=feature_dict))
-    return example, 0
-
-
-def create_tf_records(
-    dataset,
-    output_path,
-    labels,
-    back_thresh,
-    num_shards=1,
-    cropped=True,
-    by_label=False,
-):
-    output_path = Path(output_path)
-    if output_path.is_dir():
-        logging.info("Clearing dir %s", output_path)
-        for child in output_path.glob("*"):
-            if child.is_file():
-                child.unlink()
-    output_path.mkdir(parents=True, exist_ok=True)
-    samples = dataset.samples
-    # keys = list(samples.keys())
-    np.random.shuffle(samples)
-    dataset.load_db()
-    dataset.set_read_only(True)
-    db = dataset.db
-    total_num_annotations_skipped = 0
-    num_labels = len(labels)
-    # pool = multiprocessing.Pool(4)
-    logging.info("writing to output path: %s for %s samples", output_path, len(samples))
-    writers = []
-    lbl_counts = [0] * num_labels
-    # lbl_counts[l] = 0
-    logging.info("labels are %s", labels)
-
-    writers = []
-    for label in labels:
-        for i in range(num_shards):
-            if by_label:
-                safe_l = label.replace("/", "-")
-                name = f"{safe_l}-%05d-of-%05d.tfrecord" % (i, num_shards)
-            else:
-                name = f"%05d-of-%05d.tfrecord" % (i, num_shards)
-            writers.append(tf.io.TFRecordWriter(str(output_path / name)))
-        if not by_label:
-            break
-    load_first = 200
-    try:
-        count = 0
-        while len(samples) > 0:
-            local_set = samples[:load_first]
-            samples = samples[load_first:]
-            loaded = []
-            for sample in local_set:
-                data, features = get_data(sample, db)
-                # features = forest_features(db, sample.clip_id, sample.track_id)
-                if data is None:
-                    continue
+    return example
 
-                loaded.append(((data, features), sample))
-            loaded = np.array(loaded, dtype=object)
-            np.random.shuffle(loaded)
 
-            for data, sample in loaded:
-                try:
-                    logging.info("Writing out data")
-                    tf_example, num_annotations_skipped = create_tf_example(
-                        data, output_path, sample, labels, "", dataset.segment_length
-                    )
-                    total_num_annotations_skipped += num_annotations_skipped
-                    l_i = labels.index(sample.label)
-                    if by_label:
-                        wrtier = writers[
-                            num_shards * l_i + lbl_counts[l_i] % num_shards
-                        ]
-                    else:
-                        writer = writers[count % num_shards]
-                    writer.write(tf_example.SerializeToString())
-                    lbl_counts[l_i] += 1
-                    # print("saving example", [count % num_shards])
-                    count += 1
-                    if count % 100 == 0:
-                        logging.info("saved %s", count)
-                    # count += 1
-                except Exception as e:
-                    logging.error("Error saving ", exc_info=True)
-
-    except:
-        logging.error("Error saving track info", exc_info=True)
-    for writer in writers:
-        writer.close()
-
-    logging.info(
-        "Finished writing, skipped %d annotations.", total_num_annotations_skipped
-    )
-
-
-@lru_cache(maxsize=1000)
+# @lru_cache(maxsize=1000)
+# only going to get called once now per track
 def get_track_data(clip_id, track_id, db):
     background = db.get_clip_background(clip_id)
 
     track_frames = db.get_track(
-        clip_id,
-        track_id,
-        original=False,
-        channels=[TrackChannels.thermal],
+        clip_id, track_id, channels=[TrackChannels.thermal], crop=True
     )
     clip_meta = db.get_clip_meta(clip_id)
     frame_temp_median = clip_meta["frame_temp_median"]
@@ -237,41 +149,173 @@ def get_track_data(clip_id, track_id, db):
     return background, track_frames, features, frame_temp_median
 
 
-def get_data(sample, db):
+def save_data(samples, writer, labels, extra_args):
+    sample_data = get_data(samples, extra_args)
+    if sample_data is None:
+        return 0
+    saved = 0
+    try:
+        for data in sample_data:
+            tf_example = create_tf_example(
+                data[0], data[1], data[2], labels, extra_args["num_frames"]
+            )
+            writer.write(tf_example.SerializeToString())
+            saved += 1
+    except:
+        logging.error(
+            "Could not save data for %s", samples[0].source_file, exc_info=True
+        )
+    return saved
+
+
+def get_data(clip_samples, extra_args):
     # prepare the sample data for saving
+    if len(clip_samples) == 0:
+        return None
+    data = []
     crop_rectangle = tools.Rectangle(2, 2, 160 - 2 * 2, 140 - 2 * 2)
+    if clip_samples[0].source_file.stem == ".hdf5":
+        db = TrackDatabase(clip_samples[0].source_file)
+    else:
+        db = RawDatabase(clip_samples[0].source_file)
+        db.load_frames()
+        # going to redo segments to get rid of ffc segments
+
+    clip_id = clip_samples[0].clip_id
     try:
-        background, track_frames, features, frame_temp_median = get_track_data(
-            sample.clip_id, sample.track_id, db
-        )
+        background = db.get_clip_background()
+        if background is None:
+            frame_data = db.get_frames()
+            background = np.median(frame_data, axis=0)
+            del frame_data
+        clip_meta = db.get_clip_meta(extra_args.get("tag_precedence"))
+        frame_temp_median = clip_meta.frame_temp_median
+
+        # group samples by track_id
+        samples_by_track = {}
+        for s in clip_samples:
+            samples_by_track.setdefault(s.track_id, []).append(s)
+
+        for track_id in samples_by_track.keys():
+            samples = samples_by_track[track_id]
+            if clip_samples[0].source_file.stem != ".hdf5":
+                track = next(
+                    (track for track in clip_meta.tracks if track.track_id == track_id),
+                    None,
+                )
+                if track is None:
+                    logging.error(
+                        "Cannot find track %s in clip %s", track_id, clip_meta.clip_id
+                    )
+                    continue
+
+                # GP All assumes we dont have a track over multiple bins (Whcih we probably never want)
 
-        thermals = []  # np.empty(len(frames), dtype=object)
-        filtered = []  # np.empty(len(frames), dtype=object)
-        for i, frame in enumerate(track_frames):
-            if frame.frame_number not in sample.frame_indices:
-                continue
-            frame.float_arrays()
-            frame.filtered = frame.thermal - frame.region.subimage(background)
-            temp = frame_temp_median[i]
-            frame.resize_with_aspect((32, 32), crop_rectangle, keep_edge=True)
-            frame.thermal -= temp
-            np.clip(frame.thermal, a_min=0, a_max=None, out=frame.thermal)
+                track.calculate_segments(
+                    extra_args.get("segment_frame_spacing", 9),
+                    extra_args.get("segment_width", 25),
+                    extra_args.get("segment_type"),
+                    extra_args.get("segment_min_avg_mass"),
+                    max_segments=extra_args.get("max_segments"),
+                    dont_filter=extra_args.get("dont_filter_segment", False),
+                    skip_ffc=extra_args.get("skip_ffc", True),
+                    ffc_frames=clip_meta.ffc_frames,
+                )
+                samples = track.samples
+                frame_temp_median = {}
+                track_frames = []
+                for frame_i in range(
+                    track.start_frame, track.start_frame + track.num_frames
+                ):
+                    f = db.frames[frame_i]
+                    f.region = track.regions_by_frame[frame_i]
+                    frame_temp_median[frame_i] = np.median(f.thermal)
+                    f.thermal = f.region.subimage(f.thermal)
+                    f.float_arrays()
+                    track_frames.append(f)
 
-            frame.thermal, stats = imageprocessing.normalize(frame.thermal, new_max=255)
-            if not stats[0]:
-                frame.thermal = np.zeros((frame.thermal.shape))
-                # continue
-            frame.filtered, stats = imageprocessing.normalize(
-                frame.filtered, new_max=255
+            else:
+                track_frames = db.get_track(
+                    clip_id, track_id, channels=[TrackChannels.thermal], crop=True
+                )
+            logging.debug("Saving %s samples %s", track_id, len(samples))
+            used_frames = []
+
+            features = forest_features(
+                track_frames,
+                background,
+                frame_temp_median,
+                [f.region for f in track_frames],
+                normalize=True,
+                cropped=True,
             )
-            if not stats[0]:
-                frame.filtered = np.zeros((frame.filtered.shape))
 
-            filtered.append(frame.filtered)
-            thermals.append(frame.thermal)
-        thermals = np.array(thermals)
-        filtered = np.array(filtered)
+            by_frame_number = {}
+            max_diff = 0
+            min_diff = 0
+            for f in track_frames:
+                if f.region.blank or f.region.width <= 0 or f.region.height <= 0:
+                    continue
+
+                by_frame_number[f.frame_number] = (f, frame_temp_median[f.frame_number])
+                f.float_arrays()
+                diff_frame = f.thermal - f.region.subimage(background)
+                new_max = np.amax(diff_frame)
+                new_min = np.amin(diff_frame)
+                if new_min < min_diff:
+                    min_diff = new_min
+                if new_max > max_diff:
+                    max_diff = new_max
+
+            # normalize by maximum difference between background and tracked region
+            # probably only need to use difference on the frames used for this record
+            # also min_diff maybe could just be set to 0 and clip values below 0,
+            # these represent pixels whcih are cooler than the background
+            for sample in samples:
+                thermals = []  # np.empty(len(frames), dtype=object)
+                filtered = []  # np.empty(len(frames), dtype=object)
+                for frame_number in sample.frame_indices:
+                    frame, temp_median = by_frame_number[frame_number]
+                    # no need to do work twice
+                    if frame_number not in used_frames:
+                        used_frames.append(frame_number)
+                        frame.filtered = frame.thermal - frame.region.subimage(
+                            background
+                        )
+                        frame.resize_with_aspect(
+                            (32, 32), crop_rectangle, keep_edge=True
+                        )
+                        frame.thermal -= temp_median
+                        np.clip(frame.thermal, a_min=0, a_max=None, out=frame.thermal)
+
+                        frame.thermal, stats = imageprocessing.normalize(
+                            frame.thermal, new_max=255
+                        )
+                        if not stats[0]:
+                            frame.thermal = np.zeros((frame.thermal.shape))
+                            # continue
+                        # f2 = frame.filtered.copy()
+                        # frame.filtered, stats = imageprocessing.normalize(
+                        #     frame.filtered, new_max=255
+                        # )
+                        # np.clip(frame.filtered, a_min=min_diff, a_max=None, out=frame.filtered)
+
+                        frame.filtered, stats = imageprocessing.normalize(
+                            frame.filtered, min=min_diff, max=max_diff, new_max=255
+                        )
+
+                        if not stats[0]:
+                            frame.filtered = np.zeros((frame.filtered.shape))
+                        f2 = np.uint8(frame.filtered)
+                    filtered.append(frame.filtered)
+                    thermals.append(frame.thermal)
+
+                thermals = np.array(thermals)
+                filtered = np.array(filtered)
+                data.append((sample, (thermals, filtered), features))
     except:
-        logging.error("Cant get segment %s", sample, exc_info=True)
-        return None, None
-    return (thermals, filtered), features
+        logging.error(
+            "Cant get Samples for %s", clip_samples[0].source_file, exc_info=True
+        )
+        return None
+    return data
diff --git a/ml_tools/trackdatabase.py b/ml_tools/trackdatabase.py
index 3eee6b92..bcf2e98e 100644
--- a/ml_tools/trackdatabase.py
+++ b/ml_tools/trackdatabase.py
@@ -15,10 +15,10 @@
 from dateutil.parser import parse as parse_date
 from .frame import Frame, TrackChannels
 import json
-from dateutil.parser import parse as parse_date
 
 import numpy as np
 from track.region import Region
+from ml_tools.datasetstructures import TrackHeader, ClipHeader
 
 special_datasets = [
     "tag_frames",
@@ -68,11 +68,11 @@ def __init__(self, database_filename, read_only=False):
         """
 
         self.database = database_filename
-        if not os.path.exists(database_filename):
-            logging.info("Creating new database %s", database_filename)
-            f = h5py.File(database_filename, "w")
-            f.create_group("clips")
-            f.close()
+        # if not os.path.exists(database_filename):
+        #     logging.info("Creating new database %s", database_filename)
+        #     f = h5py.File(database_filename, "w")
+        #     f.create_group("clips")
+        #     f.close()
         HDF5Manager.READ_ONLY = read_only
 
     def set_read_only(self, read_only):
@@ -386,17 +386,26 @@ def get_track_meta(self, clip_id, track_id):
 
         return result
 
-    def dataset_track(self, dataset, track_id):
+    def dataset_track(self, dataset, clip_id, track_id):
         result = hdf5_attributes_dictionary(dataset)
-        preds = dataset.get("model_predictions")
-        if preds is not None:
-            result["model_predictions"] = {}
-            for model in preds:
-                model_preds = hdf5_attributes_dictionary(preds[model])
-                result["model_predictions"][model] = model_preds
-        result["id"] = track_id
-
-        return result
+        regions = dataset.get("regions")[:]
+        regions_by_frame = {}
+        for r in regions:
+            r = Region.region_from_array(r)
+            regions_by_frame[r.frame_number] = r
+        header = TrackHeader(
+            clip_id=clip_id,
+            track_id=int(track_id),
+            label=result.get("human_tag"),
+            num_frames=len(regions),
+            regions=regions_by_frame,
+            start_frame=result["start_frame"],
+            confidence=result.get("human_tag_confidence"),
+            human_tags=result.get("human_tags"),
+            source_file=self.database,
+            # frame_temp_median=frame_temp_median,
+        )
+        return header
 
     def get_track_predictions(self, clip_id, track_id):
         """
@@ -412,13 +421,13 @@ def get_track_predictions(self, clip_id, track_id):
         return None
 
     def get_clip_background(self, clip_id):
-        with HDF5Manager(self.database) as f:
-            clip = f["clips"][str(clip_id)]
-            if "background_frame" in clip:
-                return clip["background_frame"][:]
+        with HDF5Manager(self.database) as clip:
+            frames = clip["frames"]
+            if "background" in frames:
+                return frames["background"][:]
         return None
 
-    def get_clip_meta(self, clip_id):
+    def get_clip_meta(self, clip_id=None):
         """
         Gets metadata for given clip
         :param clip_id:
@@ -426,37 +435,49 @@ def get_clip_meta(self, clip_id):
         """
 
         with HDF5Manager(self.database) as f:
-            dataset = f["clips"][str(clip_id)]
+            if "clips" in f:
+                dataset = f["clips"][str(clip_id)]
+            else:
+                dataset = f
             result = hdf5_attributes_dictionary(dataset)
-            result["tracks"] = len(dataset)
-            tag_frames = dataset.get("tag_frames")
-            if tag_frames:
-                result["tag_frames"] = {}
-                for key, value in tag_frames.attrs.items():
-                    result["tag_frames"][key] = value
-
-                tag_regions = tag_frames.get("tag_regions")
-                if tag_regions is not None:
-                    result["tag_frames"]["tag_regions"] = {}
-                    for key, value in tag_regions.attrs.items():
-                        result["tag_frames"]["tag_regions"][key] = value
-        return result
+            clip_header = ClipHeader(
+                clip_id=int(result["clip_id"]),
+                station_id=result.get("station_id"),
+                source_file=self.database,
+                location=result.get("station_id"),
+                camera=result.get("device_id"),
+                frames_per_second=9,
+                events=result.get("event", ""),
+                trap=result.get("trap", ""),
+                tracks=[],
+                ffc_frames=result.get("ffc_frames"),
+                rec_time=None,
+                frame_temp_median=result.get("frame_temp_median"),
+            )
+        return clip_header
 
-    def get_clip_tracks(self, clip_id):
+    def get_clip_tracks(self, clip_id=None):
         """
         Gets metadata for given clip
         :param clip_id:
         :return:
         """
         tracks = []
+        clip_header = self.get_clip_meta()
         with HDF5Manager(self.database) as f:
-            clip = f["clips"][str(clip_id)]
+            if "clips" in f:
+                clip = f["clips"][str(clip_id)]
+            else:
+                clip = f["tracks"]
             for track_id in clip:
                 if track_id in special_datasets:
                     continue
-                track = self.dataset_track(clip[track_id], track_id)
-                tracks.append(track)
-        return tracks
+                track = self.dataset_track(
+                    clip[track_id], clip_header.clip_id, track_id
+                )
+                track.station_id = clip_header.station_id
+                clip_header.tracks.append(track)
+        return clip_header
 
     def get_tag(self, clip_id, track_id):
         with HDF5Manager(self.database) as f:
@@ -472,31 +493,21 @@ def get_frame(self, clip_id, track_id, frame, original=False):
             return frames[0]
         return None
 
-    def get_clip(
+    def get_frames(
         self,
         clip_id,
         frame_numbers=None,
         channels=None,
     ):
-        frames = []
-        with HDF5Manager(self.database) as f:
-            clip = f["clips"][str(clip_id)]
-            if "original_frames" not in clip:
-                return None
-            frames_node = clip["original_frames"]
-            if frame_numbers is None:
-                frame_numbers = []
-                for f_i in frames_node:
-                    frame_numbers.append(int(f_i))
-                frame_numbers.sort()
-            frame_iter = iter(frame_numbers)
-
-            for frame_number in frame_iter:
-                frame = frames_node[str(frame_number)][:, :]
-                frames.append(
-                    Frame.from_channels([frame], [TrackChannels.thermal], frame_number)
-                )
-        return frames
+        # frames = []
+        with HDF5Manager(self.database) as clip_node:
+            raw_frames = clip_node["frames"]["thermals"][:]
+            return raw_frames
+        #     for frame_number, frame in enumerate(raw_frames):
+        #         frames.append(
+        #             Frame.from_channels([frame], [TrackChannels.thermal], frame_number)
+        #         )
+        # return frames
 
     def get_track(
         self,
@@ -504,9 +515,9 @@ def get_track(
         track_id,
         start_frame=None,
         end_frame=None,
-        original=False,
         frame_numbers=None,
         channels=None,
+        crop=False,
     ):
         """
         Fetches a track data from database with optional slicing.
@@ -516,98 +527,38 @@ def get_track(
         :param end_frame: last frame of slice to return (exclusive).
         :return: a list of numpy arrays of shape [channels, height, width] and of type np.int16
         """
-        with HDF5Manager(self.database) as f:
-            clips = f["clips"]
-            clip_node = clips[str(clip_id)]
-            track_node = clip_node[str(track_id)]
+        with HDF5Manager(self.database) as clip_node:
+            tracks = clip_node["tracks"]
+            track_node = tracks[str(track_id)]
 
-            bounds = track_node.attrs["bounds_history"]
+            bounds = track_node["regions"]
+            track_start = int(round(track_node.attrs.get("start_frame")))
             if start_frame is None:
-                start_frame = 0
+                start_frame = track_start
             if end_frame is None:
-                end_frame = track_node.attrs["frames"]
-            track_start = track_node.attrs.get("start_frame")
+                end_frame = int(round(track_node.attrs["end_frame"])) + 1
             bad_frames = track_node.attrs.get("skipepd_frames", [])
             result = []
-            if original:
-                track_node = clip_node["original_frames"]
-            else:
-                if "cropped" in track_node:
-                    track_node = track_node["cropped"]
 
             if frame_numbers is None:
                 frame_iter = range(start_frame, end_frame)
             else:
                 frame_iter = iter(frame_numbers)
-
+            frames = clip_node["frames"]["thermals"][:]
             for frame_number in frame_iter:
-                if original:
-                    region = Region.region_from_array(bounds[frame_number])
-                    region.frame_number = frame_number + track_start
-                    frame = track_node[str(frame_number + track_start)][:, :]
-                    result.append(
-                        Frame.from_channels(
-                            [frame],
-                            [TrackChannels.thermal],
-                            frame_number + track_start,
-                            region=region,
-                        )
+                region = Region.region_from_array(bounds[frame_number - track_start])
+                # assert region.frame_number == frame_number
+                region.frame_number = frame_number
+                frame = frames[frame_number]
+                frame = region.subimage(frame)
+                result.append(
+                    Frame.from_channels(
+                        [frame],
+                        [TrackChannels.thermal],
+                        frame_number,
+                        region=region,
                     )
-                else:
-                    if frame_number in bad_frames:
-                        continue
-                    region = Region.region_from_array(bounds[frame_number])
-                    if channels is None:
-                        try:
-                            frame = track_node[str(frame_number)][:, :, :]
-                            if frame.shape[0] < 5:
-                                frame_channels = [
-                                    TrackChannels.thermal,
-                                    TrackChannels.filtered,
-                                ]
-                                if frame.shape[0] == 3:
-                                    frame_channels.append(TrackChannels.mask)
-                                f = Frame.from_channels(
-                                    frame,
-                                    frame_channels,
-                                    frame_number + track_start,
-                                    region=region,
-                                )
-                            else:
-                                f = Frame.from_array(
-                                    frame,
-                                    frame_number + track_start,
-                                    flow_clipped=True,
-                                    region=region,
-                                )
-                            result.append(f)
-                        except:
-                            logging.debug(
-                                "trying to get clip %s track %s frame %s",
-                                clip_id,
-                                track_id,
-                                frame_number + track_start,
-                                exc_info=True,
-                            )
-                    else:
-                        try:
-                            frame = track_node[str(frame_number)][channels, :, :]
-                            result.append(
-                                Frame.from_channels(
-                                    frame,
-                                    channels,
-                                    frame_number + track_start,
-                                    region=region,
-                                )
-                            )
-                        except:
-                            logging.debug(
-                                "trying to get clip %s track %s frame %s",
-                                clip_id,
-                                track_id,
-                                frame_number,
-                                exc_info=True,
-                            )
+                )
 
         return result
 
diff --git a/mldataset/makedataset.py b/mldataset/makedataset.py
new file mode 100644
index 00000000..65368843
--- /dev/null
+++ b/mldataset/makedataset.py
@@ -0,0 +1,341 @@
+"" """
+classifier-pipeline - this is a server side component that manipulates cptv
+files and to create a classification model of animals present
+Copyright (C) 2018, The Cacophony Project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program. If not, see <http://www.gnu.org/licenses/>.
+"""
+
+import os
+import logging
+import time
+from multiprocessing import Process, Queue
+import traceback
+from ml_tools import tools
+
+from track.clip import Clip, ClipStats
+import numpy as np
+import json
+from pathlib import Path
+
+import h5py
+from cptv import CPTVReader
+import yaml
+from track.cliptrackextractor import is_affected_by_ffc
+from track.track import Track
+from track.region import Region
+
+FPS = 9
+
+
+def process_job(loader, queue, out_dir, config):
+    i = 0
+    while True:
+        i += 1
+        filename = queue.get()
+        logging.info("Processing %s", filename)
+        try:
+            if filename == "DONE":
+                break
+            else:
+                loader.process_file(str(filename), out_dir, config)
+            if i % 50 == 0:
+                logging.info("%s jobs left", queue.qsize())
+        except Exception as e:
+            logging.error("Process_job error %s %s", filename, e)
+            traceback.print_exc()
+
+
+class ClipLoader:
+    def __init__(self, config):
+        self.config = config
+
+        # number of threads to use when processing jobs.
+        self.workers_threads = 9
+
+    def process_all(self, root, out_dir):
+        out_dir.mkdir(parents=True, exist_ok=True)
+        job_queue = Queue()
+        processes = []
+        for i in range(max(1, self.workers_threads)):
+            p = Process(
+                target=process_job,
+                args=(self, job_queue, out_dir, self.config),
+            )
+            processes.append(p)
+            p.start()
+
+        file_paths = []
+        #
+        for folder_path, _, files in os.walk(root):
+            for name in files:
+                if os.path.splitext(name)[1] in [".cptv"]:
+                    full_path = os.path.join(folder_path, name)
+                    file_paths.append(full_path)
+        # allows us know the order of processing
+        file_paths.sort()
+        for file_path in file_paths:
+            job_queue.put(file_path)
+
+        logging.info("Processing %d", job_queue.qsize())
+        for i in range(len(processes)):
+            job_queue.put("DONE")
+        for process in processes:
+            try:
+                process.join()
+            except KeyboardInterrupt:
+                logging.info("KeyboardInterrupt, terminating.")
+                for process in processes:
+                    process.terminate()
+                exit()
+
+    def process_file(self, filename, out_dir, config):
+        start = time.time()
+        filename = Path(filename)
+        logging.info(f"processing %s", filename)
+        metadata_file = filename.with_suffix(".txt")
+        if not metadata_file.exists():
+            logging.error("No meta data found for %s", metadata_file)
+            return
+
+        metadata = tools.load_clip_metadata(metadata_file)
+        r_id = metadata["id"]
+        out_file = out_dir / f"{r_id}.hdf5"
+        tracker_version = metadata.get("tracker_version")
+        logging.info("Tracker version is %s", tracker_version)
+        if out_file.exists() and tracker_version > 9:
+            logging.warning("Already loaded %s", filename)
+            # going to add some missing fierlds
+            with h5py.File(out_file, "a") as f:
+                if f.attrs.get("rec_time") is None:
+                    with open(filename, "rb") as cptv:
+                        reader = CPTVReader(cptv)
+                        video_start_time = reader.timestamp.astimezone(Clip.local_tz)
+                    f.attrs["rec_time"] = video_start_time.isoformat()
+            return
+        if len(metadata.get("Tracks")) == 0:
+            logging.error("No tracks found for %s", filename)
+            return
+
+        clip = Clip(config.tracking["thermal"], filename)
+        clip.load_metadata(
+            metadata,
+            config.load.tag_precedence,
+        )
+
+        with h5py.File(out_file, "w") as f:
+            try:
+                logging.info("creating clip %s", clip.get_id())
+
+                clip_id = str(clip.get_id())
+
+                clip_node = f
+                triggered_temp_thresh = None
+                camera_model = None
+                with open(clip.source_file, "rb") as f:
+                    reader = CPTVReader(f)
+                    clip.set_res(reader.x_resolution, reader.y_resolution)
+                    if clip.from_metadata:
+                        for track in clip.tracks:
+                            track.crop_regions()
+                    if reader.model:
+                        camera_model = reader.model.decode()
+                    clip.set_model(camera_model)
+
+                    # if we have the triggered motion threshold should use that
+                    # maybe even override dynamic threshold with this value
+                    if reader.motion_config:
+                        motion = yaml.safe_load(reader.motion_config)
+                        triggered_temp_thresh = motion.get("triggeredthresh")
+                        if triggered_temp_thresh:
+                            clip.temp_thresh = triggered_temp_thresh
+
+                    video_start_time = reader.timestamp.astimezone(Clip.local_tz)
+                    clip.set_video_stats(video_start_time)
+                    print("Adding rec time", video_start_time)
+                    frames = clip_node.create_group("frames")
+                    ffc_frames = []
+                    stats = ClipStats()
+                    cropped_stats = ClipStats()
+                    num_frames = 0
+                    cptv_frames = []
+                    region_adjust = 0
+                    for frame in reader:
+                        if frame.background_frame:
+                            back_node = frames.create_dataset(
+                                "background",
+                                frame.pix.shape,
+                                chunks=frame.pix.shape,
+                                compression="gzip",
+                                dtype=frame.pix.dtype,
+                            )
+                            back_node[:, :] = frame.pix
+                            # pre tracker verison 10 there was a bug where back frame was counted in region frame number
+                            #  so frame 0 is actually the background frame, no tracks shouild ever start at 0
+                            if tracker_version < 10:
+                                region_adjust = -1
+                                logging.info("Adjusting regions by %s", region_adjust)
+                            continue
+                        ffc = is_affected_by_ffc(frame)
+                        if ffc:
+                            ffc_frames.append(num_frames)
+
+                        cptv_frames.append(frame.pix)
+                        stats.add_frame(frame.pix)
+                        cropped_stats.add_frame(clip.crop_rectangle.subimage(frame.pix))
+                        num_frames += 1
+                    cptv_frames = np.uint16(cptv_frames)
+                    thermal_node = frames.create_dataset(
+                        "thermals",
+                        cptv_frames.shape,
+                        chunks=(1, *cptv_frames.shape[1:]),
+                        compression="gzip",
+                        dtype=cptv_frames.dtype,
+                    )
+
+                    thermal_node[:, :, :] = cptv_frames
+                    stats.completed()
+                    cropped_stats.completed()
+                    group_attrs = clip_node.attrs
+                    group_attrs["clip_id"] = r_id
+                    group_attrs["rec_time"] = video_start_time.isoformat()
+                    group_attrs["num_frames"] = np.uint16(num_frames)
+                    group_attrs["ffc_frames"] = np.uint16(ffc_frames)
+                    group_attrs["device_id"] = metadata["deviceId"]
+                    stationId = metadata.get("stationId", 0)
+                    group_attrs["station_id"] = stationId
+                    group_attrs["crop_rectangle"] = np.uint8(
+                        clip.crop_rectangle.to_ltrb()
+                    )
+                    group_attrs["max_temp"] = np.uint16(cropped_stats.max_temp)
+                    group_attrs["min_temp"] = np.uint16(cropped_stats.min_temp)
+                    group_attrs["mean_temp"] = np.uint16(cropped_stats.mean_temp)
+                    group_attrs["frame_temp_min"] = np.uint16(
+                        cropped_stats.frame_stats_min
+                    )
+                    group_attrs["frame_temp_max"] = np.uint16(
+                        cropped_stats.frame_stats_max
+                    )
+                    group_attrs["frame_temp_median"] = np.uint16(
+                        cropped_stats.frame_stats_median
+                    )
+                    group_attrs["frame_temp_mean"] = np.uint16(
+                        cropped_stats.frame_stats_mean
+                    )
+                    group_attrs["start_time"] = clip.video_start_time.isoformat()
+                    group_attrs["res_x"] = clip.res_x
+                    group_attrs["res_y"] = clip.res_y
+                    if camera_model is not None:
+                        group_attrs["model"] = clip.camera_model
+
+                    if triggered_temp_thresh is not None:
+                        group_attrs["temp_thresh"] = triggered_temp_thresh
+
+                    if clip.tags:
+                        clip_tags = []
+                        for track in clip.tags:
+                            if track["what"]:
+                                clip_tags.append(track["what"])
+                            elif track["detail"]:
+                                clip_tags.append(track["detail"])
+                        group_attrs["tags"] = clip_tags
+
+                tracks_group = clip_node.create_group("tracks")
+
+                tracks = metadata.get("Tracks", [])
+                for track in tracks:
+                    track_id = track["id"]
+
+                    track_group = tracks_group.create_group(str(track_id))
+
+                    node_attrs = track_group.attrs
+                    node_attrs["id"] = track_id
+                    tags = track.get("tags", [])
+                    tag = Track.get_best_human_tag(
+                        tags, self.config.load.tag_precedence, 0
+                    )
+
+                    master_tag = [
+                        t
+                        for t in tags
+                        if t.get("automatic")
+                        and not isinstance(t.get("data", ""), str)
+                        and t.get("data", {}).get("name") == "Master"
+                    ]
+                    if len(master_tag) > 0:
+                        master_tag = master_tag[0]
+                        node_attrs["ai_tag"] = master_tag["what"]
+                        node_attrs["ai_tag_confidence"] = master_tag["confidence"]
+
+                    if tag is not None:
+                        node_attrs["human_tag"] = tag["what"]
+                        node_attrs["human_tag_confidence"] = tag["confidence"]
+
+                    human_tags = [
+                        (t.get("what"), t["confidence"])
+                        for t in tags
+                        if t.get("automatic", False) != True
+                    ]
+                    if len(human_tags) > 0:
+                        node_attrs["human_tags"] = [h[0] for h in human_tags]
+                        node_attrs["human_tags_confidence"] = np.float32(
+                            [h[1] for h in human_tags]
+                        )
+
+                    start = None
+                    end = None
+
+                    prev_frame = None
+                    regions = []
+                    for i, r in enumerate(track.get("positions")):
+                        if isinstance(r, list):
+                            region = Region.region_from_array(r[1])
+                            if region.frame_number is None:
+                                if i == 0:
+                                    frame_number = round(r[0] * FPS)
+                                    region.frame_number = frame_number
+                                else:
+                                    region.frame_number = prev_frame + 1
+                        else:
+                            region = Region.region_from_json(r)
+                        if region.frame_number is None:
+                            if "frameTime" in r:
+                                if i == 0:
+                                    region.frame_number = round(r["frameTime"] * 9)
+                                else:
+                                    region.frame_number = prev_frame + 1
+                        # new_f = region.frame_number + region_adjust
+                        prev_frame = region.frame_number
+                        region.frame_number = region.frame_number + region_adjust
+                        assert region.frame_number >= 0
+                        regions.append(region.to_array())
+                        if start is None:
+                            start = region.frame_number
+                        end = region.frame_number
+                    node_attrs["start_frame"] = start
+                    node_attrs["end_frame"] = min(num_frames, end)
+
+                    region_array = np.uint16(regions)
+                    regions = track_group.create_dataset(
+                        "regions",
+                        region_array.shape,
+                        chunks=region_array.shape,
+                        compression="gzip",
+                        dtype=region_array.dtype,
+                    )
+                    regions[:, :] = region_array
+            except:
+                logging.error("Error saving file %s", filename, exc_info=True)
+                f.close()
+                out_file.unlink()
diff --git a/modelevaluate.py b/modelevaluate.py
index 25b5b557..30be6cc2 100644
--- a/modelevaluate.py
+++ b/modelevaluate.py
@@ -13,6 +13,8 @@
 import pickle
 import sys
 import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
 import json
 import pickle
 from config.config import Config
@@ -24,8 +26,9 @@
 )
 from ml_tools import tools
 from ml_tools.trackdatabase import TrackDatabase
+from ml_tools.rawdb import RawDatabase
 import tensorflow as tf
-from ml_tools.dataset import Dataset
+from ml_tools.dataset import Dataset, filter_clip, filter_track
 import pytz
 from ml_tools.datasetstructures import SegmentType
 from dateutil.parser import parse as parse_date
@@ -33,11 +36,66 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from pathlib import Path
-from ml_tools.preprocess import preprocess_ir
+from ml_tools.preprocess import preprocess_ir, preprocess_frame
+from ml_tools.frame import Frame
+from ml_tools import imageprocessing
+import cv2
+from config.loadconfig import LoadConfig
+from sklearn.metrics import confusion_matrix
 
+root_logger = logging.getLogger()
+for handler in root_logger.handlers:
+    if isinstance(handler, logging.StreamHandler) and handler.stream == sys.stderr:
+        root_logger.removeHandler(handler)
 crop_rectangle = tools.Rectangle(0, 0, 160, 120)
 PROB_THRESHOLD = 0.8
 
+land_birds = [
+    "pukeko",
+    "california quail",
+    "brown quail",
+    "black swan",
+    "quail",
+    "pheasant",
+    "penguin",
+    "duck",
+]
+
+
+def get_mappings(label_paths):
+    regroup = {}
+    for l, path in label_paths.items():
+        if l in land_birds:
+            regroup[l] = l
+            continue
+        split_path = path.split(".")
+        if len(split_path) == 1:
+            regroup[l] = l
+        elif path.startswith("all.mammal"):
+            if len(split_path) == 4:
+                regroup[l] = split_path[-2]
+            else:
+                regroup[l] = l
+        else:
+            # print("l", l, " has ", path)
+            parent = split_path[-2]
+            # print("Parent is", parent, path)
+            if parent == "kiwi" or split_path[-1] == "kiwi":
+                regroup[l] = "kiwi"
+            elif parent == "other":
+                regroup[l] = l
+
+            else:
+                if "bird." in path:
+                    regroup[l] = "bird"
+
+                elif len(split_path) > 2:
+                    regroup[l] = split_path[-3]
+                else:
+                    regroup[l] = split_path[-1]
+
+    return regroup
+
 
 def load_args():
     parser = argparse.ArgumentParser()
@@ -54,20 +112,20 @@ def load_args():
         default="test",
         help="Dataset to use train, validation, test ( Default)",
     )
+
     parser.add_argument(
-        "--confusion",
-        help="Confusion matrix filename, used if you want to save confusion matrix image",
-    )
-    parser.add_argument(
-        "--tracks", action="count", help="Evaluate whole track rather than samples"
+        "--evaluate-dir",
+        help="Evalute directory of cptv files",
     )
-    parser.add_argument("-c", "--config-file", help="Path to config file to use")
 
-    parser.add_argument("--track-id", help="Track id to evaluate from database")
+    parser.add_argument("-c", "--config-file", help="Path to config file to use")
 
-    parser.add_argument("--clip-id", help="Clip id to evaluate from database")
     parser.add_argument("-d", "--date", help="Use clips after this")
 
+    parser.add_argument(
+        "confusion",
+        help="Confusion matrix filename, used if you want to save confusion matrix image",
+    )
     args = parser.parse_args()
     if args.date:
         args.date = parse_date(args.date)
@@ -89,67 +147,347 @@ def init_logging(timestamps=False):
     )
 
 
+def filter_diffs(track_frames, background):
+    max_diff = 0
+    min_diff = 0
+    for f in track_frames:
+        r = f.region
+        if r.blank or r.width <= 0 or r.height <= 0:
+            continue
+
+        diff_frame = np.float32(f.thermal) - r.subimage(background)
+        new_max = np.amax(diff_frame)
+        new_min = np.amin(diff_frame)
+        if new_min < min_diff:
+            min_diff = new_min
+        if new_max > max_diff:
+            max_diff = new_max
+    return min_diff, max_diff
+
+
+def evalute_prod_confusion(dir, confusion_file):
+    with open("label_paths.json", "r") as f:
+        label_paths = json.load(f)
+    label_mapping = get_mappings(label_paths)
+
+    labels = [
+        "bird",
+        "bird/kiwi",
+        "cat",
+        "false-positive",
+        "hedgehog",
+        "human",
+        "insect",
+        "leporidae",
+        "mustelid",
+        "penguin",
+        "possum",
+        "rodent",
+        "sheep",
+        "vehicle",
+        "wallaby",
+        "None",
+        "unidentified",
+    ]
+    y_true = []
+    y_pred = []
+    dir = Path(dir)
+    for cptv_file in dir.glob(f"**/*cptv"):
+        meta_f = cptv_file.with_suffix(".txt")
+        meta_data = None
+        with open(meta_f, "r") as t:
+            # add in some metadata stats
+            meta_data = json.load(t)
+
+        for track in meta_data.get("Tracks", []):
+            tags = track.get("tags", [])
+            human_tags = [
+                tag.get("what")
+                for tag in tags
+                if tag.get("automatic") == False
+                # and tag.get("what", "") not in LoadConfig.EXCLUDED_TAGS
+            ]
+            human_tags = set(human_tags)
+            if len(human_tags) > 1:
+                print("Conflicting tags for ", track.get("id"), cptv_file)
+            if len(human_tags) == 0:
+                print("No humans in ", tags)
+                continue
+            human_tag = human_tags.pop()
+            human_tag = label_mapping.get(human_tag, human_tag)
+            ai_tag = [
+                tag.get("what")
+                for tag in tags
+                if tag.get("automatic") is True
+                and tag.get("data", {}).get("name") == "Inc3 RF"
+            ]
+            y_true.append(human_tag)
+            if len(ai_tag) == 0:
+                y_pred.append("None")
+            else:
+                y_pred.append(ai_tag[0])
+
+    cm = confusion_matrix(y_true, y_pred, labels=labels)
+    # Log the confusion matrix as an image summary.
+    figure = plot_confusion_matrix(cm, class_names=labels)
+    plt.savefig(confusion_file, format="png")
+
+
+def evaluate_dir(
+    model,
+    dir,
+    config,
+    confusion_file,
+):
+    with open("label_paths.json", "r") as f:
+        label_paths = json.load(f)
+    label_mapping = get_mappings(label_paths)
+    # dataset = Dataset(
+    #     dir,
+    #     "dataset",
+    #     config,
+    #     consecutive_segments=False,
+    #     label_mapping=label_mapping,
+    #     raw=True,
+    #     ext=".cptv",
+    # )
+    #
+    # tracks_loaded, total_tracks = dataset.load_clips(dont_filter_segment=True)
+    reason = {}
+    y_true = []
+    y_pred = []
+    files = list(dir.glob(f"**/*cptv"))
+    files.sort()
+    for cptv_file in files:
+        # for clip in dataset.clips:
+        clip_db = RawDatabase(cptv_file)
+        clip = clip_db.get_clip_tracks(LoadConfig.DEFAULT_GROUPS)
+        if clip is None:
+            logging.warn("No clip for %s", cptv_file)
+            continue
+
+        if filter_clip(clip, reason):
+            logging.info("Filtering %s", cptv_file)
+            continue
+        clip.tracks = [
+            track
+            for track in clip.tracks
+            if not filter_track(track, config.load.excluded_tags, reason)
+        ]
+        if len(clip.tracks) == 0:
+            logging.info("No tracks after filtering %s", cptv_file)
+            continue
+        clip_db.load_frames()
+        segment_frame_spacing = int(round(clip.frames_per_second))
+        thermal_medians = []
+        for f in clip_db.frames:
+            thermal_medians.append(np.median(f.thermal))
+        thermal_medians = np.uint16(thermal_medians)
+
+        for track in clip.tracks:
+            track.calculate_segments(
+                segment_frame_spacing * 2,
+                model.params.square_width**2,
+                segment_min_mass=10,
+                segment_type=SegmentType.ALL_SECTIONS,
+                ffc_frames=clip_db.ffc_frames,
+            )
+            frame_indices = set()
+            for sample in track.samples:
+                frame_indices.update(set(sample.frame_indices))
+                sample.remapped_label = label_mapping.get(
+                    sample.original_label, sample.original_label
+                )
+                track.label = sample.remapped_label
+            frame_indices = list(frame_indices)
+            frame_indices.sort()
+            track_frames = {}
+            # frames = []
+            for i in frame_indices:
+                f = clip_db.frames[i]
+                f.region = track.regions_by_frame[f.frame_number]
+                pre_f = preprocess_frame(
+                    f, (32, 32), f.region, clip_db.background, crop_rectangle
+                )
+                track_frames[i] = pre_f
+
+            # for f in frames:
+            #     region = track.regions_by_frame[f.frame_number]
+            #     track_frame = f.crop_by_region(region)
+            #     track_frame.region = region
+            #     track_frames[region.frame_number] = track_frame
+            #
+            # min_diff, max_diff = filter_diffs(track_frames.values(), clip_db.background)
+            # for f in track_frames.values():
+            #     f.float_arrays()
+            #     f.filtered = f.thermal - f.region.subimage(clip_db.background)
+            #     f.filtered, stats = imageprocessing.normalize(
+            #         f.filtered, min=min_diff, max=max_diff, new_max=255
+            #     )
+            #     f.thermal -= thermal_medians[f.frame_number]
+            #     np.clip(f.thermal, a_min=0, a_max=None, out=f.thermal)
+            #     f.thermal, stats = imageprocessing.normalize(f.thermal, new_max=255)
+            #     if f.thermal.size > 0:
+            #         f.resize_with_aspect(
+            #             (32, 32),
+            #             crop_rectangle,
+            #             True,
+            #         )
+
+            prediction = model.classify_track_data(
+                track.track_id, track_frames, track.samples
+            )
+            y_true.append(track.label)
+            if prediction.predicted_tag() is None:
+                y_pred.append("None")
+            else:
+                # to compare to prod add conf rules
+                if (
+                    prediction.clarity > min_tag_clarity
+                    and prediction.max_score > min_tag_confidence
+                ):
+                    y_pred.append(prediction.predicted_tag())
+                else:
+                    y_pred.append("unidentified")
+            print(
+                track,
+                "Got a prediction of",
+                y_pred[-1],
+                " should be ",
+                track.label,
+                np.round(100 * prediction.predictions),
+            )
+    model.labels.append("None")
+    model.labels.append("unidentified")
+
+    cm = confusion_matrix(y_true, y_pred, labels=model.labels)
+    # Log the confusion matrix as an image summary.
+    figure = plot_confusion_matrix(cm, class_names=model.labels)
+    plt.savefig(confusion_file, format="png")
+
+
+min_tag_clarity = 0.2
+min_tag_confidence = 0.8
+
+
+# trying to figre out differing predictions
+def test_model(model_file, weights, config):
+    import tensorflow.keras.backend as K
+
+    model = None
+    tf.random.set_seed(1)
+    model = KerasModel(train_config=config.train)
+    model.load_model(model_file.parent)
+    # model.model.save_weights(model_file.parent / "final_weights")
+    # model.build_model(dropout=0.3)
+
+    # model.model.load_weights(str(weights))
+    # print("Loading weigfhts", weights)
+    model.model.load_weights(weights)
+    model = model.model
+
+    output = model.layers[1].output
+    for new_l in model.layers[2:]:
+        print(new_l.name)
+        output = new_l(output)
+    new_model = tf.keras.models.Model(model.layers[1].input, outputs=output)
+
+    print("Saving", model_file.parent / "re_save")
+    new_model.summary()
+    new_model.save(model_file.parent / "re_save")
+    return
+    for _ in range(2):
+        test = np.ones((1, 160, 160, 3), dtype=np.float32)
+        test = test / 2.0
+        print("RUN")
+        print("RUN")
+        print("RUN")
+        print("RUN")
+
+        # for l, o in zip(model.layers[1].layers, layer_outs):
+        #     print(l.name)
+        #     for x in np.array(o).ravel():
+        #         print(x)
+        print(model.predict(test))
+        # print(layer_outs)
+    return
+
+
+def re_save(model_file, weights, config):
+    tf.random.set_seed(1)
+    model = KerasModel(train_config=config.train)
+    model.load_model(model_file.parent)
+    # model.model.save_weights(model_file.parent / "final_weights")
+    # model.build_model(dropout=0.3)
+
+    # model.model.load_weights(str(weights))
+    # print("Loading weigfhts", weights)
+    # model.model.save_weight(model_file.parent / "final_weights")
+    model.model.load_weights(weights)
+    model = model.model
+
+    output = model.layers[1].output
+    for new_l in model.layers[2:]:
+        print(new_l.name)
+        output = new_l(output)
+    new_model = tf.keras.models.Model(model.layers[1].input, outputs=output)
+
+    print("Saving", model_file.parent / "re_save")
+    new_model.summary()
+    new_model.save(model_file.parent / "re_save")
+
+
 def main():
     args = load_args()
     init_logging()
     config = Config.load_from_file(args.config_file)
+    # evalute_prod_confusion(args.evaluate_dir, args.confusion)
+    print("LOading config", args.config_file)
     weights = None
     if args.model_file:
         model_file = Path(args.model_file)
     if args.weights:
         weights = model_file / args.weights
-
+    # test_model(model_file, Path(weights), config)
+    # return
     base_dir = config.tracks_folder
+    tf.random.set_seed(1)
 
     model = KerasModel(train_config=config.train)
     model.load_model(model_file, training=False, weights=weights)
-    # model = None
-    if args.track_id or args.clip_id:
-        if args.clip_id is None:
-            logging.error("Need clip id and track id")
-            sys.exit(0)
-        db = TrackDatabase(os.path.join(config.tracks_folder, "dataset.hdf5"))
-        evaluate_db_clip(model, db, args.clip_id, args.track_id)
-        sys.exit(0)
-
-    if args.tracks:
-        # def evaluate_db_clips(model, db, classifier, config, after_date):
-
-        evaluate_db_clips(model, config, args.date, args.confusion)
-        sys.exit(0)
-    model.load_training_meta(base_dir)
-    excluded, remapped = get_excluded(model.type)
-    files = base_dir + f"/training-data/{args.dataset}"
-    dataset, _, new_labels = get_dataset(
-        files,
-        model.type,
-        model.labels,
-        batch_size=64,
-        image_size=model.params.output_dim[:2],
-        preprocess_fn=model.preprocess_fn,
-        augment=False,
-        resample=False,
-        include_features=model.params.mvm,
-        one_hot=True,
-        deterministic=True,
-        shuffle=False,
-        excluded_labels=excluded,
-        remapped_labels=remapped,
-    )
-    model.labels = new_labels
-    logging.info(
-        "Dataset loaded %s, using labels %s",
-        args.dataset,
-        model.labels,
-    )
 
-    if args.confusion:
-        if config.train.tfrecords:
-            model.confusion_tfrecords(dataset, args.confusion)
-        else:
-            model.confusion(dataset, args.confusion)
-    else:
-        model.evaluate(dataset)
+    if args.evaluate_dir:
+        evaluate_dir(model, Path(args.evaluate_dir), config, args.confusion)
+    elif args.dataset:
+        model.load_training_meta(base_dir)
+        if model.params.multi_label:
+            model.labels.append("land-bird")
+        excluded, remapped = get_excluded(model.type)
+        files = base_dir + f"/training-data/{args.dataset}"
+        dataset, _, new_labels, _ = get_dataset(
+            files,
+            model.type,
+            model.labels,
+            batch_size=64,
+            image_size=model.params.output_dim[:2],
+            preprocess_fn=model.preprocess_fn,
+            augment=False,
+            resample=False,
+            include_features=model.params.mvm,
+            one_hot=True,
+            deterministic=True,
+            shuffle=False,
+            excluded_labels=excluded,
+            remapped_labels=remapped,
+            multi_label=model.params.multi_label,
+        )
+        model.labels = new_labels
+        logging.info(
+            "Dataset loaded %s, using labels %s",
+            args.dataset,
+            model.labels,
+        )
+        model.confusion_tfrecords(dataset, args.confusion)
 
 
 def evaluate_db_clips(model, config, after_date, confusion_file="tracks-confusion"):
diff --git a/piclassifier/cptvrecorder.py b/piclassifier/cptvrecorder.py
index 6ff99a6a..ba774534 100644
--- a/piclassifier/cptvrecorder.py
+++ b/piclassifier/cptvrecorder.py
@@ -2,7 +2,7 @@
 import logging
 import os
 import yaml
-from load.cliptrackextractor import ClipTrackExtractor
+from track.cliptrackextractor import ClipTrackExtractor
 from cptv import CPTVWriter
 from cptv import Frame
 from datetime import timedelta
diff --git a/piclassifier/irrecorder.py b/piclassifier/irrecorder.py
index 8a56223b..a5b4c123 100644
--- a/piclassifier/irrecorder.py
+++ b/piclassifier/irrecorder.py
@@ -12,7 +12,7 @@
 
 from ml_tools.mpeg_creator import MPEGCreator
 from piclassifier.recorder import Recorder
-from load.cliptrackextractor import ClipTrackExtractor
+from track.cliptrackextractor import ClipTrackExtractor
 from ml_tools.logs import init_logging
 from .eventreporter import log_event
 
diff --git a/piclassifier/piclassifier.py b/piclassifier/piclassifier.py
index ae750f59..c451cf7b 100644
--- a/piclassifier/piclassifier.py
+++ b/piclassifier/piclassifier.py
@@ -9,10 +9,10 @@
 from pathlib import Path
 
 from classify.trackprediction import Predictions
-from load.clip import Clip
-from load.irtrackextractor import IRTrackExtractor
+from track.clip import Clip
+from track.irtrackextractor import IRTrackExtractor
 
-from load.cliptrackextractor import ClipTrackExtractor
+from track.cliptrackextractor import ClipTrackExtractor
 
 from ml_tools.previewer import Previewer, add_last_frame_tracking
 from ml_tools import tools
@@ -219,21 +219,21 @@ def __init__(
             self.classifier = get_interpreter(model)
 
             if self.classifier.TYPE == ForestModel.TYPE:
-                self.last_x_frames = 5 * headers.fps
-                self.frames_per_classify = self.last_x_frames
+                self.predict_from_last = 5 * headers.fps
+                self.frames_per_classify = self.predict_from_last
                 PiClassifier.SKIP_FRAMES = 30
                 # probably could be even more
 
             else:
                 # self.preprocess_fn = self.get_preprocess_fn()
 
-                self.last_x_frames = 1
+                self.predict_from_last = 1
                 self.frames_per_classify = (
                     self.classifier.params.square_width
                     * self.classifier.params.square_width
                 )
                 if self.frames_per_classify > 1:
-                    self.last_x_frames = self.frames_per_classify * 2
+                    self.predict_from_last = self.frames_per_classify * 2
 
             self.max_keep_frames = (
                 self.frames_per_classify * 2 if not preview_type else None
@@ -332,9 +332,14 @@ def new_clip(self, preview_frames):
 
     def startup_classifier(self):
         # classifies an empty frame to force loading of the model into memory
-        in_shape = self.classifier.shape()[1:]
-        p_frame = np.zeros((1, *in_shape), np.float32)
-        self.classifier.predict(p_frame)
+        num_inputs, in_shape = self.classifier.shape()
+        if num_inputs > 1:
+            zero_input = []
+            for shape in in_shape:
+                zero_input.append(np.zeros((1, *shape[1:]), np.float32))
+        else:
+            zero_input = np.zeros((1, *in_shape[1:]), np.float32)
+        self.classifier.predict(zero_input)
 
     def get_active_tracks(self):
         """
@@ -417,9 +422,10 @@ def identify_last_frame(self):
             frames, prediction, mass = self.classifier.predict_track(
                 clip,
                 track,
-                last_x_frames=self.last_x_frames,
+                predict_from_last=self.predict_from_last,
                 scale=self.track_extractor.scale,
                 frames_per_classify=self.frames_per_classify,
+                num_predictions=1,
             )
             if prediction is None:
                 track_prediction.last_frame_classified = self.clip.current_frame
diff --git a/piclassifier/piclassify.py b/piclassifier/piclassify.py
index 76bc2604..7ed53611 100644
--- a/piclassifier/piclassify.py
+++ b/piclassifier/piclassify.py
@@ -29,9 +29,9 @@
 import multiprocessing
 from cptv import Frame
 from .eventreporter import log_event
-from load.irtrackextractor import IRTrackExtractor
+from track.irtrackextractor import IRTrackExtractor
 
-from load.cliptrackextractor import ClipTrackExtractor
+from track.cliptrackextractor import ClipTrackExtractor
 
 SOCKET_NAME = "/var/run/lepton-frames"
 VOSPI_DATA_SIZE = 160
diff --git a/load/clip.py b/track/clip.py
similarity index 98%
rename from load/clip.py
rename to track/clip.py
index 33d0ec2a..d09ec72a 100644
--- a/load/clip.py
+++ b/track/clip.py
@@ -470,7 +470,7 @@ def __init__(self):
         self.average_delta = None
         self.is_static_background = None
 
-    def add_frame(self, thermal, filtered):
+    def add_frame(self, thermal, filtered=None):
         f_median = np.median(thermal)
         f_max = np.max(thermal)
         f_min = np.min(thermal)
@@ -485,12 +485,10 @@ def add_frame(self, thermal, filtered):
         if filtered is not None:
             self.filtered_sum += np.sum(np.abs(filtered))
 
-    def completed(self, num_frames, height, width):
-        if num_frames == 0:
-            return
-        total = num_frames * height * width
-        self.filtered_deviation = self.filtered_sum / float(total)
-        self.mean_temp = (height * width * sum(self.frame_stats_mean)) / float(total)
+    def completed(self):
+        if self.filtered_sum is not None:
+            self.filtered_deviation = np.mean(np.uint16(self.filtered_sum))
+        self.mean_temp = np.mean(np.uint16(self.frame_stats_mean))
 
 
 def null_safe_compare(a, b, cmp_fn):
diff --git a/load/cliptrackextractor.py b/track/cliptrackextractor.py
similarity index 98%
rename from load/cliptrackextractor.py
rename to track/cliptrackextractor.py
index 92f08b6d..44fec10a 100644
--- a/load/cliptrackextractor.py
+++ b/track/cliptrackextractor.py
@@ -131,7 +131,7 @@ def parse_clip(self, clip, process_background=False):
             self.apply_track_filtering(clip)
 
         if self.calc_stats:
-            clip.stats.completed(clip.current_frame, clip.res_y, clip.res_x)
+            clip.stats.completed()
         self._tracking_time = time.time() - start
         return True
 
diff --git a/track/framebuffer.py b/track/framebuffer.py
index 92a20fea..89da3903 100644
--- a/track/framebuffer.py
+++ b/track/framebuffer.py
@@ -82,9 +82,13 @@ def get_frame(self, frame_number):
             return self.current_frame
         elif self.cache:
             return self.cache.get_frame(frame_number)
-        if len(self.frames) > frame_number:
-            return self.frames[frame_number]
-        return None
+        if self.current_frame is None:
+            return None
+        # this supports max frames etc
+        frame_ago = self.current_frame.frame_number - frame_number
+        frame = self.get_frame_ago(frame_ago)
+        assert frame == None or frame.frame_number == frame_number
+        return frame
 
     def close_cache(self):
         if self.cache:
diff --git a/load/irtrackextractor.py b/track/irtrackextractor.py
similarity index 100%
rename from load/irtrackextractor.py
rename to track/irtrackextractor.py
diff --git a/track/region.py b/track/region.py
index 397332ae..7b2bd144 100644
--- a/track/region.py
+++ b/track/region.py
@@ -70,8 +70,10 @@ def to_array(self):
 
     @classmethod
     def region_from_array(cls, region_bounds):
-        width = region_bounds[2] - region_bounds[0]
-        height = region_bounds[3] - region_bounds[1]
+        width = int(region_bounds[2]) - region_bounds[0]
+        height = int(region_bounds[3]) - region_bounds[1]
+        height = np.uint8(max(height, 0))
+        width = np.uint8(max(width, 0))
         frame_number = None
         if len(region_bounds) > 4:
             frame_number = region_bounds[4]
diff --git a/track/track.py b/track/track.py
index 21980d8d..fe78cf05 100644
--- a/track/track.py
+++ b/track/track.py
@@ -429,16 +429,26 @@ def match(self, regions):
     def get_segments(
         self,
         ffc_frames,
-        frame_temp_median,
+        # frame_temp_median,
         segment_width,
         segment_frame_spacing=9,
         repeats=1,
         min_frames=0,
         segment_frames=None,
         segment_type=SegmentType.ALL_RANDOM,
+        from_last=None,
+        max_segments=None,
     ):
-        regions = np.array(self.bounds_history)
-        frame_temp_median = np.uint16(frame_temp_median)
+        if from_last is not None:
+            if from_last == 0:
+                return []
+            regions = np.array(self.bounds_history[-from_last:])
+            start_frame = regions[0].frame_number
+        else:
+            start_frame = self.start_frame
+            regions = np.array(self.bounds_history)
+
+        # frame_temp_median = np.uint16(frame_temp_median)
         segments = []
         if segment_frames is not None:
             mass_history = np.uint16([region.mass for region in regions])
@@ -449,33 +459,31 @@ def get_segments(
                 segment = SegmentHeader(
                     self.clip_id,
                     self._id,
-                    start_frame=self.start_frame,
+                    start_frame=start_frame,
                     frames=len(frames),
                     weight=1,
                     mass=segment_mass,
                     label=None,
                     regions=regions[relative_frames],
-                    frame_temp_median=frame_temp_median[relative_frames],
+                    # frame_temp_median=frame_temp_median[relative_frames],
                     frame_indices=frames,
-                    segment_type=segment_type,
                 )
                 segments.append(segment)
         else:
-            has_mass = any([region.mass for region in regions if region.mass > 0])
             segments, _ = get_segments(
                 self.clip_id,
                 self._id,
-                self.start_frame,
-                segment_frame_spacing,
-                segment_width,
+                start_frame,
+                segment_frame_spacing=segment_frame_spacing,
+                segment_width=segment_width,
                 regions=regions,
                 ffc_frames=ffc_frames,
                 repeats=repeats,
-                frame_temp_median=frame_temp_median,
+                # frame_temp_median=frame_temp_median,
                 min_frames=min_frames,
                 segment_frames=None,
-                ignore_mass=not has_mass,
                 segment_type=segment_type,
+                max_segments=max_segments,
             )
         return segments
 
diff --git a/track/trackextractor.py b/track/trackextractor.py
index c0d27f96..10981a5a 100644
--- a/track/trackextractor.py
+++ b/track/trackextractor.py
@@ -10,9 +10,9 @@
 import numpy as np
 
 from classify.trackprediction import Predictions
-from load.clip import Clip
-from load.cliptrackextractor import ClipTrackExtractor
-from load.irtrackextractor import IRTrackExtractor
+from track.clip import Clip
+from track.cliptrackextractor import ClipTrackExtractor
+from track.irtrackextractor import IRTrackExtractor
 from pathlib import Path
 from ml_tools import tools
 from ml_tools.previewer import Previewer
diff --git a/train/train.py b/train/train.py
index bfc9d6c4..5e6d7af5 100644
--- a/train/train.py
+++ b/train/train.py
@@ -68,7 +68,7 @@ def train_model(
             run_name=run_name + "_" + "TEST",
             weights=weights,
             resample=False,
-            rebalance=True,
+            rebalance=False,
         )
     except KeyboardInterrupt:
         pass


Test patch:
diff --git a/mldataset/testdataset.py b/mldataset/testdataset.py
new file mode 100644
index 00000000..2e86c865
--- /dev/null
+++ b/mldataset/testdataset.py
@@ -0,0 +1,278 @@
+"""
+Processes a CPTV file identifying and tracking regions of interest, and saving them in the 'trk' format.
+"""
+
+import argparse
+import os
+import sys
+from collections import namedtuple
+from pathlib import Path
+
+import h5py
+import cv2
+import numpy as np
+
+
+def parse_params():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "target",
+        default=None,
+        help="h5py file to display",
+    )
+    args = parser.parse_args()
+    args.target = Path(args.target)
+    return args
+
+
+Track = namedtuple("Track", "id regions what ai_what")
+
+
+RESCALE = 4
+from dateutil.parser import parse
+
+
+def latest_date(dir):
+    import json
+
+    latest_date = None
+    dbFiles = dir.glob("**/*.cptv")
+    for dbName in dbFiles:
+        meta_f = dbName.with_suffix(".txt")
+        with open(meta_f, "r") as f:
+            # add in some metadata stats
+            meta = json.load(f)
+        if meta.get("recordingDateTime"):
+            meta["recordingDateTime"] = parse(meta["recordingDateTime"])
+            if latest_date is None or meta["recordingDateTime"] > latest_date:
+                latest_date = meta["recordingDateTime"]
+    print("Latest date is", latest_date)
+
+
+def makecsv(dir):
+    dbFiles = dir.glob("**/*.hdf5")
+    import csv
+
+    print("Making csv")
+
+    rec_counts = {}
+    track_counts = {}
+    with open("file-desc.csv", "w") as csvout:
+        writer = csv.writer(csvout)
+        writer.writerow(["file", "track id", "human tags", "extra human tags"])
+        for dbName in dbFiles:
+            f = h5py.File(dbName, "r")
+            tracks = f["tracks"]
+            track_ids = tracks.keys()
+            has_track = False
+            for track_id in track_ids:
+                has_track = True
+                human_tags = set()
+                extra_tags = set()
+                attrs = tracks[track_id].attrs
+                if "human_tag" not in attrs:
+                    continue
+                tag = attrs["human_tag"]
+                human_tags.add(tag)
+                if "human_tags" in attrs:
+                    extra = attrs["human_tags"]
+                    for e in extra:
+                        extra_tags.add(e)
+                if tag in track_counts:
+                    track_counts[tag] += 1
+                else:
+                    track_counts[tag] = 1
+                writer.writerow([dbName, track_id, human_tags, extra_tags])
+            if not has_track:
+                writer.writerow([dbName, None, None, None])
+            human_tags = list(human_tags)
+            human_tags.sort()
+            for t in human_tags:
+                if t in rec_counts:
+                    rec_counts[t] += 1
+                else:
+                    rec_counts[t] = 1
+
+    print("REC Counts", rec_counts)
+    print("Track Counts", track_counts)
+
+
+def main():
+    args = parse_params()
+    latest_date(args.target)
+    # makecsv(args.target)
+    return
+    f = h5py.File(args.target, "r")
+    clip_attrs = f.attrs
+    print(
+        "Camera model is",
+        clip_attrs.get("model", ""),
+        " res:",
+        clip_attrs["res_x"],
+        clip_attrs["res_y"],
+        " id: ",
+        clip_attrs["clip_id"],
+        " station id: ",
+        clip_attrs["station_id"],
+        " tags: ",
+        clip_attrs.get("tags"),
+    )
+    for k in clip_attrs.keys():
+        print(f"{k}: {clip_attrs[k]}")
+    ffc_frames = clip_attrs["ffc_frames"]
+    frames = f["frames"]
+    if "background" in frames:
+        background = frames["background"]
+    else:
+        print("No background found so using first frame")
+        background = frames["thermals"][0]
+    crop_rectangle = clip_attrs["crop_rectangle"]
+
+    background = background[
+        crop_rectangle[1] : crop_rectangle[3], crop_rectangle[0] : crop_rectangle[2]
+    ]
+    norm_back = norm_image(background)
+
+    height, width, _ = norm_back.shape
+    norm_back = cv2.resize(norm_back, (width * RESCALE, height * RESCALE))
+    norm_back = cv2.putText(
+        norm_back,
+        "Background",
+        (50, 50),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        1,
+        (0, 255, 0),
+        2,
+        cv2.LINE_AA,
+    )
+
+    cv2.imshow("b", norm_back)
+    cv2.waitKey()
+    cv2.destroyAllWindows()
+
+    num_frames = clip_attrs["num_frames"]
+
+    tracks_node = f["tracks"]
+    tracks = []
+    track_ids = tracks_node.keys()
+    for track_id in track_ids:
+        track_node = tracks_node[track_id]
+        positions = track_node["regions"]
+        id = track_node.attrs["id"]
+        attrs = track_node.attrs
+        human_tag = ""
+        ai_tag = ""
+        if "human_tag" in attrs:
+            tag = attrs["human_tag"]
+            conf = attrs["human_tag_confidence"]
+            human_tag = f"{tag}-{conf}"
+
+        if "ai_tag" in attrs:
+            tag = attrs["ai_tag"]
+            conf = attrs["ai_tag_confidence"]
+            ai_tag = f"{tag}-{conf}"
+        track = Track(id, positions, human_tag, ai_tag)
+        tracks.append(track)
+        print("For track", track_id)
+        for k in attrs.keys():
+            print(f"{k}: {attrs[k]}")
+    frame_i = 0
+    for frame in frames["thermals"]:
+        frame = frame[
+            crop_rectangle[1] : crop_rectangle[3], crop_rectangle[0] : crop_rectangle[2]
+        ]
+
+        filtered = np.float32(frame) - background
+        filtered[filtered < 0] = 0
+        filtered = norm_image(filtered)
+        frame = norm_image(frame)
+
+        height, width, _ = frame.shape
+        frame = cv2.resize(frame, (width * RESCALE, height * RESCALE))
+        filtered = cv2.resize(filtered, (width * RESCALE, height * RESCALE))
+
+        if frame_i in ffc_frames:
+            frame = cv2.putText(
+                frame,
+                "FFC",
+                (50, 50),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1,
+                (0, 255, 0),
+                2,
+                cv2.LINE_AA,
+            )
+            filtered = cv2.putText(
+                filtered,
+                "FFC",
+                (50, 50),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1,
+                (0, 255, 0),
+                2,
+                cv2.LINE_AA,
+            )
+        wait = 100
+        for t in tracks:
+            for r in t.regions:
+                if r[4] == frame_i:
+                    frame = cv2.rectangle(
+                        frame,
+                        (r[0] * 4, r[1] * 4),
+                        (r[2] * 4, r[3] * 4),
+                        (0, 255, 0),
+                        1,
+                    )
+                    frame = cv2.putText(
+                        frame,
+                        f"human {t.what} \nai {t.ai_what}",
+                        (r[0] * 4, r[1] * 4 + 20),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.5,
+                        (0, 255, 0),
+                        2,
+                        cv2.LINE_AA,
+                    )
+                    filtered = cv2.rectangle(
+                        filtered,
+                        (r[0] * 4, r[1] * 4),
+                        (r[2] * 4, r[3] * 4),
+                        (0, 255, 0),
+                        1,
+                    )
+                    filtered = cv2.putText(
+                        filtered,
+                        f"human {t.what} \nai {t.ai_what}",
+                        (r[0] * 4, r[1] * 4 + 20),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.5,
+                        (0, 255, 0),
+                        2,
+                        cv2.LINE_AA,
+                    )
+                    wait = 0
+        cv2.imshow("thermal", frame)
+        cv2.imshow("filtered", filtered)
+        cv2.moveWindow("thermal", 0, 0)
+        cv2.moveWindow("filtered", width * 5, 0)
+        cv2.waitKey(wait)
+        frame_i += 1
+    print("DONE")
+    cv2.waitKey(0)
+
+
+def norm_image(image):
+    max = np.amax(image)
+    min = np.amin(image)
+    norm = np.float32(image)
+    norm = 255 * (norm - min) / (max - min)
+
+    norm = np.uint8(norm)
+    norm = norm[:, :, np.newaxis]
+    norm = np.repeat(norm, 3, axis=2)
+    return norm
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_tracking_speed.py b/tests/test_tracking_speed.py
index ed1b4613..14a5cfac 100644
--- a/tests/test_tracking_speed.py
+++ b/tests/test_tracking_speed.py
@@ -1,7 +1,7 @@
 import time
 import os
-from load.clip import Clip
-from load.cliptrackextractor import ClipTrackExtractor
+from track.clip import Clip
+from track.cliptrackextractor import ClipTrackExtractor
 from config.config import Config
 from ml_tools.previewer import Previewer