From cd819c419e2b9da1fbaf963f83a784cbcb3b01ea Mon Sep 17 00:00:00 2001 From: HonzaCuhel Date: Fri, 25 Oct 2024 14:44:49 +0200 Subject: [PATCH] Fix: correct all SlimSAM mentions --- datadreamer/dataset_annotation/slimsam_annotator.py | 2 +- examples/generate_dataset_and_train_yolo.ipynb | 2 +- ...ate_instance_segmentation_dataset_and_train_yolo.ipynb | 2 +- tests/core_tests/unittests/test_annotators.py | 8 ++++---- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/datadreamer/dataset_annotation/slimsam_annotator.py b/datadreamer/dataset_annotation/slimsam_annotator.py index 3461f35..7e4e7cd 100644 --- a/datadreamer/dataset_annotation/slimsam_annotator.py +++ b/datadreamer/dataset_annotation/slimsam_annotator.py @@ -77,7 +77,7 @@ def annotate_batch( boxes_batch: List[np.ndarray], iou_threshold: float = 0.2, ) -> List[List[List[float]]]: - """Annotates images for the task of instance segmentation using the FastSAM + """Annotates images for the task of instance segmentation using the SlimSAM model. Args: diff --git a/examples/generate_dataset_and_train_yolo.ipynb b/examples/generate_dataset_and_train_yolo.ipynb index dbbf376..9ad1359 100644 --- a/examples/generate_dataset_and_train_yolo.ipynb +++ b/examples/generate_dataset_and_train_yolo.ipynb @@ -84,7 +84,7 @@ "- `--num_objects_range`: Range of objects in a prompt. Default is 1 to 3.\n", "- `--prompt_generator`: Choose between `simple`, `lm` (Mistral-7B), `tiny` (tiny LM), and `qwen2` (Qwen2.5 LM). Default is `qwen2`.\n", "- `--image_generator`: Choose image generator, e.g., `sdxl`, `sdxl-turbo` or `sdxl-lightning`. Default is `sdxl-turbo`.\n", - "- `--image_annotator`: Specify the image annotator, like `owlv2` for object detection or `clip` for image classification or `owlv2-fastsam` for instance segmentation. Default is `owlv2`.\n", + "- `--image_annotator`: Specify the image annotator, like `owlv2` for object detection or `clip` for image classification or `owlv2-slimsam` for instance segmentation. Default is `owlv2`.\n", "- `--conf_threshold`: Confidence threshold for annotation. Default is `0.15`.\n", "- `--annotation_iou_threshold`: Intersection over Union (IoU) threshold for annotation. Default is `0.2`.\n", "- `--prompt_prefix`: Prefix to add to every image generation prompt. Default is `\"\"`.\n", diff --git a/examples/generate_instance_segmentation_dataset_and_train_yolo.ipynb b/examples/generate_instance_segmentation_dataset_and_train_yolo.ipynb index 1dbd2d5..11f9812 100644 --- a/examples/generate_instance_segmentation_dataset_and_train_yolo.ipynb +++ b/examples/generate_instance_segmentation_dataset_and_train_yolo.ipynb @@ -99,7 +99,7 @@ "- `--num_objects_range`: Range of objects in a prompt. Default is 1 to 3.\n", "- `--prompt_generator`: Choose between `simple`, `lm` (Mistral-7B), `tiny` (tiny LM), and `qwen2` (Qwen2.5 LM). Default is `qwen2`.\n", "- `--image_generator`: Choose image generator, e.g., `sdxl`, `sdxl-turbo` or `sdxl-lightning`. Default is `sdxl-turbo`.\n", - "- `--image_annotator`: Specify the image annotator, like `owlv2` for object detection or `clip` for image classification or `owlv2-fastsam` for instance segmentation. Default is `owlv2`.\n", + "- `--image_annotator`: Specify the image annotator, like `owlv2` for object detection or `clip` for image classification or `owlv2-slimsam` for instance segmentation. Default is `owlv2`.\n", "- `--conf_threshold`: Confidence threshold for annotation. Default is `0.15`.\n", "- `--annotation_iou_threshold`: Intersection over Union (IoU) threshold for annotation. Default is `0.2`.\n", "- `--prompt_prefix`: Prefix to add to every image generation prompt. Default is `\"\"`.\n", diff --git a/tests/core_tests/unittests/test_annotators.py b/tests/core_tests/unittests/test_annotators.py index 7bd024b..4e78df2 100644 --- a/tests/core_tests/unittests/test_annotators.py +++ b/tests/core_tests/unittests/test_annotators.py @@ -127,7 +127,7 @@ def _check_slimsam_annotator(device: str, size: str = "base"): not torch.cuda.is_available() or total_disk_space < 16, reason="Test requires GPU and 16GB of HDD", ) -def test_cuda_fastsam_base_annotator(): +def test_cuda_slimsam_base_annotator(): _check_slimsam_annotator("cuda") @@ -135,7 +135,7 @@ def test_cuda_fastsam_base_annotator(): total_disk_space < 16, reason="Test requires at least 16GB of HDD", ) -def test_cpu_fastsam_base_annotator(): +def test_cpu_slimsam_base_annotator(): _check_slimsam_annotator("cpu") @@ -143,7 +143,7 @@ def test_cpu_fastsam_base_annotator(): not torch.cuda.is_available() or total_disk_space < 16, reason="Test requires GPU and 16GB of HDD", ) -def test_cuda_fastsam_large_annotator(): +def test_cuda_slimsam_large_annotator(): _check_slimsam_annotator("cuda", size="large") @@ -151,5 +151,5 @@ def test_cuda_fastsam_large_annotator(): total_disk_space < 16, reason="Test requires at least 16GB of HDD", ) -def test_cpu_fastsam_large_annotator(): +def test_cpu_slimsam_large_annotator(): _check_slimsam_annotator("cpu", size="large")