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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: scribble-annotation-generator
-Version: 0.0.1
+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
@@ -24,6 +24,7 @@ Requires-Python: >=3.8
 Requires-Dist: numpy
 Requires-Dist: opencv-python
 Requires-Dist: scikit-image
+Requires-Dist: scikit-learn
 Requires-Dist: scipy
 Description-Content-Type: text/markdown
 

{scribble_annotation_generator-0.0.1 → scribble_annotation_generator-0.1.1}/pyproject.toml

@@ -31,7 +31,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 
-dependencies = ["numpy", "scipy", "scikit-image", "opencv-python"]
+dependencies = ["numpy", "scipy", "scikit-image", "opencv-python", "scikit-learn"]
 
 # Version is automatically provided by hatch-vcs
 dynamic = ["version"]

{scribble_annotation_generator-0.0.1 → scribble_annotation_generator-0.1.1}/scribble_annotation_generator/crop_field.py

@@ -327,13 +327,14 @@ def generate_sample(
 
 
 def generate_crop_field_dataset(
-    output_dir: str,
     colour_map: dict,
+    output_dir: str = None,
     num_samples: int = NUM_SAMPLES_TO_GENERATE,
     min_rows: int = 4,
     max_rows: int = 6,
 ):
-    os.makedirs(output_dir, exist_ok=True)
+    if output_dir is not None:
+        os.makedirs(output_dir, exist_ok=True)
 
     for i in range(num_samples):
         num_rows = random.randint(min_rows, max_rows)

scribble_annotation_generator-0.1.1/scribble_annotation_generator/gmm_random_generator.py

@@ -0,0 +1,654 @@
+import argparse
+import math
+import random
+from collections import Counter
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Optional, Tuple
+
+import cv2
+import numpy as np
+from sklearn.mixture import GaussianMixture
+
+import os
+
+from scribble_annotation_generator.dataset import ScribbleDataset
+from scribble_annotation_generator.cli import parse_colour_map
+from scribble_annotation_generator.utils import (
+    extract_class_masks,
+    extract_object_features,
+    generate_multiclass_scribble,
+    get_objects,
+    rgb_to_indexed,
+    unpack_feature_vector,
+)
+
+
+@dataclass
+class CountModelEmpirical:
+    """Empirical distribution over count vectors."""
+
+    counts: np.ndarray
+    probs: np.ndarray
+
+    @classmethod
+    def fit(cls, count_vectors: Iterable[np.ndarray]) -> "CountModelEmpirical":
+        tuples = [tuple(vec.tolist()) for vec in count_vectors]
+        counter = Counter(tuples)
+        keys = list(counter.keys())
+        values = np.array([counter[k] for k in keys], dtype=np.float64)
+        probs = values / max(values.sum(), 1.0)
+        counts = np.array(keys, dtype=np.int64)
+        return cls(counts=counts, probs=probs)
+
+    def sample(self, num_samples: int = 1) -> np.ndarray:
+        idx = np.random.choice(len(self.counts), size=num_samples, p=self.probs)
+        return self.counts[idx]
+
+
+class GMMRandomGenerator:
+    """Fit per-class GMMs over object geometry and an empirical count model."""
+
+    def __init__(
+        self,
+        num_classes: int,
+        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]] = {}
+        self.class_spur_counts: Dict[int, List[int]] = {}
+        self.count_models: Dict[int, Optional[CountModelEmpirical]] = (
+            {}
+        )  # Per-class count models
+
+    # ------------------------------------------------------------------
+    # Fitting
+    # ------------------------------------------------------------------
+
+    def fit(self, dataset: ScribbleDataset) -> "GMMRandomGenerator":
+        print(
+            f"[GMMRandomGenerator] Starting fit on dataset with {len(dataset.filenames)} files..."
+        )
+
+        per_class_samples: Dict[int, List[np.ndarray]] = {
+            i: [] for i in range(self.num_classes)
+        }
+        per_class_spurs: Dict[int, List[int]] = {i: [] for i in range(self.num_classes)}
+        count_vectors: List[np.ndarray] = []
+
+        for class_ids, features in self._iter_objects(dataset):
+            counts = np.zeros(self.num_classes, dtype=np.int64)
+            for class_id in class_ids:
+                counts[class_id] += 1
+            count_vectors.append(counts)
+
+            for class_id, feat in zip(class_ids, features):
+                gmm_feat, spur_count = self._to_gmm_feature(feat)
+                per_class_samples[class_id].append(gmm_feat)
+                per_class_spurs[class_id].append(spur_count)
+
+        # Fit per-class count models (conditional on class presence)
+        for class_id in range(1, self.num_classes):  # Skip background class (0)
+            # Filter count vectors to only those containing this class
+            class_specific_counts = [
+                counts for counts in count_vectors if counts[class_id] > 0
+            ]
+
+            if len(class_specific_counts) > 0:
+                self.count_models[class_id] = CountModelEmpirical.fit(
+                    class_specific_counts
+                )
+                print(
+                    f"[GMMRandomGenerator] Fitted count model for class {class_id} with {len(class_specific_counts)} observations"
+                )
+            else:
+                fallback_counts = np.zeros(self.num_classes, dtype=np.int64)
+                fallback_counts[class_id] = 1
+                self.count_models[class_id] = CountModelEmpirical.fit(
+                    [fallback_counts]
+                )
+                print(
+                    f"[GMMRandomGenerator] No samples for class {class_id}, using fallback singleton count"
+                )
+
+        fitted_classes = []
+        for class_id, samples in per_class_samples.items():
+            if len(samples) < 2:
+                continue
+            data = np.stack(samples, axis=0)
+            num_components = min(self.num_components, len(samples))
+            gmm = GaussianMixture(
+                n_components=num_components,
+                covariance_type="full",
+                random_state=self.random_state,
+            )
+            gmm.fit(data)
+            self.class_gmms[class_id] = gmm
+            fitted_classes.append(class_id)
+
+            lengths = data[:, 2]
+            curvatures = data[:, 4]
+            self.class_stats[class_id] = {
+                "length_min": float(lengths.min()),
+                "length_max": float(lengths.max()),
+                "curvature_min": float(curvatures.min()),
+                "curvature_max": float(curvatures.max()),
+            }
+
+        print(
+            f"[GMMRandomGenerator] Fitted GMMs for {len(fitted_classes)} classes: {fitted_classes}"
+        )
+        self.class_spur_counts = per_class_spurs
+        return self
+
+    # ------------------------------------------------------------------
+    # Sampling
+    # ------------------------------------------------------------------
+
+    def sample_counts(self, class_id: int, num_samples: int = 1) -> np.ndarray:
+        if class_id not in self.count_models or self.count_models[class_id] is None:
+            raise ValueError(f"Count model not available for class {class_id}")
+
+        model = self.count_models[class_id]
+        samples = []
+
+        # Keep sampling until we get samples that include the target class
+        for _ in range(num_samples):
+            while True:
+                sample = model.sample(1)[0]
+                if sample[class_id] > 0:
+                    samples.append(sample)
+                    break
+
+        return np.array(samples)
+
+    def sample_objects(self, counts: np.ndarray) -> Tuple[List[dict], np.ndarray]:
+        objects: List[dict] = []
+        classes: List[int] = []
+
+        for class_id in range(self.num_classes):
+            if class_id == 0:
+                continue
+            if class_id not in self.class_gmms:
+                continue
+            num = int(counts[class_id])
+            if num <= 0:
+                continue
+
+            samples, _ = self.class_gmms[class_id].sample(num)
+            for sample in samples:
+                obj = self._from_gmm_feature(class_id, sample)
+                objects.append(obj)
+                classes.append(class_id)
+
+        classes_arr = np.array(classes, dtype=np.int64)
+        return objects, classes_arr
+
+    def sample_image(
+        self,
+        image_shape: Tuple[int, int],
+        overlap_iters: int = 40,
+        overlap_margin: float = 0.05,
+    ) -> Tuple[np.ndarray, str]:
+        # Randomly select a class to condition on
+        available_classes = [
+            c
+            for c in range(1, self.num_classes)
+            if c in self.count_models and self.count_models[c] is not None
+        ]
+
+        if not available_classes:
+            return np.zeros(image_shape, dtype=np.uint8), "Empty image"
+
+        chosen_class = random.choice(available_classes)
+        counts = self.sample_counts(chosen_class, 1)[0]
+        objects, classes = self.sample_objects(counts)
+
+        if len(objects) == 0:
+            return np.zeros(image_shape, dtype=np.uint8), "Empty image"
+
+        objects = self._resolve_overlaps(objects, overlap_iters, overlap_margin)
+
+        image = generate_multiclass_scribble(
+            image_shape=image_shape,
+            objects=objects,
+            classes=classes,
+            colour_map=self.colour_map,
+        )
+
+        # Generate description string
+        class_counts: Dict[int, int] = {}
+        for class_id in classes:
+            class_counts[class_id] = class_counts.get(class_id, 0) + 1
+
+        description_parts = []
+        for class_id in sorted(class_counts.keys()):
+            count = class_counts[class_id]
+            class_name = self.class_names[class_id]
+            if count == 1:
+                description_parts.append(f'1 object of class "{class_name}"')
+            else:
+                description_parts.append(f'{count} objects of class "{class_name}"')
+
+        description = "; ".join(description_parts)
+        return image, description
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _iter_objects(
+        self, dataset: ScribbleDataset
+    ) -> Iterable[Tuple[List[int], List[np.ndarray]]]:
+        for filename in dataset.filenames:
+            filepath = os.path.join(dataset.data_dir, filename)
+
+            if dataset.is_rgb:
+                if dataset.colour_map is None:
+                    raise ValueError("colour_map must be provided for RGB annotations")
+                mask = cv2.imread(str(filepath), cv2.IMREAD_COLOR)
+                mask = cv2.cvtColor(mask, cv2.COLOR_BGR2RGB)
+                mask = rgb_to_indexed(mask, dataset.colour_map)
+            else:
+                mask = cv2.imread(str(filepath), cv2.IMREAD_GRAYSCALE)
+
+            class_masks = extract_class_masks(mask)
+            class_ids: List[int] = []
+            features: List[np.ndarray] = []
+
+            for class_id, class_mask in class_masks.items():
+                objects = get_objects(class_mask)
+                for obj_mask in objects:
+                    feat = extract_object_features(obj_mask)
+                    class_ids.append(int(class_id))
+                    features.append(feat)
+
+            yield class_ids, features
+
+    def _to_gmm_feature(self, feat: np.ndarray) -> Tuple[np.ndarray, int]:
+        data = unpack_feature_vector(feat)
+
+        start_x = data["start_x"]
+        start_y = data["start_y"]
+        end_x = data["end_x"]
+        end_y = data["end_y"]
+
+        center_x = (start_x + end_x) / 2.0
+        center_y = (start_y + end_y) / 2.0
+
+        dx = end_x - start_x
+        dy = end_y - start_y
+        length = math.sqrt(dx * dx + dy * dy)
+        angle = math.atan2(dy, dx)
+
+        curvature = data["curvature"]
+        num_spurs = int(round(data["num_spurs"]))
+
+        return (
+            np.array([center_x, center_y, length, angle, curvature], dtype=np.float32),
+            num_spurs,
+        )
+
+    def _from_gmm_feature(self, class_id: int, sample: np.ndarray) -> dict:
+        center_x, center_y, length, angle, curvature = sample.tolist()
+
+        stats = self.class_stats.get(class_id, {})
+        length_min = stats.get("length_min", 0.0)
+        length_max = stats.get("length_max", 2.0)
+        curvature_min = stats.get("curvature_min", -1.0)
+        curvature_max = stats.get("curvature_max", 1.0)
+
+        length = float(
+            np.clip(
+                abs(length), max(1e-4, length_min), max(length_min + 1e-4, length_max)
+            )
+        )
+        curvature = float(np.clip(curvature, curvature_min, curvature_max))
+
+        center_x = float(np.clip(center_x, -1.0, 1.0))
+        center_y = float(np.clip(center_y, -1.0, 1.0))
+
+        cos_angle = math.cos(angle)
+        sin_angle = math.sin(angle)
+
+        length = self._fit_length_to_bounds(
+            center_x, center_y, length, cos_angle, sin_angle
+        )
+
+        half = 0.5 * length
+        start_x = center_x - half * cos_angle
+        start_y = center_y - half * sin_angle
+        end_x = center_x + half * cos_angle
+        end_y = center_y + half * sin_angle
+
+        spur_list = self.class_spur_counts.get(class_id, [])
+        num_spurs = int(random.choice(spur_list)) if spur_list else 0
+
+        return {
+            "start_x": start_x,
+            "start_y": start_y,
+            "end_x": end_x,
+            "end_y": end_y,
+            "num_spurs": num_spurs,
+            "curvature": curvature,
+            "cos_angle": cos_angle,
+            "sin_angle": sin_angle,
+        }
+
+    def _fit_length_to_bounds(
+        self,
+        center_x: float,
+        center_y: float,
+        length: float,
+        cos_angle: float,
+        sin_angle: float,
+    ) -> float:
+        half = 0.5 * length
+
+        max_half = half
+        if abs(cos_angle) > 1e-6:
+            max_half = min(max_half, (1.0 - abs(center_x)) / abs(cos_angle))
+        if abs(sin_angle) > 1e-6:
+            max_half = min(max_half, (1.0 - abs(center_y)) / abs(sin_angle))
+
+        max_half = max(0.0, max_half)
+        return max(2e-4, 2.0 * max_half)
+
+    def _resolve_overlaps(
+        self,
+        objects: List[dict],
+        max_iters: int,
+        min_margin: float,
+    ) -> List[dict]:
+        centers = np.array(
+            [
+                [
+                    (obj["start_x"] + obj["end_x"]) / 2.0,
+                    (obj["start_y"] + obj["end_y"]) / 2.0,
+                ]
+                for obj in objects
+            ],
+            dtype=np.float32,
+        )
+
+        radii = np.array(
+            [
+                0.5
+                * math.sqrt(
+                    (obj["end_x"] - obj["start_x"]) ** 2
+                    + (obj["end_y"] - obj["start_y"]) ** 2
+                )
+                for obj in objects
+            ],
+            dtype=np.float32,
+        )
+
+        if len(objects) < 2:
+            return objects
+
+        # Gentle centering force strength (0 = no centering, 1 = full centering)
+        centering_strength = 0.1
+
+        # Iteratively resolve overlaps and push toward center
+        for iteration in range(max_iters):
+            overlaps_found = False
+
+            # Resolve overlaps between object pairs
+            for i in range(len(objects)):
+                for j in range(i + 1, len(objects)):
+                    # Calculate distance between centers
+                    dx = centers[j, 0] - centers[i, 0]
+                    dy = centers[j, 1] - centers[i, 1]
+                    dist = math.sqrt(dx * dx + dy * dy)
+
+                    # Calculate minimum required distance
+                    min_dist = radii[i] + radii[j] + min_margin
+
+                    if dist < min_dist and dist > 1e-6:
+                        overlaps_found = True
+
+                        # Move objects apart
+                        direction_x = dx / dist
+                        direction_y = dy / dist
+
+                        overlap = min_dist - dist
+                        displacement = overlap / 2.0 + 1e-4
+
+                        centers[i, 0] -= direction_x * displacement
+                        centers[i, 1] -= direction_y * displacement
+                        centers[j, 0] += direction_x * displacement
+                        centers[j, 1] += direction_y * displacement
+
+            # Gently push objects toward center of image
+            for i in range(len(objects)):
+                # Center of image is at (0, 0) in normalized coordinates
+                # Objects closer to edges experience stronger centering force
+                dist_from_center = math.sqrt(centers[i, 0] ** 2 + centers[i, 1] ** 2)
+                max_dist = math.sqrt(
+                    2.0
+                )  # diagonal distance in normalized [-1, 1] space
+                adaptive_strength = centering_strength * (dist_from_center / max_dist)
+
+                centers[i, 0] *= 1.0 - adaptive_strength
+                centers[i, 1] *= 1.0 - adaptive_strength
+
+            # if not overlaps_found:
+            #     break
+
+        # Update objects with new centers
+        for idx, obj in enumerate(objects):
+            center_x = float(np.clip(centers[idx, 0], -1.0, 1.0))
+            center_y = float(np.clip(centers[idx, 1], -1.0, 1.0))
+
+            cos_angle = obj["cos_angle"]
+            sin_angle = obj["sin_angle"]
+            length = (
+                0.5
+                * math.sqrt(
+                    (obj["end_x"] - obj["start_x"]) ** 2
+                    + (obj["end_y"] - obj["start_y"]) ** 2
+                )
+                * 2.0
+            )
+
+            length = self._fit_length_to_bounds(
+                center_x, center_y, length, cos_angle, sin_angle
+            )
+
+            half = 0.5 * length
+            obj["start_x"] = center_x - half * cos_angle
+            obj["start_y"] = center_y - half * sin_angle
+            obj["end_x"] = center_x + half * cos_angle
+            obj["end_y"] = center_y + half * sin_angle
+
+        return objects
+
+
+def build_parser() -> argparse.ArgumentParser:
+
+    parser = argparse.ArgumentParser(
+        description="Fit per-class GMMs and sample synthetic scribble annotations.",
+    )
+    parser.add_argument(
+        "--data-dir",
+        required=True,
+        help="Path to the training dataset directory.",
+    )
+    parser.add_argument(
+        "--num-classes",
+        type=int,
+        required=True,
+        help="Total number of classes (including background).",
+    )
+    parser.add_argument(
+        "--colour-map",
+        default=None,
+        help=(
+            "Colour map specified inline as 'R,G,B=class;...' or a path to a file "
+            "with one 'R,G,B,class' or 'R,G,B' entry per line. Required for RGB annotations."
+        ),
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="./local/gmm-inference",
+        help="Directory to write generated samples.",
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=10,
+        help="Number of images to generate.",
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=512,
+        help="Output image height.",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=512,
+        help="Output image width.",
+    )
+    parser.add_argument(
+        "--num-components",
+        type=int,
+        default=5,
+        help="Number of GMM components per class.",
+    )
+    parser.add_argument(
+        "--overlap-iters",
+        type=int,
+        default=40,
+        help="Iterations for overlap resolution.",
+    )
+    parser.add_argument(
+        "--overlap-margin",
+        type=float,
+        default=0.05,
+        help="Minimum margin between objects in normalized coordinates.",
+    )
+    parser.add_argument(
+        "--class-names",
+        type=str,
+        default=None,
+        help=("Path to file with one class name per line."),
+    )
+    return parser
+
+
+def read_class_names(filepath: str) -> List[str]:
+    with open(filepath, "r") as f:
+        lines = f.read().splitlines()
+    return lines
+
+
+def train_gmm_model(data_dir: str) -> GMMRandomGenerator:
+    """
+    Trains and returns a GMMRandomGenerator model based on the dataset located in data_dir.
+
+    Args:
+        data_dir (str): Path to the training dataset directory. This directory should contain
+            a subdirectory "segmentation" with annotation files, a "colour_map.txt" file,
+            and a "class_labelling.txt" file.
+
+    Returns:
+        GMMRandomGenerator: The trained GMMRandomGenerator model.
+    """
+
+    colour_map = parse_colour_map(os.path.join(data_dir, "colour_map.txt"))
+    num_classes = len(colour_map)
+
+    dataset = ScribbleDataset(
+        num_classes=num_classes,
+        data_dir=os.path.join(data_dir, "segmentation"),
+        colour_map=colour_map,
+    )
+
+    class_names = read_class_names(os.path.join(data_dir, "class_labelling.txt"))
+
+    generator = GMMRandomGenerator(
+        num_classes=num_classes,
+        num_components=5,
+        class_names=class_names,
+        colour_map=colour_map,
+    ).fit(dataset)
+
+    return generator
+
+
+def main(argv: Optional[List[str]] = None) -> None:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    print(f"\n{'='*70}")
+    print(f"{'GMMRandomGenerator - Scribble Annotation Generation':^70}")
+    print(f"{'='*70}\n")
+
+    print(f"Configuration:")
+    print(f"  Data directory:    {args.data_dir}")
+    print(f"  Number of classes: {args.num_classes}")
+    print(f"  GMM components:    {args.num_components}")
+    print(f"  Output directory:  {args.output_dir}")
+    print(f"  Image size:        {args.width}x{args.height}")
+    print(f"  Samples to gen:    {args.num_samples}\n")
+
+    colour_map = parse_colour_map(args.colour_map) if args.colour_map else None
+
+    print("[1/3] Loading dataset...")
+    dataset = ScribbleDataset(
+        num_classes=args.num_classes,
+        data_dir=args.data_dir,
+        colour_map=colour_map,
+    )
+    print(f"      Loaded {len(dataset.filenames)} annotation files\n")
+
+    print("[2/3] Fitting GMM models...")
+    class_names = read_class_names(args.class_names) if args.class_names else None
+
+    generator = GMMRandomGenerator(
+        num_classes=args.num_classes,
+        num_components=args.num_components,
+        class_names=class_names,
+        colour_map=colour_map,
+    ).fit(dataset)
+    print()
+
+    print("[3/3] Generating samples...")
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    for idx in range(args.num_samples):
+        image, description = generator.sample_image(
+            image_shape=(args.height, args.width),
+            overlap_iters=args.overlap_iters,
+            overlap_margin=args.overlap_margin,
+        )
+
+        output_path = os.path.join(args.output_dir, f"gmm_sample_{idx:04d}.png")
+        if image.ndim == 3:
+            image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+            cv2.imwrite(output_path, image_bgr)
+        else:
+            cv2.imwrite(output_path, image)
+
+        print(f"      Sample {idx + 1}: {description}")
+
+        if (idx + 1) % 10 == 0 or idx == args.num_samples - 1:
+            print(f"      Generated {idx + 1}/{args.num_samples} samples")
+
+    print(f"\n{'='*70}")
+    print(f"✓ Generation complete! Samples saved to: {args.output_dir}")
+    print(f"{'='*70}\n")
+
+
+if __name__ == "__main__":
+    main()