cesnet-datazoo 0.1.3__tar.gz → 0.1.4__tar.gz

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cesnet-datazoo
-Version: 0.1.3
+Version: 0.1.4
 Summary: A toolkit for large network traffic datasets
 Author-email: Jan Luxemburk <luxemburk@cesnet.cz>, Karel Hynek <hynekkar@cesnet.cz>
 Maintainer-email: Jan Luxemburk <luxemburk@cesnet.cz>, Karel Hynek <hynekkar@cesnet.cz>

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo/config.py

@@ -14,8 +14,9 @@ import yaml
 from pydantic import model_validator
 from pydantic.dataclasses import dataclass
 
-from cesnet_datazoo.constants import (PHIST_BIN_COUNT, PPI_MAX_LEN, SELECTED_TCP_FLAGS,
-                                      TCP_PPI_CHANNELS, UDP_PPI_CHANNELS)
+from cesnet_datazoo.constants import (PHIST_BIN_COUNT, PPI_MAX_LEN, QUIC_SNI_COLUMN,
+                                      SELECTED_TCP_FLAGS, TCP_PPI_CHANNELS, TLS_SNI_COLUMN,
+                                      UDP_PPI_CHANNELS)
 
 if TYPE_CHECKING:
     from cesnet_datazoo.datasets.cesnet_dataset import CesnetDataset
@@ -128,6 +129,7 @@ class DatasetConfig():
         flowstats_features_boolean: Taken from `dataset.metadata.flowstats_features_boolean`.
         flowstats_features_phist: Taken from `dataset.metadata.packet_histograms` if `use_packet_histograms` is true, otherwise an empty list.
         other_fields: Taken from `dataset.metadata.other_fields` if `return_other_fields` is true, otherwise an empty list.
+        sni_column: Database column with SNI domains, can be None for datasets without SNI domains.
 
     # Configuration options
 
@@ -343,6 +345,8 @@ class DatasetConfig():
         # Configure features
         self.flowstats_features = dataset.metadata.flowstats_features
         self.flowstats_features_boolean = dataset.metadata.flowstats_features_boolean
+        sni_column = TLS_SNI_COLUMN if dataset.metadata.protocol == Protocol.TLS else QUIC_SNI_COLUMN
+        self.sni_column = sni_column if sni_column in dataset.metadata.other_fields else None
         self.other_fields = dataset.metadata.other_fields if self.return_other_fields else []
         if self.use_packet_histograms:
             if len(dataset.metadata.packet_histograms) == 0:

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo/constants.py

@@ -39,6 +39,8 @@ DEFAULT_BACKGROUND_CLASS = "default-background"
 GOOGLE_BACKGROUND_CLASS = "google-background"
 
 # Indices
-INDICES_TABLE_POS = 0
-INDICES_INDEX_POS = 1
-INDICES_LABEL_POS = 2
+INDICES_TABLE_FIELD = "TABLE"
+INDICES_INDEX_FIELD = "INDEX"
+INDICES_APP_FIELD = "APP"
+INDICES_SNI_FIELD = "SNI"
+INDICES_DTYPE = [(INDICES_TABLE_FIELD, "int32"), (INDICES_INDEX_FIELD, "int32"), (INDICES_APP_FIELD, "int32"), (INDICES_SNI_FIELD, "U50")]

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo/datasets/cesnet_dataset.py

@@ -16,7 +16,7 @@ from torch.utils.data import BatchSampler, DataLoader, RandomSampler, Sampler, S
 from typing_extensions import assert_never
 
 from cesnet_datazoo.config import AppSelection, DataLoaderOrder, DatasetConfig, ValidationApproach
-from cesnet_datazoo.constants import (APP_COLUMN, CATEGORY_COLUMN, DATASET_SIZES, INDICES_LABEL_POS,
+from cesnet_datazoo.constants import (APP_COLUMN, CATEGORY_COLUMN, DATASET_SIZES, INDICES_APP_FIELD,
                                       SERVICEMAP_FILE, UNKNOWN_STR_LABEL)
 from cesnet_datazoo.datasets.loaders import collate_fn_simple, create_df_from_dataloader
 from cesnet_datazoo.datasets.metadata.dataset_metadata import DatasetMetadata, load_metadata
@@ -555,7 +555,7 @@ class CesnetDataset():
                 train_val_rng = get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.TRAIN_VAL_SPLIT)
                 val_data_path = dataset_config._get_train_data_path()
                 val_unknown_indices = train_unknown_indices
