Repository: luxonis/luxonis-ml (language: python)

Issue:
Misc updates
- Remove the need for the `media_exists` variable. Now we check if the copy needs to take place
- Small updates in test suite for initial clean up
- Prevent duplicate LDF creation name with same `team_id` and `dataset_name`
- `create_project` in `LabelStudioConnector` which automatically generates the correct XML/labeling config
- Pre-commit hooks


Discussion / hints:

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 00000000..ad9ab615
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,11 @@
+repos:
+  - repo: https://github.com/ambv/black
+    rev: 23.3.0
+    hooks:
+    - id: black
+      language_version: python3.8
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
+    hooks:
+    - id: no-commit-to-branch
+      args: ['--branch', 'main', '--branch', 'dev']
diff --git a/examples/COCO_people_subset.ipynb b/examples/COCO_people_subset.ipynb
index 44f967ee..e2a74d46 100644
--- a/examples/COCO_people_subset.ipynb
+++ b/examples/COCO_people_subset.ipynb
@@ -10,7 +10,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "4c06d8fc",
    "metadata": {},
    "outputs": [],
@@ -20,8 +20,11 @@
     "import cv2\n",
     "import pycocotools.mask as mask_util\n",
     "from PIL import Image, ImageDraw\n",
+    "from copy import deepcopy\n",
     "\n",
-    "from luxonis_ml.ops import *"
+    "from luxonis_ml.data import *\n",
+    "\n",
+    "BUCKET_TYPE = 'local'"
    ]
   },
   {
@@ -34,70 +37,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "5cc9ddf2",
    "metadata": {
     "scrolled": true
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Requirement already satisfied: gdown in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (4.7.1)\n",
-      "Requirement already satisfied: tqdm in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from gdown) (4.65.0)\n",
-      "Requirement already satisfied: six in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from gdown) (1.16.0)\n",
-      "Requirement already satisfied: requests[socks] in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from gdown) (2.31.0)\n",
-      "Requirement already satisfied: filelock in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from gdown) (3.12.0)\n",
-      "Requirement already satisfied: beautifulsoup4 in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from gdown) (4.12.2)\n",
-      "Requirement already satisfied: soupsieve>1.2 in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from beautifulsoup4->gdown) (2.4.1)\n",
-      "Requirement already satisfied: certifi>=2017.4.17 in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from requests[socks]->gdown) (2023.5.7)\n",
-      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from requests[socks]->gdown) (3.1.0)\n",
-      "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from requests[socks]->gdown) (1.26.16)\n",
-      "Requirement already satisfied: idna<4,>=2.5 in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from requests[socks]->gdown) (3.4)\n",
-      "Requirement already satisfied: PySocks!=1.5.7,>=1.5.6 in /home/conor/Luxonis/envs/fiftyone/lib/python3.9/site-packages (from requests[socks]->gdown) (1.7.1)\n",
-      "\u001b[33mWARNING: You are using pip version 22.0.4; however, version 23.1.2 is available.\n",
-      "You should consider upgrading via the '/home/conor/Luxonis/envs/fiftyone/bin/python3.9 -m pip install --upgrade pip' command.\u001b[0m\u001b[33m\n",
-      "\u001b[0mDownloading...\n",
-      "From: https://drive.google.com/uc?id=1XlvFK7aRmt8op6-hHkWVKIJQeDtOwoRT\n",
-      "To: /home/conor/Luxonis/tmp/luxonis-ml/data/COCO_people_subset.zip\n",
-      "100%|██████████████████████████████████████| 7.78M/7.78M [00:00<00:00, 29.9MB/s]\n",
-      "Archive:  ../data/COCO_people_subset.zip\n",
-      "  inflating: ../data/person_keypoints_val2017.json  \n",
-      "   creating: ../data/person_val2017_subset/\n",
-      "  inflating: ../data/person_val2017_subset/000000001490.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000003934.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000005060.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000003255.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000001761.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000001000.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000002431.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000002006.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000002261.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000004395.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000005001.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000000872.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000002685.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000001268.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000005037.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000002473.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000001296.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000002299.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000005193.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000003553.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000001584.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000002153.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000001353.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000004765.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000002532.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000000139.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000000785.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000000885.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000004134.jpg  \n",
-      "  inflating: ../data/person_val2017_subset/000000003156.jpg  \n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "! pip install gdown\n",
     "! gdown 1XlvFK7aRmt8op6-hHkWVKIJQeDtOwoRT -O ../data/COCO_people_subset.zip\n",
@@ -114,7 +59,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "7fa5d2b9",
    "metadata": {},
    "outputs": [],
@@ -159,7 +104,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "4404049f",
    "metadata": {},
    "outputs": [],
@@ -220,56 +165,47 @@
     "            'keypoints': keypoints,\n",
     "            'split': splits[i]\n",
     "        }\n",
-    "    })"
+    "    })\n",
+    "    \n",
+    "original_additions = deepcopy(additions)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8171a7f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "team_id = \"luxonis_team_id\"\n",
+    "team_name = \"luxonis\"\n",
+    "dataset_name = \"coco_test\"\n",
+    "\n",
+    "dataset_id = LuxonisDataset.create(\n",
+    "    team_id=team_id,\n",
+    "    team_name=team_name,\n",
+    "    dataset_name=dataset_name\n",
+    "\n",
+    ")\n",
+    "print(f\"Dataset id: {dataset_id}\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "id": "9a30216e",
    "metadata": {
-    "scrolled": true
+    "scrolled": false
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/conor/Luxonis/tmp/luxonis-ml/src/luxonis_ml/ops/dataset.py:266: UserWarning: Updating from a previously saved source!\n",
-      "  warnings.warn(f\"Updating from a previously saved source!\")\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Checking for additions or modifications...\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [00:00<00:00, 38.69it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "No new additions!\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "with LuxonisDataset('luxonis', 'coco_test', bucket_type='local') as dataset:\n",
+    "with LuxonisDataset(\n",
+    "    team_id=team_id,\n",
+    "    dataset_id=dataset_id,\n",
+    "    team_name=team_name,\n",
+    "    dataset_name=dataset_name,\n",
+    "    bucket_type=BUCKET_TYPE\n",
+    ") as dataset:\n",
     "    \n",
     "    # defines the format of the COCO data source,\n",
     "    # which in this case is just a single image per training example\n",
@@ -287,7 +223,75 @@
     "        'edges': (np.array(data['categories'][0]['skeleton'])-1).tolist()\n",
     "    }) # optional for keypoints to define the skeleton visualization\n",
     "    \n",
-    "    dataset.add(additions, note=\"adding initial data :)\")"
+    "    dataset.add(additions)\n",
+    "    dataset.create_version(note=\"Initially adding data\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "58991ba0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with LuxonisDataset(\n",
+    "    team_id=team_id,\n",
+    "    dataset_id=dataset_id,\n",
+    "    bucket_type=BUCKET_TYPE\n",
+    ") as dataset:\n",
+    "    \n",
+    "    additions = deepcopy(original_additions)\n",
+    "    \n",
+    "    # add a classification label to the dataset\n",
+    "    for addition in additions:\n",
+    "        addition['image']['class'] = 'person'\n",
+    "    \n",
+    "    dataset.add(additions)\n",
+    "    dataset.create_version(note=\"Add classification labels\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e60d5a44",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with LuxonisDataset(\n",
+    "    team_id=team_id,\n",
+    "    dataset_id=dataset_id,\n",
+    ") as dataset:\n",
+    "    \n",
+    "    print(\"Dataset version:\", dataset.version)\n",
+    "    v10 = dataset.fo_dataset.load_saved_view('version_1.0')\n",
+    "    print(\"Class attribute presence in version 1.0:\")\n",
+    "    print([sample['class'] is not None for sample in v10])\n",
+    "    v11 = dataset.fo_dataset.load_saved_view('version_1.1')\n",
+    "    print(\"Class attribute presence in version 1.1:\")\n",
+    "    print([sample['class'] is not None for sample in v11])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cdcc8595",
+   "metadata": {},
+   "source": [
+    "### Permanently delete dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2724865d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with LuxonisDataset(\n",
+    "    team_id=team_id,\n",
+    "    dataset_id=dataset_id,\n",
+    ") as dataset:\n",
+    "    \n",
+    "    dataset.delete_dataset()"
    ]
   }
  ],
@@ -307,7 +311,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16"
+   "version": "3.9.17"
   }
  },
  "nbformat": 4,
diff --git a/requirements-dev.txt b/requirements-dev.txt
new file mode 100644
index 00000000..0aabd6cc
--- /dev/null
+++ b/requirements-dev.txt
@@ -0,0 +1,1 @@
+pre-commit==3.2.1
diff --git a/requirements.txt b/requirements.txt
index 8d7f36fe..9aa30cdb 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,10 +4,7 @@ ngraph>=0.2.0
 numpy>=1.22
 pytest>=7.2.0
 seaborn>=0.12.1
-torch>=1.10.0
 tqdm>=4.64.1
-unique-names-generator>=1.0.2
-mlflow>=2.0.1
 lakefs_client>=0.90.0
 webdataset>=0.2.31
 b2>=3.6.0
@@ -15,8 +12,8 @@ opencv-python>=4.5.5.64
 PyYAML>=6.0
 pandasql>=0.7.3
 pycocotools>=2.0.6
-pytorch-lightning>=1.9.0
 boto3>=1.17.39
 sqlalchemy==1.4.46
 idx2numpy==1.2.3
 fiftyone==0.21.0
+label-studio-sdk>=0.0.28
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 8c3e703f..712b51e9 100644
--- a/setup.py
+++ b/setup.py
@@ -1,28 +1,34 @@
 import io
 
-from setuptools import setup
+from setuptools import setup, find_packages
 
-with open('requirements.txt') as f:
-    required = f.readlines()
+base_requires = open("requirements.txt").readlines()
+loader_requires = open("src/luxonis_ml/loader/requirements.txt").readlines()
+tracker_requires = open("src/luxonis_ml/tracker/requirements.txt").readlines()
+all_requires = base_requires + loader_requires + tracker_requires
 
 setup(
-    name='luxonis-ml',
-    version='0.0.1',
-    description='This package provides MLOps tools for training models for OAK devices',
+    name="luxonis-ml",
+    version="0.0.1",
+    description="This package provides MLOps tools for training models for OAK devices",
     long_description=io.open("README.md", encoding="utf-8").read(),
     long_description_content_type="text/markdown",
-    url='https://github.com/luxonis/luxonis-ml-library/',
+    url="https://github.com/luxonis/luxonis-ml-library/",
     keywords="ml ops luxonis oak camera",
-    author='Luxonis',
-    author_email='support@luxonis.com',
-    license='MIT',
-    packages=['luxonis_ml'], # https://docs.python.org/3/distutils/setupscript.html point 2.1
+    author="Luxonis",
+    author_email="support@luxonis.com",
+    license="MIT",
+    packages=find_packages(where="src"),
     package_dir={"": "src"},  # https://stackoverflow.com/a/67238346/5494277
-    package_data={"luxonis_ml": ["ops/*.py"]},
-    install_requires=required,
+    install_requires=base_requires,
+    extras_require={
+        "loader": loader_requires,
+        "tracker": tracker_requires,
+        "all": all_requires,
+    },
     include_package_data=True,
     classifiers=[
-        'License :: MIT License',
+        "License :: MIT License",
         "Programming Language :: Python",
         "Programming Language :: Python :: 3",
         "Programming Language :: Python :: 3.6",
@@ -30,7 +36,7 @@
         "Programming Language :: Python :: 3.8",
     ],
     entry_points={
-        "mlflow.request_header_provider": "unused=luxonis_ml.ops.mlflow_plugins:LuxonisRequestHeaderProvider",
-        "console_scripts": ["luxonis_ml=luxonis_ml.luxonis_ml:main"]
+        "mlflow.request_header_provider": "unused=luxonis_ml.tracker.mlflow_plugins:LuxonisRequestHeaderProvider",
+        "console_scripts": ["luxonis_ml=luxonis_ml.luxonis_ml:main"],
     },
 )
