API reference

Apply test-time augmentation (TTA) to the given image.

Args:
    image: The image to be augmented.

Returns:
    list: A list of augmented images.

Note:
    Currently, only horizontal flip is enabled. Additional transformations like
    vertical flip and color jitter are commented out but can be enabled as needed.

Initializes the CLIPAnnotator with a specific seed and device.

Args:
    seed (float): Seed for reproducibility. Defaults to 42.
    device (str): The device to run the model on. Defaults to 'cuda'.

Annotates images using the OWLv2 model.

Args:
    images: The images to be annotated.
    objects: A list of objects (text) to test against the images.
    conf_threshold (float, optional): Confidence threshold for the annotations. Defaults to 0.1.
    synonym_dict (dict, optional): Dictionary for handling synonyms in labels. Defaults to None.

Returns:
    List[List[int]]: A list of lists of labels for each image.

Releases the model and optionally empties the CUDA cache.

Args:
    empty_cuda_cache (bool, optional): Whether to empty the CUDA cache. Defaults to False.

Initializes the OWLv2Annotator with a specific seed and device.

Args:
    seed (float): Seed for reproducibility. Defaults to 42.
    device (str): The device to run the model on. Defaults to 'cuda'.

Annotates images using the OWLv2 model.

Args:
    images: The images to be annotated.
    prompts: Prompts to guide the annotation.
    conf_threshold (float, optional): Confidence threshold for the annotations. Defaults to 0.1.
    iou_threshold (float, optional): Intersection over union threshold for non-maximum suppression. Defaults to 0.2.
    use_tta (bool, optional): Flag to apply test-time augmentation. Defaults to False.
    synonym_dict (dict, optional): Dictionary for handling synonyms in labels. Defaults to None.

Returns:
    tuple: A tuple containing the final bounding boxes, scores, and labels for the annotations.

Releases the model and optionally empties the CUDA cache.

Args:
    empty_cuda_cache (bool, optional): Whether to empty the CUDA cache. Defaults to False.

A class for testing images against a set of textual objects using the CLIP model.

A class for generating images based on textual prompts, with optional CLIP model testing.

Initializes the StableDiffusionImageGenerator with the given arguments.

Generates a batch of images based on the provided prompts.

Args:
    prompts (List[str]): A list of positive prompts to guide image generation.
    negative_prompt (str): The negative prompt to avoid certain features in the image.
    prompt_objects (Optional[List[List[str]]]): Optional list of objects to be used in CLIP model testing.

Returns:
    List[Image.Image]: A list of generated images.

Releases the models and optionally empties the CUDA cache.

Initializes the StableDiffusionLightningImageGenerator with the given
arguments.

Generates a batch of images using the Stable Diffusion Lightning model based
on the provided prompts.

Args:
    prompts (List[str]): A list of positive prompts to guide image generation.
    negative_prompt (str): The negative prompt to avoid certain features in the image.
    prompt_objects (Optional[List[List[str]]]): Optional list of objects for each prompt for CLIP model testing.
    batch_size (int): The number of images to generate in each batch.

Returns:
    List[Image.Image]: A list of generated images.

Releases the model and optionally empties the CUDA cache.

Initializes the StableDiffusionTurboImageGenerator with the given
arguments.

Generates a batch of images using the Stable Diffusion Turbo model based on
the provided prompts.

Args:
    prompts (List[str]): A list of positive prompts to guide image generation.
    negative_prompt (str): The negative prompt to avoid certain features in the image.
    prompt_objects (Optional[List[List[str]]]): Optional list of objects for each prompt for CLIP model testing.
    batch_size (int): The number of images to generate in each batch.

Returns:
    List[Image.Image]: A list of generated images.

Releases the model and optionally empties the CUDA cache.

Abstract base class for prompt generation.

Abstract base class for synonym generation.

Initializes the LMPromptGenerator with class names and other settings.

Generates a list of prompts using the language model.

Args:
    prompt_texts_batch (List[str]): List of text prompts for the language model.

Returns:
    List[str]: List of generated prompts.

Generates a list of text prompts based on the class names.

Returns:
    List[str]: A list of generated prompts.

Releases the model and optionally empties the CUDA cache.

Initializes the SynonymGenerator with parameters.

Generates synonyms for a single word and returns them in a list.

Args:
    word (str): The word for which synonyms are generated.

Returns:
    List[str]: A list of generated synonyms for the word.

Releases resources and optionally empties the CUDA cache.

Args:
    empty_cuda_cache (bool): Whether to empty the CUDA cache (default is False).

Initializes the SimplePromptGenerator with class names and other settings.

Generates a list of simple prompts.

Returns:
    List[str]: A list of generated prompts in the form of "A photo of a {selected_objects}".

Generates a single simple prompt based on selected objects.

Args:
    selected_objects (List[str]): List of selected objects to include in the prompt.

Returns:
    str: A simple prompt in the form of "A photo of a {selected_objects}".

Releases resources (no action is taken in this implementation).

Initializes the LMPromptGenerator with class names and other settings.

Generates a list of prompts using the language model.

Args:
    prompt_texts_batch (List[str]): List of text prompts for the language model.

Returns:
    List[str]: List of generated prompts.

