reformatted

examples/Data_Custom_Example.ipynb

@@ -240,7 +240,9 @@
     "    {\n",
     "        \"person\": {\n",
     "            \"labels\": data[\"categories\"][0][\"keypoints\"],\n",
-    "            \"edges\": (np.array(data[\"categories\"][0][\"skeleton\"]) - 1).tolist(),\n",
+    "            \"edges\": (\n",
+    "                np.array(data[\"categories\"][0][\"skeleton\"]) - 1\n",
+    "            ).tolist(),\n",
     "        }\n",
     "    },\n",
     "    task=\"keypoints\",\n",
@@ -307,7 +309,9 @@
     "    for kp in kps:\n",
     "        kp = kp[1:].reshape(-1, 3)\n",
     "        for k in kp:\n",
-    "            cv2.circle(image, (int(k[0] * w), int(k[1] * h)), 2, (0, 255, 0), 2)\n",
+    "            cv2.circle(\n",
+    "                image, (int(k[0] * w), int(k[1] * h)), 2, (0, 255, 0), 2\n",
+    "            )\n",
     "\n",
     "    plt.imshow(image)\n",
     "    plt.axis(\"off\")  # Optional: Hide axis\n",

examples/Data_Parser_Example.ipynb

@@ -82,7 +82,9 @@
    "outputs": [],
    "source": [
     "dataset_name = \"coco_test\"\n",
-    "parser = LuxonisParser(dataset_dir, dataset_name=dataset_name, delete_existing=True)\n",
+    "parser = LuxonisParser(\n",
+    "    dataset_dir, dataset_name=dataset_name, delete_existing=True\n",
+    ")\n",
     "dataset = parser.parse(random_split=True)"
    ]
   },
@@ -125,7 +127,9 @@
     "    for kp in kps:\n",
     "        kp = kp[1:].reshape(-1, 3)\n",
     "        for k in kp:\n",
-    "            cv2.circle(image, (int(k[0] * w), int(k[1] * h)), 2, (0, 255, 0), 2)\n",
+    "            cv2.circle(\n",
+    "                image, (int(k[0] * w), int(k[1] * h)), 2, (0, 255, 0), 2\n",
+    "            )\n",
     "\n",
     "    plt.imshow(image)\n",
     "    plt.axis(\"off\")  # Optional: Hide axis\n",

examples/Embeddings_LDF_Qdrant_Example.ipynb

@@ -60,7 +60,9 @@
    "outputs": [],
    "source": [
     "# Load the data\n",
-    "data_loader = load_mnist_data(save_path=\"./data/mnist\", num_samples=640, batch_size=64)"
+    "data_loader = load_mnist_data(\n",
+    "    save_path=\"./data/mnist\", num_samples=640, batch_size=64\n",
+    ")"
    ]
   },
   {
@@ -108,7 +110,9 @@
    "outputs": [],
    "source": [
     "# Load the data\n",
-    "data_loader = load_mnist_data(save_path=\"./data/mnist\", num_samples=640, batch_size=64)"
+    "data_loader = load_mnist_data(\n",
+    "    save_path=\"./data/mnist\", num_samples=640, batch_size=64\n",
+    ")"
    ]
   },
   {
@@ -159,7 +163,9 @@
     "    and \"CUDAExecutionProvider\" in onnxruntime.get_available_providers()\n",
     "    else None\n",
     ")\n",
-    "ort_session = onnxruntime.InferenceSession(\"./data/resnet50-1.onnx\", providers=provider)\n",
+    "ort_session = onnxruntime.InferenceSession(\n",
+    "    \"./data/resnet50-1.onnx\", providers=provider\n",
+    ")\n",
     "\n",
     "# Extract embeddings from the dataset\n",
     "embeddings, labels = extract_embeddings_onnx(\n",
@@ -420,7 +426,9 @@
     "    and \"CUDAExecutionProvider\" in onnxruntime.get_available_providers()\n",
     "    else None\n",
     ")\n",
-    "ort_session = onnxruntime.InferenceSession(\"./data/resnet50-1.onnx\", providers=provider)"
+    "ort_session = onnxruntime.InferenceSession(\n",
+    "    \"./data/resnet50-1.onnx\", providers=provider\n",
+    ")"
    ]
   },
   {
@@ -446,7 +454,9 @@
     "\n",
     "# Create a collection\n",
     "qdrant_api.create_collection(\n",
-    "    collection_name=\"Mnist_LDF\", properties=[\"label\", \"image_path\"], vector_size=2048\n",
+    "    collection_name=\"Mnist_LDF\",\n",
+    "    properties=[\"label\", \"image_path\"],\n",
+    "    vector_size=2048,\n",
     ")"
    ]
   },