diff --git a/src/luxonis_ml/__init__.py b/src/luxonis_ml/__init__.py
index 058ea05c..8b137891 100644
--- a/src/luxonis_ml/__init__.py
+++ b/src/luxonis_ml/__init__.py
@@ -1,1 +1,1 @@
-from luxonis_ml.ops import *
+
diff --git a/src/luxonis_ml/data/__init__.py b/src/luxonis_ml/data/__init__.py
new file mode 100644
index 00000000..12ea9a74
--- /dev/null
+++ b/src/luxonis_ml/data/__init__.py
@@ -0,0 +1,5 @@
+from .session import LuxonisSession
+from .dataset import LuxonisDataset, LDFComponent, IType, HType, LDFTransactionType
+from .version import LuxonisVersion
+from .parsers import LuxonisParser, DatasetType
+from .labelstudio import LabelStudioConnector
diff --git a/src/luxonis_ml/data/dataset.py b/src/luxonis_ml/data/dataset.py
new file mode 100644
index 00000000..9a4e558c
--- /dev/null
+++ b/src/luxonis_ml/data/dataset.py
@@ -0,0 +1,1083 @@
+import fiftyone as fo
+import fiftyone.core.odm as foo
+import luxonis_ml.data.fiftyone_plugins as fop
+import luxonis_ml.data.utils.data_utils as data_utils
+from fiftyone import ViewField as F
+import os, subprocess, shutil
+from pathlib import Path
+import glob
+import warnings
+import cv2
+from PIL import Image
+import pickle
+import json
+import boto3
+import glob
+from tqdm import tqdm
+from enum import Enum
+import numpy as np
+from datetime import datetime
+import pymongo
+from bson.objectid import ObjectId
+from .version import LuxonisVersion
+from .utils.s3_utils import sync_from_s3, sync_to_s3, check_s3_file_existence
+
+
+class HType(Enum):
+    """Individual file type"""
+
+    IMAGE = 1
+    JSON = 2
+
+
+class IType(Enum):
+    """Image type for IMAGE HType"""
+
+    BGR = 1
+    MONO = 2
+    DISPARITY = 3
+    DEPTH = 4
+
+
+class LDFTransactionType(Enum):
+    """The type of transaction"""
+
+    END = "END"
+    ADD = "ADD"
+    UPDATE = "UPDATE"
+    DELETE = "DELETE"
+
+
+class LDFComponent:
+    htypes = [HType.IMAGE, HType.JSON]
+    itypes = [IType.BGR, IType.MONO, IType.DISPARITY, IType.DEPTH]
+
+    def __init__(self, name, htype, itype=None):
+        if htype not in self.htypes:
+            raise Exception(f"{htype} is not a valid HType!")
+        self.name = name
+        self.htype = htype
+        if htype == HType.IMAGE:
+            if itype is not None:
+                if itype not in self.itypes:
+                    raise Exception(f"{itype} is not a valid IType!")
+                self.itype = itype
+            else:
+                raise Exception("itype parameter is required for image component")
+
+
+class LDFSource:
+    types = ["real_oak", "synthetic_oak", "custom"]
+
+    def __init__(self, name, components, main_component=None, source_type=None):
+        self.name = name
+        self.components = {component.name: component for component in components}
+        if main_component is not None:
+            self.main_component = main_component
+        else:
+            self.main_component = list(self.components.keys())[
+                0
+            ]  # make first component main component by default
+        if source_type is not None:
+            self.source_type = source_type
+        else:
+            self.source_type = "custom"
+
+    def add_component(self, component):
+        self.components += component
+
+
+class LuxonisDataset:
+    """
+    Uses the LDF (Luxonis Dataset Format)
+
+    The goal is to standardize an arbitrary number of devices and sensor configurations, synthetic OAKs, and other sources
+    """
+
+    @staticmethod
+    def create(
+        team_id,
+        team_name,
+        dataset_name,
+    ):
+        conn = foo.get_db_conn()
+        res = list(
+            conn.luxonis_dataset_document.find(
+                {"$and": [{"team_id": team_id}, {"dataset_name": dataset_name}]}
+            )
+        )
+        if len(res) >= 2:
+            raise Exception("More than 1 dataset exists under this name!")
+
+        if len(res):
+            return str(res[0]["_id"])
+        else:
+            dataset_doc = fop.LuxonisDatasetDocument(
+                team_id=team_id, team_name=team_name, dataset_name=dataset_name
+            )
+            dataset_doc = dataset_doc.save(upsert=True)
+
+            return str(dataset_doc.id)
+
+    def __init__(
+        self,
+        team_id,
+        dataset_id,
+        team_name=None,
+        dataset_name=None,
+        bucket_type="local",
+        override_bucket_type=False,
+    ):
+        """
+        team name: team under which you can find all datasets
+        dataset name: name of the dataset
+        bucket_type: underlying storage for images, which can be local or an AWS bucket
+        override_bucket_type: option to change underlying storage from saved setting in DB
+        """
+
+        self.conn = foo.get_db_conn()
+
+        self.team_id = team_id
+        self.dataset_id = dataset_id
+        self.team_name = team_name
+        self.dataset_name = dataset_name
+        self.full_name = f"{self.team_id}-{self.dataset_id}"
+        self.bucket_type = bucket_type
+        self.bucket_choices = ["local", "aws"]
+        self.override_bucket_type = override_bucket_type
+        if self.bucket_type not in self.bucket_choices:
+            raise Exception(f"Bucket type {self.bucket_type} is not supported!")
+
+        credentials_cache_file = (
+            f"{str(Path.home())}/.cache/luxonis_ml/credentials.json"
+        )
+        if os.path.exists(credentials_cache_file):
+            with open(credentials_cache_file) as file:
+                self.creds = json.load(file)
+        else:
+            self.creds = {}
+
+        self._init_path()
+        self.version = 0.0
+        self.tasks = ["class", "boxes", "segmentation", "keypoints"]
+        self.compute_heatmaps = True  # TODO: could make this configurable
+
+    def __enter__(self):
+        if self.full_name in fo.list_datasets():
+            self.fo_dataset = fo.load_dataset(self.full_name)
+        else:
+            self.fo_dataset = fo.Dataset(self.full_name)
+            self.conn.luxonis_dataset_document.update_one
+        self.fo_dataset.persistent = True
+
+        res = list(
+            self.conn.luxonis_dataset_document.find(
+                {
+                    "$and": [
+                        {"team_id": self.team_id},
+                        {"_id": ObjectId(self.dataset_id)},
+                    ]
+                }
+            )
+        )
+
+        if len(res):
+            assert len(res) == 1
+            self.dataset_doc = fop.LuxonisDatasetDocument.objects.get(
+                team_id=self.team_id, id=ObjectId(self.dataset_id)
+            )
+
+            tmp_bucket_type = self.bucket_type
+
+            self._doc_to_class()
+
+            if self.override_bucket_type:
+                self.bucket_type = tmp_bucket_type
+
+            self._init_path()
+
+        else:
+            raise Exception("Dataset not found!")
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._class_to_doc()
+        self.dataset_doc.save(upsert=True)
+
+    def _doc_to_class(self):
+        if hasattr(self.dataset_doc, "path") and self.dataset_doc.path is not None:
+            self.path = self.dataset_doc.path
+        if (
+            hasattr(self.dataset_doc, "bucket_type")
+            and self.dataset_doc.bucket_type is not None
+        ):
+            self.bucket_type = self.dataset_doc.bucket_type
+        if (
+            hasattr(self.dataset_doc, "current_version")
+            and self.dataset_doc.current_version is not None
+        ):
+            self.version = self.dataset_doc.current_version
+
+        doc = list(self._get_source())
+        if len(doc):
+            doc = doc[0]
+            components = [
+                LDFComponent(name, HType(htype), IType(itype))
+                for name, htype, itype in list(
+                    zip(
+                        doc["component_names"],
+                        doc["component_htypes"],
+                        doc["component_itypes"],
+                    )
+                )
+            ]
+            main_component = components[0].name
+            self.source = LDFSource(
+                name=doc["name"],
+                components=components,
+                main_component=main_component,
+                source_type=doc["source_type"],
+            )
+        else:
+            self.source = None
+
+    def _class_to_doc(self):
+        self.dataset_doc.fo_dataset_id = self.fo_dataset._doc.id
+        self.dataset_doc.path = self.path
+        self.dataset_doc.bucket_type = self.bucket_type
+        self.dataset_doc.current_version = self.version
+
+    def _get_credentials(self, key):
+        if key in self.creds.keys():
+            return self.creds[key]
+        else:
+            return os.environ[key]
+
+    def _get_source(self):
+        return self.conn.luxonis_source_document.find(
+            {"$and": [{"_luxonis_dataset_id": self.dataset_doc.id}]}
+        ).limit(1)
+
+    def _init_boto3_client(self):
+        self.client = boto3.client(
+            "s3",
+            endpoint_url=self._get_credentials("AWS_S3_ENDPOINT_URL"),
+            aws_access_key_id=self._get_credentials("AWS_ACCESS_KEY_ID"),
+            aws_secret_access_key=self._get_credentials("AWS_SECRET_ACCESS_KEY"),
+        )
+        self.bucket = self.path.split("//")[1].split("/")[0]
+        self.bucket_path = self.path.split(self.bucket + "/")[-1]
+        # self._create_directory(self.bucket_path)
+
+    def _init_path(self):
+        if self.bucket_type == "local":
+            self.path = f"{str(Path.home())}/.cache/luxonis_ml/data/{self.team_id}/datasets/{self.dataset_id}"
+            os.makedirs(self.path, exist_ok=True)
+        elif self.bucket_type == "aws":
+            self.path = f"s3://{self._get_credentials('AWS_BUCKET')}/{self.team_id}/datasets/{self.dataset_id}"
+            self._init_boto3_client()
+
+    def _create_directory(self, path, clear_contents=False):
+        if self.bucket_type == "local":
+            os.makedirs(path, exist_ok=True)
+        elif self.bucket_type == "aws":
+            resp = self.client.list_objects(
+                Bucket=self.bucket,
+                Prefix=f"{self.bucket_path}/{path}/",
+                Delimiter="/",
+                MaxKeys=1,
+            )
+            if "Contents" not in resp:
+                self.client.put_object(
+                    Bucket=self.bucket, Key=f"{self.bucket_path}/{path}/"
+                )
+            elif clear_contents:
+                for content in resp["Contents"]:
+                    self.client.delete_object(Bucket=self.bucket, Key=content["Key"])
+
+    def _save_version(self, samples, note):
+        version = LuxonisVersion(self, samples=samples, note=note)
+        samples = version.get_samples()
+
+        version_view = self.fo_dataset[samples]
+        self.fo_dataset.save_view(f"version_{self.version}", version_view)
+
+    def create_source(
+        self,
+        name=None,
+        oak_d_default=False,
+        # calibration=None,
+        # calibration_mapping=None,
+        custom_components=None,
+    ):
+        """
+        name (str): optional rename for standard sources, required for custom sources
+        calibration (dict): DepthAI calibration file for real device (TODO: or Unity devices)
+        calibration_mapping (dict): for calibration files from non-standard products,
+            maps camera ID to a LDFComponent for the sensor and ID pairs to disparities
+        custom_components (list): list of LDFComponents for a custom source
+
+        TODO: more features with calibration
+        """
+
+        if oak_d_default:
+            if name is None:
+                name = "oak_d"
+            components = [
+                LDFComponent("left", htype=HType.IMAGE, itype=IType.MONO),
+                LDFComponent("right", htype=HType.IMAGE, itype=IType.MONO),
+                LDFComponent("color", htype=HType.IMAGE, itype=IType.BGR),
+                LDFComponent("disparity", htype=HType.IMAGE, itype=IType.DISPARITY),
+                LDFComponent("depth", htype=HType.IMAGE, itype=IType.DEPTH),
+            ]
+
+            source = LDFSource(
+                name, components, main_component="color", source_type="real_oak"
+            )
+
+        elif name is not None and custom_components is not None:
+            source = LDFSource(name, custom_components)
+        else:
+            raise Exception("name and custom_components required for custom sources")
+
+        if self.dataset_doc.id is None:
+            self._class_to_doc()
+            self.dataset_doc.save(safe=True, upsert=True)
+        res = self._get_source()
+        if len(list(res)):  # exists
+            warnings.warn(f"Updating from a previously saved source!")
+            self.conn.luxonis_source_document.delete_many(
+                {"_luxonis_dataset_id": self.dataset_doc.id}
+            )
+        luxonis_dataset_id = self.dataset_doc.id
+        source_doc = fop.LuxonisSourceDocument(
+            luxonis_dataset_id=luxonis_dataset_id,
+            name=source.name,
+            source_type=source.source_type,
+            component_names=[c.name for _, c in source.components.items()],
+            component_htypes=[c.htype.value for _, c in source.components.items()],
+            component_itypes=[c.itype.value for _, c in source.components.items()],
+        )
+        source_doc.save(upsert=True)
+
+        self.source = source
+
+        self.fo_dataset.add_group_field(
+            self.source.name, default=self.source.main_component
+        )
+
+        return source
+
+    def launch_app(self):
+        session = fo.launch_app(dataset)
+
+    def set_classes(self, classes, task=None):
+        if task is not None:
+            if task not in self.tasks:
+                raise Exception(f"Task {task} is not a supported task")
+            self.fo_dataset.classes[task] = classes
+        else:
+            for task in self.tasks:
+                self.fo_dataset.classes[task] = classes
+
+    def set_mask_targets(self, mask_targets):
+        if 0 in mask_targets.keys():
+            raise Exception("Cannot set 0! This is assumed to be the background class")
+        self.fo_dataset.mask_targets = {"segmentation": mask_targets}
+
+    def set_skeleton(self, skeleton):
+        """
+        class_name (str): name of class to add fo skeleton for
+        skeleton (dict):
+            labels (list): list of strings for what each keypoint represents
+            egdes (list): list of length length-2 lists of keypoint indices denoting
+                          how keypoints are connected
+
+        NOTE: this only supports one class, which seems to be a fiftyone limitation
+        """
+
+        self.fo_dataset.skeletons.update(
+            {
+                "keypoints": fo.KeypointSkeleton(
+                    labels=skeleton["labels"], edges=skeleton["edges"]
+                )
+            }
+        )
+        self.fo_dataset.save()
+
+    def sync_from_cloud(self):
+        if self.bucket_type == "local":
+            print("This is a local dataset! Cannot sync")
+        else:
+            sync_from_s3(
+                # non_streaming_dir=f"{str(Path.home())}/.cache/luxonis_ml/data/{self.team_id}/datasets/{self.dataset_id}",
+                non_streaming_dir=f"{str(Path.home())}/.cache/luxonis_ml/data",
+                bucket=self.bucket,
+                bucket_dir=os.path.join(self.team_id, "datasets", self.dataset_id),
+                endpoint_url=self._get_credentials("AWS_S3_ENDPOINT_URL"),
+            )
+
+    def get_classes(self):
+        classes = set()
+        classes_by_task = {}
+        for task in self.fo_dataset.classes:
+            if len(self.fo_dataset.classes[task]):
+                classes_by_task[task] = self.fo_dataset.classes[task]
+                for cls in self.fo_dataset.classes[task]:
+                    classes.add(cls)
+        mask_classes = set()
+        if "segmentation" in self.fo_dataset.mask_targets.keys():
+            for key in self.fo_dataset.mask_targets["segmentation"]:
+                mask_classes.add(self.fo_dataset.mask_targets["segmentation"][key])
+        classes_by_task["segmentation"] = list(mask_classes)
+        classes = list(classes.union(mask_classes))
+        classes.sort()
+
+        return classes, classes_by_task
+
+    def get_classes_count(self):
+        """Returns dictionary with number of occurances for each class. If no class label is present returns empty dict."""
+        try:
+            count_dict = self.fo_dataset.count_values("class.classifications.label")
+        except:
+            warnings.warn(
+                "No 'class' label present in the dataset. Returning empty dictionary."
+            )
+            count_dict = {}
+        return count_dict
+
+    def delete_dataset(self):
+        """
+        Deletes the entire dataset, aside from the images
+        """
+
+        res = list(
+            self.conn.luxonis_dataset_document.find(
+                {
+                    "$and": [
+                        {"team_id": self.team_id},
+                        {"_id": ObjectId(self.dataset_id)},
+                    ]
+                }
+            )
+        )
+        if len(res):
+            ldf_doc = res[0]
+
+            self.conn.luxonis_source_document.delete_many(
+                {"_luxonis_dataset_id": ldf_doc["_id"]}
+            )
+            self.conn.version_document.delete_many(
+                {"_dataset_id": ldf_doc["_dataset_id"]}
+            )
+            self.conn.transaction_document.delete_many({"_dataset_id": ldf_doc["_id"]})
+            self.conn.luxonis_dataset_document.delete_many(
+                {
+                    "$and": [
+                        {"team_id": self.team_id},
+                        {"_id": ObjectId(self.dataset_id)},
+                    ]
+                }
+            )
+        fo.delete_dataset(self.full_name)
+
+    def _make_transaction(
+        self, action, sample_id=None, field=None, value=None, component=None
+    ):
+        transaction_doc = fop.TransactionDocument(
+            dataset_id=self.dataset_doc.id,
+            created_at=datetime.utcnow(),
+            executed=False,
+            action=action.value,
+            sample_id=sample_id,
+            field=field,
+            value={"value": value},
+            component=component,
+            version=-1,  # encodes no version yet assigned
+        )
+        transaction_doc = transaction_doc.save(upsert=True)
+
+        return transaction_doc.id
+
+    def _check_transactions(self, for_versioning=False):
+        if for_versioning:
+            attribute, value = "version", -1
+        else:
+            attribute, value = "executed", False
+
+        transactions = list(
+            self.conn.transaction_document.find(
+                {"_dataset_id": self.dataset_doc.id}
+            ).sort("created_at", pymongo.ASCENDING)
+        )
+        if len(transactions):
+            if transactions[-1]["action"] != LDFTransactionType.END.value:
+                i = -1
+                while (
+                    i != -len(transactions) - 1 and transactions[i][attribute] == value
+                ):
+                    self.conn.transaction_document.delete_many(
+                        {"_id": transactions[i]["_id"]}
+                    )
+                    i -= 1
+                return None
+            else:
+                i = -2
+                while (
+                    i != -len(transactions) - 1 and transactions[i][attribute] == value
+                ):
+                    i -= 1
+                return transactions[i + 1 :]
+        else:
+            return None
+
+    def _execute_transaction(self, tid):
+        self.conn.transaction_document.update_one(
+            {"_id": tid}, {"$set": {"executed": True}}
+        )
+
+    def _version_transaction(self, tid):
+        self.conn.transaction_document.update_one(
+            {"_id": tid}, {"$set": {"version": self.version}}
+        )
+
+    def _incr_version(self, media_change, field_change):
+        if media_change:  # major change
+            self.version += 1
+        if field_change:  # minor change
+            self.version += 0.1
+        self.version = round(self.version, 1)
+
+    def _add_filter(self, additions):
+        """
+        Filters out any additions to the dataset already existing
+        """
+
+        source = self.source
+        components = self.source.components
+
+        latest_view = self.fo_dataset.match(
+            (F("latest") == True) | (F("version") == -1)
+        )
+        filepaths = {component_name: [] for component_name in components}
+        for component_name in self.fo_dataset.group_slices:
+            self.fo_dataset.group_slice = component_name
+            paths = [sample["filepath"] for sample in latest_view]
+            filepaths[component_name] = list(
+                np.array(
+                    [f"{path.split('/')[-2]}/{path.split('/')[-1]}" for path in paths]
+                )
+            )
+
+        print("Checking for additions or modifications...")
+
+        transaction_to_additions = {}
+        media_change = False
+        field_change = False
+
+        for i, addition in tqdm(enumerate(additions), total=len(additions)):
+            # change the filepath for all components
+            for component_name in addition.keys():
+                filepath = addition[component_name]["filepath"]
+                additions[i][component_name]["_old_filepath"] = filepath
+                if not data_utils.is_modified_filepath(self, filepath):
+                    additions[i][component_name][
+                        "_new_image_name"
+                    ] = data_utils.generate_hashname(filepath)
+                granule = data_utils.get_granule(filepath, addition, component_name)
+                new_filepath = f"/{self.team_id}/datasets/{self.dataset_id}/{component_name}/{granule}"
+                additions[i][component_name]["filepath"] = new_filepath
+
+            group = fo.Group(name=source.name)
+
+            # check for ADD or UPDATE cases in the dataset
+            for component_name in addition.keys():
+                filepath = addition[component_name]["filepath"]
+                granule = data_utils.get_granule(filepath, addition, component_name)
+                new_filepath = f"/{self.team_id}/datasets/{self.dataset_id}/{component_name}/{granule}"
+                candidate = f"{component_name}/{granule}"
+                if candidate not in filepaths[component_name]:
+                    # ADD case
+                    media_change = True
+                    additions[i][component_name]["_group"] = group
+                    tid = self._make_transaction(
+                        LDFTransactionType.ADD,
+                        sample_id=None,
+                        field="filepath",
+                        value=new_filepath,
+                        component=component_name,
+                    )
+                    tid = str(tid)
+                    transaction_to_additions[tid] = i
+                else:
+                    # check for UPDATE
+                    self.fo_dataset.group_slice = component_name
+                    sample_view = self.fo_dataset.match(F("filepath") == new_filepath)
+                    # find the most up to date sample
+                    max_version = -np.inf
+                    for sample in sample_view:
+                        if sample.version == -1:
+                            latest_sample = sample
+                            max_version = sample.version
+                            break
+                        if sample.version > max_version:
+                            latest_sample = sample
+                            max_version = sample.version
+
+                    changes = data_utils.check_fields(
+                        self, latest_sample, addition, component_name
+                    )
+                    for change in changes:
+                        field, value = list(change.items())[0]
+                        field_change = True
+                        tid = self._make_transaction(
+                            LDFTransactionType.UPDATE,
+                            sample_id=latest_sample._id,
+                            field=field,
+                            value=value,
+                            component=component_name,
+                        )
+
+        if media_change or field_change:
+            self._make_transaction(LDFTransactionType.END)
+
+            return transaction_to_additions, media_change, field_change
+        else:
+            return None
+
+    def _add_extract(self, additions, from_bucket):
+        """
+        Filters out any additions to the dataset already existing
+        """
+
+        print("Extracting dataset media...")
+
+        components = self.source.components
+        if self.bucket_type == "local":
+            local_cache = f"{str(Path.home())}/.cache/luxonis_ml/data/{self.team_id}/datasets/{self.dataset_id}"
+        elif self.bucket_type == "aws":
+            local_cache = f"{str(Path.home())}/.cache/luxonis_ml/tmp"
+        os.makedirs(local_cache, exist_ok=True)
+        for component_name in components:
+            os.makedirs(f"{local_cache}/{component_name}", exist_ok=True)
+
+        sync = False
+
+        for i, addition in tqdm(enumerate(additions), total=len(additions)):
+            add_heatmaps = {}
+            for component_name in addition.keys():
+                if component_name not in components.keys():
+                    raise Exception(
+                        f"Component {component_name} is not present in source"
+                    )
+                component = addition[component_name]
+                if "filepath" not in component.keys():
+                    raise Exception("Must specify filepath for every component!")
+
+                filepath = component["_old_filepath"]
+                granule = data_utils.get_granule(filepath, addition, component_name)
+                local_path = f"{local_cache}/{component_name}/{granule}"
+
+                if self.bucket_type == "local":
+                    if os.path.exists(local_path):
+                        continue
+                elif self.bucket_type == "aws":
+                    s3_path = component["filepath"][1:]
+                    if check_s3_file_existence(
+                        self.bucket,
+                        s3_path,
+                        self._get_credentials("AWS_S3_ENDPOINT_URL"),
+                    ):
+                        continue
+                    sync = True
+
+                if from_bucket:
+                    old_prefix = filepath.split(f"s3://{self.bucket}/")[-1]
+                    new_prefix = f"{self.bucket_path}/{component_name}/{granule}"
+                    self.client.copy_object(
+                        Bucket=self.bucket,
+                        Key=new_prefix,
+                        CopySource={"Bucket": self.bucket, "Key": old_prefix},
+                    )
+                elif not data_utils.is_modified_filepath(self, filepath):
+                    cmd = f"cp {filepath} {local_path}"
+                    subprocess.check_output(cmd, shell=True)
+                else:
+                    pass  # TODO: some logging here on not actually extracting data due to a path error, most likely due to memory issues
+
+                if (
+                    self.compute_heatmaps
+                    and not from_bucket
+                    and (
+                        components[component_name].itype == IType.DISPARITY
+                        or components[component_name].itype == IType.DEPTH
+                    )
+                ):
+                    heatmap_component = f"{component_name}_heatmap"
+                    os.makedirs(f"{local_cache}/{heatmap_component}", exist_ok=True)
+                    im = cv2.imread(filepath, cv2.IMREAD_UNCHANGED)
+                    if im.dtype == np.uint16:
+                        im = im / 8
+                    heatmap = (im / 96 * 255).astype(np.uint8)
+                    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+                    cv2.imwrite(f"{local_cache}/{heatmap_component}/{granule}", heatmap)
+                    new_filepath = f"/{self.team_id}/datasets/{self.dataset_id}/{heatmap_component}/{granule}"
+                    add_heatmaps[heatmap_component] = new_filepath
+
+            if self.compute_heatmaps and not from_bucket and sync:
+                for heatmap_component in add_heatmaps:
+                    additions[i][heatmap_component] = {}
+                    additions[i][heatmap_component]["filepath"] = add_heatmaps[
+                        heatmap_component
+                    ]
+
+        if self.bucket_type == "aws" and not from_bucket:
+            sync_to_s3(
+                bucket=self.bucket,
+                s3_dir=self.bucket_path,
+                local_dir=local_cache,
+                endpoint_url=self._get_credentials("AWS_S3_ENDPOINT_URL"),
+            )
+
+        if self.bucket_type != "local":
+            shutil.rmtree(local_cache)
+
+        return additions
+
+    def _add_execute(self, additions=None, transaction_to_additions=None):
+        source = self.source
+        samples = []
+
+        transactions = self._check_transactions()
+
+        if transactions is None:
+            raise Exception("There are no changes to the dataset to execute!")
+
+        for transaction in transactions:
+            if transaction["action"] == LDFTransactionType.ADD.value:
+                if additions is None or transaction_to_additions is None:
+                    raise Exception(
+                        "additions and transaction_to_additions required for adding data"
+                    )
+
+                if transaction_to_additions.get(str(transaction["_id"])) is None:
+                    # TODO: find a better way to handle those
+                    print(f"Cannot find transaction with id: {str(transaction['_id'])}")
+                    print(f"Seems like a failed transaction?")
+                    continue
+
+                addition = additions[transaction_to_additions[str(transaction["_id"])]]
+                component_name = transaction["component"]
+                component = addition[component_name]
+                group = component["_group"]
+                sample = fo.Sample(
+                    filepath=component["filepath"], version=self.version, latest=True
+                )
+                sample[source.name] = group.element(component_name)
+
+                for ann in component.keys():
+                    if ann == "class" and component["class"] is not None:
+                        sample["class"] = data_utils.construct_class_label(
+                            self, component["class"]
+                        )
+                    elif ann == "boxes" and component["boxes"] is not None:
+                        sample["boxes"] = data_utils.construct_boxes_label(
+                            self, component["boxes"]
+                        )
+                    elif (
+                        ann == "segmentation" and component["segmentation"] is not None
+                    ):
+                        sample[
+                            "segmentation"
+                        ] = data_utils.construct_segmentation_label(
+                            self, component["segmentation"]
+                        )
+                    elif ann == "keypoints" and component["keypoints"] is not None:
+                        sample["keypoints"] = data_utils.construct_keypoints_label(
+                            self, component["keypoints"]
+                        )
+                    elif ann.startswith("_"):
+                        continue  # ignore temporary attributes
+                    else:
+                        sample[ann] = component[ann]
+
+                if (
+                    self.compute_heatmaps
+                    and f"{component_name}_heatmap" in addition.keys()
+                ):
+                    sample["heatmap"] = fo.Heatmap(
+                        map_path=addition[f"{component_name}_heatmap"]["filepath"]
+                    )
+
+                sample["tid"] = transaction["_id"].binary.hex()
+                sample["latest"] = False
+                sample["version"] = -1.0
+                samples.append(sample)
+
+            elif transaction["action"] == LDFTransactionType.UPDATE.value:
+                self.fo_dataset.group_slice = transaction["component"]
+                old_sample = self.fo_dataset[transaction["sample_id"]]
+
+                # check if there is already a working example with filepath and version as -1
+                sample = self.fo_dataset.match(
+                    (F("filepath") == old_sample.filepath) & (F("version") == -1)
+                )
+                if len(sample):
+                    # update the existing sample
+                    for sample in sample:
+                        break
+                    if transaction["field"] in self.tasks:
+                        if transaction["field"] == "class":
+                            sample[
+                                transaction["field"]
+                            ] = data_utils.construct_class_label(
+                                self, transaction["value"]["value"]
+                            )
+                        elif transaction["field"] == "boxes":
+                            sample[
+                                transaction["field"]
+                            ] = data_utils.construct_boxes_label(
+                                self, transaction["value"]["value"]
+                            )
+                        elif transaction["field"] == "segmentation":
+                            sample[
+                                transaction["field"]
+                            ] = data_utils.construct_segmentation_label(
+                                self, transaction["value"]["value"]
+                            )
+                        elif transaction["field"] == "keypoints":
+                            sample[
+                                transaction["field"]
+                            ] = data_utils.construct_keypoints_label(
+                                self, transaction["value"]["value"]
+                            )
+                    else:
+                        sample[transaction["field"]] = transaction["value"]["value"]
+                    sample.save()
+                else:
+                    # create a new sample which is a copy of an old sample
+                    sample_dict = old_sample.to_dict()
+                    sample = fo.Sample.from_dict(sample_dict)
+                    if transaction["field"] in self.tasks:
+                        if transaction["field"] == "class":
+                            sample[
+                                transaction["field"]
+                            ] = data_utils.construct_class_label(
+                                self, transaction["value"]["value"]
+                            )
+                        elif transaction["field"] == "boxes":
+                            sample[
+                                transaction["field"]
+                            ] = data_utils.construct_boxes_label(
+                                self, transaction["value"]["value"]
+                            )
+                        elif transaction["field"] == "segmentation":
+                            sample[
+                                transaction["field"]
+                            ] = data_utils.construct_segmentation_label(
+                                self, transaction["value"]["value"]
+                            )
+                        elif transaction["field"] == "keypoints":
+                            sample[
+                                transaction["field"]
+                            ] = data_utils.construct_keypoints_label(
+                                self, transaction["value"]["value"]
+                            )
+                    else:
+                        sample[transaction["field"]] = transaction["value"]["value"]
+
+                    sample["tid"] = transaction["_id"].binary.hex()
+                    sample["latest"] = False
+                    sample["version"] = -1.0
+                    self.fo_dataset.add_sample(sample)
+
+            self._execute_transaction(transaction["_id"])
+
+        if len(samples):
+            self.fo_dataset.add_samples(samples)
+
+    def _add_defaults(self, additions):
+        for i, addition in tqdm(enumerate(additions), total=len(additions)):
+            for component_name in addition.keys():
+                if "split" not in addition[component_name].keys():
+                    addition[component_name]["split"] = "train"
+
+    def add(self, additions, from_bucket=False, add_defaults=True):
+        """
+        Function to add data and automatically log transactions
+
+        additions: a list of dictionaries describing each Voxel51 sample
+            Each dict contains keys of components
+            The value of each component can contain the following keys:
+                filepath (media) : path to image, video, point cloud, or voxel
+                -----
+                class              : name of class for an entire image (Image Classification)
+                boxes              : list of Nx5 numpy arrays [class, xmin, ymin, width, height] (Object Detection)
+                segmentation       : numpy array where pixels correspond to integer values of classes
+                keypoints          : list of classes and (x,y) keypoints
+        from_bucket: True if adding images to a cloud bucket instead of locally
+        add_defaults: Add default values, such as split. Useful when calling the method the first time. Best if set to false on updates.
+        """
+
+        if from_bucket and self.bucket_type == "local":
+            raise Exception("from_bucket must be False for local dataset!")
+
+        self._check_transactions()  # will clear transactions any interrupted transactions
+
+        # TODO: add a try/except to gracefully handle any errors
+
+        # add defaults before
+        if add_defaults:
+            self._add_defaults(additions)
+
+        filter_result = self._add_filter(additions)
+        if filter_result is None:
+            print("No additions or modifications")
+            return
+        else:
+            transaction_to_additions, media_change, field_change = filter_result
+
+        additions = self._add_extract(additions, from_bucket)
+
+        self._add_execute(additions, transaction_to_additions)
+
+    def delete(self, deletions):
+        """
+        Function to delete data by sample ID
+
+        deletions: a list of sample IDs as strings
+        """
+
+        # TODO: add a try/except to gracefully handle any errors
+
+        self._check_transactions()  # will clear transactions any interrupted transactions
+
+        for delete_id in deletions:
+            # assume we want to delete all components in a group
+            sample = self.fo_dataset[delete_id]
+            gid = sample[self.source.name]["id"]
+            group = self.fo_dataset.get_group(gid)
+            for component_name in group:
+                sample = self.fo_dataset[group[component_name]["id"]]
+                tid = self._make_transaction(
+                    LDFTransactionType.DELETE,
+                    sample_id=sample.id,
+                    component=component_name,
+                )
+                sample.latest = False
+                sample.save()
+                self._execute_transaction(tid)
+
+        tid = self._make_transaction(LDFTransactionType.END)
+        self._execute_transaction(tid)
+
+    def create_version(self, note):
+        def get_current_sample(transaction):
+            sample = self.fo_dataset.match(F("tid") == transaction["_id"].binary.hex())
+            if not len(sample):
+                return None
+            for sample in sample:
+                break
+            return sample
+
+        def get_previous_sample(transaction):
+            sample = self.fo_dataset[transaction["sample_id"]]
+            return sample
+
+        # TODO: exception handling
+
+        transactions = self._check_transactions(for_versioning=True)
+        if transactions is None:
+            raise Exception("There are no changes to the dataset to version!")
+
+        contains_add = (
+            True
+            if np.sum(
+                [t["action"] == LDFTransactionType.ADD.value for t in transactions]
+            )
+            else False
+        )
+        contains_update = (
+            True
+            if np.sum(
+                [t["action"] == LDFTransactionType.UPDATE.value for t in transactions]
+            )
+            else False
+        )
+        contains_delete = (
+            True
+            if np.sum(
+                [t["action"] == LDFTransactionType.DELETE.value for t in transactions]
+            )
+            else False
+        )
+        media_change = contains_add or contains_delete
+        field_change = contains_update
+        self._incr_version(media_change, field_change)
+
+        if not media_change and not field_change:
+            # TODO: this could just be some [INFO] or [DEBUG] information later
+            print("No changes to version!")
+            return
+
+        add_samples = set()
+        deprecate_samples = set()
+        for transaction in transactions:
+            if transaction["action"] != LDFTransactionType.END.value:
+                self.fo_dataset.group_slice = transaction["component"]
+
+            if transaction["action"] == LDFTransactionType.ADD.value:
+                sample = get_current_sample(transaction)
+                if sample is None:
+                    # TODO: Find a better way to handle those.
+                    print(
+                        f"Sample is none for transaction ADD with tid {transaction['_id'].binary.hex()}. Skipping..."
+                    )
+                    self._version_transaction(transaction["_id"])
+                    continue
+                add_samples.add(sample["id"])
+
+            elif transaction["action"] == LDFTransactionType.UPDATE.value:
+                sample = get_current_sample(transaction)
+                if sample is None:  # a new sample which has had both ADD and UPDATE
+                    # sample = get_previous_sample(transaction)
+                    # TODO: Find a better way to handle those.
+                    print(f"Sample is none, skipping and versioning transatction")
+                    self._version_transaction(transaction["_id"])
+                    continue
+                else:
+                    add_samples.add(sample["id"])
+                    prev_sample = get_previous_sample(transaction)
+                    deprecate_samples.add(prev_sample["id"])
+
+            self._version_transaction(transaction["_id"])
+
+        for sample_id in add_samples:
+            sample = self.fo_dataset[sample_id]
+            sample["latest"] = True
+            sample["version"] = self.version
+            sample.save()
+        for sample_id in deprecate_samples:
+            sample = self.fo_dataset[sample_id]
+            sample["latest"] = False
+            sample.save()
+
+        version_samples = []
+        for component_name in self.source.components:
+            self.fo_dataset.group_slice = component_name
+            latest_view = self.fo_dataset.match(F("latest") == True)
+            version_samples += [sample.id for sample in latest_view]
+        self._save_version(version_samples, note)
+
+    def create_view(self, name, expr, version=None):
+        if version is None:
+            version = self.version
+
+        version_view = self.fo_dataset.load_saved_view(f"version_{version}")
+
+        view = version_view.match(expr)
+
+        self.fo_dataset.save_view(name, view)
+
+        return view
diff --git a/src/luxonis_ml/ops/dataset_recognition.py b/src/luxonis_ml/data/dataset_recognition.py
similarity index 67%
rename from src/luxonis_ml/ops/dataset_recognition.py
rename to src/luxonis_ml/data/dataset_recognition.py
index 7d70ddfd..26bbbe4d 100644
--- a/src/luxonis_ml/ops/dataset_recognition.py
+++ b/src/luxonis_ml/data/dataset_recognition.py
@@ -4,6 +4,7 @@
 from yaml.scanner import ScannerError
 from pathlib import Path
 
+
 # TODO: place for common parser args
 def recognize(dataset_path: str) -> str:
     """
@@ -12,19 +13,27 @@ def recognize(dataset_path: str) -> str:
     NOTE: Dataset type checking is done by some significant property of the dataset (has to contain json file, yaml file,..).
     """
 
-    dataset_path = dataset_path if dataset_path[-1] != '/' else dataset_path[:-1]
+    dataset_path = dataset_path if dataset_path[-1] != "/" else dataset_path[:-1]
 
     if not os.path.isdir(dataset_path):
         raise Exception("Invalid path name - not a directory.")
 
-    file_dir_list = [f"{dataset_path}/{dir_file}" for dir_file in os.listdir(dataset_path)] 
-    json_files = list(filter(lambda file_dir: file_dir.split(".")[-1] == "json", file_dir_list))
-    yaml_files = list(filter(lambda file_dir: file_dir.split(".")[-1] == "yaml", file_dir_list))
+    file_dir_list = [
+        f"{dataset_path}/{dir_file}" for dir_file in os.listdir(dataset_path)
+    ]
+    json_files = list(
+        filter(lambda file_dir: file_dir.split(".")[-1] == "json", file_dir_list)
+    )
+    yaml_files = list(
+        filter(lambda file_dir: file_dir.split(".")[-1] == "yaml", file_dir_list)
+    )
     dirs = list(filter(lambda file_dir: os.path.isdir(file_dir), file_dir_list))
 
     if json_files and dirs:
         if len(json_files) > 1:
-            raise Exception("Possible COCO dataset but multiple json files present - possible ambiguity.")
+            raise Exception(
+                "Possible COCO dataset but multiple json files present - possible ambiguity."
+            )
         json_file = json_files[0]
 
         try:
@@ -38,7 +47,9 @@ def recognize(dataset_path: str) -> str:
 
     if yaml_files and dirs:
         if len(yaml_files) > 1:
-            raise Exception("Possible YOLO dataset but multiple yaml files present - possible ambiguity.")
+            raise Exception(
+                "Possible YOLO dataset but multiple yaml files present - possible ambiguity."
+            )
         yaml_file = yaml_files[0]
         try:
             yaml.safe_load(Path(yaml_file).read_text())
diff --git a/src/luxonis_ml/fiftyone_plugins/__init__.py b/src/luxonis_ml/data/fiftyone_plugins/__init__.py
similarity index 100%
rename from src/luxonis_ml/fiftyone_plugins/__init__.py
rename to src/luxonis_ml/data/fiftyone_plugins/__init__.py
diff --git a/src/luxonis_ml/fiftyone_plugins/database.py b/src/luxonis_ml/data/fiftyone_plugins/database.py
similarity index 63%
rename from src/luxonis_ml/fiftyone_plugins/database.py
rename to src/luxonis_ml/data/fiftyone_plugins/database.py
index fde1eeae..877e1162 100644
--- a/src/luxonis_ml/fiftyone_plugins/database.py
+++ b/src/luxonis_ml/data/fiftyone_plugins/database.py
@@ -4,23 +4,27 @@
     ObjectIdField,
     StringField,
     FloatField,
-    IntField
+    IntField,
+    DictField,
+    BooleanField,
 )
 
 from fiftyone.core.odm.document import Document
 from fiftyone.core.odm.dataset import DatasetDocument
 
+
 class LuxonisDatasetDocument(Document):
     """Backing document for LuxonisDataset."""
 