-                train_labels = train_indices[:, INDICES_LABEL_POS]
+                train_labels = train_indices[INDICES_APP_FIELD]
                 if dataset_config.train_dates_weigths is not None:
                     assert dataset_config.val_known_size != "all"
                     # When weight sampling is used, val_known_size is kept but the resulting train size can be smaller due to no enough samples in some train dates
@@ -626,6 +626,7 @@ class CesnetDataset():
                 flowstats_features_boolean=dataset_config.flowstats_features_boolean,
                 flowstats_features_phist=dataset_config.flowstats_features_phist,
                 other_fields=self.dataset_config.other_fields,
+                sni_column=self.dataset_config.sni_column,
                 ppi_channels=dataset_config.get_ppi_channels(),
                 ppi_transform=dataset_config.ppi_transform,
                 flowstats_transform=dataset_config.flowstats_transform,
@@ -644,6 +645,7 @@ class CesnetDataset():
                 flowstats_features_boolean=dataset_config.flowstats_features_boolean,
                 flowstats_features_phist=dataset_config.flowstats_features_phist,
                 other_fields=self.dataset_config.other_fields,
+                sni_column=self.dataset_config.sni_column,
                 ppi_channels=dataset_config.get_ppi_channels(),
                 ppi_transform=dataset_config.ppi_transform,
                 flowstats_transform=dataset_config.flowstats_transform,
@@ -664,6 +666,7 @@ class CesnetDataset():
                 flowstats_features_boolean=dataset_config.flowstats_features_boolean,
                 flowstats_features_phist=dataset_config.flowstats_features_phist,
                 other_fields=self.dataset_config.other_fields,
+                sni_column=self.dataset_config.sni_column,
                 ppi_channels=dataset_config.get_ppi_channels(),
                 ppi_transform=dataset_config.ppi_transform,
                 flowstats_transform=dataset_config.flowstats_transform,

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo/datasets/metadata/dataset_metadata.py

@@ -39,5 +39,5 @@ class DatasetMetadata():
 metadata_df = pd.read_csv(os.path.join(os.path.dirname(__file__), "metadata.csv"), index_col="Name", keep_default_na=False)
 def load_metadata(dataset_name: str) -> DatasetMetadata:
     d = metadata_df.loc[dataset_name].to_dict()
-    d = {k.replace(" ", "_").lower(): v for k, v in d.items()}
+    d = {k.replace(" ", "_").lower(): v for k, v in d.items()} # type: ignore
     return DatasetMetadata(**d)

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo/datasets/statistics.py