examples/Embeddings_LDF_Weaviate_Example.ipynb

@@ -71,7 +71,9 @@
    "outputs": [],
    "source": [
     "# Load the data\n",
-    "data_loader = load_mnist_data(save_path=\"./data/mnist\", num_samples=640, batch_size=64)"
+    "data_loader = load_mnist_data(\n",
+    "    save_path=\"./data/mnist\", num_samples=640, batch_size=64\n",
+    ")"
    ]
   },
   {
@@ -119,7 +121,9 @@
    "outputs": [],
    "source": [
     "# Load the data\n",
-    "data_loader = load_mnist_data(save_path=\"./data/mnist\", num_samples=640, batch_size=64)"
+    "data_loader = load_mnist_data(\n",
+    "    save_path=\"./data/mnist\", num_samples=640, batch_size=64\n",
+    ")"
    ]
   },
   {
@@ -170,7 +174,9 @@
     "    and \"CUDAExecutionProvider\" in onnxruntime.get_available_providers()\n",
     "    else None\n",
     ")\n",
-    "ort_session = onnxruntime.InferenceSession(\"./data/resnet50-1.onnx\", providers=provider)\n",
+    "ort_session = onnxruntime.InferenceSession(\n",
+    "    \"./data/resnet50-1.onnx\", providers=provider\n",
+    ")\n",
     "\n",
     "# Extract embeddings from the dataset\n",
     "embeddings, labels = extract_embeddings_onnx(\n",
@@ -232,7 +238,9 @@
     "# Insert the embeddings into the collection\n",
     "uuids = [str(uuid.uuid5(uuid.NAMESPACE_DNS, str(e))) for e in embeddings]\n",
     "label_list_dict = [{\"label\": label} for label in labels]\n",
-    "weaviate_api.insert_embeddings(uuids, embeddings, label_list_dict, batch_size=50)"
+    "weaviate_api.insert_embeddings(\n",
+    "    uuids, embeddings, label_list_dict, batch_size=50\n",
+    ")"
    ]
   },
   {
@@ -254,8 +262,12 @@
    ],
    "source": [
     "# Search for the nearest neighbors\n",
-    "search_uuids, scores = weaviate_api.search_similar_embeddings(embeddings[0], top_k=5)\n",
-    "payloads = weaviate_api.retrieve_payloads_by_ids(search_uuids, properties=[\"label\"])\n",
+    "search_uuids, scores = weaviate_api.search_similar_embeddings(\n",
+    "    embeddings[0], top_k=5\n",
+    ")\n",
+    "payloads = weaviate_api.retrieve_payloads_by_ids(\n",
+    "    search_uuids, properties=[\"label\"]\n",
+    ")\n",
     "\n",
     "# Print the search results\n",
     "for u, p in zip(search_uuids, payloads):\n",
@@ -414,7 +426,9 @@
     "    and \"CUDAExecutionProvider\" in onnxruntime.get_available_providers()\n",
     "    else None\n",
     ")\n",
-    "ort_session = onnxruntime.InferenceSession(\"./data/resnet50-1.onnx\", providers=provider)"
+    "ort_session = onnxruntime.InferenceSession(\n",
+    "    \"./data/resnet50-1.onnx\", providers=provider\n",
+    ")"
    ]
   },
   {
@@ -555,8 +569,12 @@
    ],
    "source": [
     "# Search for the nearest neighbors\n",
-    "search_uuids, scores = weaviate_api.search_similar_embeddings(first_emb, top_k=5)\n",
-    "payloads = weaviate_api.retrieve_payloads_by_ids(search_uuids, properties=[\"label\"])\n",
+    "search_uuids, scores = weaviate_api.search_similar_embeddings(\n",
+    "    first_emb, top_k=5\n",
+    ")\n",
+    "payloads = weaviate_api.retrieve_payloads_by_ids(\n",
+    "    search_uuids, properties=[\"label\"]\n",
+    ")\n",
     "\n",
     "# Print the search results\n",
     "for u, p, s in zip(search_uuids, payloads, scores):\n",
@@ -743,7 +761,9 @@
    "source": [
     "# Setup Weaviate\n",
     "weaviate_api = WeaviateAPI(\"http://localhost:8080\")\n",
-    "weaviate_api.create_collection(collection_name=\"Mnist_LFS\", properties=[\"image_path\"])"
+    "weaviate_api.create_collection(\n",
+    "    collection_name=\"Mnist_LFS\", properties=[\"image_path\"]\n",
+    ")"
    ]
   },
   {
@@ -784,7 +804,9 @@
    ],
    "source": [
     "# Search for the nearest neighbors\n",
-    "search_uuids, scores = weaviate_api.search_similar_embeddings(embeddings[0], top_k=5)\n",
+    "search_uuids, scores = weaviate_api.search_similar_embeddings(\n",
+    "    embeddings[0], top_k=5\n",
+    ")\n",
     "\n",
     "# Print the search results\n",
     "for u, s in zip(search_uuids, scores):\n",

examples/Embeddings_Processing_Example.ipynb

@@ -103,7 +103,9 @@
    "source": [
     "desired_size = int(len(embeddings) * 0.05)\n",
     "# desired_size = 10\n",
-    "selected_image_indices = find_representative_kmedoids(similarity_matrix, desired_size)\n",
+    "selected_image_indices = find_representative_kmedoids(\n",
+    "    similarity_matrix, desired_size\n",
+    ")\n",
     "# selected_image_indices = find_representative_greedy_qdrant(qdrant_client, desired_size, 0, \"mnist3\")"
    ]
   },