-    dataset_id = ObjectIdField(db_field="_dataset_id", required=True)
-    dataset_id_str = StringField(required=True)
-    team_name = StringField(required=True)
-    dataset_name = StringField(required=True)
-    path = StringField(required=True)
-    bucket_type = StringField(required=True)
+    team_id = StringField(required=True)
+    team_name = StringField()
+    dataset_name = StringField()
+    fo_dataset_id = ObjectIdField(db_field="_dataset_id")
+    path = StringField()
+    bucket_type = StringField()
     current_version = FloatField()
 
+
 class LuxonisSourceDocument(Document):
     """Backing document for LuxonisDataset sources."""
 
@@ -31,13 +35,28 @@ class LuxonisSourceDocument(Document):
     component_htypes = ListField(IntField())
     component_itypes = ListField(IntField())
 
+
 class VersionDocument(Document):
     """Backing document for LuxonisDataset versions."""
 
-    id = ObjectIdField()
+    # id = ObjectIdField()
     number = FloatField(required=True)
     dataset_id = ObjectIdField(db_field="_dataset_id", required=True)
     dataset_id_str = StringField(required=True)
     created_at = DateTimeField(required=True)
     samples = ListField(StringField())
     note = StringField()
+
+
+class TransactionDocument(Document):
+    """Backing document for LuxonisDataset transactions."""
+
+    dataset_id = ObjectIdField(db_field="_dataset_id", required=True)
+    created_at = DateTimeField(required=True)
+    executed = BooleanField(required=True)
+    action = StringField(required=True)
+    sample_id = ObjectIdField()
+    field = StringField()
+    value = DictField()
+    component = StringField()
+    version = FloatField()
diff --git a/src/luxonis_ml/data/labelstudio.py b/src/luxonis_ml/data/labelstudio.py
new file mode 100644
index 00000000..6252d09b
--- /dev/null
+++ b/src/luxonis_ml/data/labelstudio.py
@@ -0,0 +1,227 @@
+from .dataset import LuxonisDataset
+from label_studio_sdk import Client, Project
+
+# from label_studio_converter import brush
+from typing import List, Union
+from pycocotools import mask as maskUtils
+import numpy as np
+import xml.etree.ElementTree as ET
+
+
+class InputStream:
+    def __init__(self, data):
+        self.data = data
+        self.i = 0
+
+    def read(self, size):
+        out = self.data[self.i : self.i + size]
+        self.i += size
+        return int(out, 2)
+
+
+def access_bit(data, num):
+    """from bytes array to bits by num position"""
+    base = int(num // 8)
+    shift = 7 - int(num % 8)
+    return (data[base] & (1 << shift)) >> shift
+
+
+def bytes2bit(data):
+    """get bit string from bytes data"""
+    return "".join([str(access_bit(data, i)) for i in range(len(data) * 8)])
+
+
+def rle_to_mask(rle: List[int], height: int, width: int) -> np.array:
+    """
+    Converts rle to image mask
+    Args:
+        rle: your long rle
+        height: original_height
+        width: original_width
+
+    Returns: np.array
+    """
+
+    rle_input = InputStream(bytes2bit(rle))
+
+    num = rle_input.read(32)
+    word_size = rle_input.read(5) + 1
+    rle_sizes = [rle_input.read(4) + 1 for _ in range(4)]
+    # print('RLE params:', num, 'values,', word_size, 'word_size,', rle_sizes, 'rle_sizes')
+
+    i = 0
+    out = np.zeros(num, dtype=np.uint8)
+    while i < num:
+        x = rle_input.read(1)
+        j = i + 1 + rle_input.read(rle_sizes[rle_input.read(2)])
+        if x:
+            val = rle_input.read(word_size)
+            out[i:j] = val
+            i = j
+        else:
+            while i < j:
+                val = rle_input.read(word_size)
+                out[i] = val
+                i += 1
+
+    image = np.reshape(out, [height, width, 4])[:, :, 3]
+    return image
+
+
+def generate_random_color():
+    return [np.random.randint(0, 255) for _ in range(3)]
+
+
+class LabelStudioConnector:
+    def __init__(self, dataset: LuxonisDataset) -> None:
+        self.dataset = dataset
+        self.client = Client(
+            self.dataset._get_credentials("LABELSTUDIO_URL"),
+            self.dataset._get_credentials("LABELSTUDIO_KEY"),
+        )
+        self.client.check_connection()
+
+    def get_projects(self):
+        return self.client.list_projects()
+
+    def push(
+        self,
+        project_id: int,
+        sample_ids: List,
+    ) -> None:
+        # TODO: add option to push annotations for cases such as pre-annotation
+
+        tasks = []
+        project = Project.get_from_id(self.client, project_id)
+
+        for sample_id in sample_ids:
+            sample = self.dataset.fo_dataset[sample_id]
+            filepath = sample["filepath"]
+
+            if self.dataset.bucket_type == "aws":
+                tasks.append({"image": f"s3://{self.dataset.bucket}{filepath}"})
+            else:
+                raise NotImplementedError()
+
+        project.import_tasks(tasks)
+
+    def pull(self, project_id: int) -> None:
+        additions = []
+        project = Project.get_from_id(self.client, project_id)
+
+        for task in project.export_tasks():
+            filepath = task["data"]["image"]
+            component_name = filepath.split("/")[
+                -2
+            ]  # TODO: maybe update this, as it's a bit hacky
+            addition = {component_name: {"filepath": filepath}}
+
+            cls = []
+            boxes = []
+            segmentation = None
+            seg_name_to_class = {
+                v: k
+                for k, v in self.dataset.fo_dataset.mask_targets["segmentation"].items()
+            }
+            keypoints = []
+
+            for annotation in task["annotations"]:
+                annotation = annotation["result"]
+                for ann in annotation:
+                    if ann["type"] == "choices":
+                        cls += ann["value"]["choices"]
+                    if ann["type"] == "rectanglelabels":
+                        # ih, iw = ann['original_height'], ann['original_width']
+                        x, y = ann["value"]["x"] / 100, ann["value"]["y"] / 100
+                        w, h = ann["value"]["width"] / 100, ann["value"]["height"] / 100
+                        box_cls = ann["value"]["rectanglelabels"][0]
+                        boxes.append([box_cls, x, y, w, h])
+                    if ann["type"] == "brushlabels":
+                        mask = rle_to_mask(
+                            ann["value"]["rle"],
+                            ann["original_height"],
+                            ann["original_width"],
+                        )
+                        mask_cls = ann["value"]["brushlabels"][0]
+                        if segmentation is None:
+                            segmentation = np.zeros(
+                                (ann["original_height"], ann["original_width"])
+                            )
+                        segmentation[mask > 128] = seg_name_to_class[mask_cls]
+                    if ann["type"] == "keypointlabels":
+                        pass
+
+            if len(cls):
+                addition[component_name]["class"] = cls
+            if len(boxes):
+                addition[component_name]["boxes"] = boxes
+            if segmentation is not None:
+                addition[component_name]["segmentation"] = segmentation
+            if len(keypoints):
+                addition[component_name]["keypoints"] = keypoints
+
+            additions.append(addition)
+
+        self.dataset.add(additions, media_exists=True)
+
+    def create_project(self, project_name: str):
+        root = ET.Element("View")
+
+        image = ET.SubElement(root, "Image")
+        image.set("name", "image")
+        image.set("value", "$image")
+        image.set("zoom", "true")
+
+        classes = self.dataset.fo_dataset.classes
+
+        if "class" in classes and len(classes["class"]):
+            choices = ET.SubElement(root, "Choices")
+            choices.set("name", "choice")
+            choices.set("toName", "image")
+            for label in classes["class"]:
+                choice = ET.SubElement(choices, "Choice")
+                choice.set("value", label)
+
+        if "boxes" in classes and len(classes["boxes"]):
+            inst = ET.SubElement(root, "RectangleLabels")
+            inst.set("name", "label")
+            inst.set("toName", "image")
+            for label in classes["boxes"]:
+                r, g, b = generate_random_color()
+                lbl = ET.SubElement(inst, "Label")
+                lbl.set("value", label)
+                lbl.set("background", f"rgba({r},{b},{g},0.7)")
+
+        if "segmentation" in classes and len(classes["segmentation"]):
+            inst = ET.SubElement(root, "BrushLabels")
+            inst.set("name", "tag")
+            inst.set("toName", "image")
+            for label in classes["segmentation"]:
+                r, g, b = generate_random_color()
+                lbl = ET.SubElement(inst, "Label")
+                lbl.set("value", label)
+                lbl.set("background", f"rgba({r},{b},{g},0.7)")
+
+        if "keypoints" in classes and len(classes["keypoints"]):
+            inst = ET.SubElement(root, "KeyPointLabels")
+            inst.set("name", "kp-1")
+            inst.set("toName", "image")
+            for label in classes["keypoints"]:
+                r, g, b = generate_random_color()
+                lbl = ET.SubElement(inst, "Label")
+                lbl.set("value", label)
+                lbl.set("background", f"rgba({r},{b},{g},0.7)")
+
+        tree = ET.ElementTree(root)
+        tree.write("tmp.xml", encoding="utf-8", xml_declaration=False)
+
+        with open("tmp.xml") as file:
+            label_config_str = file.read()
+
+        project = Project(
+            self.dataset._get_credentials("LABELSTUDIO_URL"),
+            self.dataset._get_credentials("LABELSTUDIO_KEY"),
+        )
+        project.start_project(title=project_name, label_config=label_config_str)
+
+        return project.id
diff --git a/src/luxonis_ml/ops/parsers.py b/src/luxonis_ml/data/parsers.py
similarity index 72%
rename from src/luxonis_ml/ops/parsers.py
rename to src/luxonis_ml/data/parsers.py
index fe834bf3..08d5908e 100644
--- a/src/luxonis_ml/ops/parsers.py
+++ b/src/luxonis_ml/data/parsers.py
@@ -1,5 +1,3 @@
-#!/usr/bin/env python3
-
 import json, yaml
 import os
 import numpy as np
@@ -9,14 +7,16 @@
 import warnings
 import random
 import shutil
-# from luxonis_ml.ops.dataset_type import DatasetType as dt
+
+# from luxonis_ml.data.dataset_type import DatasetType as dt
 from enum import Enum
 import threading
 from threading import Thread
 import xml.etree.ElementTree as ET
-from luxonis_ml.ops import *
+from luxonis_ml.data import *
 import csv
 
+
 class DatasetType(Enum):
     LDF = "LDF"
     S3DIR = "S3DIR"
@@ -34,11 +34,13 @@ class DatasetType(Enum):
     NUMPY = "numpy"
     UNKNOWN = "unknown"
 
+
 """
 Functions used with LuxonisDataset to convert other data formats to LDF
 """
-class LuxonisParser:
 
+
+class LuxonisParser:
     def __init__(self):
         self.percentage = 0
         self.error_message = None
@@ -51,12 +53,15 @@ def inner(*args, **kwargs):
             args[0].parsing_in_progress = True
 
             bucket_type = "aws" if args[1][2] == DatasetType.S3DIR else "local"
-            with LuxonisDataset(args[1][0], args[1][1], bucket_type=bucket_type) as dataset:
+            with LuxonisDataset(
+                args[1][0], args[1][1], bucket_type=bucket_type
+            ) as dataset:
                 print("setting component")
                 custom_components = [
-                    LDFComponent(name="image",
-                                htype=HType.IMAGE, # data component type
-                                itype=IType.BGR # image type
+                    LDFComponent(
+                        name="image",
+                        htype=HType.IMAGE,  # data component type
+                        itype=IType.BGR,  # image type
                     )
                 ]
                 # Just set source to some default name, as of right now, I think that parser users will care about its name
@@ -73,18 +78,18 @@ def inner(*args, **kwargs):
 
             args[0].parsing_in_progress = False
             args[0].percentage = 100.0
+
         return inner
 
     @parsing_wrapper
     def from_s3_directory_format(
-            self,
-            dataset,
-            image_dir,
-            splits=None,
-            dataset_size=None,
-            override_main_component=None
-        ):
-
+        self,
+        dataset,
+        image_dir,
+        splits=None,
+        dataset_size=None,
+        override_main_component=None,
+    ):
         ## define main component
         if override_main_component is not None:
             component_name = override_main_component
@@ -92,41 +97,40 @@ def from_s3_directory_format(
             component_name = dataset.source.main_component
 
         ## initialize additions
-        contents = dataset.client.list_objects(Bucket=dataset.bucket, Prefix=image_dir)['Contents']
+        contents = dataset.client.list_objects(Bucket=dataset.bucket, Prefix=image_dir)[
+            "Contents"
+        ]
 
         if splits is None:
             idxs = np.arange(len(contents))
             np.random.shuffle(idxs)
-            i1 = round(0.8*len(contents))
-            i2 = round(0.9*len(contents))
-            splits = {contents[idxs[n]]['Key']:'train' for n in range(0,i1)}
-            val = {contents[idxs[n]]['Key']:'val' for n in range(i1,i2)}
-            test = {contents[idxs[n]]['Key']:'test' for n in range(i2,len(contents))}
+            i1 = round(0.8 * len(contents))
+            i2 = round(0.9 * len(contents))
+            splits = {contents[idxs[n]]["Key"]: "train" for n in range(0, i1)}
+            val = {contents[idxs[n]]["Key"]: "val" for n in range(i1, i2)}
+            test = {contents[idxs[n]]["Key"]: "test" for n in range(i2, len(contents))}
             splits.update(val)
             splits.update(test)
 
         additions = []
         for metadata in contents:
-            filepath = metadata['Key']
-            additions.append({
-                component_name: {
-                    'filepath': filepath,
-                    'split': splits[filepath]
-                }
-            })
+            filepath = metadata["Key"]
+            additions.append(
+                {component_name: {"filepath": filepath, "split": splits[filepath]}}
+            )
 
         dataset.add(additions, note="S3 directory parser", from_bucket=True)
 
     @parsing_wrapper
     def from_coco_format(
-            self,
-            dataset,
-            image_dir,
-            annotation_path,
-            split,
-            dataset_size=None,
-            override_main_component=None
-        ):
+        self,
+        dataset,
+        image_dir,
+        annotation_path,
+        split,
+        dataset_size=None,
+        override_main_component=None,
+    ):
         """
         Constructs a LDF dataset from a COCO type dataset.
         Arguments:
@@ -150,56 +154,76 @@ def from_coco_format(
         ## load data
         with open(annotation_path) as file:
             coco = json.load(file)
-        coco_images = coco['images']
-        coco_annotations = coco['annotations']
-        coco_categories = coco['categories']
-        categories = {cat['id']:cat['name'] for cat in coco_categories}
-        keypoint_definitions = {cat['id']:{'keypoints':cat['keypoints'],'skeleton':cat['skeleton']} \
-                                for cat in coco_categories if 'keypoints' in cat.keys() and 'skeleton' in cat.keys()}
+        coco_images = coco["images"]
+        coco_annotations = coco["annotations"]
+        coco_categories = coco["categories"]
+        categories = {cat["id"]: cat["name"] for cat in coco_categories}
+        keypoint_definitions = {
+            cat["id"]: {"keypoints": cat["keypoints"], "skeleton": cat["skeleton"]}
+            for cat in coco_categories
+            if "keypoints" in cat.keys() and "skeleton" in cat.keys()
+        }
 
         ## dataset construction loop
         iter = tqdm(coco_images)
         for image in iter:
-            self.percentage = round((iter.n/iter.total)*100, 2)
+            self.percentage = round((iter.n / iter.total) * 100, 2)
 
             new_ann = {component_name: {"annotations": []}}
 
-            image_name = image['file_name']
-            image_id = image['id']
+            image_name = image["file_name"]
+            image_id = image["id"]
             image_path = os.path.join(image_dir, image_name)
             if os.path.exists(image_path):
                 image = cv2.imread(image_path)
             else:
                 warnings.warn(f"skipping {image_path} as it does no exist!")
 
-            annotations = [ann for ann in coco_annotations if ann['image_id'] == image_id]
+            annotations = [
+                ann for ann in coco_annotations if ann["image_id"] == image_id
+            ]
             for annotation in annotations:
                 new_ann_instance = {}
 
-                cat_id = annotation['category_id']
+                cat_id = annotation["category_id"]
                 class_name = categories[cat_id]
-                new_ann_instance['class_name'] = class_name
+                new_ann_instance["class_name"] = class_name
 
                 class_id = dataset._add_class(new_ann_instance)
-                new_ann_instance['class'] = class_id
+                new_ann_instance["class"] = class_id
 
                 if cat_id in keypoint_definitions.keys():
-                    dataset._add_keypoint_definition(class_id, keypoint_definitions[cat_id])
-
-                if 'segmentation' in annotation.keys():
-                    segmentation = annotation['segmentation']
-                    if isinstance(segmentation, list) and len(segmentation) > 0: # polygon format
-                        new_ann_instance['segmentation'] = {'type':'polygon', 'task':'instance', 'data':segmentation}
-                    elif isinstance(segmentation, dict) and len(segmentation['counts']) > 0:
-                        new_ann_instance['segmentation'] = {'type':'mask', 'task':'instance', 'data':segmentation}
-
-                if 'bbox' in annotation.keys():
-                    new_ann_instance['bbox'] = annotation['bbox']
+                    dataset._add_keypoint_definition(
+                        class_id, keypoint_definitions[cat_id]
+                    )
 
-                if 'keypoints' in annotation.keys():
-                    new_ann_instance['keypoints'] = annotation['keypoints']
+                if "segmentation" in annotation.keys():
+                    segmentation = annotation["segmentation"]
+                    if (
+                        isinstance(segmentation, list) and len(segmentation) > 0
+                    ):  # polygon format
+                        new_ann_instance["segmentation"] = {
+                            "type": "polygon",
+                            "task": "instance",
+                            "data": segmentation,
+                        }
+                    elif (
+                        isinstance(segmentation, dict)
+                        and len(segmentation["counts"]) > 0
+                    ):
+                        new_ann_instance["segmentation"] = {
+                            "type": "mask",
+                            "task": "instance",
+                            "data": segmentation,
+                        }
+
+                if "bbox" in annotation.keys():
+                    new_ann_instance["bbox"] = annotation["bbox"]
+
+                if "keypoints" in annotation.keys():
+                    new_ann_instance["keypoints"] = annotation["keypoints"]
 
-                new_ann[component_name]['annotations'].append(new_ann_instance)
+                new_ann[component_name]["annotations"].append(new_ann_instance)
 
             ## add instance to the dataset
             dataset.add_data(
@@ -207,7 +231,7 @@ def from_coco_format(
             )
 
             ## limit dataset size
-            if dataset_size is not None and iter.n+1 >= dataset_size:
+            if dataset_size is not None and iter.n + 1 >= dataset_size:
                 break
 
         ## Convert to webdataset
@@ -216,15 +240,14 @@ def from_coco_format(
 
     @parsing_wrapper
     def from_voc_format(
-            self,
-            dataset,
-            image_dir,
-            xml_annotation_files_paths,
-            split,
-            dataset_size=None,
-            override_main_component=None
-        ):
-
+        self,
+        dataset,
+        image_dir,
+        xml_annotation_files_paths,
+        split,
+        dataset_size=None,
+        override_main_component=None,
+    ):
         """
         Constructs a LDF dataset from a VOC type dataset.
         Arguments:
@@ -250,7 +273,7 @@ def from_voc_format(
         ## dataset construction loop
         iter = tqdm(xml_annotation_files_paths)
         for xml_annotation_file_path in iter:
-            self.percentage = round((iter.n/iter.total)*100, 2)
+            self.percentage = round((iter.n / iter.total) * 100, 2)
 
             new_ann = {component_name: {"annotations": []}}
 
@@ -273,18 +296,23 @@ def from_voc_format(
                         if info.tag == "name":
                             new_ann_instance["class_name"] = info.text
                             class_id = dataset._add_class(new_ann_instance)
-                            new_ann_instance['class'] = class_id
+                            new_ann_instance["class"] = class_id
                         if info.tag == "bndbox":
                             for point in info:
-                                if point.tag=="xmin":
+                                if point.tag == "xmin":
                                     bbox_xmin = int(point.text)
-                                if point.tag=="ymin":
+                                if point.tag == "ymin":
                                     bbox_ymin = int(point.text)
-                                if point.tag=="xmax":
+                                if point.tag == "xmax":
                                     bbox_xmax = int(point.text)
-                                if point.tag=="ymax":
+                                if point.tag == "ymax":
                                     bbox_ymax = int(point.text)
-                            coco_bbox = [bbox_xmin, bbox_ymin, bbox_xmax-bbox_xmin, bbox_ymax-bbox_ymin] #x_min, y_min, width, height
+                            coco_bbox = [
+                                bbox_xmin,
+                                bbox_ymin,
+                                bbox_xmax - bbox_xmin,
+                                bbox_ymax - bbox_ymin,
+                            ]  # x_min, y_min, width, height
                             new_ann_instance["bbox"] = coco_bbox
 
                     new_ann[component_name]["annotations"].append(new_ann_instance)
@@ -295,7 +323,7 @@ def from_voc_format(
             )
 
             ## limit dataset size
-            if dataset_size is not None and iter.n+1 >= dataset_size:
+            if dataset_size is not None and iter.n + 1 >= dataset_size:
                 break
 
         ## Convert to webdataset
@@ -304,15 +332,15 @@ def from_voc_format(
 
     @parsing_wrapper
     def from_yolo4_format(
-            self,
-            dataset,
-            image_dir,
-            txt_annotations_file_path,
-            classes_txt_file_path,
-            split,
-            dataset_size=None,
-            override_main_component=None
-        ):
+        self,
+        dataset,
+        image_dir,
+        txt_annotations_file_path,
+        classes_txt_file_path,
+        split,
+        dataset_size=None,
+        override_main_component=None,
+    ):
         """
         Constructs a LDF dataset from a YOLO4 type dataset.
         Arguments:
@@ -336,20 +364,20 @@ def from_yolo4_format(
 
         ## get class names
         if not os.path.exists(classes_txt_file_path):
-            raise RuntimeError('classes file path non-existent.')
-        with open(classes_txt_file_path, 'r', encoding='utf-8-sig') as text:
+            raise RuntimeError("classes file path non-existent.")
+        with open(classes_txt_file_path, "r", encoding="utf-8-sig") as text:
             classes = text.readlines()
             classes = [class_name.replace("\n", "") for class_name in classes]
 
         ## dataset construction loop
-        with open(txt_annotations_file_path, 'r', encoding='utf-8-sig') as text:
+        with open(txt_annotations_file_path, "r", encoding="utf-8-sig") as text:
             iter = tqdm(text.readlines())
             for line in iter:
-                self.percentage = round((iter.n/iter.total)*100, 2)
+                self.percentage = round((iter.n / iter.total) * 100, 2)
 
                 new_ann = {component_name: {"annotations": []}}
 
-                image_name = line.split(' ')[0]
+                image_name = line.split(" ")[0]
                 image_path = os.path.join(image_dir, image_name)
                 if os.path.exists(image_path):
                     image = cv2.imread(image_path)
@@ -357,26 +385,35 @@ def from_yolo4_format(
                     warnings.warn(f"skipping {image_path} as it does no exist!")
                     continue
 
-                for annotation in line.split(' ')[1:]:
+                for annotation in line.split(" ")[1:]:
                     new_ann_instance = {}
 
-                    bbox_xmin, bbox_ymin, bbox_xmax, bbox_ymax, class_idx = [int(x) for x in annotation.split(',')]
+                    bbox_xmin, bbox_ymin, bbox_xmax, bbox_ymax, class_idx = [
+                        int(x) for x in annotation.split(",")
+                    ]
 
                     new_ann_instance["class_name"] = classes[class_idx]
                     class_id = dataset._add_class(new_ann_instance)
-                    new_ann_instance['class'] = class_id
-
-                    coco_bbox_format = [bbox_xmin, bbox_ymin, bbox_xmax-bbox_xmin, bbox_ymax-bbox_ymin] #x_min, y_min, width, height
+                    new_ann_instance["class"] = class_id
+
+                    coco_bbox_format = [
+                        bbox_xmin,
+                        bbox_ymin,
+                        bbox_xmax - bbox_xmin,
+                        bbox_ymax - bbox_ymin,
+                    ]  # x_min, y_min, width, height
                     new_ann_instance["bbox"] = coco_bbox_format
                     new_ann[component_name]["annotations"].append(new_ann_instance)
 
                 ## add to dataset
                 dataset.add_data(
-                    self.source_name, {component_name: image, "json": new_ann}, split=split
+                    self.source_name,
+                    {component_name: image, "json": new_ann},
+                    split=split,
                 )
 
                 ## limit dataset size
-                if dataset_size is not None and iter.n+1 >= dataset_size:
+                if dataset_size is not None and iter.n + 1 >= dataset_size:
                     break
 
         ## Convert to webdataset
@@ -385,14 +422,14 @@ def from_yolo4_format(
 
     @parsing_wrapper
     def from_yolo5_format(
-            self,
-            dataset,
-            image_dir,
-            #txt_annotation_files_paths, # list of paths to the text annotation files where each line encodes a bounding box
-            split,
-            dataset_size=None,
-            override_main_component=None
-        ):
+        self,
+        dataset,
+        image_dir,
+        # txt_annotation_files_paths, # list of paths to the text annotation files where each line encodes a bounding box
+        split,
+        dataset_size=None,
+        override_main_component=None,
+    ):
         """
         Constructs a LDF dataset from a YOLO5 type dataset.
         Arguments:
@@ -415,21 +452,23 @@ def from_yolo5_format(
             component_name = dataset.sources[self.source_name].main_component
 
         ## get class names
-        root_path = os.path.split(os.path.split(image_dir)[0])[0] #go two levels up
-        yaml_files = [fname for fname in os.listdir(root_path) if fname.endswith('.yaml')]
+        root_path = os.path.split(os.path.split(image_dir)[0])[0]  # go two levels up
+        yaml_files = [
+            fname for fname in os.listdir(root_path) if fname.endswith(".yaml")
+        ]
         if len(yaml_files) > 1:
-            raise RuntimeError('Multiple YAML files - possible ambiguity')
+            raise RuntimeError("Multiple YAML files - possible ambiguity")
         else:
             yaml_path = yaml_files[0]
-            yolo = yaml.safe_load(Path(os.path.join(root_path,yaml_path)).read_text())
-            classes = yolo['names']
-            if yolo['nc'] != len(classes):
+            yolo = yaml.safe_load(Path(os.path.join(root_path, yaml_path)).read_text())
+            classes = yolo["names"]
+            if yolo["nc"] != len(classes):
                 raise Exception(f"nc in YOLO YAML file does not match names!")
 
         ## dataset construction loop
         iter = tqdm(os.listdir(image_dir))
         for image_name in iter:
-            self.percentage = round((iter.n/iter.total)*100, 2)
+            self.percentage = round((iter.n / iter.total) * 100, 2)
 
             image_path = os.path.join(image_dir, image_name)
             if os.path.exists(image_path):
@@ -438,40 +477,50 @@ def from_yolo5_format(
                 warnings.warn(f"skipping {image_path} as it does no exist!")
                 continue
 
-            ext = image_path.split('.')[-1]
-            label_path = image_path.replace('/images/', '/labels/').replace(f".{ext}", '.txt')
+            ext = image_path.split(".")[-1]
+            label_path = image_path.replace("/images/", "/labels/").replace(
+                f".{ext}", ".txt"
+            )
             if not os.path.exists(label_path):
-                warnings.warn(f"Skipping image {image_path} - label {label_path} not found!")
+                warnings.warn(
+                    f"Skipping image {image_path} - label {label_path} not found!"
+                )
                 continue
 
             # add YOLO data
-            new_ann = {component_name: {'annotations': []}}
+            new_ann = {component_name: {"annotations": []}}
 
             h, w = image.shape[0], image.shape[1]
 
             with open(label_path) as file:
                 lines = file.readlines()
                 rows = len(lines)
-                data = np.array([float(num) for line in lines for num in line.replace('\n','').split(' ')])
+                data = np.array(
+                    [
+                        float(num)
+                        for line in lines
+                        for num in line.replace("\n", "").split(" ")
+                    ]
+                )
                 data = data.reshape(rows, -1).tolist()
 
             for row in data:
                 new_ann_instance = {}
                 yolo_class_id = int(row[0])
                 class_name = classes[yolo_class_id]
-                new_ann_instance['class_name'] = class_name
+                new_ann_instance["class_name"] = class_name
                 class_id = dataset._add_class(new_ann_instance)
-                new_ann_instance['class'] = class_id
+                new_ann_instance["class"] = class_id
 
                 # bounding box
 
-                bbox_xcenter = row[1]*w
-                bbox_ycenter = row[2]*h
+                bbox_xcenter = row[1] * w
+                bbox_ycenter = row[2] * h
 
-                #bbox_xmin = annotation["coordinates"]["x"]
-                #bbox_ymin = annotation["coordinates"]["y"]
-                bbox_width = row[3]*w
-                bbox_height = row[4]*h
+                # bbox_xmin = annotation["coordinates"]["x"]
+                # bbox_ymin = annotation["coordinates"]["y"]
+                bbox_width = row[3] * w
+                bbox_height = row[4] * h
 
                 bbox_xmin = bbox_xcenter - bbox_width / 2
                 bbox_ymin = bbox_ycenter - bbox_height / 2
@@ -480,10 +529,9 @@ def from_yolo5_format(
                 new_ann_instance["bbox"] = coco_bbox_format
                 new_ann[component_name]["annotations"].append(new_ann_instance)
 
-
                 #############################
-                #new_ann_instance['bbox'] = [, , , ]
-                #new_ann[component_name]['annotations'].append(new_ann_instance)
+                # new_ann_instance['bbox'] = [, , , ]
+                # new_ann[component_name]['annotations'].append(new_ann_instance)
 
             ## add to dataset
             dataset.add_data(
@@ -491,7 +539,7 @@ def from_yolo5_format(
             )
 
             ## limit dataset size
-            if dataset_size is not None and iter.n+1 >= dataset_size:
+            if dataset_size is not None and iter.n + 1 >= dataset_size:
                 break
 
         ## Convert to webdataset
@@ -500,14 +548,14 @@ def from_yolo5_format(
 
     @parsing_wrapper
     def from_tfodc_format(
-            self,
-            dataset,
-            image_dir,
-            csv_file_path,
-            split,
-            dataset_size=None,
-            override_main_component=None
-        ):
+        self,
+        dataset,
+        image_dir,
+        csv_file_path,
+        split,
+        dataset_size=None,
+        override_main_component=None,
+    ):
         """
         Constructs a LDF dataset from a TFObjectDetectionCSV type dataset.
         Arguments:
@@ -532,9 +580,8 @@ def from_tfodc_format(
 
         ## extract annotations for each image
         annotations = {}
-        reader = csv.reader(open(csv_file_path),delimiter=',')
+        reader = csv.reader(open(csv_file_path), delimiter=",")
         for n, row in enumerate(reader):
-
             if n == 0:
                 idx_fname = row.index("filename")
                 idx_class = row.index("class")
@@ -558,7 +605,7 @@ def from_tfodc_format(
         ## dataset construction loop
         iter = tqdm(annotations.keys())
         for image_name in iter:
-            self.percentage = round((iter.n/iter.total)*100, 2)
+            self.percentage = round((iter.n / iter.total) * 100, 2)
 
             image_path = os.path.join(image_dir, image_name)
             if os.path.exists(image_path):
@@ -576,9 +623,14 @@ def from_tfodc_format(
 
                 new_ann_instance["class_name"] = class_name
                 class_id = dataset._add_class(new_ann_instance)
-                new_ann_instance['class'] = class_id
-
-                coco_bbox = [bbox_xmin, bbox_ymin, bbox_xmax-bbox_xmin, bbox_ymax-bbox_ymin] #x_min, y_min, width, height
+                new_ann_instance["class"] = class_id
+
+                coco_bbox = [
+                    bbox_xmin,
+                    bbox_ymin,
+                    bbox_xmax - bbox_xmin,
+                    bbox_ymax - bbox_ymin,
+                ]  # x_min, y_min, width, height
                 new_ann_instance["bbox"] = coco_bbox
                 new_ann[component_name]["annotations"].append(new_ann_instance)
 
@@ -588,7 +640,7 @@ def from_tfodc_format(
             )
 
             ## limit dataset size
-            if dataset_size is not None and iter.n+1 >= dataset_size:
+            if dataset_size is not None and iter.n + 1 >= dataset_size:
                 break
 
         ## Convert to webdataset
@@ -597,14 +649,14 @@ def from_tfodc_format(
 
     @parsing_wrapper
     def from_cml_format(
-            self,
-            dataset,
-            image_dir,
-            annotation_path,
-            split,
-            dataset_size=None,
-            override_main_component=None
-        ):
+        self,
+        dataset,
+        image_dir,
+        annotation_path,
+        split,
+        dataset_size=None,
+        override_main_component=None,
+    ):
         """
         Constructs a LDF dataset from a CreateML type dataset.
         Arguments:
@@ -634,7 +686,7 @@ def from_cml_format(
         ## dataset construction loop
         iter = tqdm(cml_annotations)
         for annotations_instance in iter:
-            self.percentage = round((iter.n/iter.total)*100, 2)
+            self.percentage = round((iter.n / iter.total) * 100, 2)
 
             image_name = annotations_instance["image"]
             image_path = os.path.join(image_dir, image_name)
@@ -648,7 +700,6 @@ def from_cml_format(
 
             annotations = annotations_instance["annotations"]
             for annotation in annotations:
-
                 class_name = annotation["label"]
                 bbox_xcenter = annotation["coordinates"]["x"]
                 bbox_ycenter = annotation["coordinates"]["y"]
@@ -663,7 +714,7 @@ def from_cml_format(
 
                 new_ann_instance["class_name"] = class_name
                 class_id = dataset._add_class(new_ann_instance)
-                new_ann_instance['class'] = class_id
+                new_ann_instance["class"] = class_id
 
                 coco_bbox_format = [bbox_xmin, bbox_ymin, bbox_width, bbox_height]
                 new_ann_instance["bbox"] = coco_bbox_format
@@ -675,7 +726,7 @@ def from_cml_format(
             )
 
             ## limit dataset size
-            if dataset_size is not None and iter.n+1 >= dataset_size:
+            if dataset_size is not None and iter.n + 1 >= dataset_size:
                 break
 
         ## Convert to webdataset
@@ -684,14 +735,14 @@ def from_cml_format(
 
     @parsing_wrapper
     def from_numpy_format(
-            self,
-            dataset,
-            images,
-            labels,
-            split,
-            dataset_size=None,
-            override_main_component=None
-        ):
+        self,
+        dataset,
+        images,
+        labels,
+        split,
+        dataset_size=None,
+        override_main_component=None,
+    ):
         """
         Constructs a LDF dataset from data provided in numpy arrays.
         Arguments:
@@ -720,7 +771,7 @@ def from_numpy_format(
         ## dataset construction loop
         iter = tqdm(zip(images, labels))
         for image, label in iter:
-            self.percentage = round((iter.n/iter.total)*100, 2)
+            self.percentage = round((iter.n / iter.total) * 100, 2)
 
             new_ann = {component_name: {"annotations": []}}
             new_ann_instance = {}
@@ -736,13 +787,13 @@ def from_numpy_format(
         dataset.to_webdataset(split, query)
 
     def train_test_split_image_classification_directory_tree(
-            self,
-            directory_tree_path,
-            destination_path1,
-            destination_path2,
-            split_proportion,
-            folders_to_ignore=[]
-        ):
+        self,
+        directory_tree_path,
+        destination_path1,
+        destination_path2,
+        split_proportion,
+        folders_to_ignore=[],
+    ):
         """
         Split directory tree into two parts. This is useful as some classification directory tree format datasets
         (e.g. Caltech101) do not separately provide data for training, validation and testing.
@@ -757,13 +808,12 @@ def train_test_split_image_classification_directory_tree(
         """
 
         for folder_name in os.listdir(directory_tree_path):
-
             if folder_name in folders_to_ignore:
                 continue
 
             image_names = os.listdir(os.path.join(directory_tree_path, folder_name))
             random.shuffle(image_names)
-            split_idx = int(len(image_names)*split_proportion)
+            split_idx = int(len(image_names) * split_proportion)
             image_names1 = image_names[:split_idx]
             image_names2 = image_names[split_idx:]
 
@@ -771,21 +821,25 @@ def train_test_split_image_classification_directory_tree(
             os.mkdir(os.path.join(destination_path2, folder_name))
 
             for image_name in image_names1:
-                shutil.copyfile(src=os.path.join(directory_tree_path, folder_name, image_name),
-                                dst=os.path.join(destination_path1, folder_name, image_name))
+                shutil.copyfile(
+                    src=os.path.join(directory_tree_path, folder_name, image_name),
+                    dst=os.path.join(destination_path1, folder_name, image_name),
+                )
             for image_name in image_names2:
-                shutil.copyfile(src=os.path.join(directory_tree_path, folder_name, image_name),
-                                dst=os.path.join(destination_path2, folder_name, image_name))
+                shutil.copyfile(
+                    src=os.path.join(directory_tree_path, folder_name, image_name),
+                    dst=os.path.join(destination_path2, folder_name, image_name),
+                )
 
     @parsing_wrapper
     def from_image_classification_directory_tree_format(
-            self,
-            dataset,
-            class_folders_paths,
-            split,
-            dataset_size=None,
-            override_main_component=None
-        ):
+        self,
+        dataset,
+        class_folders_paths,
+        split,
+        dataset_size=None,
+        override_main_component=None,
+    ):
         """
         Constructs a LDF dataset from a directory tree whose subfolders define image classes.
         Arguments:
@@ -806,7 +860,7 @@ def from_image_classification_directory_tree_format(
             component_name = dataset.sources[self.source_name].main_component
 
         if len(class_folders_paths) == 0:
-            raise RuntimeError('Directory tree is empty')
+            raise RuntimeError("Directory tree is empty")
 
         ## dataset construction loop
         iter1 = tqdm(class_folders_paths)
@@ -814,7 +868,9 @@ def from_image_classification_directory_tree_format(
             class_folder_name = os.path.split(class_folder_path)[-1]
             iter2 = tqdm(os.listdir(class_folder_path))
             for image_name in iter2:
-                self.percentage = round((iter1.n/iter1.total)*100, 2) + round(((iter2.n/iter2.total)/iter1.total)*100, 2)
+                self.percentage = round((iter1.n / iter1.total) * 100, 2) + round(
+                    ((iter2.n / iter2.total) / iter1.total) * 100, 2
+                )
 
                 image_path = os.path.join(class_folder_path, image_name)
                 if os.path.exists(image_path):
@@ -831,11 +887,13 @@ def from_image_classification_directory_tree_format(
 
                 ## add data to the provided LDF dataset instance
                 dataset.add_data(
-                    self.source_name, {component_name: image, "json": new_ann}, split=split
+                    self.source_name,
+                    {component_name: image, "json": new_ann},
+                    split=split,
                 )
 
                 ## limit dataset size
-                if dataset_size is not None and iter.n+1 >= dataset_size:
+                if dataset_size is not None and iter.n + 1 >= dataset_size:
                     break
 
         ## Convert to webdataset
@@ -844,15 +902,15 @@ def from_image_classification_directory_tree_format(
 
     @parsing_wrapper
     def from_image_classification_with_text_annotations_format(
-            self,
-            dataset,
-            image_dir,
-            info_file_path,
-            split,
-            delimiter=" ",
-            dataset_size=None,
-            override_main_component=None
-        ):
+        self,
+        dataset,
+        image_dir,
+        info_file_path,
+        split,
+        delimiter=" ",
+        dataset_size=None,
+        override_main_component=None,
+    ):
         """
         Constructs a LDF dataset based on image paths and labels from text annotations.
         Arguments:
@@ -875,18 +933,20 @@ def from_image_classification_with_text_annotations_format(
             component_name = dataset.sources[self.source_name].main_component
 
         if not os.path.exists(info_file_path):
-            raise RuntimeError('Info file path non-existent.')
+            raise RuntimeError("Info file path non-existent.")
 
         ## dataset construction loop
         with open(info_file_path) as f:
             lines = f.readlines()
             iter = tqdm(lines)
             for line in iter:
-                self.percentage = round((iter.n/iter.total)*100, 2)
+                self.percentage = round((iter.n / iter.total) * 100, 2)
                 try:
                     image_path, label = line.split(delimiter)
                 except:
-                    raise RuntimeError('Unable to split the info file based on the provided delimiter.')
+                    raise RuntimeError(
+                        "Unable to split the info file based on the provided delimiter."
+                    )
 
                 label = label.strip()
 
@@ -904,11 +964,13 @@ def from_image_classification_with_text_annotations_format(
 
                 ## add data to the provided LDF dataset instance
                 dataset.add_data(
-                    self.source_name, {component_name: image, "json": new_ann}, split=split
+                    self.source_name,
+                    {component_name: image, "json": new_ann},
+                    split=split,
                 )
 
                 ## limit dataset size
-                if dataset_size is not None and iter.n+1 >= dataset_size:
+                if dataset_size is not None and iter.n + 1 >= dataset_size:
                     break
 
         # Convert to webdataset
@@ -938,22 +1000,28 @@ def get_error_message(self):
     def get_parsing_in_progress(self):
         return self.parsing_in_progress
 
-    def parse_to_dataset(self, dataset_type, team_name, dataset_name, *args, new_thread = False, **kwargs):
+    def parse_to_dataset(
+        self, dataset_type, team_name, dataset_name, *args, new_thread=False, **kwargs
+    ):
         # Order of args passed in parsing functions has to be (self, dataset_info) + other args
 
         dataset_info = (team_name, dataset_name, dataset_type)
         if not new_thread:
-            LuxonisParser.DATASET_TYPE_TO_FUNCTION[dataset_type](self, dataset_info, *args, **kwargs)
+            LuxonisParser.DATASET_TYPE_TO_FUNCTION[dataset_type](
+                self, dataset_info, *args, **kwargs
+            )
         else:
+
             def thread_exception_hook(args):
                 self.error_message = str(args.exc_value)
                 self.parsing_in_progress = False
+
             threading.excepthook = thread_exception_hook
 
             self.thread = threading.Thread(
                 target=Parser.DATASET_TYPE_TO_FUNCTION[dataset_type],
                 args=(self, dataset_info) + args,
                 kwargs=kwargs,
-                daemon=True
+                daemon=True,
             )
             self.thread.start()
diff --git a/src/luxonis_ml/ops/session.py b/src/luxonis_ml/data/session.py
similarity index 88%
rename from src/luxonis_ml/ops/session.py
rename to src/luxonis_ml/data/session.py
index c9d850af..70cb0467 100644
--- a/src/luxonis_ml/ops/session.py
+++ b/src/luxonis_ml/data/session.py
@@ -3,12 +3,15 @@
 from pathlib import Path
 import fiftyone as fo
 
+
 class SessionThread(threading.Thread):
     def __init__(self):
         super(SessionThread, self).__init__()
 
     def start(self, cmd):
-        self.process = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True)
+        self.process = subprocess.Popen(
+            cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True
+        )
 
     def stop(self):
         # print(self.process)
@@ -21,13 +24,15 @@ def stop(self):
         return_code = self.process.wait()
         print(f"Process terminated with exit code {return_code}")
 
-class LuxonisSession:
 
+class LuxonisSession:
     def __init__(self):
         self.mount_path = f"{str(Path.home())}/.luxonis_mount"
         os.makedirs(self.mount_path, exist_ok=True)
 
-        credentials_cache_file = f'{str(Path.home())}/.cache/luxonis_ml/credentials.json'
+        credentials_cache_file = (
+            f"{str(Path.home())}/.cache/luxonis_ml/credentials.json"
+        )
         if os.path.exists(credentials_cache_file):
             with open(credentials_cache_file) as file:
                 self.creds = json.load(file)
@@ -44,7 +49,6 @@ def _get_cmd(self):
         return f"s3fs luxonis-test-bucket {self.mount_path} -o passwd_file={str(Path.home())}/.passwd-s3fs -o curldbg -o url={self._get_credentials('AWS_S3_ENDPOINT_URL')} -o use_path_request_style"
 
     def _app_worker(self, dataset):
-
         app_session = fo.launch_app(dataset.fo_dataset)
 
         while True:
diff --git a/src/luxonis_ml/data/utils/__init__.py b/src/luxonis_ml/data/utils/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/luxonis_ml/data/utils/data_utils.py b/src/luxonis_ml/data/utils/data_utils.py
new file mode 100644
index 00000000..b0e22864
--- /dev/null
+++ b/src/luxonis_ml/data/utils/data_utils.py
@@ -0,0 +1,285 @@
+import numpy as np
+import hashlib
+import fiftyone as fo
+import os
+import subprocess
+import uuid
+
+
+def get_granule(filepath, addition, component_name):
+    granule = filepath.split("/")[-1]
+    if "_new_image_name" in addition[component_name].keys():
+        filepath = filepath.replace(
+            granule, addition[component_name]["_new_image_name"]
+        )
+        granule = addition[component_name]["_new_image_name"]
+    return granule
+
+
+def check_media(dataset, filepath, mount_path, component_name, granule, from_bucket):
+    source = dataset.source
+    if dataset.bucket_type == "local":
+        local_path = f"{dataset.path}/{mount_path}"
+        md5 = hashlib.md5()
+        with open(filepath, "rb") as file:
+            for chunk in iter(lambda: file.read(8192), b""):
+                md5.update(chunk)
+        old_checksum = md5.hexdigest()
+    elif dataset.bucket_type == "aws":
+        prefix = f"{dataset.bucket_path}/{component_name}/{granule}"
+        resp = dataset.client.list_objects(
+            Bucket=dataset.bucket, Prefix=prefix, Delimiter="/", MaxKeys=1
+        )
+        exists = "Contents" in resp
+        if not exists:
+            raise Exception("Issue with dataset! Sample file does not exist on S3!")
+        old_checksum = resp["Contents"][0]["ETag"].strip('"')
+
+    if from_bucket:
+        if dataset.bucket_type == "aws":
+            resp = dataset.client.list_objects(
+                Bucket=dataset.bucket, Prefix=filepath, Delimiter="/", MaxKeys=1
+            )
+            exists = "Contents" in resp
+            if not exists:
+                raise Exception(f"File {filepath} not found on S3!")
+            new_checksum = resp["Contents"][0]["ETag"].strip('"')
+    else:
+        md5 = hashlib.md5()
+        with open(filepath, "rb") as file:
+            for chunk in iter(lambda: file.read(8192), b""):
+                md5.update(chunk)
+        new_checksum = md5.hexdigest()
+    if new_checksum != old_checksum:
+        if dataset.bucket_type == "local":
+            new_path = f"{dataset.path}/{dataset.team}/datasets/{dataset.name}/archive/{source.name}/{component_name}/{granule}"
+            os.makedirs(new_path, exist_ok=True)
+            cmd = f"cp {local_path} {new_path}"
+            subprocess.check_output(cmd, shell=True)
+        elif dataset.bucket_type == "aws":
+            new_prefix = f"{dataset.bucket_path}/archive/{source.name}/{component_name}/{granule}"
+            dataset.client.copy_object(
+                Bucket=dataset.bucket,
+                Key=new_prefix,
+                CopySource={"Bucket": dataset.bucket, "Key": prefix},
+            )
+        # update existing datasat sample to point to this new_prefix
+        # limitation: all old versions point to the latest archive only (TODO: fix)
+        for sample in sample_view:
+            sample.filepath = mount_path.replace(
+                f"/{dataset.team}/{dataset.name}/",
+                f"/{dataset.team}/{dataset.name}/archive/",
+            )
+            sample.save()
+
+        return False
+    else:
+        return True
+
+
+def check_classification(val1, val2):
+    if isinstance(val1, str):
+        # prevent zip from taking letters only
+        val1 = [val1]
+
+    # TODO: Should be updated to work properly with multi-label classification
+    # This includes the bottom for loop which seems to return only the first value where dict is not identical.
+
+    if val2 is None and val1 is not None:
+        val1 = fo.Classification(label=val1[0]).to_dict()
+        return [{"class": val1}]
+
+    for val1, val2 in list(zip(val1, val2["classifications"])):
+        val1 = fo.Classification(label=val1).to_dict()
+        if len(val1.keys()) == len(val2.keys()):
+            for key in val1:
+                if not key.startswith("_"):
+                    if val1[key] != val2[key]:
+                        return [{"class": val1}]
+        else:
+            return [{"class": val1}]
+    return []
+
+
+def check_boxes(dataset, val1, val2):
+    if val2 is None and val1 is not None:
+        return [{"boxes": val1}]
+
+    if len(val1) == len(val2["detections"]):
+        for val1, val2 in list(zip(val1, val2["detections"])):
+            if isinstance(val1[0], str) and val2["label"] != val1[0]:
+                return [{"boxes": val1}]
+            if (
+                not isinstance(val1[0], str)
+                and val2["label"] != dataset.fo_dataset.classes["boxes"][val1[0]]
+            ):
+                val1[0] = dataset.fo_dataset.classes["boxes"][val1[0]]
+                return [{"boxes": val1}]
+            for c1, c2 in list(zip(val1[1:], val2["bounding_box"])):
+                if abs(c1 - c2) > 1e-8:
+                    return [{"boxes": val1}]
+    else:
+        return [{"boxes": val1}]
+    return []
+
+
+def check_segmentation(val1, val2):
+    if (val1.shape != val2.shape) or (np.linalg.norm(val1 - val2) > 1e-8):
+        return [{"segmentation": val1}]
+    return []
+
+
+def check_keypoints(dataset, val1, val2):
+    if len(val1) == len(val2["keypoints"]):
+        for val1, val2 in list(zip(val1, val2["keypoints"])):
+            if isinstance(val1[0], str) and val2["label"] != val1[0]:
+                return [{"keypoints": val1}]
+            if (
+                not isinstance(val1[0], str)
+                and val2["label"] != dataset.fo_dataset.classes["keypoints"][val1[0]]
+            ):
+                val1[0] = dataset.fo_dataset.classes["keypoints"][val1[0]]
+                return [{"keypoints": val1}]
+            for c1, c2 in list(zip(val1[1], val2["points"])):
+                if not (np.isnan(c1[0]) and isinstance(c2[0], dict)):
+                    if not np.isnan(c1[0]) and not isinstance(c2[0], dict):
+                        if c1[0] != c2[0] or c1[1] != c2[1]:
+                            return [{"keypoints": val1}]
+                    else:
+                        return [{"keypoints": val1}]
+
+    else:
+        return [{"keypoints": val1}]
+    return []
+
+
+def check_fields(dataset, latest_sample, addition, component_name):
+    changes = []
+
+    ignore_fields_match = set(
+        [
+            dataset.source.name,
+            "version",
+            "metadata",
+            "latest",
+            "tags",
+            "tid",
+            "_group",
+            "_old_filepath",
+            "_new_image_name",
+        ]
+    )
+    ignore_fields_check = set(
+        ["filepath", "_group", "_old_filepath", "_new_image_name"]
+    )
+
+    sample_dict = latest_sample.to_dict()
+    f1 = set(addition[component_name].keys())
+    f2 = set(sample_dict.keys())
+    new = f1 - f1.intersection(f2) - ignore_fields_match
+    missing = f2 - f1.intersection(f2) - ignore_fields_match
+    if len(new):
+        for new_field in list(new):
+            changes.append({new_field: addition[component_name][new_field]})
+    if len(missing):
+        for missing_field in list(missing):
+            changes.append({missing_field: None})
+
+    check_fields = list(f1.intersection(f2) - ignore_fields_check)
+
+    for field in check_fields:
+        val1 = addition[component_name][field]
+        val2 = sample_dict[field]
+
+        if field in dataset.tasks:
+            if field == "class":
+                changes += check_classification(val1, val2)
+            elif field == "boxes":
+                changes += check_boxes(dataset, val1, val2)
+            elif field == "segmentation":
+                val2 = latest_sample.segmentation.mask
+                changes += check_segmentation(val1, val2)
+            elif field == "keypoints":
+                changes += check_keypoints(dataset, val1, val2)
+            else:
+                raise NotImplementedError()
+        else:
+            if not val1 == val2:
+                changes.append({field: val1})
+
+    return changes
+
+
+def construct_class_label(dataset, classes):
+    if not isinstance(classes, list):  # fix for only one class
+        classes = [classes]
+
+    classifications = []
+    for cls in classes:
+        label = None
+        if isinstance(cls, str):
+            label = cls
+        elif isinstance(cls, dict):
+            label = cls.get("label")
+        else:
+            label = dataset.fo_dataset.classes["class"][int(cls)]
+
+        classifications.append(fo.Classification(label=label))
+
+    return fo.Classifications(classifications=classifications)
+
+
+def construct_boxes_label(dataset, boxes):
+    if not isinstance(boxes[0], list):  # fix for only one box without a nested list
+        boxes = [boxes]
+    return fo.Detections(
+        detections=[
+            fo.Detection(
+                label=box[0]
+                if isinstance(box[0], str)
+                else dataset.fo_dataset.classes["boxes"][int(box[0])],
+                bounding_box=box[1:5],
+            )
+            for box in boxes
+        ]
+    )
+
+
+def construct_segmentation_label(dataset, mask):
+    if isinstance(mask, list):
+        mask = np.array(mask)
+    return fo.Segmentation(mask=mask)
+
+
+def construct_keypoints_label(dataset, kps):
+    if not isinstance(kps[0], list):  # fix for only one kp without a nested list
+        kps = [kps]
+    return fo.Keypoints(
+        keypoints=[
+            fo.Keypoint(
+                label=kp[0]
+                if isinstance(kp[0], str)
+                else dataset.fo_dataset.classes["keypoints"][int(kp[0])],
+                points=kp[1],
+            )
+            for kp in kps
+        ]
+    )
+
+
+def generate_hashname(filepath):
+    # Read the contents of the file
+    with open(filepath, "rb") as file:
+        file_contents = file.read()
+
+    # Generate the UUID5 based on the file contents and the NAMESPACE_URL
+    file_hash_uuid = uuid.uuid5(uuid.NAMESPACE_URL, file_contents.hex())
+
+    return str(file_hash_uuid) + os.path.splitext(filepath)[1]
+
+
+def is_modified_filepath(dataset, filepath):
+    if filepath.startswith(f"/{dataset.team_id}"):
+        return True
+    else:
+        return False
diff --git a/src/luxonis_ml/data/utils/s3_utils.py b/src/luxonis_ml/data/utils/s3_utils.py
new file mode 100644
index 00000000..b8cf8ea0
--- /dev/null
+++ b/src/luxonis_ml/data/utils/s3_utils.py
@@ -0,0 +1,109 @@
+import os
+import boto3
+from botocore.exceptions import NoCredentialsError, ClientError
+import datetime
+from concurrent.futures import ThreadPoolExecutor
+import hashlib
+
+
+def get_file_hash(local_file):
+    local_file_hash = None
+    with open(local_file, "rb") as f:
+        local_file_hash = hashlib.md5(f.read()).hexdigest()
+    return local_file_hash
+
+
+def download_file(bucket, s3_file, local_dir):
+    try:
+        local_file_path = os.path.join(local_dir, s3_file.key)
+        if (
+            os.path.exists(local_file_path)
+            and get_file_hash(local_file_path) == s3_file.e_tag[1:-1]
+        ):
+            print(f"File {local_file_path} already exists and is up to date.")
+        else:
+            bucket.download_file(s3_file.key, local_file_path)
+            print(f"Downloaded {s3_file.key} to {local_file_path}")
+    except Exception as e:
+        print(f"Failed to download {s3_file.key}. Reason: {e}")
+
+
+def sync_from_s3(non_streaming_dir, bucket, bucket_dir, endpoint_url=None):
+    os.makedirs(non_streaming_dir, exist_ok=True)
+
+    print("Syncing from cloud...")
+    s3 = boto3.resource("s3", endpoint_url=endpoint_url)
+    bucket = s3.Bucket(bucket)
+    try:
+        with ThreadPoolExecutor(
+            min(os.cpu_count(), int(os.getenv("MAX_S3_CONCURRENCY", 4)))
+        ) as executor:
+            for s3_file in bucket.objects.filter(Prefix=bucket_dir):
+                if not os.path.exists(
+                    os.path.dirname(non_streaming_dir + "/" + s3_file.key)
+                ):
+                    os.makedirs(os.path.dirname(non_streaming_dir + "/" + s3_file.key))
+                executor.submit(download_file, bucket, s3_file, non_streaming_dir)
+
+    except NoCredentialsError:
+        print("No AWS credentials found")
+    except Exception as e:
+        print("Unable to download files. Reason:", e)
+
+
+def upload_file(bucket, local_file, s3_file):
+    try:
+        # Check if S3 object already exists
+        try:
+            s3_object = bucket.Object(s3_file)
+            s3_object.load()
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "404":  # The object does not exist.
+                s3_object = None
+            else:  # Something else has gone wrong.
+                raise RuntimeError(f"Error while validating the object on S3.")
+
+        local_file_hash = get_file_hash(local_file)
+
+        if s3_object is None or s3_object.e_tag[1:-1] != local_file_hash:
+            bucket.upload_file(local_file, s3_file)
+            print(f"Uploaded {local_file} to {s3_file}")
+        else:
+            print(f"File {s3_file} is already up to date.")
+    except Exception as e:
+        print(f"Failed to upload {local_file}. Reason: {e}")
+
+
+def sync_to_s3(bucket, s3_dir, local_dir, endpoint_url=None):
+    print("Syncing to cloud...")
+    s3 = boto3.resource("s3", endpoint_url=endpoint_url)
+    bucket = s3.Bucket(bucket)
+
+    try:
+        files_to_upload = [
+            os.path.join(dp, f) for dp, _, fn in os.walk(local_dir) for f in fn
+        ]
+
+        with ThreadPoolExecutor() as executor:
+            for local_file_path in files_to_upload:
+                s3_dest = os.path.join(
+                    s3_dir, os.path.relpath(local_file_path, local_dir)
+                )
+                executor.submit(upload_file, bucket, local_file_path, s3_dest)
+
+    except NoCredentialsError:
+        print("No AWS credentials found")
+    except Exception as e:
+        print("Unable to upload files. Reason:", e)
+
+
+def check_s3_file_existence(bucket, s3_file, endpoint_url=None):
+    s3 = boto3.resource("s3", endpoint_url=endpoint_url)
+    bucket = s3.Bucket(bucket)
+
+    try:
+        s3_object = bucket.Object(s3_file)
+        s3_object.load()
+        return True
+    except ClientError as e:
+        return False
diff --git a/src/luxonis_ml/ops/version.py b/src/luxonis_ml/data/version.py
similarity index 58%
rename from src/luxonis_ml/ops/version.py
rename to src/luxonis_ml/data/version.py
index 136f88db..42f354a5 100644
--- a/src/luxonis_ml/ops/version.py
+++ b/src/luxonis_ml/data/version.py
@@ -1,10 +1,9 @@
 import fiftyone.core.odm as foo
-import luxonis_ml.fiftyone_plugins as fop
-from bson import ObjectId
+import luxonis_ml.data.fiftyone_plugins as fop
 from datetime import datetime
 
-class LuxonisVersion:
 
+class LuxonisVersion:
     def __init__(self, dataset, version_number=None, samples=None, note=None):
         """
         dataset: LuxonisDataset instance
@@ -14,38 +13,47 @@ def __init__(self, dataset, version_number=None, samples=None, note=None):
 
         self.conn = foo.get_db_conn()
 
-        self.version_number = dataset.version if version_number is None else version_number
+        self.version_number = (
+            dataset.version if version_number is None else version_number
+        )
         dataset_id = dataset.fo_dataset._doc.id
         self.dataset_id_str = dataset_id.binary.hex()
 
         if not self._exists():
-
             if samples is None:
                 raise Exception("samples must be defined for creating a new version")
 
-            _id = ObjectId()
-
             version_doc = fop.VersionDocument(
-                id=_id,
                 number=self.version_number,
                 dataset_id=dataset_id,
                 dataset_id_str=self.dataset_id_str,
                 created_at=datetime.utcnow(),
                 samples=samples,
-                note=note
+                note=note,
             )
 
             version_doc.save(upsert=True)
-        else:
-            print("this already exists")
 
         self.doc = self._find()
 
     def _exists(self):
-        return bool(list(self.conn.version_document.find({ "$and": [{"dataset_id_str": self.dataset_id_str}, {"number": self.version_number}] }).limit(1)))
+        return bool(
+            list(
+                self.conn.version_document.find(
+                    {
+                        "$and": [
+                            {"dataset_id_str": self.dataset_id_str},
+                            {"number": self.version_number},
+                        ]
+                    }
+                ).limit(1)
+            )
+        )
 
     def _find(self):
-        return self.conn.version_document.find({"dataset_id_str": self.dataset_id_str, "number": self.version_number})[0]
+        return self.conn.version_document.find(
+            {"dataset_id_str": self.dataset_id_str, "number": self.version_number}
+        )[0]
 
     def get_samples(self):
-        return self.doc['samples']
+        return self.doc["samples"]
diff --git a/src/luxonis_ml/loader/__init__.py b/src/luxonis_ml/loader/__init__.py
new file mode 100644
index 00000000..c071f444
--- /dev/null
+++ b/src/luxonis_ml/loader/__init__.py
@@ -0,0 +1,2 @@
+from .loader import LuxonisLoader
+from .augmentations import Augmentations, TrainAugmentations, ValAugmentations
diff --git a/src/luxonis_ml/loader/augmentations.py b/src/luxonis_ml/loader/augmentations.py
new file mode 100644
index 00000000..374c6f58
--- /dev/null
+++ b/src/luxonis_ml/loader/augmentations.py
@@ -0,0 +1,230 @@
+import albumentations as A
+from albumentations import *
+from albumentations.pytorch import ToTensorV2
+import numpy as np
+import cv2
+import torch
+
+
+class Augmentations:
+    def __init__(self, train_rgb: bool = True):
+        """Base class for creating Augmentations object"""
+        self.train_rgb = train_rgb
+
+    def _parse_cfg(
+        self, image_size: list, augmentations: dict, keep_aspect_ratio: bool = True
+    ):
+        """Parses provided config and returns Albumentations Compose object"""
+        image_size = image_size
+
+        # Always perform Resize
+        if keep_aspect_ratio:
+            resize = A.Sequential(
+                [
+                    A.LongestMaxSize(max_size=max(image_size), interpolation=1),
+                    A.PadIfNeeded(
+                        min_height=image_size[0],
+                        min_width=image_size[1],
+                        border_mode=0,
+                        value=(0, 0, 0),
+                    ),
+                ],
+                p=1,
+            )
+        else:
+            resize = A.Resize(image_size[0], image_size[1])
+
+        all_augs = [resize]
+        if augmentations:
+            for aug in augmentations:
+                all_augs.append(eval(aug["name"])(**aug.get("params", {})))
+        all_augs.append(ToTensorV2())
+
+        return A.Compose(
+            all_augs,
+            bbox_params=A.BboxParams(format="coco", label_fields=["bbox_classes"]),
+            keypoint_params=A.KeypointParams(
+                format="xy", label_fields=["keypoints_classes"], remove_invisible=False
+            ),
+        )
+
+    def __call__(self, data: tuple):
+        """Performs augmentations on provided data"""
+        img, anno_dict = data
+        present_annotations = anno_dict.keys()
+        img_in = img.numpy()
+
+        # albumentations expects with RGB image with HWC format
+        img_in = np.transpose(img_in, (1, 2, 0))
+        img_in = cv2.cvtColor(img_in, cv2.COLOR_BGR2RGB)
+
+        classes = anno_dict.get("class", torch.zeros(1))
+
+        seg = anno_dict.get("segmentation", torch.zeros((1, *img_in.shape[:-1])))
+        masks = [m.numpy() for m in seg]
+
+        # COCO format in albumentations is [x,y,w,h] non-normalized
+        bboxes = anno_dict.get("bbox", torch.zeros((0, 5)))
+        ih, iw, _ = img_in.shape
+        bboxes_points = bboxes[:, 1:]
+        bboxes_points[:, 0::2] *= iw
+        bboxes_points[:, 1::2] *= ih
+        bboxes_points = check_bboxes(bboxes_points)
+        bbox_classes = bboxes[:, 0]
+
+        # albumentations expects "list" of keypoints e.g. [(x,y),(x,y),(x,y),(x,y)]
+        keypoints = anno_dict.get("keypoints", torch.zeros((1, 3 + 1)))
+        keypoints_classes = keypoints[:, 0]
+        keypoints_flat = torch.reshape(keypoints[:, 1:], (-1, 3))
+        keypoints_points = keypoints_flat[:, :2]
+        keypoints_points[:, 0] *= iw
+        keypoints_points[:, 1] *= ih
+        keypoints_visibility = keypoints_flat[:, 2]
+
+        transformed = self.transform(
+            image=img_in,
+            masks=masks,
+            bboxes=bboxes_points,
+            bbox_classes=bbox_classes,
+            keypoints=keypoints_points,
+            keypoints_classes=keypoints_visibility,  # not object class, but per-kp class
+        )
+
+        (
+            transformed_image,
+            out_bboxes,
+            transformed_mask,
+            final_keypoints,
+        ) = post_augment_process(
+            transformed, keypoints, keypoints_classes, use_rgb=self.train_rgb
+        )
+
+        out_annotations = create_out_annotations(
+            present_annotations,
+            classes=classes,
+            bboxes=out_bboxes,
+            masks=transformed_mask,
+            keypoints=final_keypoints,
+        )
+
+        return transformed_image, out_annotations
+
+
+class TrainAugmentations(Augmentations):
+    def __init__(
+        self,
+        image_size: list,
+        augmentations: dict,
+        train_rgb: bool = True,
+        keep_aspect_ratio: bool = True,
+    ):
+        """Class for train augmentations"""
+        super().__init__(train_rgb=train_rgb)
+        self.transform = self._parse_cfg(
+            image_size=image_size,
+            augmentations=augmentations,
+            keep_aspect_ratio=keep_aspect_ratio,
+        )
+
+
+class ValAugmentations(Augmentations):
+    def __init__(
+        self,
+        image_size: list,
+        augmentations: dict,
+        train_rgb: bool = True,
+        keep_aspect_ratio: bool = True,
+    ):
+        """Class for val augmentations"""
+        super().__init__(train_rgb=train_rgb)
+        self.transform = self._parse_cfg(
+            image_size=image_size,
+            augmentations=[a for a in augmentations if a["name"] == "Normalize"],
+            keep_aspect_ratio=keep_aspect_ratio,
+        )
+
+
+def post_augment_process(
+    transformed: dict,
+    keypoints: torch.Tensor,
+    keypoints_classes: np.array,
+    use_rgb: bool = True,
+):
+    """Post process augmentation outputs to prepare for training"""
+    transformed_image = transformed["image"]
+    if not use_rgb:
+        transformed_image = transformed_image.flip(-3)
+
+    transformed_mask = torch.stack(transformed["masks"])  # stack list of masks
+
+    transformed_bboxes = torch.tensor(transformed["bboxes"])
+    transformed_bbox_classes = torch.tensor(transformed["bbox_classes"])
+    # merge bboxes and classes back together
+    transformed_bbox_classes = torch.unsqueeze(transformed_bbox_classes, dim=-1)
+    out_bboxes = torch.cat((transformed_bbox_classes, transformed_bboxes), dim=1)
+
+    out_bboxes[:, 1::2] /= transformed_image.shape[2]
+    out_bboxes[:, 2::2] /= transformed_image.shape[1]
+
+    transformed_keypoints = torch.tensor(transformed["keypoints"])
+    transformed_keypoints_classes = torch.tensor(transformed["keypoints_classes"])
+    # merge keypoints and classes back together
+    transformed_keypoints_classes = torch.unsqueeze(
+        transformed_keypoints_classes, dim=-1
+    )
+    out_keypoints = torch.cat(
+        (transformed_keypoints, transformed_keypoints_classes), dim=1
+    )
+
+    out_keypoints = torch.reshape(out_keypoints, (-1, 3))
+    out_keypoints = mark_invisible_keypoints(out_keypoints, transformed_image)
+    out_keypoints[..., 0] /= transformed_image.shape[2]
+    out_keypoints[..., 1] /= transformed_image.shape[1]
+    out_keypoints = torch.reshape(
+        out_keypoints, (keypoints.shape[0], keypoints.shape[1] - 1)
+    )
+
+    final_keypoints = torch.zeros_like(keypoints)
+    final_keypoints[:, 1:] = out_keypoints
+    final_keypoints[:, 0] = keypoints_classes
+
+    return transformed_image, out_bboxes, transformed_mask, final_keypoints
+
+
+def mark_invisible_keypoints(keypoints: torch.Tensor, image: np.array):
+    """Mark invisible keypoints with label == 0"""
+    _, h, w = image.shape
+    for kp in keypoints:
+        if not (0 <= kp[0] < w and 0 <= kp[1] < h):
+            kp[2] = 0
+    return keypoints
+
+
+def check_bboxes(bboxes: torch.Tensor):
+    """Check bbox annotations and correct those with width or height 0"""
+    for i in range(bboxes.shape[0]):
+        if bboxes[i, 2] == 0:
+            bboxes[i, 2] = 1
+        if bboxes[i, 3] == 0:
+            bboxes[i, 3] = 1
+    return bboxes
+
+
+def create_out_annotations(
+    present_annotations: list,
+    classes: torch.Tensor,
+    bboxes: torch.Tensor,
+    masks: torch.Tensor,
+    keypoints: torch.Tensor,
+):
+    """Create dictionary of output annotations"""
+    out_annotations = {}
+    if "class" in present_annotations:
+        out_annotations["class"] = classes
+    if "bbox" in present_annotations:
+        out_annotations["bbox"] = bboxes
+    if "segmentation" in present_annotations:
+        out_annotations["segmentation"] = masks
+    if "keypoints" in present_annotations:
+        out_annotations["keypoints"] = keypoints
+    return out_annotations
diff --git a/src/luxonis_ml/ops/loader.py b/src/luxonis_ml/loader/loader.py
similarity index 66%
rename from src/luxonis_ml/ops/loader.py
rename to src/luxonis_ml/loader/loader.py
index a566369a..a3c4c738 100644
--- a/src/luxonis_ml/ops/loader.py
+++ b/src/luxonis_ml/loader/loader.py
@@ -1,16 +1,12 @@
-#!/usr/bin/env python3
-
 from pathlib import Path
 import cv2
 import numpy as np
-import glob, os, subprocess
 import torch
-from .dataset import HType, IType
 from fiftyone import ViewField as F
 
-class LuxonisLoader(torch.utils.data.Dataset):
 
-    def __init__(self, dataset, view='train', stream=False, augmentations=None):
+class LuxonisLoader(torch.utils.data.Dataset):
+    def __init__(self, dataset, view="train", stream=False, augmentations=None):
         """
         LuxonisDataset dataset: LuxonisDataset to use
         str view: either a saved view from a query or the name of a split
@@ -18,8 +14,11 @@ def __init__(self, dataset, view='train', stream=False, augmentations=None):
         """
 
         self.dataset = dataset
-        if view in ['train', 'val', 'test']:
-            self.samples = dataset.fo_dataset.match(
+        if view in ["train", "val", "test"]:
+            version_view = self.dataset.fo_dataset.load_saved_view(
+                f"version_{dataset.version}"
+            )
+            self.samples = version_view.match(
                 (F("latest") == True) & (F("split") == view)
             )
         else:
@@ -27,78 +26,86 @@ def __init__(self, dataset, view='train', stream=False, augmentations=None):
 
         # this will automatically use the main component in a group
         # if another component is desired, you must set dataset.fo_dataset.group_slice = 'other component'
+        self.ids = self.samples.values("id")
         self.paths = self.samples.values("filepath")
         self.stream = stream
         self.classes, self.classes_by_task = self.dataset.get_classes()
         self.nc = len(self.classes)
-        self.ns = len(self.dataset.fo_dataset.mask_targets.keys())
-        if 'keypoints' in self.dataset.fo_dataset.skeletons.keys():
-            self.nk = len(self.dataset.fo_dataset.skeletons['keypoints']['labels'])
+        self.ns = len(
+            self.dataset.fo_dataset.mask_targets.get("segmentation", {}).keys()
+        )
+        if "keypoints" in self.dataset.fo_dataset.skeletons.keys():
+            self.nk = len(self.dataset.fo_dataset.skeletons["keypoints"]["labels"])
         else:
             self.nk = 0
         self.augmentations = augmentations
 
-        if not self.stream and self.dataset.bucket_type != 'local':
+        if not self.stream and self.dataset.bucket_type != "local":
             self.dataset.sync_from_cloud()
 
+        self.dataset.fo_dataset.group_slice = self.dataset.source.main_component
+
     def __len__(self):
-        return len(self.paths)
+        return len(self.ids)
 
     def __getitem__(self, idx):
+        sample_id = self.ids[idx]
         path = self.paths[idx]
-        sample = self.samples[path]
-        if self.stream and self.dataset.bucket_type != 'local':
-            img_path = f'{str(Path.home())}/.luxonis_mount/{path}'
+        sample = self.dataset.fo_dataset[sample_id]
+        if self.stream and self.dataset.bucket_type != "local":
+            img_path = f"{str(Path.home())}/.luxonis_mount/{path}"
         else:
-            img_path = f'{str(Path.home())}/.cache/luxonis_ml/data/{path}'
-        img = np.transpose(cv2.imread(img_path), (2,0,1))
+            img_path = f"{str(Path.home())}/.cache/luxonis_ml/data/{path}"
+        img = np.transpose(cv2.imread(img_path), (2, 0, 1))
         _, ih, iw = img.shape
 
-        keys = sample.to_dict().keys()
-
         classify = np.zeros(self.nc)
-        bboxes = np.zeros((0,5))
+        bboxes = np.zeros((0, 5))
         seg = np.zeros((self.ns, ih, iw))
-        keypoints = np.zeros((0, self.nk*3+1))
+        keypoints = np.zeros((0, self.nk * 3 + 1))
         present_annotations = set()
 
         anno_dict = {}
-        if "classification" in keys:
-            cls = self.classes.index(sample.classification)
-            classify[cls] = classify[cls] + 1
+        if "class" in sample and sample["class"] is not None:
+            classes = sample["class"]["classifications"]
+            for cls in classes:
+                cls = self.classes.index(cls.label)
+                classify[cls] = classify[cls] + 1
             present_annotations.add("class")
-        if "boxes" in keys:
+        if "boxes" in sample and sample["boxes"] is not None:
             detections = sample.boxes.detections
             for det in detections:
-                box = np.array([
-                    self.classes.index(det.label),
-                    det.bounding_box[0],
-                    det.bounding_box[1],
-                    det.bounding_box[2],
-                    det.bounding_box[3],
-                ]).reshape(1,5)
+                box = np.array(
+                    [
+                        self.classes.index(det.label),
+                        det.bounding_box[0],
+                        det.bounding_box[1],
+                        det.bounding_box[2],
+                        det.bounding_box[3],
+                    ]
+                ).reshape(1, 5)
                 bboxes = np.append(bboxes, box, axis=0)
             present_annotations.add("bbox")
-        if "segmentation" in keys:
+        if "segmentation" in sample and sample["segmentation"] is not None:
             mask = sample.segmentation.mask
             for key in np.unique(mask):
                 if key != 0:
-                    seg[int(key)-1,...] = mask==key
+                    seg[int(key) - 1, ...] = mask == key
             present_annotations.add("segmentation")
-        if "keypoints" in keys:
+        if "keypoints" in sample and sample["keypoints"] is not None:
             for kps in sample.keypoints.keypoints:
                 cls = self.classes.index(kps.label)
-                pnts = np.array(kps.points).reshape((-1,2)).astype(np.float32)
+                pnts = np.array(kps.points).reshape((-1, 2)).astype(np.float32)
                 kps = np.zeros((len(pnts), 3))
                 nan_key = np.isnan(pnts[:, 0])
                 kps[~nan_key, 2] = 2
                 kps[:, :2] = pnts
-                kps[nan_key, :2] = 0 # use 0 instead of NaN
+                kps[nan_key, :2] = 0  # use 0 instead of NaN
                 kps = kps.flatten()
                 nk = len(kps)
                 kps = np.concatenate([[cls], kps])
-                points = np.zeros((1, self.nk*3+1))
-                points[0,:nk+1] = kps
+                points = np.zeros((1, self.nk * 3 + 1))
+                points[0, : nk + 1] = kps
                 keypoints = np.append(keypoints, points, axis=0)
 
             present_annotations.add("keypoints")
@@ -155,7 +162,7 @@ def collate_fn(batch):
             keypoint_annos = [anno["keypoints"] for anno in anno_dicts]
             label_keypoints = []
             for i, points in enumerate(keypoint_annos):
-                l_kps = torch.zeros((points.shape[0], points.shape[1]+1))
+                l_kps = torch.zeros((points.shape[0], points.shape[1] + 1))
                 l_kps[:, 0] = i  # add target image index for build_targets()
                 l_kps[:, 1:] = points
                 label_keypoints.append(l_kps)
diff --git a/src/luxonis_ml/loader/requirements.txt b/src/luxonis_ml/loader/requirements.txt
new file mode 100644
index 00000000..a0bc976c
--- /dev/null
+++ b/src/luxonis_ml/loader/requirements.txt
@@ -0,0 +1,1 @@
+torch>=1.12.1
\ No newline at end of file
diff --git a/src/luxonis_ml/luxonis_ml.py b/src/luxonis_ml/luxonis_ml.py
index cc770c27..7b63b043 100644
--- a/src/luxonis_ml/luxonis_ml.py
+++ b/src/luxonis_ml/luxonis_ml.py
@@ -1,20 +1,22 @@
 #!/usr/bin/env python3
 
-from luxonis_ml.ops import *
+from luxonis_ml.data import *
 import argparse
-import sys, os
-import warnings
+import os
 import json
 from pathlib import Path
 
+
 def _config():
     AWS_BUCKET = input("AWS Bucket: ")
     AWS_ACCESS_KEY_ID = input("AWS Access Key: ")
     AWS_SECRET_ACCESS_KEY = input("AWS Secret Access Key: ")
     AWS_S3_ENDPOINT_URL = input("AWS Endpoint URL: ")
     MONGO_URI = input("MongoDB URI: ")
+    LABELSTUDIO_URL = input("label-studio URL: ")
+    LABELSTUDIO_KEY = input("label-studio API Key: ")
 
-    cache_dir = f'{str(Path.home())}/.cache/luxonis_ml'
+    cache_dir = f"{str(Path.home())}/.cache/luxonis_ml"
     if not os.path.exists(cache_dir):
         os.mkdir(cache_dir)
 
@@ -24,44 +26,54 @@ def _config():
         "AWS_ACCESS_KEY_ID": AWS_ACCESS_KEY_ID,
         "AWS_SECRET_ACCESS_KEY": AWS_SECRET_ACCESS_KEY,
         "AWS_S3_ENDPOINT_URL": AWS_S3_ENDPOINT_URL,
-        "MONGO_URI": MONGO_URI
+        "MONGO_URI": MONGO_URI,
+        "LABELSTUDIO_URL": LABELSTUDIO_URL,
+        "LABELSTUDIO_KEY": LABELSTUDIO_KEY,
     }
 
-    with open(cache_file, 'w') as file:
+    with open(cache_file, "w") as file:
         json.dump(credentials, file)
 
-    fo_config = {
-        "database_dir": "null",
-        "database_uri": MONGO_URI
-    }
+    fo_config = {"database_dir": "null", "database_uri": MONGO_URI}
     fo_config_file = f"{str(Path.home())}/.fiftyone/config.json"
 
-    with open(fo_config_file, 'w') as file:
+    with open(fo_config_file, "w") as file:
         json.dump(fo_config, file)
 
     print(f"Credentials saved to {cache_file}!")
 
+
 def _dataset_sync(args):
     with LuxonisDataset(args.team, args.dataset_name) as dataset:
         dataset.sync_from_cloud()
 
-def main():
 
+def main():
     parser = argparse.ArgumentParser()
 
-    subparsers = parser.add_subparsers(dest='main')
-    parser_config = subparsers.add_parser('config', help='Configure API keys for luxonis_ml')
-    parser_dataset = subparsers.add_parser('dataset', help='Dataset programs to work with LDF')
-
-    dataset_subparsers = parser_dataset.add_subparsers(dest='dataset')
-
-    parser_dataset_sync = dataset_subparsers.add_parser('sync', help='Sync dataset media from cloud')
-    parser_dataset_sync.add_argument('-t', '--team', type=str, help='Team name', default=None)
-    parser_dataset_sync.add_argument('-n', '--dataset_name', type=str, help='Name of dataset', default=None)
+    subparsers = parser.add_subparsers(dest="main")
+    parser_config = subparsers.add_parser(
+        "config", help="Configure API keys for luxonis_ml"
+    )
+    parser_dataset = subparsers.add_parser(
+        "dataset", help="Dataset programs to work with LDF"
+    )
+
+    dataset_subparsers = parser_dataset.add_subparsers(dest="dataset")
+
+    parser_dataset_sync = dataset_subparsers.add_parser(
+        "sync", help="Sync dataset media from cloud"
+    )
+    parser_dataset_sync.add_argument(
+        "-t", "--team", type=str, help="Team name", default=None
+    )
+    parser_dataset_sync.add_argument(
+        "-n", "--dataset_name", type=str, help="Name of dataset", default=None
+    )
 
     args = parser.parse_args()
 
-    if args.main == 'config':
+    if args.main == "config":
         _config()
-    elif args.dataset == 'sync':
+    elif args.dataset == "sync":
         _dataset_sync(args)
diff --git a/src/luxonis_ml/ops/__init__.py b/src/luxonis_ml/ops/__init__.py
deleted file mode 100644
index c412fd9d..00000000
--- a/src/luxonis_ml/ops/__init__.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from .session import LuxonisSession
-from .tracker import LuxonisTracker, LuxonisTrackerPL
-from .dataset import LuxonisDataset, LDFComponent, IType, HType
-from .loader import LuxonisLoader
-from .version import LuxonisVersion
-from .parsers import LuxonisParser, DatasetType
diff --git a/src/luxonis_ml/ops/dataset.py b/src/luxonis_ml/ops/dataset.py
deleted file mode 100644
index 06a7c3c7..00000000
--- a/src/luxonis_ml/ops/dataset.py
+++ /dev/null
@@ -1,654 +0,0 @@
-#!/usr/bin/env python3
-
-import fiftyone as fo
-import fiftyone.core.odm as foo
-import luxonis_ml.fiftyone_plugins as fop
-import luxonis_ml.ops.utils.data_utils as data_utils
-from fiftyone import ViewField as F
-import os, subprocess, shutil
-from pathlib import Path
-import glob
-import warnings
-import cv2
-from PIL import Image
-import pickle
-import json
-from pathlib import Path
-import boto3
-import glob
-from tqdm import tqdm
-from enum import Enum
-import numpy as np
-from .version import LuxonisVersion
-
-class HType(Enum):
-    """ Individual file type """
-    IMAGE = 1
-    JSON = 2
-
-class IType(Enum):
-    """ Image type for IMAGE HType """
-    BGR = 1
-    MONO = 2
-    DISPARITY = 3
-    DEPTH = 4
-
-class LDFComponent:
-    htypes = [HType.IMAGE, HType.JSON]
-    itypes = [IType.BGR, IType.MONO, IType.DISPARITY, IType.DEPTH]
-
-    def __init__(self, name, htype, itype=None):
-        if htype not in self.htypes:
-            raise Exception(f"{htype} is not a valid HType!")
-        self.name = name
-        self.htype = htype
-        if htype == HType.IMAGE:
-            if itype is not None:
-                if itype not in self.itypes:
-                    raise Exception(f"{itype} is not a valid IType!")
-                self.itype = itype
-            else:
-                raise Exception("itype parameter is required for image component")
-
-class LDFSource:
-    types = ['real_oak', 'synthetic_oak', 'custom']
-
-    def __init__(self, name, components, main_component=None, source_type=None):
-        self.name = name
-        self.components = {component.name:component for component in components}
-        if main_component is not None:
-            self.main_component = main_component
-        else:
-            self.main_component = list(self.components.keys())[0] # make first component main component by default
-        if source_type is not None:
-            self.source_type = source_type
-        else:
-            self.source_type = 'custom'
-
-    def add_component(self, component):
-        self.components += component
-
-class LuxonisDataset:
-    """
-    Uses the LDF (Luxonis Dataset Format)
-
-    The goal is to standardize an arbitrary number of devices and sensor configurations, synthetic OAKs, and other sources
-    """
-
-    def __init__(self, team_name, dataset_name, bucket_type='local', override_bucket_type=False):
-        """
-        team name: team under which you can find all datasets
-        dataset name: name of the dataset
-        bucket_type: underlying storage for images, which can be local or an AWS bucket
-        override_bucket_type: option to change underlying storage from saved setting in DB
-        """
-
-        self.conn = foo.get_db_conn()
-
-        self.team = team_name
-        self.name = dataset_name
-        self.full_name = f"{self.team}-{self.name}"
-        self.bucket_type = bucket_type
-        self.bucket_choices = ['local', 'aws']
-        self.override_bucket_type = override_bucket_type
-        if self.bucket_type not in self.bucket_choices:
-            raise Exception(f"Bucket type {self.bucket_type} is not supported!")
-
-        credentials_cache_file = f'{str(Path.home())}/.cache/luxonis_ml/credentials.json'
-        if os.path.exists(credentials_cache_file):
-            with open(credentials_cache_file) as file:
-                self.creds = json.load(file)
-        else:
-            self.creds = {}
-
-        self._init_path()
-        self.version = -1
-        self.tasks = ['class', 'boxes', 'segmentation', 'keypoints']
-
-    def __enter__(self):
-
-        if self.full_name in fo.list_datasets():
-            self.fo_dataset = fo.load_dataset(self.full_name)
-        else:
-            self.fo_dataset = fo.Dataset(self.full_name)
-        self.fo_dataset.persistent = True
-
-        res = list(self.conn.luxonis_dataset_document.find(
-            { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
-        ))
-
-        if len(res): # already exists
-            assert len(res) == 1
-            self.dataset_doc = fop.LuxonisDatasetDocument.objects.get(
-                team_name=self.team,
-                dataset_name=self.name
-            )
-
-            tmp_bucket_type = self.bucket_type
-
-            self._doc_to_class()
-
-            if self.override_bucket_type:
-                self.bucket_type = tmp_bucket_type
-
-            self._init_path()
-
-        else: # create new
-            dataset_id = self.fo_dataset._doc.id
-            dataset_id_str = dataset_id.binary.hex()
-            self.dataset_doc = fop.LuxonisDatasetDocument(
-                dataset_id=dataset_id,
-                dataset_id_str=dataset_id_str,
-                team_name=self.team,
-                dataset_name=self.name,
-                path=self.path,
-                bucket_type='aws',
-                current_version=self.version,
-            )
-
-            self._init_path()
-            self.source = None # assuming a single dataset can only have one source
-
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-
-        self._class_to_doc()
-        self.dataset_doc.save(safe=True, upsert=True)
-
-    def _doc_to_class(self):
-
-        self.path = self.dataset_doc.path
-        self.bucket_type = self.dataset_doc.bucket_type
-        self.version = self.dataset_doc.current_version
-
-        doc = list(self._get_source())
-        if len(doc):
-            doc = doc[0]
-            components = [
-                LDFComponent(name, HType(htype), IType(itype)) \
-                for name, htype, itype in list(zip(doc['component_names'], doc['component_htypes'], doc['component_itypes']))
-            ]
-            main_component = components[0].name
-            self.source = LDFSource(
-                name=doc['name'],
-                components=components,
-                main_component=main_component,
-                source_type=doc['source_type']
-            )
-        else:
-            self.source = None
-
-    def _class_to_doc(self):
-
-        self.dataset_doc.path = self.path
-        self.dataset_doc.bucket_type = self.bucket_type
-        self.dataset_doc.current_version = self.version
-
-    def _get_credentials(self, key):
-        if key in self.creds.keys():
-            return self.creds[key]
-        else:
-            return os.environ[key]
-
-    def _get_source(self):
-        return self.conn.luxonis_source_document.find(
-            { "$and": [{"_luxonis_dataset_id": self.dataset_doc.id}] }
-        ).limit(1)
-
-    def _init_boto3_client(self):
-        self.client = boto3.client('s3',
-                        endpoint_url = self._get_credentials('AWS_S3_ENDPOINT_URL'),
-                        aws_access_key_id = self._get_credentials('AWS_ACCESS_KEY_ID'),
-                        aws_secret_access_key = self._get_credentials('AWS_SECRET_ACCESS_KEY')
-                      )
-        self.bucket = self.path.split('//')[1].split('/')[0]
-        self.bucket_path = self.path.split(self.bucket+'/')[-1]
-        # self._create_directory(self.bucket_path)
-
-    def _init_path(self):
-        if self.bucket_type == 'local':
-            self.path = f"{str(Path.home())}/.cache/luxonis_ml/data/{self.team}/datasets/{self.name}"
-            os.makedirs(self.path, exist_ok=True)
-        elif self.bucket_type == 'aws':
-            self.path = f"s3://{self._get_credentials('AWS_BUCKET')}/{self.team}/datasets/{self.name}"
-            self._init_boto3_client()
-
-    def _create_directory(self, path, clear_contents=False):
-        if self.bucket_type == 'local':
-            os.makedirs(path, exist_ok=True)
-        elif self.bucket_type == 'aws':
-            resp = self.client.list_objects(Bucket=self.bucket, Prefix=f"{self.bucket_path}/{path}/", Delimiter='/', MaxKeys=1)
-            if 'Contents' not in resp:
-                self.client.put_object(Bucket=self.bucket, Key=f"{self.bucket_path}/{path}/")
-            elif clear_contents:
-                for content in resp['Contents']:
-                    self.client.delete_object(Bucket=self.bucket, Key=content['Key'])
-
-    def _save_version(self, samples, note):
-
-        version = LuxonisVersion(self, samples=samples, note=note)
-        samples = version.get_samples()
-
-        version_view = self.fo_dataset[samples]
-        self.fo_dataset.save_view(f"version_{self.version}", version_view)
-
-    def create_source(self,
-        name=None,
-        oak_d_default=False,
-        # calibration=None,
-        # calibration_mapping=None,
-        custom_components=None):
-        """
-        name (str): optional rename for standard sources, required for custom sources
-        calibration (dict): DepthAI calibration file for real device (TODO: or Unity devices)
-        calibration_mapping (dict): for calibration files from non-standard products,
-            maps camera ID to a LDFComponent for the sensor and ID pairs to disparities
-        custom_components (list): list of LDFComponents for a custom source
-
-        TODO: more features with calibration
-        """
-
-        if oak_d_default:
-            if name is None: name = "oak_d"
-            components = [
-                LDFComponent('left', htype=HType.IMAGE, itype=IType.MONO),
-                LDFComponent('right', htype=HType.IMAGE, itype=IType.MONO),
-                LDFComponent('color', htype=HType.IMAGE, itype=IType.BGR),
-                LDFComponent('disparity', htype=HType.IMAGE, itype=IType.DISPARITY),
-                LDFComponent('depth', htype=HType.IMAGE, itype=IType.DEPTH)
-            ]
-
-            source = LDFSource(name, components, main_component='color', source_type='real_oak')
-
-        elif name is not None and custom_components is not None:
-            source = LDFSource(name, custom_components)
-        else:
-            raise Exception("name and custom_components required for custom sources")
-
-        if self.dataset_doc.id is None:
-            self._class_to_doc()
-            self.dataset_doc.save(safe=True, upsert=True)
-        res = self._get_source()
-        if len(list(res)): # exists
-            warnings.warn(f"Updating from a previously saved source!")
-            self.conn.luxonis_source_document.delete_many(
-                { "_luxonis_dataset_id": self.dataset_doc.id }
-            )
-        luxonis_dataset_id = self.dataset_doc.id
-        source_doc = fop.LuxonisSourceDocument(
-            luxonis_dataset_id=luxonis_dataset_id,
-            name=source.name,
-            source_type=source.source_type,
-            component_names=[c.name for _,c in source.components.items()],
-            component_htypes=[c.htype.value for _,c in source.components.items()],
-            component_itypes=[c.itype.value for _,c in source.components.items()],
-        )
-        source_doc.save(upsert=True)
-
-        self.source = source
-
-        self.fo_dataset.add_group_field(self.source.name, default=self.source.main_component)
-
-        return source
-
-    def launch_app(self):
-
-        session = fo.launch_app(dataset)
-
-    def set_classes(self, classes, task=None):
-        if task is not None:
-            if task not in self.tasks:
-                raise Exception(f'Task {task} is not a supported task')
-            self.fo_dataset.classes[task] = classes
-        else:
-            for task in self.tasks:
-                self.fo_dataset.classes[task] = classes
-
-    def set_mask_targets(self, mask_targets):
-        if 0 in mask_targets.keys():
-            raise Exception('Cannot set 0! This is assumed to be the background class')
-        self.fo_dataset.mask_targets = {
-            'segmentation': mask_targets
-        }
-
-    def set_skeleton(self, skeleton):
-        """
-        class_name (str): name of class to add fo skeleton for
-        skeleton (dict):
-            labels (list): list of strings for what each keypoint represents
-            egdes (list): list of length length-2 lists of keypoint indices denoting
-                          how keypoints are connected
-
-        NOTE: this only supports one class, which seems to be a fiftyone limitation
-        """
-
-        self.fo_dataset.skeletons.update({
-            'keypoints': fo.KeypointSkeleton(
-                labels=skeleton['labels'],
-                edges=skeleton['edges']
-            )
-        })
-        self.fo_dataset.save()
-
-    def sync_from_cloud(self):
-
-        self.non_streaming_dir = f"{str(Path.home())}/.cache/luxonis_ml/data/{self.team}/datasets/{self.name}"
-        os.makedirs(self.non_streaming_dir, exist_ok=True)
-
-        if self.bucket_type == 'local':
-            print("This is a local dataset! Cannot sync")
-        elif self.bucket_type == 'aws':
-            print("Syncing from cloud...")
-            cmd = f"aws s3 sync s3://{self.bucket}/{self.team}/datasets/{self.name} \
-                    {self.non_streaming_dir} \
-                    --endpoint-url={self._get_credentials('AWS_S3_ENDPOINT_URL')}"
-            process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
-            while True:
-                output = process.stdout.readline()
-                if output == b'' and process.poll() is not None:
-                    break
-                if output:
-                    print(output.strip().decode())
-
-    def get_classes(self):
-
-        classes = set()
-        classes_by_task = {}
-        for task in self.fo_dataset.classes:
-            if len(self.fo_dataset.classes[task]):
-                classes_by_task[task] = self.fo_dataset.classes[task]
-                for cls in self.fo_dataset.classes[task]:
-                    classes.add(cls)
-        mask_classes = set()
-        if 'segmentation' in self.fo_dataset.mask_targets.keys():
-            for key in self.fo_dataset.mask_targets['segmentation']:
-                mask_classes.add(self.fo_dataset.mask_targets['segmentation'][key])
-        classes_by_task['segmentation'] = list(mask_classes)
-        classes = list(classes.union(mask_classes))
-        classes.sort()
-
-        return classes, classes_by_task
-
-    def delete(self):
-        """
-        Deletes the entire dataset, aside from the images
-        """
-
-        res = list(self.conn.luxonis_dataset_document.find(
-            { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
-        ))
-        if len(res):
-            ldf_doc = res[0]
-
-            self.conn.luxonis_source_document.delete_many(
-                { "_luxonis_dataset_id": ldf_doc["_id"] }
-            )
-            self.conn.version_document.delete_many(
-                { "_dataset_id": ldf_doc["_dataset_id"] }
-            )
-            self.conn.luxonis_dataset_document.delete_many(
-                { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
-            )
-        fo.delete_dataset(self.full_name)
-
-    def _incr_version(self, version_code):
-        if version_code == 0: # major change
-            self.version += 1
-        elif version_code == 1: # minor change
-            self.version += 0.1
-
-    def _add_filter(self, additions, from_bucket, skip_media_check):
-        """
-        Filters out any additions to the dataset already existing
-        """
-
-        source = self.source
-        components = self.source.components
-
-        latest_view = self.fo_dataset.match(
-            F("latest") == True
-        )
-        filepaths = [sample['filepath'] for sample in latest_view]
-        filepaths = np.array([f"{path.split('/')[-2]}/{path.split('/')[-1]}" for path in filepaths])
-        updated = np.zeros(len(filepaths))
-
-        print("Checking for additions or modifications...")
-
-        filtered = []
-
-        media_change = False
-        field_change = False
-
-        for i, addition in tqdm(enumerate(additions), total=len(additions)):
-            for component_name in addition.keys():
-                filepath = addition[component_name]['filepath']
-                granule = data_utils.get_granule(filepath, addition, component_name)
-                candidate = f"{component_name}/{granule}"
-                if candidate not in list(filepaths):
-                    # ADD case
-                    media_change = True
-                    filtered.append(addition)
-                    break
-                else:
-                    # check for UPDATE
-                    mount_path = f"/{self.team}/datasets/{self.name}/{component_name}/{granule}"
-                    sample_view = self.fo_dataset.match( # TODO: could probably make this more efficient
-                        F("filepath") == mount_path
-                    )
-                    max_version = -1
-                    for sample in sample_view:
-                        if sample.version > max_version:
-                            latest_sample = sample
-
-                    if not skip_media_check:
-                        check = data_utils.check_media(self, filepath, mount_path, component_name, granule, from_bucket)
-                        if not check:
-                            media_change = True
-                            filtered.append(addition)
-                            if len(updated): updated[filepaths==candidate] = 1
-                            break
-
-                    check = data_utils.check_fields(
-                        self,
-                        latest_sample,
-                        addition,
-                        component_name
-                    )
-                    if not check:
-                        field_change = True
-                        filtered.append(addition)
-                        if len(updated): updated[filepaths==candidate] = 1
-                        break
-
-        additions = filtered
-        if len(additions) == 0:
-            print('No new additions!')
-            return None
-        else:
-            version_samples = []
-            for i, sample in tqdm(enumerate(latest_view), total=len(latest_view)):
-                if updated[i]:
-                    sample.latest = False
-                    sample.save()
-                else:
-                    version_samples.append(sample.id)
-            if media_change:
-                version_code = 0
-            elif field_change:
-                version_code = 1
-            else:
-                version_code = -1
-            return additions, version_samples, version_code
-
-    def _add_extract(self, additions, from_bucket, compute_heatmaps):
-        """
-        Filters out any additions to the dataset already existing
-        """
-
-        print("Extractig dataset media...")
-
-        components = self.source.components
-        if self.bucket_type == 'local':
-            local_cache = f'{str(Path.home())}/.cache/luxonis_ml/data/{self.team}/datasets/{self.name}'
-        elif self.bucket_type == 'aws':
-            local_cache = f'{str(Path.home())}/.cache/luxonis_ml/tmp'
-        os.makedirs(local_cache, exist_ok=True)
-        for component_name in components:
-            os.makedirs(f'{local_cache}/{component_name}', exist_ok=True)
-
-        for i, addition in tqdm(enumerate(additions), total=len(additions)):
-
-            add_heatmaps = {}
-            for component_name in addition.keys():
-                if component_name not in components.keys():
-                    raise Exception(f"Component {component_name} is not present in source")
-                component = addition[component_name]
-                if 'filepath' not in component.keys():
-                    raise Exception("Must specify filepath for every component!")
-
-                filepath = component['filepath']
-                granule = data_utils.get_granule(filepath, addition, component_name)
-                new_filepath = f"/{self.team}/datasets/{self.name}/{component_name}/{granule}"
-                additions[i][component_name]['filepath'] = new_filepath
-
-                if from_bucket:
-                    new_prefix = f'{self.bucket_path}/{component_name}/{granule}'
-                    self.client.copy_object(
-                        Bucket=self.bucket,
-                        Key=new_prefix,
-                        CopySource={'Bucket': self.bucket, 'Key': filepath}
-                    )
-                else:
-                    cmd = f"cp {filepath} {local_cache}/{component_name}/{granule}"
-                    subprocess.check_output(cmd, shell=True)
-
-                if compute_heatmaps and not from_bucket and \
-                (components[component_name].itype == IType.DISPARITY or \
-                components[component_name].itype == IType.DEPTH):
-
-                    heatmap_component = f"{component_name}_heatmap"
-                    os.makedirs(f'{local_cache}/{heatmap_component}', exist_ok=True)
-                    im = cv2.imread(filepath, cv2.IMREAD_UNCHANGED)
-                    if im.dtype == np.uint16:
-                        im = im/8
-                    heatmap = (im/96*255).astype(np.uint8)
-                    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-                    cv2.imwrite(f'{local_cache}/{heatmap_component}/{granule}', heatmap)
-                    new_filepath = f"/{self.team}/datasets/{self.name}/{heatmap_component}/{granule}"
-                    add_heatmaps[heatmap_component] = new_filepath
-
-            if compute_heatmaps and not from_bucket:
-                for heatmap_component in add_heatmaps:
-                    additions[i][heatmap_component] = {}
-                    additions[i][heatmap_component]['filepath'] = add_heatmaps[heatmap_component]
-
-        if self.bucket_type == 'aws' and not from_bucket:
-            print("Syncing to S3 bucket...")
-            cmd = f"aws s3 sync {local_cache} \
-                    s3://{self.bucket}/{self.team}/datasets/{self.name} \
-                    --endpoint-url={self._get_credentials('AWS_S3_ENDPOINT_URL')}"
-            process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
-            while True:
-                output = process.stdout.readline()
-                if output == b'' and process.poll() is not None:
-                    break
-                if output:
-                    print(output.strip().decode())
-
-        if self.bucket_type != 'local':
-            shutil.rmtree(local_cache)
-
-        return additions
-
-    def _add_fiftyone(self, additions, version_samples, compute_heatmaps):
-        source = self.source
-        samples = []
-        for i, addition in enumerate(additions):
-            component_names = [component_name for component_name in addition.keys() if not component_name.endswith('_heatmap')]
-            group = fo.Group(name=source.name)
-            for component_name in component_names:
-
-                component = addition[component_name]
-
-                sample = fo.Sample(filepath=component['filepath'], version=self.version, latest=True)
-                sample[source.name] = group.element(component_name)
-
-                for ann in component.keys():
-                    if ann == 'class':
-                        sample['class'] = fo.Classification(label=component['class'])
-                    elif ann == 'boxes':
-                        sample['boxes'] = fo.Detections(detections=[
-                            fo.Detection(
-                                label=box[0] if isinstance(box[0], str) else self.fo_dataset.classes['boxes'][int(box[0])],
-                                bounding_box=box[1:5]
-                            )
-                            for box in component['boxes']
-                        ])
-                    elif ann == 'segmentation':
-                        sample['segmentation'] = fo.Segmentation(mask=component['segmentation'])
-                    elif ann == 'keypoints':
-                        sample['keypoints'] = fo.Keypoints(keypoints=[
-                            fo.Keypoint(
-                                label=kp[0] if isinstance(kp[0], str) else self.fo_dataset.classes['keypoints'][int(kp[0])],
-                                points=kp[1]
-                            )
-                            for kp in component['keypoints']
-                        ])
-                    elif ann == 'new_image_name':
-                        continue # ignore this as an attribute
-                    else:
-                        sample[ann] = component[ann]
-
-                if 'split' not in component.keys():
-                    sample['split'] = 'train' # default split
-
-                if compute_heatmaps and f'{component_name}_heatmap' in addition.keys():
-                    sample['heatmap'] = fo.Heatmap(map_path=addition[f'{component_name}_heatmap']['filepath'])
-
-                samples.append(sample)
-
-        new_ids = self.fo_dataset.add_samples(samples)
-        version_samples += new_ids
-
-        return version_samples
-
-    def add(
-        self,
-        additions,
-        note,
-        from_bucket=False,
-        compute_heatmaps=True,
-        skip_media_check=False
-    ):
-        """
-        Function to add data and automatically version the data
-
-        dataset: LuxonisDataset instance
-        dict: a list of dictionaries describing each Voxel51 sample
-            Each dict contains keys of components
-            The value of each component can contain the following keys:
-                filepath (media) : path to image, video, point cloud, or voxel
-                -----
-                class              : name of class for an entire image (Image Classification)
-                boxes              : list of Nx6 numpy arrays [class, xmin, ymin, width, height] (Object Detection)
-                segmentation       : numpy array where pixels correspond to integer values of classes
-                keypoints          : list of classes and (x,y) keypoints
-        """
-
-        if from_bucket and self.bucket_type == 'local':
-            raise Exception('from_bucket must be False for local dataset!')
-
-        filter_result = self._add_filter(additions, from_bucket, skip_media_check)
-        if filter_result is None:
-            return
-        else:
-            additions, version_samples, version_code = filter_result
-
-        self._incr_version(version_code)
-
-        additions = self._add_extract(additions, from_bucket, compute_heatmaps)
-
-        version_samples = self._add_fiftyone(additions, version_samples, compute_heatmaps)
-
-        self._save_version(version_samples, note)
diff --git a/src/luxonis_ml/ops/utils/data_utils.py b/src/luxonis_ml/ops/utils/data_utils.py
deleted file mode 100644
index 872e810d..00000000
--- a/src/luxonis_ml/ops/utils/data_utils.py
+++ /dev/null
@@ -1,154 +0,0 @@
-import numpy as np
-import hashlib
-
-def get_granule(filepath, addition, component_name):
-    granule = filepath.split('/')[-1]
-    if 'new_image_name' in addition[component_name].keys():
-        filepath = filepath.replace(granule, addition[component_name]['new_image_name'])
-        granule = addition[component_name]['new_image_name']
-    return granule
-
-def check_media(dataset, filepath, mount_path, component_name, granule, from_bucket):
-    source = dataset.source
-    if dataset.bucket_type == 'local':
-        local_path = f'{dataset.path}/{mount_path}'
-        md5 = hashlib.md5()
-        with open(filepath, 'rb') as file:
-            for chunk in iter(lambda: file.read(8192), b''):
-                md5.update(chunk)
-        old_checksum = md5.hexdigest()
-    elif dataset.bucket_type == 'aws':
-        prefix = f'{dataset.bucket_path}/{component_name}/{granule}'
-        resp = dataset.client.list_objects(Bucket=dataset.bucket, Prefix=prefix, Delimiter='/', MaxKeys=1)
-        exists = 'Contents' in resp
-        if not exists:
-            raise Exception("Issue with dataset! Sample file does not exist on S3!")
-        old_checksum = resp['Contents'][0]['ETag'].strip('"')
-
-    if from_bucket:
-        if dataset.bucket_type == 'aws':
-            resp = dataset.client.list_objects(Bucket=dataset.bucket, Prefix=filepath, Delimiter='/', MaxKeys=1)
-            exists = 'Contents' in resp
-            if not exists:
-                raise Exception(f"File {filepath} not found on S3!")
-            new_checksum = resp['Contents'][0]['ETag'].strip('"')
-    else:
-        md5 = hashlib.md5()
-        with open(filepath, 'rb') as file:
-            for chunk in iter(lambda: file.read(8192), b''):
-                md5.update(chunk)
-        new_checksum = md5.hexdigest()
-    if new_checksum != old_checksum:
-        if dataset.bucket_type == 'local':
-            new_path = f'{dataset.path}/{dataset.team}/datasets/{dataset.name}/archive/{source.name}/{component_name}/{granule}'
-            os.makedirs(new_path, exist_ok=True)
-            cmd = f'cp {local_path} {new_path}'
-            subprocess.check_output(cmd, shell=True)
-        elif dataset.bucket_type == 'aws':
-            new_prefix = f'{dataset.bucket_path}/archive/{source.name}/{component_name}/{granule}'
-            dataset.client.copy_object(
-                Bucket=dataset.bucket,
-                Key=new_prefix,
-                CopySource={'Bucket': dataset.bucket, 'Key': prefix}
-            )
-        # update existing datasat sample to point to this new_prefix
-        # limitation: all old versions point to the latest archive only (TODO: fix)
-        for sample in sample_view:
-            sample.filepath = mount_path.replace(
-                f'/{dataset.team}/{dataset.name}/',
-                f'/{dataset.team}/{dataset.name}/archive/'
-            )
-            sample.save()
-
-        return False
-    else:
-        return True
-
-def check_classification(val1, val2):
-    val1 = fo.Classification(label=val1).to_dict()
-    if len(val1.keys()) == len(val2.keys()):
-        for key in val1:
-            if not key.startswith('_'):
-                if val1[key] != val2[key]:
-                    return False
-    else:
-        return False
-    return True
-
-def check_boxes(val1, val2):
-    if len(val1) == len(val2['detections']):
-        for val1, val2 in list(zip(val1, val2['detections'])):
-            if isinstance(val1[0], str) and val2['label'] != val1[0]:
-                return False
-            if not isinstance(val1[0], str) and val2['label'] != dataset.fo_dataset.classes['boxes'][val1[0]]:
-                return False
-            for c1, c2 in list(zip(val1[1:], val2['bounding_box'])):
-                if abs(c1-c2) > 1e-8:
-                    return False
-    else:
-        return False
-    return True
-
-def check_segmentation(val1, val2):
-    if (val1.shape != val2.shape) or (np.linalg.norm(val1-val2) > 1e-8):
-        return False
-    return True
-
-def check_keypoints(val1, val2):
-    if len(val1) == len(val2['keypoints']):
-        for val1, val2 in list(zip(val1, val2['keypoints'])):
-            if isinstance(val1[0], str) and val2['label'] != val1[0]:
-                return False
-            if not isinstance(val1[0], str) and val2['label'] != dataset.fo_dataset.classes['keypoints'][val1[0]]:
-                return False
-            for c1, c2 in list(zip(val1[1], val2['points'])):
-                if not (np.isnan(c1[0]) and isinstance(c2[0], dict)):
-                    if not np.isnan(c1[0]) and not isinstance(c2[0], dict):
-                        if c1[0] != c2[0] or c1[1] != c2[1]:
-                            return False
-                    else:
-                        return False
-
-    else:
-        return False
-    return True
-
-def check_fields(dataset, latest_sample, addition, component_name):
-
-    ignore_fields_match = set([dataset.source.name, 'version', 'metadata', 'latest', 'tags', 'split'])
-    ignore_fields_check = set(['filepath'])
-
-    sample_dict = latest_sample.to_dict()
-    f1 = set(addition[component_name].keys())
-    f2 = set(sample_dict.keys())
-    new = f1 - f1.intersection(f2)
-    missing = f2 - f1.intersection(f2) - ignore_fields_match
-    if len(new) or len(missing):
-        return False
-
-    check_fields = list(f1.intersection(f2) - ignore_fields_check)
-    for field in check_fields:
-        val1 = addition[component_name][field]
-        val2 = sample_dict[field]
-
-        if field in dataset.tasks:
-            if field == 'class':
-                check = check_classification(val1, val2)
-                if not check: return False
-            elif field == 'boxes':
-                check = check_boxes(val1, val2)
-                if not check: return False
-            elif field == 'segmentation':
-                val2 = latest_sample.segmentation.mask
-                check = check_segmentation(val1, val2)
-                if not check: return False
-            elif field == 'keypoints':
-                check = check_keypoints(val1, val2)
-                if not check: return False
-            else:
-                raise NotImplementedError()
-        else:
-            if not val1 == val2:
-                return False
-
-    return True
diff --git a/src/luxonis_ml/tracker/__init__.py b/src/luxonis_ml/tracker/__init__.py
new file mode 100644
index 00000000..f7624149
--- /dev/null
+++ b/src/luxonis_ml/tracker/__init__.py
@@ -0,0 +1,2 @@
+from .tracker import LuxonisTracker, LuxonisTrackerPL
+from .mlflow_plugins import LuxonisRequestHeaderProvider
diff --git a/src/luxonis_ml/ops/mlflow_plugins.py b/src/luxonis_ml/tracker/mlflow_plugins.py
similarity index 54%
rename from src/luxonis_ml/ops/mlflow_plugins.py
rename to src/luxonis_ml/tracker/mlflow_plugins.py
index adb4f85e..9bbbda8a 100644
--- a/src/luxonis_ml/ops/mlflow_plugins.py
+++ b/src/luxonis_ml/tracker/mlflow_plugins.py
@@ -1,11 +1,12 @@
 import os
 from mlflow import __version__
-from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider
+from mlflow.tracking.request_header.abstract_request_header_provider import (
+    RequestHeaderProvider,
+)
 
 _USER_AGENT = "User-Agent"
-_DEFAULT_HEADERS = {
-    _USER_AGENT: "mlflow-python-client/%s" % __version__
-}
+_DEFAULT_HEADERS = {_USER_AGENT: "mlflow-python-client/%s" % __version__}
+
 
 class LuxonisRequestHeaderProvider(RequestHeaderProvider):
     def in_context(self):
@@ -13,6 +14,6 @@ def in_context(self):
 
     def request_headers(self):
         headers = dict(**_DEFAULT_HEADERS)
-        headers['CF-Access-Client-Id'] = os.environ.get('MLFLOW_CLOUDFLARE_ID')
-        headers['CF-Access-Client-Secret'] = os.environ.get('MLFLOW_CLOUDFLARE_SECRET')
+        headers["CF-Access-Client-Id"] = os.environ.get("MLFLOW_CLOUDFLARE_ID")
+        headers["CF-Access-Client-Secret"] = os.environ.get("MLFLOW_CLOUDFLARE_SECRET")
         return headers
diff --git a/src/luxonis_ml/tracker/requirements.txt b/src/luxonis_ml/tracker/requirements.txt
new file mode 100644
index 00000000..d7806a1a
--- /dev/null
+++ b/src/luxonis_ml/tracker/requirements.txt
@@ -0,0 +1,3 @@
+unique-names-generator>=1.0.2
+pytorch-lightning>=2.0.0
+mlflow>=2.0.1
\ No newline at end of file
diff --git a/src/luxonis_ml/ops/tracker.py b/src/luxonis_ml/tracker/tracker.py
similarity index 66%
rename from src/luxonis_ml/ops/tracker.py
rename to src/luxonis_ml/tracker/tracker.py
index 6b4fecf6..d018b3e3 100644
--- a/src/luxonis_ml/ops/tracker.py
+++ b/src/luxonis_ml/tracker/tracker.py
@@ -4,9 +4,6 @@
 from pathlib import Path
 import glob
 from unique_names_generator import get_random_name
-from PIL import Image
-import cv2
-import numpy as np
 
 from pytorch_lightning.utilities import rank_zero_only
 from pytorch_lightning.loggers.logger import Logger as plLogger
@@ -20,25 +17,35 @@ class LuxonisTracker:
         - MLFlow
     """
 
-    def __init__(self,
-                 project_name,
-                 run_name=None,
-                 hyperparameter_config=None,
-                 save_directory='runs',
-                 is_tensorboard=False,
-                 is_wandb=False,
-                 is_mlflow=False,
-                 is_sweep=False,
-                 wandb_entity=None,
-                 mlflow_tracking_uri=None):
-
+    def __init__(
+        self,
+        project_name=None,
+        project_id=None,
+        run_name=None,
+        run_id=None,
+        hyperparameter_config=None,
+        save_directory="runs",
+        is_tensorboard=False,
+        is_wandb=False,
+        is_mlflow=False,
+        is_sweep=False,
+        wandb_entity=None,
+        mlflow_tracking_uri=None,
+    ):
         self.project_name = project_name
+        self.project_id = project_id
         self.save_directory = save_directory
         self.is_tensorboard = is_tensorboard
         self.is_wandb = is_wandb
         self.is_mlflow = is_mlflow
         self.config = hyperparameter_config
 
+        self.run_id = run_id  # if using MLFlow then it will continue previous run
+
+        if is_wandb or is_mlflow:
+            if self.project_name is None and self.project_id is None:
+                raise Exception("Either project_name or project_id must be specified!")
+
         if self.is_wandb and wandb_entity is None:
             raise Exception("Must specify wandb_entity when using wandb!")
         else:
@@ -48,53 +55,63 @@ def __init__(self,
         else:
             self.mlflow_tracking_uri = mlflow_tracking_uri
 
-        if not (
-            self.is_tensorboard or \
-            self.is_wandb or \
-            self.is_mlflow
-        ):
+        if not (self.is_tensorboard or self.is_wandb or self.is_mlflow):
             raise Exception("At least one integration must be used!")
 
         self.run_name = self._get_run_name() if run_name is None else run_name
-        Path(f"{self.save_directory}/{self.run_name}").mkdir(parents=True, exist_ok=True)
+        Path(f"{self.save_directory}/{self.run_name}").mkdir(
+            parents=True, exist_ok=True
+        )
 
         if self.is_tensorboard:
             from torch.utils.tensorboard import SummaryWriter
-            log_dir=f"{save_directory}/tensorboard_logs/{self.run_name}"
-            self.tensorboard_logger = SummaryWriter(
-                log_dir=log_dir
-            )
+
+            log_dir = f"{save_directory}/tensorboard_logs/{self.run_name}"
+            self.tensorboard_logger = SummaryWriter(log_dir=log_dir)
             self.tensorboard_logger.add_hparams(
                 self.config,
-                {'placeholder_metric': 0}, # placeholder metric is needed due to this issue: https://github.com/tensorflow/tensorboard/issues/5476
+                {
+                    "placeholder_metric": 0
+                },  # placeholder metric is needed due to this issue: https://github.com/tensorflow/tensorboard/issues/5476
             )
 
         if self.is_wandb:
             import wandb
+
             self.wandb = wandb
 
             log_dir = f"{save_directory}/wandb_logs"
             Path(log_dir).mkdir(parents=True, exist_ok=True)
 
             self.wandb.init(
-                project=project_name,
+                project=self.project_name
+                if self.project_name != None
+                else self.project_id,
                 entity=self.wandb_entity,
                 dir=log_dir,
                 name=self.run_name,
-                config=self.config
+                config=self.config,
             )
 
         if self.is_mlflow:
             import mlflow
+
             self.mlflow = mlflow
 
             self.artifacts_dir = f"{self.save_directory}/{self.run_name}/artifacts"
             Path(self.artifacts_dir).mkdir(parents=True, exist_ok=True)
 
             self.mlflow.set_tracking_uri(self.mlflow_tracking_uri)
-            self.mlflow.set_experiment(self.project_name)
+
+            if self.project_id is not None:
+                self.project_name = None
+            self.mlflow.set_experiment(
+                experiment_name=self.project_name, experiment_id=self.project_id
+            )
+
+            # if self.run_id == None then create new run, else use alredy created one
             self.mlflow.start_run(
-                run_name=self.run_name, nested=is_sweep
+                run_id=self.run_id, run_name=self.run_name, nested=self.is_sweep
             )
             self.mlflow.log_params(self.config)
 
@@ -103,13 +120,13 @@ def _get_next_run_number(self):
         log_dirs = glob.glob(f"{self.save_directory}/*")
         log_dirs = [path.split("/")[-1] for path in log_dirs if os.path.isdir(path)]
         # find the numbers based on the naming convention
-        nums = [path.split('-')[0] for path in log_dirs]
+        nums = [path.split("-")[0] for path in log_dirs]
         nums = [int(num) for num in nums if num.isnumeric()]
 
         if len(nums) == 0:
             return 0
         else:
-            return max(nums)+1
+            return max(nums) + 1
 
     def _get_run_name(self):
         name_without_number = get_random_name(separator="-", style="lowercase")
@@ -138,33 +155,36 @@ def log_image(self, name, img, step):
         """
 
         if self.is_tensorboard:
-            self.tensorboard_logger.add_image(name, img, dataformats='HWC')
+            self.tensorboard_logger.add_image(name, img, dataformats="HWC")
 
         if self.is_wandb:
             wandb_image = self.wandb.Image(img, caption=name)
             self.wandb.log({name: wandb_image})
 
         if self.is_mlflow:
-            self.mlflow.log_image(
-                img, f"{name}_{step}.png"
-            )
+            self.mlflow.log_image(img, f"{name}_{step}.png")
+
 
 class LuxonisTrackerPL(plLogger):
-    def __init__(self,
-                 project_name,
-                 run_name=None,
-                 save_directory='runs',
-                 is_tensorboard=False,
-                 is_wandb=False,
-                 is_mlflow=False,
-                 is_sweep=False,
-                 wandb_entity=None,
-                 mlflow_tracking_uri=None,
-                 rank=0):             
-        
+    def __init__(
+        self,
+        project_name=None,
+        project_id=None,
+        run_name=None,
+        run_id=None,
+        save_directory="runs",
+        is_tensorboard=False,
+        is_wandb=False,
+        is_mlflow=False,
+        is_sweep=False,
+        wandb_entity=None,
+        mlflow_tracking_uri=None,
+        rank=0,
+    ):
         plLogger.__init__(self)
 
         self.project_name = project_name
+        self.project_id = project_id
         self.save_directory = save_directory
         self.is_tensorboard = is_tensorboard
         self.is_wandb = is_wandb
@@ -172,6 +192,12 @@ def __init__(self,
         self.is_sweep = is_sweep
         self.rank = rank
 
+        self.run_id = run_id  # if using MLFlow then it will continue previous run
+
+        if is_wandb or is_mlflow:
+            if self.project_name is None and self.project_id is None:
+                raise Exception("Either project_name or project_id must be specified!")
+
         if self.is_wandb and wandb_entity is None:
             raise Exception("Must specify wandb_entity when using wandb!")
         else:
@@ -181,27 +207,23 @@ def __init__(self,
         else:
             self.mlflow_tracking_uri = mlflow_tracking_uri
 
-        if not (
-            self.is_tensorboard or \
-            self.is_wandb or \
-            self.is_mlflow
-        ):
+        if not (self.is_tensorboard or self.is_wandb or self.is_mlflow):
             raise Exception("At least one integration must be used!")
-        
 
         self._experiment = None
 
         if run_name:
             self.run_name = run_name
         else:
-            # create new directory if rank==0 else return newest run 
+            # create new directory if rank==0 else return newest run
             if rank == 0:
                 self.run_name = self._get_run_name()
             else:
                 self.run_name = self._get_latest_run_name()
 
-        Path(f"{self.save_directory}/{self.run_name}").mkdir(parents=True, exist_ok=True)
-
+        Path(f"{self.save_directory}/{self.run_name}").mkdir(
+            parents=True, exist_ok=True
+        )
 
     @property
     def name(self):
@@ -216,58 +238,68 @@ def version(self):
     def experiment(self):
         if self._experiment is not None:
             return self._experiment
-        
+
         self._experiment = {}
 
         if self.is_tensorboard:
             from torch.utils.tensorboard import SummaryWriter
-            log_dir=f"{self.save_directory}/tensorboard_logs/{self.run_name}"
-            self._experiment["tensorboard"] = SummaryWriter(
-                log_dir=log_dir
-            )
+
+            log_dir = f"{self.save_directory}/tensorboard_logs/{self.run_name}"
+            self._experiment["tensorboard"] = SummaryWriter(log_dir=log_dir)
 
         if self.is_wandb:
             import wandb
+
             self._experiment["wandb"] = wandb
 
             log_dir = f"{self.save_directory}/wandb_logs"
             Path(log_dir).mkdir(parents=True, exist_ok=True)
 
             self._experiment["wandb"].init(
-                project=self.project_name,
+                project=self.project_name
+                if self.project_name != None
+                else self.project_id,
                 entity=self.wandb_entity,
                 dir=log_dir,
                 name=self.run_name,
-                #config=self.config # TODO: this also?
+                # config=self.config # TODO: this also?
             )
 
         if self.is_mlflow:
             import mlflow
+
             self._experiment["mlflow"] = mlflow
 
             self.artifacts_dir = f"{self.save_directory}/{self.run_name}/artifacts"
             Path(self.artifacts_dir).mkdir(parents=True, exist_ok=True)
 
-            self._experiment["mlflow"].set_tracking_uri("https://ml-track.luxonis.com")
-            self._experiment["mlflow"].set_experiment(self.project_name)
+            self._experiment["mlflow"].set_tracking_uri(self.mlflow_tracking_uri)
+
+            if self.project_id is not None:
+                self.project_name = None
+            self._experiment["mlflow"].set_experiment(
+                experiment_name=self.project_name, experiment_id=self.project_id
+            )
+
+            # if self.run_id == None then create new run, else use alredy created one
             self._experiment["mlflow"].start_run(
-                run_name=self.run_name, nested=self.is_sweep
+                run_id=self.run_id, run_name=self.run_name, nested=self.is_sweep
             )
-        
+
         return self._experiment
-    
+
     def _get_next_run_number(self):
         # find all directories that should be runs
         log_dirs = glob.glob(f"{self.save_directory}/*")
         log_dirs = [path.split("/")[-1] for path in log_dirs if os.path.isdir(path)]
         # find the numbers based on the naming convention
-        nums = [path.split('-')[0] for path in log_dirs]
+        nums = [path.split("-")[0] for path in log_dirs]
         nums = [int(num) for num in nums if num.isnumeric()]
 
         if len(nums) == 0:
             return 0
         else:
-            return max(nums)+1
+            return max(nums) + 1
 
     def _get_run_name(self):
         name_without_number = get_random_name(separator="-", style="lowercase")
@@ -279,21 +311,30 @@ def _get_latest_run_name(self):
         log_dirs = glob.glob(f"{self.save_directory}/*")
         log_dirs = [path for path in log_dirs if os.path.isdir(path)]
         # find run names based on the naming convention and sort them by last modified time
-        runs = [l.replace(f"{self.save_directory}/","") for l in log_dirs if l.split("-")[0].isnumeric()]
-        runs.sort(key = lambda x: os.path.getmtime(os.path.join(self.save_directory, x)), reverse=True)
+        runs = [
+            l.replace(f"{self.save_directory}/", "")
+            for l in log_dirs
+            if l.split("-")[0].isnumeric()
+        ]
+        runs.sort(
+            key=lambda x: os.path.getmtime(os.path.join(self.save_directory, x)),
+            reverse=True,
+        )
         return runs[0]
 
     @rank_zero_only
     def log_hyperparams(self, params):
-        """ Log hypeparameter dictionary
+        """Log hypeparameter dictionary
 
         Args:
-            params (dict): Dict of hyperparameters key and value pairs 
+            params (dict): Dict of hyperparameters key and value pairs
         """
         if self.is_tensorboard:
             self.experiment["tensorboard"].add_hparams(
                 params,
-                {'placeholder_metric': 0}, # placeholder metric is needed due to this issue: https://github.com/tensorflow/tensorboard/issues/5476
+                {
+                    "placeholder_metric": 0
+                },  # placeholder metric is needed due to this issue: https://github.com/tensorflow/tensorboard/issues/5476
             )
         if self.is_wandb:
             self.experiment["wandb"].config.update(params)
@@ -302,7 +343,7 @@ def log_hyperparams(self, params):
 
     @rank_zero_only
     def log_metrics(self, metrics, step):
-        """ Log metric dictionary
+        """Log metric dictionary
 
         Args:
             metrics (dict): Dict of metric key and value pairs
@@ -318,7 +359,7 @@ def log_metrics(self, metrics, step):
 
     @rank_zero_only
     def log_images(self, imgs, step):
-        """ Log imgs dictionary
+        """Log imgs dictionary
 
         Args:
             imgs (dict): Dict of image key and value pairs where key is image caption and value is img data
@@ -329,7 +370,7 @@ def log_images(self, imgs, step):
 
     @rank_zero_only
     def log_image(self, caption, img, step):
-        """ Log one image with a caption
+        """Log one image with a caption
 
         Args:
             caption (str): Image caption
@@ -338,7 +379,9 @@ def log_image(self, caption, img, step):
         """
 
         if self.is_tensorboard:
-            self.experiment["tensorboard"].add_image(caption, img, step, dataformats='HWC')
+            self.experiment["tensorboard"].add_image(
+                caption, img, step, dataformats="HWC"
+            )
 
         if self.is_wandb:
             wandb_image = self._experiment["wandb"].Image(img, caption=caption)
@@ -346,11 +389,7 @@ def log_image(self, caption, img, step):
             self._experiment["wandb"].log({caption: wandb_image})
 
         if self.is_mlflow:
-            ## TODO: has yet to be tested!
-            name = "name_of_choice"
-            self.mlflow.log_image(
-                img, f"{name}_{step}.png"
-            )
+            self._experiment["mlflow"].log_image(img, f"{caption}_{step}.png")
 
     @rank_zero_only
     def save(self):
@@ -361,4 +400,4 @@ def save(self):
     def finalize(self, status):
         # Optional. Any code that needs to be run after training
         # finishes goes here
-        pass
\ No newline at end of file
+        pass


Test patch:
diff --git a/tests/test.py b/tests/test.py
index 65319f8c..b1446772 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -1,17 +1,30 @@
 import unittest
-import subprocess, shutil, os
+import subprocess, shutil, os, glob, json, time
+import numpy as np
+import cv2
+from PIL import Image, ImageDraw
 from pathlib import Path
 import fiftyone.core.odm as foo
-from luxonis_ml.ops import *
+from luxonis_ml.data import *
+from copy import deepcopy
+from fiftyone import ViewField as F
+from bson.objectid import ObjectId
 
 unittest.TestLoader.sortTestMethodsUsing = None
 
-class LuxonisDatasetTester(unittest.TestCase):
+"""
+NOTE: Depsite the name, this is integration testing not unit testing.
+The test cases are meant to happen in a certain order to check how
+the state of the LDF changes with each step.
+"""
+
 
+class LuxonisDatasetTester(unittest.TestCase):
     @classmethod
     def setUpClass(self):
-        self.team = "unittest"
-        self.name = "coco"
+        self.team_id = "d7625eef-ad99-4019-af95-ffa5ebd48e3c"
+        self.team_name = "unittest"
+        self.dataset_name = "coco"
 
         self.coco_images_path = "../data/person_val2017_subset"
         self.coco_annotation_path = "../data/person_keypoints_val2017.json"
@@ -21,7 +34,9 @@ def setUpClass(self):
             os.remove(self.coco_annotation_path)
 
         print("Downloading data from GDrive...")
-        cmd = "gdown 1XlvFK7aRmt8op6-hHkWVKIJQeDtOwoRT -O ../data/COCO_people_subset.zip"
+        cmd = (
+            "gdown 1XlvFK7aRmt8op6-hHkWVKIJQeDtOwoRT -O ../data/COCO_people_subset.zip"
+        )
         subprocess.check_output(cmd, shell=True)
 
         print("Extracting data...")
@@ -30,143 +45,662 @@ def setUpClass(self):
 
         self.conn = foo.get_db_conn()
         # if dataset already exists, delete it
-        # could be problematic if .delete() breaks
-        res = list(self.conn.luxonis_dataset_document.find(
-            { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
-        ))
+        res = list(
+            self.conn.luxonis_dataset_document.find(
+                {
+                    "$and": [
+                        {"team_id": self.team_id},
+                        {"dataset_name": self.dataset_name},
+                    ]
+                }
+            )
+        )
         if len(res):
-            with LuxonisDataset(self.team, self.name) as dataset:
-                dataset.delete()
+            with LuxonisDataset(self.team_id, res[0]["_id"]) as dataset:
+                dataset.delete_dataset()
 
-    def test_local_init(self):
+        # get COCO data for testing
+        img_dir = "../data/person_val2017_subset"
+        annot_file = "../data/person_keypoints_val2017.json"
+
+        # get paths to images sorted by number
+        im_paths = glob.glob(img_dir + "/*.jpg")
+        nums = np.array([int(path.split("/")[-1].split(".")[0]) for path in im_paths])
+        idxs = np.argsort(nums)
+        im_paths = list(np.array(im_paths)[idxs])
+
+        # load
+        with open(annot_file) as file:
+            data = json.load(file)
+        imgs = data["images"]
+        anns = data["annotations"]
+        self.categories = data["categories"]
+
+        # create some artificial splits
+        splits = ["train" for _ in range(20)] + ["val" for _ in range(10)]
+        self.additions = []  # list of additional training examples
+
+        for i, path in enumerate(im_paths):
+            # find annotations matching the COCO image
+            gran = path.split("/")[-1]
+            img = [img for img in imgs if img["file_name"] == gran][0]
+            img_id = img["id"]
+            img_anns = [ann for ann in anns if ann["image_id"] == img_id]
+
+            # load the image
+            im = cv2.imread(path)
+            height, width, _ = im.shape
+
+            # initialize annotations for LDF
+            mask = np.zeros(
+                (height, width)
+            )  # segmentation mask is always a HxW numpy array
+            boxes = (
+                []
+            )  # bounding boxes are a list of [class, x, y, width, height] of the box
+            keypoints = []  # keypoints are a list of classes and (x,y) points
+
+            # create the additions
+            for ann in img_anns:
+                # COCO-specific conversion for segmentation
+                seg = ann["segmentation"]
+                if isinstance(seg, list):
+                    for s in seg:
+                        poly = np.array(s).reshape(-1, 2)
+                        poly = [(poly[i, 0], poly[i, 1]) for i in range(len(poly))]
+                        m = Image.new("L", (width, height), 0)
+                        ImageDraw.Draw(m).polygon(poly, outline=1, fill=1)
+                        m = np.array(m)
+                        mask[m == 1] = 1
+                # COCO-specific conversion for bounding boxes
+                x, y, w, h = ann["bbox"]
+                boxes.append(["person", x / width, y / height, w / width, h / height])
+                # COCO-specific conversion for keypoints
+                kps = np.array(ann["keypoints"]).reshape(-1, 3)
+                keypoint = []
+                for kp in kps:
+                    if kp[2] == 0:
+                        keypoint.append((float("nan"), float("nan")))
+                    else:
+                        keypoint.append((kp[0] / width, kp[1] / height))
+                keypoints.append(["person", keypoint])
+            data = {
+                "filepath": path,
+                "segmentation": mask,
+                "boxes": boxes,
+                "keypoints": keypoints,
+                "split": splits[i],
+            }
+            if i % 2 == 0:
+                A_dict = data
+            else:
+                B_dict = data
+                self.additions.append({"A": A_dict, "B": B_dict})
 
-        with LuxonisDataset(self.team, self.name) as dataset:
+        self.dataset_id = LuxonisDataset.create(
+            self.team_id, self.team_name, self.dataset_name
+        )
+
+        print("Testing", self.dataset_id)
+
+    def test_local_init(self):
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
             pass
 
         curr = self.conn.luxonis_dataset_document.find(
-            { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
+            {"$and": [{"team_id": self.team_id}, {"_id": ObjectId(self.dataset_id)}]}
         )
         res = list(curr)
         self.assertGreater(len(res), 0, "Document not created")
         self.assertEqual(len(res), 1, "Multiple documents created")
         res = res[0]
-        self.assertEqual(res['team_name'], 'unittest', "Wrong team name")
-        self.assertEqual(res['dataset_name'], 'coco', "Wrong dataset name")
-        self.assertEqual(res['current_version'], 0., "Version initialize failure")
-        self.assertEqual(res['path'], f'{Path.home()}/.cache/luxonis_ml/data/unittest/datasets/coco', "Dataset path failure")
-        self.assertEqual(res['bucket_type'], 'local', "Default bucket type is not local")
-
-        curr = self.conn.datasets.find(
-            { "name": f"{self.team}-{self.name}" }
+        self.assertEqual(res["team_id"], self.team_id, "Wrong team ID")
+        self.assertEqual(res["_id"], ObjectId(self.dataset_id), "Wrong dataset ID")
+        self.assertEqual(res["team_name"], "unittest", "Wrong team name")
+        self.assertEqual(res["dataset_name"], "coco", "Wrong dataset name")
+        self.assertEqual(res["current_version"], 0.0, "Version initialize failure")
+        self.assertEqual(
+            res["path"],
+            f"{Path.home()}/.cache/luxonis_ml/data/{self.team_id}/datasets/{self.dataset_id}",
+            "Dataset path failure",
+        )
+        self.assertEqual(
+            res["bucket_type"], "local", "Default bucket type is not local"
         )
+
+        curr = self.conn.datasets.find({"name": f"{self.team_id}-{self.dataset_id}"})
         res2 = list(curr)
         self.assertGreater(len(res2), 0, "Underlying fo dataset not found")
         self.assertEqual(len(res2), 1, "Duplicate underlying fo datasets")
         res2 = res2[0]
-        self.assertEqual(res['_dataset_id'], res2['_id'], "Luxonis dataset does not reference fo dataset")
+        self.assertEqual(
+            res["_dataset_id"],
+            res2["_id"],
+            "Luxonis dataset does not reference fo dataset",
+        )
 
-        with LuxonisDataset(self.team, self.name, bucket_type='aws') as dataset:
+        with LuxonisDataset(
+            self.team_id, self.dataset_id, bucket_type="aws"
+        ) as dataset:
             dataset.version = 5
 
         curr = self.conn.luxonis_dataset_document.find(
-            { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
+            {"$and": [{"team_id": self.team_id}, {"_id": ObjectId(self.dataset_id)}]}
         )
         res = list(curr)
         self.assertGreater(len(res), 0, "Document not created")
         self.assertEqual(len(res), 1, "Multiple documents created")
         res = res[0]
-        self.assertEqual(res['current_version'], 5., "Update version field fail")
-        self.assertEqual(res['bucket_type'], 'local', "Default override_bucket_type arg fail")
+        self.assertEqual(res["current_version"], 5.0, "Update version field fail")
+        self.assertEqual(
+            res["bucket_type"], "local", "Default override_bucket_type arg fail"
+        )
 
-    def test_aws_init(self):
+        with LuxonisDataset(
+            self.team_id, self.dataset_id, bucket_type="aws"
+        ) as dataset:
+            dataset.version = 0  # set back to 0
 
-        with LuxonisDataset(self.team, self.name, bucket_type='aws', override_bucket_type=True) as dataset:
+    def test_aws_init(self):
+        with LuxonisDataset(
+            self.team_id, self.dataset_id, bucket_type="aws", override_bucket_type=True
+        ) as dataset:
             pass
 
         curr = self.conn.luxonis_dataset_document.find(
-            { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
+            {"$and": [{"team_id": self.team_id}, {"_id": ObjectId(self.dataset_id)}]}
         )
         res = list(curr)
         self.assertGreater(len(res), 0, "Document not created")
         self.assertEqual(len(res), 1, "Multiple documents created")
         res = res[0]
-        self.assertEqual(res['bucket_type'], 'aws', "Default override_bucket_type arg fail")
+        self.assertEqual(
+            res["bucket_type"], "aws", "Default override_bucket_type arg fail"
+        )
 
     def test_source(self):
-
-        with LuxonisDataset(self.team, self.name) as dataset:
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
             dataset.create_source(
-                'test_source',
-                custom_components = [
-                    LDFComponent('A', HType.IMAGE, IType.BGR),
-                    LDFComponent('B', HType.IMAGE, IType.DEPTH),
-                ]
+                "test_source",
+                custom_components=[
+                    LDFComponent("A", HType.IMAGE, IType.BGR),
+                    LDFComponent("B", HType.IMAGE, IType.DEPTH),
+                ],
             )
 
         curr = self.conn.luxonis_dataset_document.find(
-            { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
+            {"$and": [{"team_id": self.team_id}, {"_id": ObjectId(self.dataset_id)}]}
         )
         res = list(curr)[0]
         curr = self.conn.luxonis_source_document.find(
-            { "_luxonis_dataset_id": res["_id"] }
+            {"_luxonis_dataset_id": res["_id"]}
         )
         res2 = list(curr)
         self.assertGreater(len(res2), 0, "Source document not created")
         self.assertEqual(len(res2), 1, "Multiple source documents created")
         res2 = res2[0]
-        self.assertEqual(res2['name'], 'test_source', "Wrong source name")
-        self.assertEqual(res2['source_type'], 'custom', "Wrong source type")
-        self.assertEqual(res2['component_names'], ['A','B'], "Wrong component names")
-        self.assertEqual(res2['component_htypes'], [1,1], "Wrong HType")
-        self.assertEqual(res2['component_itypes'], [1,4], "Wrong IType")
+        self.assertEqual(res2["name"], "test_source", "Wrong source name")
+        self.assertEqual(res2["source_type"], "custom", "Wrong source type")
+        self.assertEqual(res2["component_names"], ["A", "B"], "Wrong component names")
+        self.assertEqual(res2["component_htypes"], [1, 1], "Wrong HType")
+        self.assertEqual(res2["component_itypes"], [1, 4], "Wrong IType")
 
-        with LuxonisDataset(self.team, self.name) as dataset:
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
             dataset.create_source(
-                'test_source',
-                custom_components = [
-                    LDFComponent('A', HType.IMAGE, IType.BGR),
-                    LDFComponent('B', HType.IMAGE, IType.BGR),
-                ]
+                "test_source",
+                custom_components=[
+                    LDFComponent("A", HType.IMAGE, IType.BGR),
+                    LDFComponent("B", HType.IMAGE, IType.BGR),
+                ],
             )
 
         curr = self.conn.luxonis_source_document.find(
-            { "_luxonis_dataset_id": res["_id"] }
+            {"_luxonis_dataset_id": res["_id"]}
         )
         res2 = list(curr)[0]
-        self.assertEqual(res2['component_itypes'], [1,1], "Wrong IType after changing source")
+        self.assertEqual(
+            res2["component_itypes"], [1, 1], "Wrong IType after changing source"
+        )
+
+    def test_transactions(self):
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
+            dataset.set_classes(
+                ["person", "orange", "pear"]
+            )  # needed for classification and detection
+            dataset.set_mask_targets(
+                {1: "person", 2: "orange", 3: "pear"}
+            )  # needed for segmentation
+            dataset.set_skeleton(
+                {
+                    "labels": self.categories[0]["keypoints"],
+                    "edges": (np.array(self.categories[0]["skeleton"]) - 1).tolist(),
+                }
+            )  # optional for keypoints to define the skeleton visualization
+
+            result = dataset._check_transactions()
+            self.assertEqual(result, None, "Empty transactions fail")
+
+            dataset._make_transaction(LDFTransactionType.ADD)
+            result = dataset._check_transactions()
+            self.assertEqual(result, None, "Missing END transaction fail")
+            time.sleep(0.5)
+            curr = self.conn.transaction_document.find(
+                {"_dataset_id": dataset.dataset_doc.id}
+            )
+            res = list(curr)
+            self.assertEqual(len(res), 0, "Faulty transaction not deleted")
+
+            dataset._make_transaction(LDFTransactionType.UPDATE)
+            dataset._make_transaction(LDFTransactionType.DELETE)
+            dataset._make_transaction(LDFTransactionType.END)
+            time.sleep(0.5)
+            result = dataset._check_transactions()
+            self.assertEqual(len(result), 3, "With END transaction fail")
+
+            # cleanup
+            self.conn.transaction_document.delete_many(
+                {"_dataset_id": dataset.dataset_doc.id}
+            )
+
+    def test_add(self):
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
+            dataset.create_source(
+                "test_source",
+                custom_components=[
+                    LDFComponent("A", HType.IMAGE, IType.BGR),
+                    LDFComponent("B", HType.IMAGE, IType.BGR),
+                ],
+            )
+
+            transaction_to_additions, media_change, field_change = dataset._add_filter(
+                self.additions
+            )
+
+            self.assertEqual(
+                len(transaction_to_additions),
+                30,
+                "Wrong number of entries in transaction_to_additions",
+            )
+            self.assertEqual(media_change, True, "media_change failed")
+            self.assertEqual(field_change, False, "field_change failed")
+
+            curr = self.conn.transaction_document.find(
+                {"_dataset_id": dataset.dataset_doc.id}
+            )
+            res = list(curr)
+            self.assertEqual(len(res), 31, "Wrong number of saved transactions")
+            num_adds = np.sum([True if t["action"] == "ADD" else False for t in res])
+            num_updates = np.sum(
+                [True if t["action"] == "UPDATE" else False for t in res]
+            )
+            num_deletes = np.sum(
+                [True if t["action"] == "DELETE" else False for t in res]
+            )
+            num_ends = np.sum([True if t["action"] == "END" else False for t in res])
+            self.assertEqual(num_adds, 30, "Wrong number of ADDs")
+            self.assertEqual(num_updates, 0, "Wrong number of UPDATEs")
+            self.assertEqual(num_deletes, 0, "Wrong number of DELETEs")
+            self.assertEqual(num_ends, 1, "Wrong number of ENDs")
+
+            # Execution
+            dataset._add_execute(self.additions, transaction_to_additions)
+
+            self.assertEqual(len(dataset.fo_dataset), 15, "Not added to dataset")
+            curr = self.conn.transaction_document.find(
+                {"_dataset_id": dataset.dataset_doc.id}
+            )
+            res = list(curr)
+            num_executed = np.sum(
+                [True if t["executed"] == True else False for t in res]
+            )
+            self.assertEqual(num_executed, 31, "Wrong number of executed transactions")
+
+    def test_version_1(self):
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
+            dataset.create_version(note="test version 1")
+            self.assertEqual(dataset.version, 1.0, "Wrong data version incr")
+
+            num_latest, num_version = 0, 0
+            dataset.fo_dataset.group_slice = "A"
+            for sample in dataset.fo_dataset:
+                if sample.latest:
+                    num_latest += 1
+                if sample.version == 1.0:
+                    num_version += 1
+            self.assertEqual(
+                num_latest, 15, "Not all versioned samples are latest for A"
+            )
+            self.assertEqual(
+                num_version,
+                15,
+                "Not all versioned samples have the right version for A",
+            )
+
+            num_latest, num_version = 0, 0
+            dataset.fo_dataset.group_slice = "B"
+            for sample in dataset.fo_dataset:
+                if sample.latest:
+                    num_latest += 1
+                if sample.version == 1.0:
+                    num_version += 1
+            self.assertEqual(
+                num_latest, 15, "Not all versioned samples are latest for B"
+            )
+            self.assertEqual(
+                num_version,
+                15,
+                "Not all versioned samples have the right version for B",
+            )
+
+            curr = self.conn.version_document.find(
+                {"dataset_id_str": dataset.dataset_doc.fo_dataset_id}
+            )
+            res = list(curr)
+            self.assertEqual(len(res), 1, "Number of saved versions")
+
+    def test_modify(self):
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
+            transaction_to_additions = []
+            a = deepcopy(self.additions[-1])
+            a["A"]["class"] = ["person"]  # adding a new field (classification)
+            a["A"]["weather"] = "sunny"  # adding a new field
+            del a["A"]["split"]  # removing a field
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(media_change, False, "media_change failed")
+            self.assertEqual(field_change, True, "field_change failed")
+
+            res = dataset._check_transactions()
+            self.assertEqual(len(res), 4, "Correct number of new transactions")
+            num_adds = np.sum([True if t["action"] == "ADD" else False for t in res])
+            num_updates = np.sum(
+                [True if t["action"] == "UPDATE" else False for t in res]
+            )
+            num_deletes = np.sum(
+                [True if t["action"] == "DELETE" else False for t in res]
+            )
+            num_ends = np.sum([True if t["action"] == "END" else False for t in res])
+            num_executed = np.sum(
+                [True if t["executed"] == True else False for t in res]
+            )
+            self.assertEqual(num_adds, 0, "Wrong number of ADDs")
+            self.assertEqual(num_updates, 3, "Wrong number of UPDATEs")
+            self.assertEqual(num_deletes, 0, "Wrong number of DELETEs")
+            self.assertEqual(num_ends, 1, "Wrong number of ENDs")
+            self.assertEqual(
+                num_executed, 0, "Some transactions are unexpectedly executed"
+            )
+
+            # non-annotation field change
+            a = deepcopy(self.additions[-1])
+            a["A"]["weather"] = "stormy"
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(field_change, True, "Non-annotation field update")
+            # classification change
+            a = deepcopy(self.additions[-1])
+            a["A"]["class"] = ["orange", "pear"]
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(field_change, True, "Classification update")
+            # boxes change
+            a = deepcopy(self.additions[-1])
+            a["A"]["boxes"][0][0] = "orange"  # test class change with str
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(field_change, True, "Boxes update (change class str)")
+            a = deepcopy(self.additions[-1])
+            a["A"]["boxes"][0][0] = 1  # test class change with int
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(field_change, True, "Boxes update (change class int)")
+            a = deepcopy(self.additions[-1])
+            a["A"]["boxes"][0][1] += 0.1  # test coordinate change
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(field_change, True, "Boxes update (change coordinate)")
+            # segmentation change
+            a = deepcopy(self.additions[-1])
+            a["A"]["segmentation"] = np.zeros(a["A"]["segmentation"].shape)
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(field_change, True, "Segmentation update")
+            # keypoints change
+            a = deepcopy(self.additions[-1])
+            a["A"]["keypoints"][0][0] = "orange"  # test class change with str
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(field_change, True, "Keypoints update (change class str)")
+            a = deepcopy(self.additions[-1])
+            a["A"]["keypoints"][0][0] = 1  # test class change with int
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(field_change, True, "Keypoints update (change class int)")
+            a = deepcopy(self.additions[-1])
+            a["A"]["keypoints"][0][1][0] = (
+                a["A"]["keypoints"][0][1][0][0] + 0.1,
+                a["A"]["keypoints"][0][1][0][1],
+            )
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(
+                field_change,
+                True,
+                "Keypoints update (change coordinate non-NaN->non-NaN)",
+            )
+            a = deepcopy(self.additions[-1])
+            a["B"]["keypoints"][0][1][8] = (
+                a["A"]["keypoints"][0][1][0][0] + 0.1,
+                a["A"]["keypoints"][0][1][0][1],
+            )
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(
+                field_change, True, "Keypoints update (change coordinate NaN->non-Nan)"
+            )
+            a = deepcopy(self.additions[-1])
+            a["A"]["keypoints"][0][1][0] = (np.nan, np.nan)
+            tta, media_change, field_change = dataset._add_filter([a])
+            transaction_to_additions += tta
+            self.assertEqual(
+                field_change, True, "Keypoints update (change coordinate non-NaN->NaN)"
+            )
+
+            dataset._add_execute()
+            dataset.fo_dataset.group_slice = "A"
+            query = dataset.fo_dataset.match(
+                (F("latest") == False) & (F("version") == -1)
+            )
+            self.assertEqual(
+                len(query), 1, "Only one change in A when executing modifications"
+            )
+            for sample in query:
+                break
+            self.assertEqual(
+                sample["weather"], "stormy", "Executed change to added field"
+            )
+            self.assertEqual(sample["split"], None, "Executed change to added field")
+            self.assertEqual(
+                sample["class"]["classifications"][0]["label"],
+                "orange",
+                "Executed change to classification label",
+            )
+            # TODO: better box test
+            self.assertEqual(
+                np.sum(sample["segmentation"]["mask"]),
+                0.0,
+                "Executed change to segmentation mask",
+            )
+            # TODO: better keypoint test
+
+            dataset.fo_dataset.group_slice = "B"
+            query = dataset.fo_dataset.match(
+                (F("latest") == False) & (F("version") == -1)
+            )
+            self.assertEqual(
+                len(query), 1, "Only one change in B when executing modifications"
+            )
+            for sample in query:
+                break
+            # TODO: keypoint test
+
+    def test_version_2(self):
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
+            dataset.create_version(note="test version 2")
+            self.assertEqual(dataset.version, 1.1, "Wrong data version incr")
+
+            num_latest, num_version = 0, 0
+            dataset.fo_dataset.group_slice = "A"
+            for sample in dataset.fo_dataset:
+                if sample.latest:
+                    num_latest += 1
+                if sample.version == 1.1:
+                    num_version += 1
+            self.assertEqual(
+                num_latest, 15, "Not all versioned samples are latest for A"
+            )
+            self.assertEqual(
+                num_version, 1, "Not all versioned samples have the right version for A"
+            )
+
+            num_latest, num_version = 0, 0
+            dataset.fo_dataset.group_slice = "B"
+            for sample in dataset.fo_dataset:
+                if sample.latest:
+                    num_latest += 1
+                if sample.version == 1.1:
+                    num_version += 1
+            self.assertEqual(
+                num_latest, 15, "Not all versioned samples are latest for B"
+            )
+            self.assertEqual(
+                num_version, 1, "Not all versioned samples have the right version for B"
+            )
+
+            curr = self.conn.version_document.find(
+                {"dataset_id_str": dataset.dataset_doc.fo_dataset_id}
+            )
+            res = list(curr)
+            self.assertEqual(len(res), 2, "Number of saved versions")
+            set1 = set(res[0]["samples"])
+            set2 = set(res[1]["samples"])
+            set1_unique = set1.union(set2) - set2
+            set2_unique = set1.union(set2) - set1
+            self.assertEqual(len(set1_unique), 2, "One deprecated sample")
+            self.assertEqual(len(set2_unique), 2, "One added sample")
 
     def test_delete(self):
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
+            for sample in dataset.fo_dataset:
+                break
+            dataset.delete([sample.id])
 
+            res = dataset._check_transactions(
+                for_versioning=True
+            )  # will show the latest 3 transactions
+            self.assertEqual(len(res), 3, "Wrong number of transactions for delete")
+            num_executed = np.sum(
+                [True if t["executed"] == True else False for t in res]
+            )
+            num_adds = np.sum([True if t["action"] == "ADD" else False for t in res])
+            num_updates = np.sum(
+                [True if t["action"] == "UPDATE" else False for t in res]
+            )
+            num_deletes = np.sum(
+                [True if t["action"] == "DELETE" else False for t in res]
+            )
+            num_ends = np.sum([True if t["action"] == "END" else False for t in res])
+            self.assertEqual(num_executed, 3, "Wrong number of executed transactions")
+            self.assertEqual(num_adds, 0, "Wrong number of ADD transactions")
+            self.assertEqual(num_updates, 0, "Wrong number of UPDATE transactions")
+            self.assertEqual(num_deletes, 2, "Wrong number of DELETE transactions")
+            self.assertEqual(num_ends, 1, "Wrong number of END transactions")
+
+            dres = [t for t in res if t["action"] == "DELETE"]
+            for t in dres:
+                sample = dataset.fo_dataset[t["sample_id"]]
+                self.assertEqual(sample.latest, False, "Delete latest!=False")
+
+    def test_version_3(self):
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
+            dataset.create_version(note="deleted some samples")
+            self.assertEqual(dataset.version, 2.1, "Wrong data version incr")
+
+            num_latest, num_version = 0, 0
+            dataset.fo_dataset.group_slice = "A"
+            for sample in dataset.fo_dataset:
+                if sample.latest:
+                    num_latest += 1
+                if sample.version == 2.1:
+                    num_version += 1
+            self.assertEqual(
+                num_latest, 14, "Not all versioned samples are latest for A"
+            )
+            self.assertEqual(
+                num_version, 0, "Not all versioned samples have the right version for A"
+            )
+
+            num_latest, num_version = 0, 0
+            dataset.fo_dataset.group_slice = "B"
+            for sample in dataset.fo_dataset:
+                if sample.latest:
+                    num_latest += 1
+                if sample.version == 2.1:
+                    num_version += 1
+            self.assertEqual(
+                num_latest, 14, "Not all versioned samples are latest for B"
+            )
+            self.assertEqual(
+                num_version, 0, "Not all versioned samples have the right version for B"
+            )
+
+            curr = self.conn.version_document.find(
+                {"dataset_id_str": dataset.dataset_doc.fo_dataset_id}
+            )
+            res = list(curr)
+            self.assertEqual(len(res), 3, "Number of saved versions")
+
+    def test_delete_dataset(self):
         curr = self.conn.luxonis_dataset_document.find(
-            { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
+            {"$and": [{"team_id": self.team_id}, {"_id": ObjectId(self.dataset_id)}]}
         )
         res = list(curr)[0]
         old_id = res["_id"]
 
-        with LuxonisDataset(self.team, self.name) as dataset:
-            dataset.delete()
+        with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
+            dataset.delete_dataset()
 
         curr = self.conn.luxonis_dataset_document.find(
-            { "$and": [{"team_name": self.team}, {"dataset_name": self.name}] }
+            {"$and": [{"team_id": self.team_id}, {"_id": ObjectId(self.dataset_id)}]}
         )
         res = list(curr)
         self.assertEqual(len(res), 0, "Document not deleted")
-        # TODO: test if source is deleted
-        curr = self.conn.luxonis_source_document.find(
-            { "_luxonis_dataset_id": old_id }
-        )
+
+        curr = self.conn.luxonis_source_document.find({"_luxonis_dataset_id": old_id})
         res2 = list(curr)
         self.assertEqual(len(res2), 0, "Source document not deleted")
 
         # TODO: test dataset version documents are deleted
+        # TODO: test dataset transaction documents are deleted
+
+    # @classmethod
+    # def tearDownClass(self):
+    #     with LuxonisDataset(self.team_id, self.dataset_id) as dataset:
+    #         dataset.delete_dataset()
 
-if __name__ == "__main__":
 
-    # test_suite = unittest.TestLoader().loadTestsFromTestCase(LuxonisDatasetTester)
+if __name__ == "__main__":
     suite = unittest.TestSuite()
-    suite.addTest(LuxonisDatasetTester('test_local_init'))
-    suite.addTest(LuxonisDatasetTester('test_aws_init'))
-    suite.addTest(LuxonisDatasetTester('test_source'))
-    suite.addTest(LuxonisDatasetTester('test_delete'))
+    suite.addTest(LuxonisDatasetTester("test_local_init"))
+    suite.addTest(LuxonisDatasetTester("test_aws_init"))
+    suite.addTest(LuxonisDatasetTester("test_source"))
+    suite.addTest(LuxonisDatasetTester("test_transactions"))
+    suite.addTest(LuxonisDatasetTester("test_add"))
+    suite.addTest(LuxonisDatasetTester("test_version_1"))
+    suite.addTest(LuxonisDatasetTester("test_modify"))
+    suite.addTest(LuxonisDatasetTester("test_version_2"))
+    suite.addTest(LuxonisDatasetTester("test_delete"))
+    suite.addTest(LuxonisDatasetTester("test_version_3"))
+    suite.addTest(LuxonisDatasetTester("test_delete_dataset"))
     runner = unittest.TextTestRunner()
     runner.run(suite)