@@ -92,7 +92,7 @@ def compute_dataset_statistics(database_path: str,
                                     tables_paths=table_paths,
                                     indices=None,
                                     disabled_apps=disabled_apps,
-                                    return_all_fields=True,
+                                    return_raw_fields=True,
                                     flowstats_features=[],
                                     flowstats_features_boolean=[],
                                     flowstats_features_phist=[],

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo/metrics/classification_report.py

@@ -1,13 +1,12 @@
 import numpy as np
 from sklearn.metrics import accuracy_score, precision_recall_fscore_support
 
-from cesnet_datazoo.metrics.provider_metrics import (per_app_provider_metrics,
-                                                       provider_accuracies)
+from cesnet_datazoo.metrics.provider_metrics import per_app_provider_metrics, provider_accuracies
 from cesnet_datazoo.utils.class_info import ClassInfo
 
 
-def better_classification_report(y_true: np.ndarray, y_pred: np.ndarray, cm: np.ndarray, labels: list[int], class_info: ClassInfo, digits: int = 2, zero_division: int = 0) -> tuple[str, dict[str, float]]:
-    p, r, f1, s  = precision_recall_fscore_support(y_true, y_pred,
+def better_classification_report(test_labels: np.ndarray, preds: np.ndarray, cm: np.ndarray, labels: list[int], class_info: ClassInfo, digits: int = 2, zero_division: int = 0) -> tuple[str, dict[str, float]]:
+    p, r, f1, s  = precision_recall_fscore_support(test_labels, preds,
                                                    labels=labels,
                                                    zero_division=zero_division)
     sc_p, sc_r, sc_f1 = per_app_provider_metrics(cm, class_info=class_info)
@@ -46,20 +45,20 @@ def better_classification_report(y_true: np.ndarray, y_pred: np.ndarray, cm: np.
     report += headers_fmt.format("", *headers_avg, width=width)
     report += row_fmt_avg.format("macro avg", *row_avg, width=width, digits=digits)
 
-    acc = accuracy_score(y_true, y_pred)
-    provider_acc, failed_provider_acc = provider_accuracies(y_true, y_pred, class_info=class_info)
+    acc = accuracy_score(test_labels, preds)
+    provider_acc, failed_provider_acc = provider_accuracies(test_labels, preds, class_info=class_info)
 
     row_fmt_acc = "{:>{width}} {:>15} {:>15} {:>7.{digits}f}\n"
     report += row_fmt_acc.format("acc", "", "", acc, width=width, digits=digits)
     report += row_fmt_acc.format("provider acc", "", "", provider_acc, width=width, digits=digits)
     report += row_fmt_acc.format("failed provider acc", "", "", failed_provider_acc, width=width, digits=digits)
     metrics = {
-        "Test/Accuracy": acc,
-        "Test/Provider Accuracy": provider_acc,
-        "Test/Failed Provider Accuracy": failed_provider_acc,
-        "Test/Fscore": avg_f1,
-        "Test/Provider Fscore": avg_sc_f1,
-        "Test/Recall": avg_r,
-        "Test/Provider Recall": avg_sc_r,
+        "test/acc": acc,
+        "test/provider-acc": provider_acc,
+        "test/failed-provider-acc": failed_provider_acc,
+        "test/fscore": avg_f1,
+        "test/provider-fscore": avg_sc_f1,
+        "test/recall": avg_r,
+        "test/provider-recall": avg_sc_r,
     }
     return report, metrics

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo/metrics/provider_metrics.py

@@ -3,13 +3,13 @@ import numpy as np
 from cesnet_datazoo.utils.class_info import ClassInfo
 
 
-def provider_accuracies(y_true: np.ndarray, y_pred: np.ndarray, class_info: ClassInfo) -> tuple[float, float]:
+def provider_accuracies(true_labels: np.ndarray, preds: np.ndarray, class_info: ClassInfo) -> tuple[float, float]:
     provider_mapping_arr = np.array(list(class_info.provider_mapping.values()))
-    y_true_sc = provider_mapping_arr[y_true]
-    y_pred_sc = provider_mapping_arr[y_pred]
-    mistakes = y_true != y_pred
-    provider_acc = (y_true_sc == y_pred_sc).sum() / len(y_true_sc)
-    failed_provider_acc = (y_true_sc[mistakes] == y_pred_sc[mistakes]).sum() / mistakes.sum()
+    true_labels_provider = provider_mapping_arr[true_labels]
+    preds_provider = provider_mapping_arr[preds]
+    mistakes = true_labels != preds
+    provider_acc = (true_labels_provider == preds_provider).sum() / len(true_labels_provider)
+    failed_provider_acc = (true_labels_provider[mistakes] == preds_provider[mistakes]).sum() / mistakes.sum()
     return provider_acc, failed_provider_acc
 
 def per_app_provider_metrics(cm, class_info: ClassInfo):

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo/pytables_data/indices_setup.py

@@ -1,6 +1,7 @@
 import dataclasses
 import logging
 import os
+import time
 import warnings
 from collections import namedtuple
 from enum import Enum
@@ -9,7 +10,8 @@ import numpy as np
 import pandas as pd
 
 from cesnet_datazoo.config import DatasetConfig
-from cesnet_datazoo.constants import INDICES_INDEX_POS, INDICES_LABEL_POS, INDICES_TABLE_POS
+from cesnet_datazoo.constants import (INDICES_APP_FIELD, INDICES_DTYPE, INDICES_INDEX_FIELD,
+                                      INDICES_TABLE_FIELD)
 from cesnet_datazoo.pytables_data.pytables_dataset import init_test_indices, init_train_indices
 from cesnet_datazoo.utils.fileutils import yaml_dump, yaml_load
 from cesnet_datazoo.utils.random import RandomizedSection, get_fresh_random_generator
@@ -21,8 +23,8 @@ IndicesTuple = namedtuple("IndicesTuple", ["train_indices", "val_known_indices",
 
 
 def sort_indices(indices: np.ndarray) -> np.ndarray:
-    idxs = np.argsort(indices[:, INDICES_INDEX_POS])
-    res = idxs[np.argsort(indices[idxs, INDICES_TABLE_POS], kind="stable")]
+    idxs = np.argsort(indices[INDICES_INDEX_FIELD])
+    res = idxs[np.argsort(indices[idxs][INDICES_TABLE_FIELD], kind="stable")]
     return indices[res]
 
 def subset_and_sort_indices(dataset_config: DatasetConfig, dataset_indices: IndicesTuple) -> IndicesTuple:
@@ -61,7 +63,7 @@ def subset_and_sort_indices(dataset_config: DatasetConfig, dataset_indices: Indi
 
 def date_weight_sample_train_indices(dataset_config: DatasetConfig, train_indices: np.ndarray, num_samples: int) -> np.ndarray:
     rng = get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.DATE_WEIGHT_SAMPLING)
-    indices_per_date = [train_indices[train_indices[:, INDICES_TABLE_POS] == i] for i in np.unique(train_indices[:, INDICES_TABLE_POS])]
+    indices_per_date = [train_indices[train_indices[INDICES_TABLE_FIELD] == i] for i in np.unique(train_indices[INDICES_TABLE_FIELD])]
     weights = np.array(dataset_config.train_dates_weigths)
     weights = weights / weights.sum()
     samples_per_date = np.ceil((weights * (num_samples))).astype(int)
@@ -77,7 +79,7 @@ def date_weight_sample_train_indices(dataset_config: DatasetConfig, train_indice
     return sampled_train_indices
 
 def indices_to_app_counts(indices: np.ndarray, tables_app_enum: dict[int, str]) -> pd.Series:
-    app_counts = pd.Series(indices[:, INDICES_LABEL_POS]).value_counts()
+    app_counts = pd.Series(indices[INDICES_APP_FIELD]).value_counts()
     app_counts.index = app_counts.index.map(lambda x: tables_app_enum[x])
     return app_counts
 
@@ -99,23 +101,26 @@ def init_or_load_train_indices(dataset_config: DatasetConfig, tables_app_enum: d
     init_train_data(train_data_path)
     if not os.path.isfile(os.path.join(train_data_path, TRAIN_DATA_PARAMS_FILE)):
         log.info("Processing train indices")
+        start_time = time.time()
         train_data_params = dataset_config._get_train_data_params()
         train_known_indices, train_unknown_indices, known_apps, unknown_apps = init_train_indices(train_data_params=train_data_params,
                                                                                                   database_path=dataset_config.database_path,
                                                                                                   tables_app_enum=tables_app_enum,
+                                                                                                  sni_column=dataset_config.sni_column,
                                                                                                   servicemap=servicemap,
                                                                                                   rng=get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.INIT_TRAIN_INDICES))
         if not disable_indices_cache:
             yaml_dump({k: str(v) if isinstance(v, Enum) else list(v) if isinstance(v, tuple) else v for k, v in dataclasses.asdict(train_data_params).items()}, os.path.join(train_data_path, TRAIN_DATA_PARAMS_FILE))
             yaml_dump(known_apps, os.path.join(train_data_path, "known_apps.yaml"))
             yaml_dump(unknown_apps, os.path.join(train_data_path, "unknown_apps.yaml"))
-            np.save(os.path.join(train_data_path, "train_known_indices.npy"), train_known_indices)
-            np.save(os.path.join(train_data_path, "train_unknown_indices.npy"), train_unknown_indices)
+            np.savez_compressed(os.path.join(train_data_path, "train_indices.npz"), train_known_indices=train_known_indices, train_unknown_indices=train_unknown_indices)
+        log.info(f"Processing indices took {time.time() - start_time:.2f} seconds")
     else:
         known_apps = yaml_load(os.path.join(train_data_path, "known_apps.yaml"))
         unknown_apps = yaml_load(os.path.join(train_data_path, "unknown_apps.yaml"))
-        train_known_indices = np.load(os.path.join(train_data_path, "train_known_indices.npy"))
-        train_unknown_indices = np.load(os.path.join(train_data_path, "train_unknown_indices.npy"))
+        loaded = np.load(os.path.join(train_data_path, "train_indices.npz"))
+        train_known_indices = loaded["train_known_indices"]
+        train_unknown_indices = loaded["train_unknown_indices"]
     return train_known_indices, train_unknown_indices, known_apps, unknown_apps
 
 def init_or_load_val_indices(dataset_config: DatasetConfig, known_apps: list[str], unknown_apps: list[str], tables_app_enum: dict[int, str], disable_indices_cache: bool) -> tuple[np.ndarray, np.ndarray, str]:
@@ -123,17 +128,20 @@ def init_or_load_val_indices(dataset_config: DatasetConfig, known_apps: list[str
     init_test_data(val_data_path)
     if not os.path.isfile(os.path.join(val_data_path, TEST_DATA_PARAMS_FILE)):
         log.info("Processing validation indices")
+        start_time = time.time()
         val_known_indices, val_unknown_indices = init_test_indices(test_data_params=val_data_params,
                                                                    database_path=dataset_config.database_path,
                                                                    tables_app_enum=tables_app_enum,
+                                                                   sni_column=dataset_config.sni_column,
                                                                    rng=get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.INIT_VAL_INIDICES))
         if not disable_indices_cache:
             yaml_dump(dataclasses.asdict(val_data_params), os.path.join(val_data_path, TEST_DATA_PARAMS_FILE))
-            np.save(os.path.join(val_data_path, "val_known_indices.npy"), val_known_indices)
-            np.save(os.path.join(val_data_path, "val_unknown_indices.npy"), val_unknown_indices)
+            np.savez_compressed(os.path.join(val_data_path, "val_indices.npz"), val_known_indices=val_known_indices, val_unknown_indices=val_unknown_indices)
+        log.info(f"Processing indices took {time.time() - start_time:.2f} seconds")
     else:
-        val_known_indices = np.load(os.path.join(val_data_path, "val_known_indices.npy"))
-        val_unknown_indices = np.load(os.path.join(val_data_path, "val_unknown_indices.npy"))
+        loaded = np.load(os.path.join(val_data_path, "val_indices.npz"))
+        val_known_indices = loaded["val_known_indices"]
+        val_unknown_indices = loaded["val_unknown_indices"]
     return val_known_indices, val_unknown_indices, val_data_path
 
 def init_or_load_test_indices(dataset_config: DatasetConfig, known_apps: list[str], unknown_apps: list[str], tables_app_enum: dict[int, str], disable_indices_cache: bool) -> tuple[np.ndarray, np.ndarray, str]:
@@ -141,17 +149,20 @@ def init_or_load_test_indices(dataset_config: DatasetConfig, known_apps: list[st
     init_test_data(test_data_path)
     if not os.path.isfile(os.path.join(test_data_path, TEST_DATA_PARAMS_FILE)):
         log.info("Processing test indices")
+        start_time = time.time()
         test_known_indices, test_unknown_indices = init_test_indices(test_data_params=test_data_params,
                                                                      database_path=dataset_config.database_path,
                                                                      tables_app_enum=tables_app_enum,
+                                                                     sni_column=dataset_config.sni_column,
                                                                      rng=get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.INIT_TEST_INDICES))
         if not disable_indices_cache:
             yaml_dump(dataclasses.asdict(test_data_params), os.path.join(test_data_path, TEST_DATA_PARAMS_FILE))
-            np.save(os.path.join(test_data_path, "test_known_indices.npy"), test_known_indices)
-            np.save(os.path.join(test_data_path, "test_unknown_indices.npy"), test_unknown_indices)
+            np.savez_compressed(os.path.join(test_data_path, "test_indices.npz"), test_known_indices=test_known_indices, test_unknown_indices=test_unknown_indices)
+        log.info(f"Processing indices took {time.time() - start_time:.2f} seconds")
     else:
-        test_known_indices = np.load(os.path.join(test_data_path, "test_known_indices.npy"))
-        test_unknown_indices = np.load(os.path.join(test_data_path, "test_unknown_indices.npy"))
+        loaded = np.load(os.path.join(test_data_path, "test_indices.npz"))
+        test_known_indices = loaded["test_known_indices"]
+        test_unknown_indices = loaded["test_unknown_indices"]
     return test_known_indices, test_unknown_indices, test_data_path
 
 def init_train_data(train_data_path: str):
@@ -164,4 +175,4 @@ def init_test_data(test_data_path: str):
     os.makedirs(os.path.join(test_data_path, "preload"), exist_ok=True)
 
 def no_indices() -> np.ndarray:
-    return np.zeros((0,3), dtype=np.int64)
+    return np.empty(shape=(0,), dtype=INDICES_DTYPE)

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo/pytables_data/pytables_dataset.py

@@ -16,8 +16,8 @@ from typing_extensions import assert_never
 
 from cesnet_datazoo.config import (AppSelection, MinTrainSamplesCheck, TestDataParams,
                                    TrainDataParams)
-from cesnet_datazoo.constants import (APP_COLUMN, INDICES_INDEX_POS, INDICES_TABLE_POS, PPI_COLUMN,
-                                      QUIC_SNI_COLUMN, TLS_SNI_COLUMN)
+from cesnet_datazoo.constants import (APP_COLUMN, INDICES_DTYPE, INDICES_INDEX_FIELD,
+                                      INDICES_TABLE_FIELD, PPI_COLUMN)
 from cesnet_datazoo.pytables_data.apps_split import (is_background_app,
                                                      split_apps_topx_with_provider_groups)
 
@@ -36,12 +36,13 @@ class PyTablesDataset(Dataset):
                  flowstats_features_phist: list[str],
                  other_fields: list[str],
                  ppi_channels: list[int],
+                 sni_column: Optional[str] = None,
                  ppi_transform: Optional[Callable] = None,
                  flowstats_transform: Optional[Callable] = None,
                  flowstats_phist_transform: Optional[Callable] = None,
                  target_transform: Optional[Callable] = None,
                  return_tensors: bool = False,
-                 return_all_fields: bool = False,
+                 return_raw_fields: bool = False,
                  preload: bool = False,
                  preload_blob: Optional[str] = None,
                  disabled_apps: Optional[list[str]] = None,):
@@ -60,14 +61,14 @@ class PyTablesDataset(Dataset):
         self.flowstats_features_boolean = flowstats_features_boolean
         self.flowstats_features_phist = flowstats_features_phist
         self.other_fields = other_fields
+        self.sni_column = sni_column
         self.ppi_channels = ppi_channels
         self.ppi_transform = ppi_transform
         self.flowstats_transform = flowstats_transform
         self.flowstats_phist_transform = flowstats_phist_transform
         self.target_transform = target_transform
         self.return_tensors = return_tensors
-        self.return_all_fields = return_all_fields
-        self.sni_column = TLS_SNI_COLUMN if TLS_SNI_COLUMN in self.other_fields else QUIC_SNI_COLUMN if QUIC_SNI_COLUMN in self.other_fields else None
+        self.return_raw_fields = return_raw_fields
 
         self.preload = preload
         self.preload_blob = preload_blob
@@ -78,7 +79,7 @@ class PyTablesDataset(Dataset):
             batch_data = self.data[batch_idx]
         else:
             batch_data = load_data_from_tables(tables=self.tables, indices=self.indices[batch_idx], data_dtype=self.data_dtype)
-        if self.return_all_fields:
+        if self.return_raw_fields:
             return (batch_data, batch_idx)
 
         # Prepare data
@@ -173,16 +174,21 @@ def worker_init_fn(worker_id):
     dataset = worker_info.dataset
     dataset.pytables_worker_init(worker_id)
 
-def init_train_indices(train_data_params: TrainDataParams, database_path: str, tables_app_enum: dict[int, str], servicemap: pd.DataFrame, rng: np.random.RandomState) -> tuple[np.ndarray, np.ndarray, list[str], list[str]]:
+def init_train_indices(train_data_params: TrainDataParams, database_path: str, tables_app_enum: dict[int, str], sni_column: Optional[str], servicemap: pd.DataFrame, rng: np.random.RandomState) -> tuple[np.ndarray, np.ndarray, list[str], list[str]]:
     database, train_tables = load_database(database_path, tables_paths=train_data_params.train_tables_paths)
     inverted_tables_app_enum = {v: k for k, v in tables_app_enum.items()}
-    all_app_labels = {}
+    all_labels = {}
+    all_sni_domains = {}
     app_counts = pd.Series(dtype="int64")
     start_time = time.time()
     for i, table_path in enumerate(train_data_params.train_tables_paths):
-        all_app_labels[i] = train_tables[i].read(field=APP_COLUMN)
-        log.info(f"Reading app column for table {table_path} took {time.time() - start_time:.2f} seconds"); start_time = time.time()
-        app_counts = app_counts.add(pd.Series(all_app_labels[i]).value_counts(), fill_value=0)
+        all_labels[i] = train_tables[i].read(field=APP_COLUMN)
+        if sni_column is not None:
+            all_sni_domains[i] = train_tables[i].read(field=sni_column)
+        else:
+            all_sni_domains[i] = np.full_like(all_labels[i], "", dtype="U1")
+        log.info(f"Reading app and SNI columns for table {table_path} took {time.time() - start_time:.2f} seconds"); start_time = time.time()
+        app_counts = app_counts.add(pd.Series(all_labels[i]).value_counts(), fill_value=0)
     database.close()
     # Handle disabled apps and apps with less than min_samples_per_app samples
     if len(train_data_params.disabled_apps) > 0:
@@ -202,8 +208,9 @@ def init_train_indices(train_data_params: TrainDataParams, database_path: str, t
     # Base indices are indices of samples that are not disabled and have enough samples
     base_indices = {}
     for i, table_path in enumerate(train_data_params.train_tables_paths):
-        base_indices[i] = np.nonzero(np.isin(all_app_labels[i], disabled_apps_ids, invert=True))[0]
-    base_labels = {table_id: arr[base_indices[table_id]] for table_id, arr in all_app_labels.items()}
+        base_indices[i] = np.nonzero(np.isin(all_labels[i], disabled_apps_ids, invert=True))[0]
+    base_labels = {table_id: arr[base_indices[table_id]] for table_id, arr in all_labels.items()}
+    base_sni_domains = {table_id: arr[base_indices[table_id]] for table_id, arr in all_sni_domains.items()}
     # Apps selection
     if train_data_params.apps_selection != AppSelection.FIXED:
         app_counts = app_counts[[app for app in app_counts.index.tolist() if app not in disabled_apps_ids]]
@@ -230,26 +237,38 @@ def init_train_indices(train_data_params: TrainDataParams, database_path: str, t
     known_apps_ids = [inverted_tables_app_enum[app] for app in known_apps]
     unknown_apps_ids = [inverted_tables_app_enum[app] for app in unknown_apps]
 
-    train_known_indices, train_unknown_indices = convert_dict_indices(base_indices=base_indices, base_labels=base_labels, known_apps_ids=known_apps_ids, unknown_apps_ids=unknown_apps_ids)
+    train_known_indices, train_unknown_indices = convert_dict_indices(base_indices=base_indices,
+                                                                      base_labels=base_labels,
+                                                                      base_sni_domains=base_sni_domains,
+                                                                      known_apps_ids=known_apps_ids,
+                                                                      unknown_apps_ids=unknown_apps_ids)
     rng.shuffle(train_known_indices)
     rng.shuffle(train_unknown_indices)
-    log.info(f"Processing indices took {time.time() - start_time:.2f} seconds"); start_time = time.time()
     return train_known_indices, train_unknown_indices, known_apps, unknown_apps
 
-def init_test_indices(test_data_params: TestDataParams, database_path: str, tables_app_enum: dict[int, str], rng: np.random.RandomState) -> tuple[np.ndarray, np.ndarray]:
+def init_test_indices(test_data_params: TestDataParams, database_path: str, tables_app_enum: dict[int, str], sni_column: Optional[str], rng: np.random.RandomState) -> tuple[np.ndarray, np.ndarray]:
     database, test_tables = load_database(database_path, tables_paths=test_data_params.test_tables_paths)
     inverted_tables_app_enum = {v: k for k, v in tables_app_enum.items()}
     base_labels = {}
+    base_sni_domains = {}
     base_indices = {}
     start_time = time.time()
     for i, table_path in enumerate(test_data_params.test_tables_paths):
         base_labels[i] = test_tables[i].read(field=APP_COLUMN)
-        log.info(f"Reading app column for table {table_path} took {time.time() - start_time:.2f} seconds"); start_time = time.time()
+        if sni_column is not None:
+            base_sni_domains[i] = test_tables[i].read(field=sni_column)
+        else:
+            base_sni_domains[i] = np.full_like(base_labels[i], "", dtype="U1")
+        log.info(f"Reading app and SNI columns for table {table_path} took {time.time() - start_time:.2f} seconds"); start_time = time.time()
         base_indices[i] = np.arange(len(test_tables[i]))
     database.close()
     known_apps_ids = [inverted_tables_app_enum[app] for app in test_data_params.known_apps]
     unknown_apps_ids = [inverted_tables_app_enum[app] for app in test_data_params.unknown_apps]
-    test_known_indices, test_unknown_indices = convert_dict_indices(base_indices=base_indices, base_labels=base_labels, known_apps_ids=known_apps_ids, unknown_apps_ids=unknown_apps_ids)
+    test_known_indices, test_unknown_indices = convert_dict_indices(base_indices=base_indices,
+                                                                    base_labels=base_labels,
+                                                                    base_sni_domains=base_sni_domains,
+                                                                    known_apps_ids=known_apps_ids,
+                                                                    unknown_apps_ids=unknown_apps_ids)
     rng.shuffle(test_known_indices)
     rng.shuffle(test_unknown_indices)
     log.info(f"Processing indices took {time.time() - start_time:.2f} seconds"); start_time = time.time()
@@ -271,28 +290,32 @@ def list_all_tables(database_path: str) -> list[str]:
     with tb.open_file(database_path, mode="r") as database:
         return list(map(lambda x: x._v_pathname, iter(database.get_node(f"/flows"))))
 
-def convert_dict_indices(base_indices: dict[int, np.ndarray], base_labels: dict[int, np.ndarray], known_apps_ids: list[int], unknown_apps_ids: list[int]) -> tuple[np.ndarray, np.ndarray]:
+def convert_dict_indices(base_indices: dict[int, np.ndarray], base_labels: dict[int, np.ndarray], base_sni_domains: dict[int, np.ndarray], known_apps_ids: list[int], unknown_apps_ids: list[int]) -> tuple[np.ndarray, np.ndarray]:
     is_known = {table_id: np.isin(table_arr, known_apps_ids) for table_id, table_arr in base_labels.items()}
     is_unknown = {table_id: np.isin(table_arr, unknown_apps_ids) for table_id, table_arr in base_labels.items()}
     known_indices_dict = {table_id: table_arr[is_known[table_id]] for table_id, table_arr in base_indices.items()}
     unknown_indices_dict = {table_id: table_arr[is_unknown[table_id]] for table_id, table_arr in base_indices.items()}
     known_labels_dict = {table_id: table_arr[is_known[table_id]] for table_id, table_arr in base_labels.items()}
     unknown_labels_dict = {table_id: table_arr[is_unknown[table_id]] for table_id, table_arr in base_labels.items()}
-    known_indices = np.column_stack((
+    known_sni_domains_dict = {table_id: table_arr[is_known[table_id]] for table_id, table_arr in base_sni_domains.items()}
+    unknown_sni_domains_dict = {table_id: table_arr[is_unknown[table_id]] for table_id, table_arr in base_sni_domains.items()}
+    known_indices = np.array(list(zip(
         np.concatenate([[table_id] * table_arr.sum() for table_id, table_arr in is_known.items()]),
         np.concatenate(list(known_indices_dict.values())),
-        np.concatenate(list(known_labels_dict.values()))))
-    unknown_indices = np.column_stack((
+        np.concatenate(list(known_labels_dict.values())),
+        np.concatenate(list(known_sni_domains_dict.values())))), dtype=INDICES_DTYPE)
+    unknown_indices = np.array(list(zip(
         np.concatenate([[table_id] * table_arr.sum() for table_id, table_arr in is_unknown.items()]),
         np.concatenate(list(unknown_indices_dict.values())),
-        np.concatenate(list(unknown_labels_dict.values()))))
+        np.concatenate(list(unknown_labels_dict.values())),
+        np.concatenate(list(unknown_sni_domains_dict.values())))), dtype=INDICES_DTYPE)
     return known_indices, unknown_indices
 
 def load_data_from_tables(tables, indices: np.ndarray, data_dtype: np.dtype) -> np.ndarray:
-    sorted_indices = indices[indices[:, INDICES_TABLE_POS].argsort(kind="stable")]
-    unique_tables, split_bounderies = np.unique(sorted_indices[:, INDICES_TABLE_POS], return_index=True)
+    sorted_indices = indices[indices[INDICES_TABLE_FIELD].argsort(kind="stable")]
+    unique_tables, split_bounderies = np.unique(sorted_indices[INDICES_TABLE_FIELD], return_index=True)
     indices_per_table = np.split(sorted_indices, split_bounderies[1:])
     data = np.zeros(len(indices), dtype=data_dtype)
     for table_id, table_indices in zip(unique_tables, indices_per_table):
-        data[np.where(indices[:, INDICES_TABLE_POS] == table_id)[0]] = tables[table_id].read_coordinates(table_indices[:, INDICES_INDEX_POS])
+        data[np.where(indices[INDICES_TABLE_FIELD] == table_id)[0]] = tables[table_id].read_coordinates(table_indices[INDICES_INDEX_FIELD])
     return data

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/cesnet_datazoo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cesnet-datazoo
-Version: 0.1.3
+Version: 0.1.4
 Summary: A toolkit for large network traffic datasets
 Author-email: Jan Luxemburk <luxemburk@cesnet.cz>, Karel Hynek <hynekkar@cesnet.cz>
 Maintainer-email: Jan Luxemburk <luxemburk@cesnet.cz>, Karel Hynek <hynekkar@cesnet.cz>

{cesnet-datazoo-0.1.3 → cesnet_datazoo-0.1.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cesnet-datazoo"
-version = "0.1.3"
+version = "0.1.4"
 authors = [
   {name = "Jan Luxemburk", email = "luxemburk@cesnet.cz"},
   {name = "Karel Hynek", email = "hynekkar@cesnet.cz"},