scribble-annotation-generator 0.1.0__tar.gz → 0.1.1__tar.gz

{scribble_annotation_generator-0.1.0 → scribble_annotation_generator-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: scribble-annotation-generator
-Version: 0.1.0
+Version: 0.1.1
 Summary: Programmatically generate semi-realistic synthetic scribble annotations based on statistics from existing scribble datasets
 Project-URL: Homepage, https://github.com/alexsenden/scribble-annotation-generator
 Project-URL: Repository, https://github.com/alexsenden/scribble-annotation-generator

{scribble_annotation_generator-0.1.0 → scribble_annotation_generator-0.1.1}/scribble_annotation_generator/gmm_random_generator.py

@@ -54,11 +54,13 @@ class GMMRandomGenerator:
         num_components: int = 5,
         random_state: int = 0,
         class_names: Optional[List[str]] = None,
+        colour_map: Optional[dict[tuple[int, int, int], int]] = None,
     ) -> None:
         self.num_classes = num_classes
         self.num_components = num_components
         self.random_state = random_state
         self.class_names = class_names or [f"class_{i}" for i in range(num_classes)]
+        self.colour_map = colour_map
 
         self.class_gmms: Dict[int, GaussianMixture] = {}
         self.class_stats: Dict[int, Dict[str, float]] = {}
@@ -193,7 +195,6 @@ class GMMRandomGenerator:
     def sample_image(
         self,
         image_shape: Tuple[int, int],
-        colour_map: Optional[dict[tuple[int, int, int], int]] = None,
         overlap_iters: int = 40,
         overlap_margin: float = 0.05,
     ) -> Tuple[np.ndarray, str]:
@@ -220,7 +221,7 @@ class GMMRandomGenerator:
             image_shape=image_shape,
             objects=objects,
             classes=classes,
-            colour_map=colour_map,
+            colour_map=self.colour_map,
         )
 
         # Generate description string
@@ -579,6 +580,7 @@ def train_gmm_model(data_dir: str) -> GMMRandomGenerator:
         num_classes=num_classes,
         num_components=5,
         class_names=class_names,
+        colour_map=colour_map,
     ).fit(dataset)
 
     return generator
@@ -617,6 +619,7 @@ def main(argv: Optional[List[str]] = None) -> None:
         num_classes=args.num_classes,
         num_components=args.num_components,
         class_names=class_names,
+        colour_map=colour_map,
     ).fit(dataset)
     print()
 
@@ -626,7 +629,6 @@ def main(argv: Optional[List[str]] = None) -> None:
     for idx in range(args.num_samples):
         image, description = generator.sample_image(
             image_shape=(args.height, args.width),
-            colour_map=colour_map,
             overlap_iters=args.overlap_iters,
             overlap_margin=args.overlap_margin,
         )