@@ -356,7 +358,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "mis_img_paths = qdrant_api.retrieve_payloads_by_ids(missing_img_uuids, [\"image_path\"])\n",
+    "mis_img_paths = qdrant_api.retrieve_payloads_by_ids(\n",
+    "    missing_img_uuids, [\"image_path\"]\n",
+    ")\n",
     "mis_img_paths = [x[\"image_path\"] for x in mis_img_paths]"
    ]
   },
@@ -462,10 +466,14 @@
     "# find representative images\n",
     "desired_size = int(len(embeddings) * 0.05)\n",
     "similarity_matrix = calculate_similarity_matrix(embeddings)\n",
-    "selected_image_indices = find_representative_kmedoids(similarity_matrix, desired_size)\n",
+    "selected_image_indices = find_representative_kmedoids(\n",
+    "    similarity_matrix, desired_size\n",
+    ")\n",
     "\n",
     "selcted_ids = np.array(ids)[selected_image_indices].tolist()\n",
-    "represent_img_paths = w_api.retrieve_payloads_by_ids(selcted_ids, [\"image_path\"])\n",
+    "represent_img_paths = w_api.retrieve_payloads_by_ids(\n",
+    "    selcted_ids, [\"image_path\"]\n",
+    ")\n",
     "represent_img_paths = [x[\"image_path\"] for x in represent_img_paths]\n",
     "\n",
     "# plot\n",
@@ -626,7 +634,9 @@
     "mis_ix, new_y = find_mismatches_centroids(X, y)\n",
     "\n",
     "missing_img_uuids = np.array(ids)[mis_ix].tolist()\n",
-    "mis_img_paths = w_api.retrieve_payloads_by_ids(missing_img_uuids, [\"image_path\"])\n",
+    "mis_img_paths = w_api.retrieve_payloads_by_ids(\n",
+    "    missing_img_uuids, [\"image_path\"]\n",
+    ")\n",
     "mis_img_paths = [x[\"image_path\"] for x in mis_img_paths]\n",
     "\n",
     "# plot\n",

examples/utils/data_utils.py

@@ -27,6 +27,7 @@
 Note: This loader is particularly useful when you want to use MNIST data with models that were
 pre-trained on datasets like ImageNet and expect 3-channel RGB input.
 """
+
 import torch
 import torchvision
 import torchvision.transforms as transforms
@@ -35,15 +36,18 @@
 def mnist_transformations() -> transforms.Compose:
     """Returns composed transformations for the MNIST dataset.
 