Initializes the SynonymGenerator with parameters.

Generates synonyms for a single word and returns them in a list.

Args:
    word (str): The word for which synonyms are generated.

Returns:
    List[str]: A list of generated synonyms for the word.

Releases resources (no action is taken in this implementation).

Convert boxes with shape [n, 4] from [x, y, w, h] to [x1, y1, x2, y2] where x1y1
is top-left, x2y2=bottom-right.

Runs Non-Maximum Suppression (NMS) on inference results.
This code is borrowed from: https://github.com/ultralytics/yolov5/blob/47233e1698b89fc437a4fb9463c815e9171be955/utils/general.py#L775
Args:
    prediction: (tensor), with shape [N, 5 + num_classes], N is the number of bboxes.
    conf_thres: (float) confidence threshold.
    iou_thres: (float) iou threshold.
    classes: (None or list[int]), if a list is provided, nms only keep the classes you provide.
    agnostic: (bool), when it is set to True, we do class-independent nms, otherwise, different class would do nms respectively.
    multi_label: (bool), when it is set to True, one box can have multi labels, otherwise, one box only huave one label.
    max_det:(int), max number of output bboxes.

Returns:
     list of detections, echo item is one tensor with shape (num_boxes, 6), 6 is for [xyxy, conf, cls].

Converts a dataset into another format.

Args:
- dataset_dir (str): The directory where the source dataset is located.
- output_dir (str): The directory where the processed dataset should be saved.
- split_ratios (list of float): The ratios to split the data into training, validation, and test sets.
- copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True.


No return value.

Reads annotations from a JSON file located at the specified path.

Args:
- annotation_path (str): The path to the JSON file containing annotations.

Returns:
- dict: A dictionary containing the data loaded from the JSON file.

Splits the list of images into training, validation, and test sets.

Args:
- images (list of str): A list of image paths.
- split_ratios (list of float): The ratios to split the data into training, validation, and test sets.
- shuffle (bool, optional): Whether to shuffle the list of images. Defaults to True.

Returns:
- list of str: A list of image paths for the training set.
- list of str: A list of image paths for the validation set.
- list of str: A list of image paths for the test set.

Converts a dataset into a COCO format.

Args:
- dataset_dir (str): The directory where the source dataset is located.
- output_dir (str): The directory where the processed dataset should be saved.
- split_ratios (list of float): The ratios to split the data into training, validation, and test sets.
- copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True.

No return value.

Processes the data by dividing it into training and validation sets, and
saves the images and labels in COCO format.

Args:
- data (dict): The dictionary containing image annotations.
- image_dir (str): The directory where the source images are located.
- output_dir (str): The base directory where the processed data will be saved.
- split_ratios (float): The ratio to split the data into training, validation, and test sets.
- copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True.


No return value.

Saves the labels to a JSON file.

Args:
- dataset_output_dir (str): The directory where the labels should be saved.
- images_info (list of dict): A list of dictionaries containing image information.
- annotations (list of dict): A list of dictionaries containing annotation information.
- class_names (list of str): A list of class names.

No return value.

Converts a dataset into a LuxonisDataset format.

Args:
- dataset_dir (str): The directory where the source dataset is located.
- output_dir (str): The directory where the processed dataset should be saved.
- split_ratios (list of float): The ratios to split the data into training, validation, and test sets.
- copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True.

No return value.

Converts a dataset into a format suitable for single-label classification.

Args:
- dataset_dir (str): The directory where the source dataset is located.
- output_dir (str): The directory where the processed dataset should be saved.
- split_ratios (list of float): The ratios to split the data into training, validation, and test sets.
- copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True.

No return value.

Processes the data by removing images with multiple labels, then dividing it
into training and validation sets, and saves the images with single labels.

Args:
- data (dict): The dictionary containing image annotations.
- image_dir (str): The directory where the source images are located.
- output_dir (str): The base directory where the processed data will be saved.
- split_ratios (float): The ratio to split the data into training, validation, and test sets.
- copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True.

No return value.

Converts a dataset into a format suitable for training with YOLO, including
creating training and validation splits.

Args:
- dataset_dir (str): The directory where the source dataset is located.
- output_dir (str): The directory where the processed dataset should be saved.
- split_ratios (list of float): The ratios to split the data into training, validation, and test sets.
- copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True.

No return value.

Converts bounding box coordinates to YOLO format.

Args:
- box (list of float): A list containing the bounding box coordinates [x_min, y_min, x_max, y_max].
- image_width (int): The width of the image.
- image_height (int): The height of the image.

Returns:
- list of float: A list containing the bounding box in YOLO format [x_center, y_center, width, height].

Processes the data by dividing it into training and validation sets, and
saves the images and labels in YOLO format.

Args:
- data (dict): The dictionary containing image annotations.
- image_dir (str): The directory where the source images are located.
- output_dir (str): The base directory where the processed data will be saved.
- split_ratios (float): The ratio to split the data into training, validation, and test sets.
- copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True.


No return value.

Creates a YAML file for dataset configuration, specifying paths and class
names.

Args:
- root_dir (str): The root directory where the dataset is located.
- class_names (list of str): A list of class names.

No return value.