API reference
0.2.1
package
datadreamer
package
package
package
package
package
variable
package
datadreamer.dataset_annotation
module
aimv2_annotator
This file uses pre-trained model derived from Apple's software, provided under the Apple Sample Code License license. The license is available at: https://developer.apple.com/support/downloads/terms/apple-sample-code/Apple-Sample-Code-License.pdf In addition, this file and other parts of the repository are licensed under the Apache 2.0 License. By using this file, you agree to comply with the terms of both licenses.
module
module
module
module
module
module
module
class
AIMv2Annotator
A class for image annotation using the AIMv2 model, specializing in image classification. Attributes: model (AutoModel): The AIMv2 model for image-text similarity evaluation. processor (AutoProcessor): The processor for preparing inputs to the AIMv2 model. device (str): The device on which the model will run ('cuda' for GPU, 'cpu' for CPU). size (str): The size of the AIMv2 model to use ('base' or 'large'). Methods: _init_processor(): Initializes the AIMv2 processor. _init_model(): Initializes the AIMv2 model. annotate_batch(image, prompts, conf_threshold, use_tta, synonym_dict): Annotates the given image with bounding boxes and labels. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.class
CLIPAnnotator
A class for image annotation using the CLIP model, specializing in image classification. Attributes: model (CLIPModel): The CLIP model for image-text similarity evaluation. processor (CLIPProcessor): The processor for preparing inputs to the CLIP model. device (str): The device on which the model will run ('cuda' for GPU, 'cpu' for CPU). size (str): The size of the CLIP model to use ('base' or 'large'). Methods: _init_processor(): Initializes the CLIP processor. _init_model(): Initializes the CLIP model. annotate_batch(image, prompts, conf_threshold, use_tta, synonym_dict): Annotates the given image with bounding boxes and labels. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.class
ImgClassificationAnnotator
Base class for image classification annotators using transformers models. Attributes: model: The model for image-text similarity evaluation. processor: The processor for preparing inputs to the model. device (str): The device on which the model will run ('cuda' for GPU, 'cpu' for CPU). size (str): The size of the model to use ('base' or 'large'). Methods: _init_processor(): Initializes the processor. _init_model(): Initializes the model. annotate_batch(image, prompts, conf_threshold, use_tta, synonym_dict): Annotates the given image with bounding boxes and labels. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.class
BaseAnnotator
Abstract base class for creating annotators. Attributes: seed (float): A seed value to ensure reproducibility in annotation processes. task_definition (TaskList): An enumeration of the task type. Default is OBJECT_DETECTION, which can be overridden by subclasses for specific tasks. Methods: annotate_batch(): Abstract method to be implemented by subclasses. It should contain the logic for performing annotation based on the task definition. release(): Abstract method to be implemented by subclasses. It should contain the logic for releasing the resources used by the annotator.
class
class
OWLv2Annotator
A class for image annotation using the OWLv2 model, specializing in object detection. Attributes: model (Owlv2ForObjectDetection): The OWLv2 model for object detection. processor (Owlv2Processor): The processor for the OWLv2 model. device (str): The device on which the model will run ('cuda' for GPU, 'cpu' for CPU). size (str): The size of the OWLv2 model to use ('base' or 'large'). Methods: _init_model(): Initializes the OWLv2 model. _init_processor(): Initializes the processor for the OWLv2 model. annotate_batch(image, prompts, conf_threshold, use_tta, synonym_dict): Annotates the given image with bounding boxes and labels. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.class
SAM2Annotator
A class for image annotation using the SAM2.1 model, specializing in instance segmentation. Attributes: model (SAM2ImagePredictor): The SAM2.1 model for instance segmentation. device (str): The device on which the model will run ('cuda' for GPU, 'cpu' for CPU). size (str): The size of the SAM model to use ('base' or 'large'). Methods: _init_model(): Initializes the SAM2.1 model. annotate_batch(image, prompts, conf_threshold, use_tta, synonym_dict): Annotates the given image with bounding boxes and labels. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.class
SlimSAMAnnotator
A class for image annotation using the SlimSAM model, specializing in instance segmentation. Attributes: model (SAM): The SAM model for instance segmentation. processor (SamProcessor): The processor for the SAM model. device (str): The device on which the model will run ('cuda' for GPU, 'cpu' for CPU). size (str): The size of the SAM model to use ('base' or 'large'). Methods: _init_model(): Initializes the SAM model. _init_processor(): Initializes the processor for the SAM model. annotate_batch(image, prompts, conf_threshold, use_tta, synonym_dict): Annotates the given image with bounding boxes and labels. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.module
datadreamer.dataset_annotation.utils
function
apply_tta(image) -> List[transforms.Compose]: List[transforms.Compose]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.
function
function
function
mask_to_polygon(mask: np.ndarray, epsilon_ratio: float = 0.001) -> List[List[int]]: List[List[int]]Converts a binary mask to a smoothed polygon. Args: mask: The binary mask to be converted. epsilon_ratio: Controls the smoothing level. Higher values mean more smoothing. Returns: List: A list of vertices of the polygon.
class
datadreamer.dataset_annotation.AIMv2Annotator(datadreamer.dataset_annotation.ImgClassificationAnnotator)
class
datadreamer.dataset_annotation.CLIPAnnotator(datadreamer.dataset_annotation.ImgClassificationAnnotator)
class
datadreamer.dataset_annotation.ImgClassificationAnnotator(datadreamer.dataset_annotation.BaseAnnotator)
method
__init__(self, seed: float = 42, device: str = 'cuda', size: str = 'base')Initializes the image classification annotator. Args: seed (float): Seed for reproducibility. Defaults to 42. device (str): The device to run the model on. Defaults to 'cuda'. size (str): The model size to use.
variable
variable
variable
variable
method
annotate_batch(self, images: List
[
PIL.Image.Image
], objects: List
[
str
], conf_threshold: float = 0.1, synonym_dict: Dict
[
str
,
List
[
str
]
]
|
None = None) -> List[np.ndarray]: List[np.ndarray]Annotates images using the CLIP 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[np.ndarray]: A list of the annotations for each image.
method
release(self, empty_cuda_cache: bool = False)Releases the model and optionally empties the CUDA cache. Args: empty_cuda_cache (bool, optional): Whether to empty the CUDA cache. Defaults to False.
class
datadreamer.dataset_annotation.BaseAnnotator(abc.ABC)
class
datadreamer.dataset_annotation.TaskList(enum.Enum)
class
datadreamer.dataset_annotation.OWLv2Annotator(datadreamer.dataset_annotation.BaseAnnotator)
method
__init__(self, seed: float = 42, device: str = 'cuda', size: str = 'base')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'.
variable
variable
variable
variable
method
annotate_batch(self, images: List
[
PIL.Image.Image
], prompts: List
[
str
]
|
List
[
PIL.Image.Image
], conf_threshold: float = 0.1, iou_threshold: float = 0.2, use_tta: bool = False, synonym_dict: Dict
[
str
,
List
[
str
]
]
|
None = None) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]: Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]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.
method
release(self, empty_cuda_cache: bool = False)Releases the model and optionally empties the CUDA cache. Args: empty_cuda_cache (bool, optional): Whether to empty the CUDA cache. Defaults to False.
class
datadreamer.dataset_annotation.SAM2Annotator(datadreamer.dataset_annotation.BaseAnnotator)
method
__init__(self, seed: float = 42, device: str = 'cuda', size: str = 'base', mask_format: str = 'rle')Initializes the SAM2Annotator. Args: seed (float): Seed for reproducibility. Defaults to 42. device (str): The device to run the model on. Defaults to 'cuda'. size (str): The size of the SAM2.1 model to use ('base' or 'large'). mask_format (str): The format of the output masks ('rle' or 'polyline').variable
variable
variable
variable
variable
method
annotate_batch(self, images: List
[
PIL.Image.Image
], boxes_batch: List
[
np.ndarray
], conf_threshold: float = 0.2) -> List[List[List[float]]]|List[List[dict]]: List[List[List[float]]]|List[List[dict]]Annotates images for the task of instance segmentation using the SAM2.1 model. Args: images: The images to be annotated. boxes_batch: The bounding boxes of found objects. conf_threshold (float, optional): Confidence threshold for the annotations. Defaults to 0.2. Returns: List: A list containing masks.
method
release(self, empty_cuda_cache: bool = False)Releases the model and optionally empties the CUDA cache. Args: empty_cuda_cache (bool, optional): Whether to empty the CUDA cache. Defaults to False.
class
datadreamer.dataset_annotation.SlimSAMAnnotator(datadreamer.dataset_annotation.BaseAnnotator)
method
__init__(self, seed: float = 42, device: str = 'cuda', size: str = 'base', mask_format: str = 'rle')Initializes the SlimSAMAnnotator. Args: seed (float): Seed for reproducibility. Defaults to 42. device (str): The device to run the model on. Defaults to 'cuda'. size (str): The size of the SAM model to use ('base' or 'large'). mask_format (str): The format of the output masks ('rle' or 'polyline').variable
variable
variable
variable
variable
method
annotate_batch(self, images: List
[
PIL.Image.Image
], boxes_batch: List
[
np.ndarray
], conf_threshold: float = 0.2) -> List[List[List[float]]]: List[List[List[float]]]Annotates images for the task of instance segmentation using the SlimSAM model. Args: images: The images to be annotated. boxes_batch: The bounding boxes of found objects. conf_threshold (float, optional): Confidence threshold for the annotations. Defaults to 0.2. Returns: List: A list containing the final segment masks represented as a polygon.
method
release(self, empty_cuda_cache: bool = False)Releases the model and optionally empties the CUDA cache. Args: empty_cuda_cache (bool, optional): Whether to empty the CUDA cache. Defaults to False.
package
datadreamer.image_generation
module
module
module
module
module
module
class
StableDiffusionImageGenerator
A subclass of ImageGenerator that uses the Stable Diffusion model for image generation. Attributes: base (DiffusionPipeline): The base Stable Diffusion model for initial image generation. refiner (DiffusionPipeline): The refiner Stable Diffusion model for enhancing generated images. base_processor (Compel): Processor for the base model. refiner_processor (Compel): Processor for the refiner model. Methods: _init_gen_model(): Initializes the generative models for image generation. _init_processor(): Initializes the processors for the models. generate_images_batch(prompts, negative_prompt, prompt_objects): Generates a batch of images based on the provided prompts. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.
class
StableDiffusionLightningImageGenerator
A subclass of ImageGenerator specifically designed to use the Stable Diffusion Lightning model for faster image generation. Attributes: pipe (StableDiffusionXLPipeline): The Stable Diffusion Lightning model for image generation. Methods: _init_gen_model(): Initializes the Stable Diffusion Lightning model. _init_compel(): Initializes the Compel model for text prompt weighting. generate_images_batch(prompts, negative_prompt, prompt_objects): Generates a batch of images based on the provided prompts. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.
class
StableDiffusionTurboImageGenerator
A subclass of ImageGenerator specifically designed to use the Stable Diffusion Turbo model for faster image generation. Attributes: base (AutoPipelineForText2Image): The Stable Diffusion Turbo model for image generation. Methods: _init_gen_model(): Initializes the Stable Diffusion Turbo model. generate_images_batch(prompts, negative_prompt, prompt_objects): Generates a batch of images based on the provided prompts. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.
class
Shuttle3DiffusionImageGenerator
A subclass of ImageGenerator specifically designed to use the Shuttle 3 Diffusion model for faster image generation. Attributes: pipe (DiffusionPipeline): The Shuttle 3 Diffusion model for image generation. Methods: _init_gen_model(): Initializes the Shuttle 3 Diffusion model. generate_images_batch(prompts, negative_prompt, prompt_objects): Generates a batch of images based on the provided prompts. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.
module
datadreamer.image_generation.clip_image_tester
class
ClipImageTester
A class for testing images against a set of textual objects using the CLIP model. Attributes: clip (CLIPModel): The CLIP model for image-text similarity evaluation. clip_processor (CLIPProcessor): The processor for preparing inputs to the CLIP model. device (str): The device on which the model will run ('cuda' for GPU, 'cpu' for CPU). Methods: test_image(image, objects, conf_threshold): Tests the given image against a list of objects. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.class
datadreamer.image_generation.clip_image_tester.ClipImageTester
method
__init__(self, device: str = 'cuda')Initializes the ClipImageTester with the CLIP model and processor.
variable
variable
variable
method
test_image(self, image: Image.Image, objects: List
[
str
], conf_threshold: float = 0.05) -> Tuple[bool, torch.Tensor, int]: Tuple[bool, torch.Tensor, int]Tests the generated image against a set of objects using the CLIP model. Args: image (Image.Image): The image to be tested. objects (List[str]): A list of objects (text) to test against the image. conf_threshold (float): Confidence threshold for considering an object as present. Returns: tuple: A tuple containing a boolean indicating if the image passes the test, the probabilities of the objects, and the number of objects that passed the test.
method
test_images_batch(self, images: List
[
Image.Image
], objects: List
[
List
[
str
]
], conf_threshold: float = 0.05) -> Tuple[List[bool], List[torch.Tensor], List[int]]: Tuple[List[bool], List[torch.Tensor], List[int]]Tests the generated images against a set of objects using the CLIP model. Args: images (List[Image.Image]): The images to be tested. objects (List[List[str]]): A list of objects (text) to test against the images. conf_threshold (float, optional): Confidence threshold for considering an object as present. Defaults to 0.05. Returns: Tuple[List[bool], List[torch.Tensor], List[int]]: A tuple containing a list of booleans indicating if the images pass the test, a list of probabilities of the objects, and a list of the number of objects that passed the test.
method
release(self, empty_cuda_cache = False)Releases the model and optionally empties the CUDA cache. Args: empty_cuda_cache (bool, optional): Whether to empty the CUDA cache. Defaults to False.
module
datadreamer.image_generation.image_generator
class
ImageGenerator
A class for generating images based on textual prompts, with optional CLIP model testing. Attributes: prompt_prefix (str): Optional prefix to add to every prompt. prompt_suffix (str): Optional suffix to add to every prompt, e.g., for adding details like resolution. negative_prompt (str): A string of negative prompts to guide the generation away from certain features. use_clip_image_tester (bool): Flag to use CLIP model testing for generated images. image_tester_patience (int): The number of attempts to generate an image that passes CLIP testing. batch_size (int): The number of images to generate in each batch. seed (float): Seed for reproducibility. clip_image_tester (ClipImageTester): Instance of ClipImageTester if use_clip_image_tester is True. device (str): The device on which the model will run ('cuda' for GPU, 'cpu' for CPU). Methods: set_seed(seed): Sets the seed for random number generators. generate_images(prompts, prompt_objects): Generates images based on provided prompts and optional object prompts. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache. (Abstract method) generate_images_batch(prompts, negative_prompt, prompt_objects): Generates a batch of images based on the provided prompts. Abstract method) Note: The actual model for image generation needs to be defined in the subclass.class
datadreamer.image_generation.image_generator.ImageGenerator
method
variable
variable
variable
variable
variable
variable
variable
variable
variable
static method
ImageGenerator.set_seed(seed: int)Sets the seed for random number generators in Python and PyTorch. Args: seed (int): The seed value to set.
method
generate_images(self, prompts: Union
[
str
,
List
[
str
]
], prompt_objects: Optional
[
List
[
List
[
str
]
]
] = None) -> List[Image.Image]: List[Image.Image]Generates images based on the provided prompts and optional object prompts. Args: prompts (Union[str, List[str]]): Single prompt or a list of prompts to guide the image generation. prompt_objects (Optional[List[List[str]]]): Optional list of objects for each prompt for CLIP model testing. Yields: List[Image.Image]: A batch of generated images.
method
release(self, empty_cuda_cache = False)Releases resources and optionally empties the CUDA cache.
method
generate_images_batch(self, prompts: List
[
str
], negative_prompt: str, prompt_objects: Optional
[
List
[
List
[
str
]
]
] = None) -> List[Image.Image]: List[Image.Image]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.
class
datadreamer.image_generation.StableDiffusionImageGenerator(datadreamer.image_generation.image_generator.ImageGenerator)
method
__init__(self, args, kwargs)Initializes the StableDiffusionImageGenerator with the given arguments.
variable
variable
variable
variable
method
generate_images_batch(self, prompts: List
[
str
], negative_prompt: str, prompt_objects: Optional
[
List
[
List
[
str
]
]
] = None) -> List[Image.Image]: List[Image.Image]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.
method
release(self, empty_cuda_cache = False)Releases the models and optionally empties the CUDA cache.
class
datadreamer.image_generation.StableDiffusionLightningImageGenerator(datadreamer.image_generation.image_generator.ImageGenerator)
method
__init__(self, args, kwargs)Initializes the StableDiffusionLightningImageGenerator with the given arguments.
variable
variable
method
generate_images_batch(self, prompts: List
[
str
], negative_prompt: str, prompt_objects: Optional
[
List
[
List
[
str
]
]
] = None) -> List[Image.Image]: List[Image.Image]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. Returns: List[Image.Image]: A list of generated images.
method
release(self, empty_cuda_cache = False)Releases the model and optionally empties the CUDA cache.
class
datadreamer.image_generation.StableDiffusionTurboImageGenerator(datadreamer.image_generation.image_generator.ImageGenerator)
method
__init__(self, args, kwargs)Initializes the StableDiffusionTurboImageGenerator with the given arguments.
variable
variable
method
generate_images_batch(self, prompts: List
[
str
], negative_prompt: str, prompt_objects: Optional
[
List
[
List
[
str
]
]
] = None, batch_size: int = 1) -> List[Image.Image]: List[Image.Image]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.
method
release(self, empty_cuda_cache = False)Releases the model and optionally empties the CUDA cache.
class
datadreamer.image_generation.Shuttle3DiffusionImageGenerator(datadreamer.image_generation.image_generator.ImageGenerator)
method
__init__(self, args, kwargs)Initializes the Shuttle3DiffusionImageGenerator with the given arguments.
variable
method
generate_images_batch(self, prompts: List
[
str
], negative_prompt: str, prompt_objects: Optional
[
List
[
List
[
str
]
]
] = None) -> List[Image.Image]: List[Image.Image]Generates a batch of images using the Shuttle 3 Diffusion 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. Returns: List[Image.Image]: A list of generated images.
method
release(self, empty_cuda_cache = False)Releases the model and optionally empties the CUDA cache.
package
datadreamer.pipelines
module
datadreamer.pipelines.generate_dataset_from_scratch
variable
variable
variable
variable
variable
variable
variable
function
function
function
package
datadreamer.prompt_generation
module
module
module
module
module
module
module
module
module
class
LMPromptGenerator
A language model-based prompt generator class, extending PromptGenerator. Attributes: device (str): Device to run the language model on ('cuda' for GPU, 'cpu' for CPU). num_objects_range (List[int]): Range for number of objects in a single image. model (AutoModelForCausalLM): The pre-trained causal language model for generating prompts. tokenizer (AutoTokenizer): The tokenizer for the pre-trained language model. pipeline (pipeline): The HuggingFace pipeline for generating text. Methods: _init_lang_model(): Initializes the language model and tokenizer. _remove_incomplete_sentence(text): Removes incomplete sentences from the generated prompt. _create_lm_prompt_text(selected_objects): Creates a text prompt for the language model. _create_lm_prompt_text_batch(selected_objects_batch): Creates a batch of text prompts for the language model. _postprocess_prompt(prompt): Post-processes the generated prompt. _test_prompt(prompt, selected_objects): Tests if the generated prompt is valid. generate_prompts_batch(prompt_texts_batch): Generates a batch of prompts using the language model. generate_prompts(): Generates a list of prompts based on the class names. release(empty_cuda_cache): Releases resources and optionally empties the CUDA cache.class
LMSynonymGenerator
Synonym generator that generates synonyms for a list of words using a language model. Args: synonyms_number (int): Number of synonyms to generate for each word. seed (Optional[float]): Seed for randomization. device (str): Device for model inference (default is "cuda"). Methods: _init_lang_model(): Initializes the language model and tokenizer. _generate_synonyms(prompt_text): Generates synonyms based on a given prompt text. _extract_synonyms(text): Extracts synonyms from a text containing synonyms. _create_prompt_text(word): Creates a prompt text for generating synonyms for a given word. generate_synonyms(word): Generates synonyms for a single word and returns them in a list. release(empty_cuda_cache): Releases resources (no action is taken in this implementation).
class
ProfanityFilter
Class for filtering bad words from texts and checking if texts are safe. Attributes: device (str): Device to run the language model on ('cuda' for GPU, 'cpu' for CPU). use_lm (bool): Whether to use a language model for checking text safety. seed (Optional[float]): Seed for randomization. model (AutoModelForCausalLM): The pre-trained causal language model for checking text safety. tokenizer (AutoTokenizer): The tokenizer for the pre-trained language model. Methods: set_seed(seed): Sets the random seed for consistent prompt generation. _init_lang_model(): Initializes the language model and tokenizer. _contains_bad_words(texts): Checks if a list of texts contain bad words. _check_lm_safety(text): Checks if a text is safe using a language model. is_safe(classes): Checks if a list of classes is safe. release(empty_cuda_cache): Releases the model and optionally empties the CUDA cache.class
Qwen2LMPromptGenerator
A language model-based prompt generator class, extending PromptGenerator. Attributes: device (str): Device to run the language model on ('cuda' for GPU, 'cpu' for CPU). model (AutoModelForCausalLM): The pre-trained causal language model for generating prompts. tokenizer (AutoTokenizer): The tokenizer for the pre-trained language model. pipeline (pipeline): The HuggingFace pipeline for generating text. Methods: _init_lang_model(): Initializes the language model and tokenizer. _remove_caption_sentences(text): Removes caption sentences from the generated prompt. _create_lm_prompt_text(selected_objects): Creates a text prompt for the language model. _postprocess_prompt(prompt): Post-processes the generated prompt. generate_prompts_batch(prompt_texts_batch): Generates a batch of prompts using the language model.class
SimplePromptGenerator
Prompt generator that creates simple prompts for text generation tasks. Args: class_names (List[str]): List of class names or objects for prompt generation. prompts_number (int): Number of prompts to generate. num_objects_range (Optional[List[int]]): Range for the number of objects to include in prompts. seed (Optional[float]): Seed for randomization. Methods: generate_prompts(): Generates a list of simple prompts. generate_prompt(selected_objects): Generates a single simple prompt based on selected objects. release(empty_cuda_cache): Releases resources (no action is taken in this implementation).
class
TinyLlamaLMPromptGenerator
A language model-based prompt generator class, extending PromptGenerator. Attributes: device (str): Device to run the language model on ('cuda' for GPU, 'cpu' for CPU). model (AutoModelForCausalLM): The pre-trained causal language model for generating prompts. tokenizer (AutoTokenizer): The tokenizer for the pre-trained language model. pipeline (pipeline): The HuggingFace pipeline for generating text. Methods: _init_lang_model(): Initializes the language model and tokenizer. _remove_caption_sentences(text): Removes caption sentences from the generated prompt. _create_lm_prompt_text(selected_objects): Creates a text prompt for the language model. _postprocess_prompt(prompt): Post-processes the generated prompt. generate_prompts_batch(prompt_texts_batch): Generates a batch of prompts using the language model.class
WordNetSynonymGenerator
Synonym generator that generates synonyms for a list of words using WordNet. Args: synonyms_number (int): Number of synonyms to generate for each word. seed (Optional[float]): Seed for randomization. device (str): Device to run the prompt generator on ('cuda' for GPU, 'cpu' for CPU). Methods: generate_synonyms(word): Generates synonyms for a single word and returns them in a list.module
datadreamer.prompt_generation.prompt_generator
class
PromptGenerator
Abstract base class for prompt generation. Attributes: class_names (List[str]): List of class names or objects for prompt generation. prompts_number (int): Number of prompts to generate. num_objects_range (Optional[List[int]]): Range for the number of objects to include in prompts. seed (Optional[float]): Seed for randomization. device (str): Device to run the prompt generator on ('cuda' for GPU, 'cpu' for CPU). quantization (str): Quantization type for the prompt generator. Methods: set_seed(seed): Sets the random seed for consistent prompt generation. save_prompts(prompts, save_path): Saves generated prompts to a JSON file. generate_prompts(): Abstract method to generate prompts (must be implemented in subclasses). release(empty_cuda_cache): Abstract method to release resources (must be implemented in subclasses).class
datadreamer.prompt_generation.prompt_generator.PromptGenerator(abc.ABC)
method
__init__(self, class_names: List
[
str
], prompts_number: int = 10, num_objects_range: Optional
[
List
[
int
]
] = None, batch_size: int = 1, seed: Optional
[
float
] = None, device: str = 'cuda', quantization: Optional
[
Literal
[
'
none
'
,
'
4bit
'
]
] = 'none')Initializes the PromptGenerator with class names and other settings.
variable
variable
variable
variable
variable
variable
variable
static method
PromptGenerator.set_seed(seed: int)Sets the random seed for consistent prompt generation. Args: seed (int): The random seed.
method
save_prompts(self, prompts: List
[
str
], save_path: str)Saves generated prompts to a JSON file. Args: prompts (List[str]): List of generated prompts. save_path (str): Path to save the prompts JSON file.
method
generate_prompts(self) -> List[str]: List[str]Abstract method to generate prompts (must be implemented in subclasses). Returns: List[str]: A list of generated prompts.
method
release(self, empty_cuda_cache = False)Abstract method to release resources (must be implemented in subclasses).
module
datadreamer.prompt_generation.synonym_generator
class
SynonymGenerator
Abstract base class for synonym generation. Attributes: synonyms_number (int): Number of synonyms to generate for each word. seed (Optional[float]): Seed for randomization. device (str): Device to run the prompt generator on ('cuda' for GPU, 'cpu' for CPU). Methods: _init_lang_model(): Initializes the language model and tokenizer. _generate_synonyms(prompt_text): Generates synonyms based on a given prompt text. _extract_synonyms(text): Extracts synonyms from a text containing synonyms. _create_prompt_text(word): Creates a prompt text for generating synonyms for a given word. generate_synonyms_for_list(words): Generates synonyms for a list of words and returns them in a dictionary. generate_synonyms(word): Generates synonyms for a single word and returns them in a list. save_synonyms(synonyms, save_path): Saves the generated synonyms to a JSON file. release(empty_cuda_cache): Releases resources (no action is taken in this implementation).class
datadreamer.prompt_generation.synonym_generator.SynonymGenerator(abc.ABC)
method
__init__(self, synonyms_number: int = 3, seed: Optional
[
float
] = 42, device: str = 'cuda')Initializes the SynonymGenerator with parameters.
variable
variable
variable
method
generate_synonyms_for_list(self, words: List
[
str
]) -> Dict: DictGenerates synonyms for a list of words and returns them in a dictionary. Args: words (List[str]): List of words for which synonyms are generated. Returns: dict: A dictionary where each word is associated with a list of its most common synonyms.
method
save_synonyms(self, synonyms: Dict, save_path: str)Saves the generated synonyms to a JSON file. Args: synonyms: The synonyms to save (typically a dictionary). save_path (str): The path to the JSON file where synonyms will be saved.
method
generate_synonyms(self, word: str) -> List[str]: List[str]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.
method
release(self, empty_cuda_cache = False)Abstract method to release resources (must be implemented in subclasses).
class
datadreamer.prompt_generation.LMPromptGenerator(datadreamer.prompt_generation.prompt_generator.PromptGenerator)
method
__init__(self, class_names: List
[
str
], prompts_number: int = 10, num_objects_range: Optional
[
List
[
int
]
] = None, batch_size: int = 1, seed: Optional
[
float
] = 42, device: str = 'cuda', quantization: Optional
[
Literal
[
'
none
'
,
'
4bit
'
]
] = 'none')Initializes the LMPromptGenerator with class names and other settings.
variable
variable
variable
method
generate_prompts_batch(self, prompt_texts_batch: List
[
str
]) -> List[str]: List[str]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.
method
generate_prompts(self) -> List[str]: List[str]Generates a list of text prompts based on the class names. Returns: List[str]: A list of generated prompts.
method
release(self, empty_cuda_cache = False)Releases the model and optionally empties the CUDA cache.
class
datadreamer.prompt_generation.LMSynonymGenerator(datadreamer.prompt_generation.synonym_generator.SynonymGenerator)
method
__init__(self, synonyms_number: int = 3, seed: Optional
[
float
] = 42, device: str = 'cuda')Initializes the SynonymGenerator with parameters.
variable
variable
variable
method
generate_synonyms(self, word: str) -> List[str]: List[str]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.
method
release(self, empty_cuda_cache = False)Releases resources and optionally empties the CUDA cache. Args: empty_cuda_cache (bool): Whether to empty the CUDA cache (default is False).
class
datadreamer.prompt_generation.ProfanityFilter
constant
method
__init__(self, device: str = 'cuda', use_lm: bool = False, seed: Optional
[
float
] = 42)Initializes the ProfanityFilter with parameters.
variable
variable
variable
variable
variable
static method
ProfanityFilter.set_seed(seed: int)Sets the random seed for consistent prompt generation. Args: seed (int): The random seed.
method
is_safe(self, classes: List
[
str
]) -> bool: boolChecks if a list of classes is safe. Args: classes (List[str]): List of classes to check for safety. Returns: bool: True if the classes are safe, False otherwise.
method
release(self, empty_cuda_cache = False)Releases the model and optionally empties the CUDA cache.
class
datadreamer.prompt_generation.Qwen2LMPromptGenerator(datadreamer.prompt_generation.LMPromptGenerator)
method
__init__(self, class_names: List
[
str
], prompts_number: int = 10, num_objects_range: Optional
[
List
[
int
]
] = None, batch_size: int = 1, seed: Optional
[
float
] = 42, device: str = 'cuda', quantization: Optional
[
Literal
[
'
none
'
,
'
4bit
'
]
] = 'none')Initializes the LMPromptGenerator with class names and other settings.
method
generate_prompts_batch(self, prompt_texts_batch: List
[
str
]) -> List[str]: List[str]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.
class
datadreamer.prompt_generation.SimplePromptGenerator(datadreamer.prompt_generation.prompt_generator.PromptGenerator)
method
__init__(self, args, kwargs)Initializes the SimplePromptGenerator with class names and other settings.
method
generate_prompts(self) -> List[str]: List[str]Generates a list of simple prompts. Returns: List[str]: A list of generated prompts in the form of "A photo of a {selected_objects}".method
generate_prompt(self, selected_objects: List
[
str
]) -> str: strGenerates 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}".method
release(self, empty_cuda_cache = False)Releases resources (no action is taken in this implementation).
class
datadreamer.prompt_generation.TinyLlamaLMPromptGenerator(datadreamer.prompt_generation.LMPromptGenerator)
method
__init__(self, class_names: List
[
str
], prompts_number: int = 10, num_objects_range: Optional
[
List
[
int
]
] = None, batch_size: int = 1, seed: Optional
[
float
] = 42, device: str = 'cuda', quantization: Optional
[
Literal
[
'
none
'
,
'
4bit
'
]
] = 'none')Initializes the LMPromptGenerator with class names and other settings.
method
generate_prompts_batch(self, prompt_texts_batch: List
[
str
]) -> List[str]: List[str]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.
class
datadreamer.prompt_generation.WordNetSynonymGenerator(datadreamer.prompt_generation.synonym_generator.SynonymGenerator)
method
__init__(self, synonyms_number: int = 3, seed: Optional
[
float
] = 42, device: str = 'cuda')Initializes the SynonymGenerator with parameters.
method
generate_synonyms(self, word: str) -> List[str]: List[str]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.
method
release(self, empty_cuda_cache: bool = False)Releases resources (no action is taken in this implementation).
package
datadreamer.utils
module
module
module
module
module
module
module
module
module
module
module
module
class
BaseConverter
Abstract base class for converter.
class
COCOConverter
Class for converting a dataset to COCO format. Format: dataset_dir ├── train │ ├── data │ │ ├── 0.jpg │ │ ├── 1.jpg │ ├── labels.json ├── validation │ ├── data │ ├── labels.json ├── test │ ├── data │ ├── labels.json
class
class
LuxonisDatasetConverter
Class for converting a dataset to LuxonisDataset format.
class
SingleLabelClsConverter
Class for converting a dataset for single-label classification task. NOTE: The number of images after conversion may be different from the number of images in the original dataset, as images with zero or more than one labels are removed. Format: dataset_dir ├── train │ ├── class_1 │ │ ├── image1.jpg │ │ ├── image2.jpg │ ├── class_2 │ │ ├── image3.jpg │ │ ├── image4.jpg ├── val │ ├── class_1 │ ├── class_2 ├── test │ ├── class_1 │ ├── class_2
class
VOCConverter
Class for converting a dataset to PASCAL VOC format. Format: dataset_dir ├── Annotations/ # XML files with image annotations (bounding boxes, labels, segmentation) | ├── 0.xml | ├── 1.xml | ├── ... ├── ImageSets/ │ ├── Main/ # Main image splits (train, val, test, etc.) │ ├── train.txt │ ├── val.txt │ ├── test.txt ├── JPEGImages/ # Images in JPG format | ├── 0.jpg | ├── 1.jpg | ├── ... ├── SegmentationClass/ # Segmentation masks (with class labels) | ├── 0.png | ├── 1.png | ├── ... ├── SegmentationObject/ # Segmentation masks (object-wise masks)labels.json | ├── 0.png | ├── 1.png | ├── ...
class
YOLOConverter
Class for converting a dataset to YOLO format. Format: dataset_dir ├── train │ ├── images │ │ ├── 0.jpg │ │ ├── 1.jpg │ ├── labels │ │ ├── 0.txt │ │ ├── 1.txt ├── val │ ├── images │ ├── labels ├── test │ ├── images │ ├── labels
module
datadreamer.utils.bad_words
constant
module
datadreamer.utils.convert_dataset
function
convert_dataset(input_dir: str, output_dir: str, dataset_format: str, split_ratios: List
[
float
], dataset_plugin: Optional
[
str
] = None, dataset_name: Optional
[
str
] = None, is_instance_segmentation: bool = False, keep_unlabeled_images: bool = False, copy_files: bool = True, seed: int = 42)Converts a dataset from one format to another. Args: input_dir (str): Directory containing the images and annotations. output_dir (str): Directory where the processed dataset will be saved. dataset_format (str): Format of the dataset. Can be 'yolo', 'coco', 'voc', 'luxonis-dataset', or 'cls-single'. split_ratios (lis of float): List of ratios for train, val, and test splits. dataset_plugin (str, optional): Plugin for Luxonis dataset. Defaults to None. dataset_name (str, optional): Name of the Luxonis dataset. Defaults to None. is_instance_segmentation (bool, optional): Whether the dataset is for instance segmentation. Defaults to False. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the files to the output directory. Defaults to True. seed (int, optional): Random seed. Defaults to 42. No return value.
function
module
datadreamer.utils.dataset_utils
function
save_annotations_to_json(image_paths, labels_list, boxes_list = None, masks_list = None, class_names = None, save_dir = None, file_name = 'annotations.json')Saves annotations to a JSON file. Args: image_paths (list): List of image paths. labels_list (list): List of labels. boxes_list (list, optional): List of bounding boxes. Defaults to None. masks_list (list, optional): List of instance segmentation masks. Defaults to None. class_names (list, optional): List of class names. Defaults to None. save_dir (str, optional): Directory to save the JSON file. Defaults to None. file_name (str, optional): Name of the JSON file. Defaults to 'annotations.json'. No return value.
module
datadreamer.utils.merge_raw_datasets
function
merge_datasets(input_dirs: List
[
str
], output_dir: str, copy_files: bool = True)Merges multiple raw datasets into a single dataset. Args: input_dirs (List[str]): A list of input directories containing raw datasets. output_dir (str): The output directory where the merged dataset will be saved. copy_files (bool, optional): Whether to copy the files from the input directories to the output directory. Defaults to True. No return value.
function
module
datadreamer.utils.nms
function
xywh2xyxy(x)Convert boxes with shape [n, 4] from [x, y, w, h] to [x1, y1, x2, y2] where x1y1 is top-left, x2y2=bottom-right.
function
non_max_suppression(prediction, conf_thres = 0.25, iou_thres = 0.45, classes = None, agnostic = False, multi_label = False, max_det = 300) -> List[np.ndarray]: List[np.ndarray]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].
class
datadreamer.utils.BaseConverter(abc.ABC)
method
method
convert(self, dataset_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False, copy_files: bool = True)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. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True. No return value.
static method
BaseConverter.read_annotations(annotation_path: str) -> Dict: DictReads 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.
static method
BaseConverter.make_splits(images: List
[
str
], split_ratios: List
[
float
], shuffle: bool = True) -> Tuple[List, List, List]: Tuple[List, List, List]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.
class
datadreamer.utils.COCOConverter(datadreamer.utils.BaseConverter)
method
variable
method
convert(self, dataset_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False, copy_files: bool = True)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. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True. No return value.
method
process_data(self, data: Dict, image_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False, copy_files: bool = True)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 (list of float): The ratios to split the data into training, validation, and test sets. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True. No return value.
method
save_labels(self, dataset_output_dir: str, images_info: List
[
Dict
], annotations: List
[
Dict
], class_names: List
[
str
])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.
class
datadreamer.utils.Config(luxonis_ml.utils.LuxonisConfig)
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
variable
class
datadreamer.utils.LuxonisDatasetConverter(datadreamer.utils.BaseConverter)
method
variable
variable
variable
method
convert(self, dataset_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False, copy_files: bool = True)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. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True. No return value.
method
process_data(self, data: Dict, dataset_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False)Processes the data into LuxonisDataset format. Args: data (dict): The data to process. 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. No return value.
class
datadreamer.utils.SingleLabelClsConverter(datadreamer.utils.BaseConverter)
method
method
convert(self, dataset_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False, copy_files: bool = True)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. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True. No return value.
method
process_data(self, data: Dict, image_dir: str, output_dir: str, split_ratios: List
[
float
], copy_files: bool = True)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.
class
datadreamer.utils.VOCConverter(datadreamer.utils.BaseConverter)
method
variable
method
convert(self, dataset_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False, copy_files: bool = True)Convert a dataset to PASCAL VOC 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. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True. No return value.
method
process_data(self, data: Dict, image_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False, copy_files: bool = True)Process the data, create XML annotations and handle dataset splitting. Args: data (dict): The dictionary containing image annotations. image_dir (str): The directory where the source images are 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. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True. No return value.
method
create_image_sets(self, image_sets_dir: str, train_images: List
[
str
], val_images: List
[
str
], test_images: List
[
str
])Create text files for each split (train, val, test). Args: image_sets_dir (str): The directory where the image sets will be saved. train_images (list of str): The list of image names for the training set. val_images (list of str): The list of image names for the validation set. test_images (list of str): The list of image names for the test set. No return value.
method
create_xml(self, annotation: Dict, folder_name: str, image_name: str, width: int, height: int, class_names: List
[
str
]) -> Element: ElementGenerate XML annotation for a single image. Args: annotation (dict): The dictionary containing image annotations. folder_name (str): The name of the folder where the image is located. image_name (str): The name of the image file. width (int): The width of the image. height (int): The height of the image. class_names (list of str): The list of class names. Returns: Element: The root element of the XML annotation.
method
create_segmentation_masks(self, annotation: Dict, image_name: str, width: int, height: int) -> Tuple[np.ndarray, np.ndarray]: Tuple[np.ndarray, np.ndarray]Create and save segmentation masks (class and object). Args: annotation (dict): The dictionary containing image annotations. image_name (str): The name of the image file. width (int): The width of the image. height (int): The height of the image. Returns: tuple of np.ndarray: A tuple containing the class mask and object mask.
class
datadreamer.utils.YOLOConverter(datadreamer.utils.BaseConverter)
method
variable
method
convert(self, dataset_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False, copy_files: bool = True)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. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True. No return value.
method
convert_to_yolo_format(self, box: List
[
float
], image_width: int, image_height: int) -> List[float]: List[float]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].
method
convert_masks_to_yolo_format(self, masks: List
[
List
[
float
]
]
|
Dict, w: int, h: int) -> List[float]: List[float]Converts masks to YOLO format. Args: masks (list of list of float): A list containing the masks. w (int): The width of the image. h (int): The height of the image. Returns: list of float: A list containing the masks in YOLO format.
method
process_data(self, data: Dict, image_dir: str, output_dir: str, split_ratios: List
[
float
], keep_unlabeled_images: bool = False, copy_files: bool = True)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. keep_unlabeled_images (bool, optional): Whether to keep images with no annotations. Defaults to False. copy_files (bool, optional): Whether to copy the source files to the output directory, otherwise move them. Defaults to True. No return value.
method
create_data_yaml(self, root_dir: str, class_names: List
[
str
])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.