-    Transforms the images from 1 channel grayscale to 3 channels RGB and resizes them.
+    Transforms the images from 1 channel grayscale to 3 channels RGB and
+    resizes them.
     """
     return transforms.Compose(
         [
             transforms.Grayscale(num_output_channels=3),
             transforms.Lambda(lambda x: x.convert("RGB")),
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+            ),
         ]
     )
 
@@ -70,7 +74,9 @@ def load_mnist_data(
     )
 
     # If num_samples is set to -1, use the entire dataset
-    num_samples = min(num_samples, len(dataset)) if num_samples != -1 else len(dataset)
+    num_samples = (
+        min(num_samples, len(dataset)) if num_samples != -1 else len(dataset)
+    )
 
     # Create a subset of the dataset using Subset class
     subset = torch.utils.data.Subset(dataset, torch.arange(num_samples))

examples/utils/torch_utils.py

@@ -11,8 +11,8 @@
 
 # PyTorch and ONNX model loading and exporting functions
 def load_model_resnet50(discard_last_layer: bool) -> nn.Module:
-    """Load a pre-trained ResNet-50 model with the last fully connected layer
-    removed."""
+    """Load a pre-trained ResNet-50 model with the last fully connected
+    layer removed."""
     model = models.resnet50(weights=resnet.ResNet50_Weights.IMAGENET1K_V1)
     if discard_last_layer:
         model = nn.Sequential(
@@ -83,7 +83,9 @@ def save_embeddings(
     torch.save(labels, save_path + "labels.pth")
 
 
-def load_embeddings(save_path: str = "./") -> Tuple[torch.Tensor, torch.Tensor]:
+def load_embeddings(
+    save_path: str = "./",
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """Load embeddings and labels tensors from the specified path."""
     embeddings = torch.load(save_path + "embeddings.pth")
     labels = torch.load(save_path + "labels.pth")
@@ -99,18 +101,21 @@ def generate_new_embeddings(
     emb_batch_size: int = 64,
     transform: transforms.Compose = None,
 ):
-    """Generate embeddings for new images using a given ONNX runtime session.
+    """Generate embeddings for new images using a given ONNX runtime
+    session.
 
     @type img_paths: List[str]
     @param img_paths: List of image paths for new images.
     @type ort_session: L{InferenceSession}
     @param ort_session: ONNX runtime session.
     @type output_layer_name: str
-    @param output_layer_name: Name of the output layer in the ONNX model.
+    @param output_layer_name: Name of the output layer in the ONNX
+        model.
     @type emb_batch_size: int
     @param emb_batch_size: Batch size for generating embeddings.
     @type transform: torchvision.transforms
-    @param transform: Optional torchvision transform for preprocessing images.
+    @param transform: Optional torchvision transform for preprocessing
+        images.
     @rtype: List[List[float]]
     @return: List of embeddings for the new images.
     """
@@ -141,7 +146,9 @@ def generate_new_embeddings(
         batch_tensor = torch.stack(batch_tensors).cuda()
 
         # Run the ONNX model on the batch
-        ort_inputs = {ort_session.get_inputs()[0].name: batch_tensor.cpu().numpy()}
+        ort_inputs = {
+            ort_session.get_inputs()[0].name: batch_tensor.cpu().numpy()
+        }
         ort_outputs = ort_session.run([output_layer_name], ort_inputs)
 
         # Append the embeddings from the batch to the new_embeddings list

luxonis_ml/data/__main__.py

@@ -59,7 +59,9 @@ def get_dataset_info(name: str) -> Tuple[int, List[str], List[str]]:
 def print_info(name: str) -> None:
     dataset = LuxonisDataset(name)
     _, classes = dataset.get_classes()
-    table = Table(title="Classes", box=rich.box.ROUNDED, row_styles=["yellow", "cyan"])
+    table = Table(
+        title="Classes", box=rich.box.ROUNDED, row_styles=["yellow", "cyan"]
+    )
     table.add_column("Task", header_style="magenta i", max_width=30)
     table.add_column("Class Names", header_style="magenta i", max_width=50)
     for task, c in classes.items():