scribble-annotation-generator 0.0.1__py3-none-any.whl

scribble_annotation_generator/cli.py

@@ -0,0 +1,195 @@
+import argparse
+import os
+from typing import Dict, Tuple
+
+from scribble_annotation_generator.crop_field import (
+    NUM_SAMPLES_TO_GENERATE,
+    generate_crop_field_dataset,
+)
+from scribble_annotation_generator.nn import train_and_infer
+
+
+def parse_colour_map(value: str) -> Dict[Tuple[int, int, int], int]:
+    """Parse a colour map from an inline string or a file path."""
+
+    def _validate_rgb(rgb: Tuple[int, int, int]) -> Tuple[int, int, int]:
+        r, g, b = rgb
+        for channel in (r, g, b):
+            if channel < 0 or channel > 255:
+                raise ValueError("RGB values must be between 0 and 255")
+        return rgb
+
+    mapping: Dict[Tuple[int, int, int], int] = {}
+
+    if os.path.isfile(value):
+        with open(value, "r", encoding="utf-8") as handle:
+            for idx, line in enumerate(handle):
+                stripped = line.strip()
+                if not stripped:
+                    continue
+                parts = [part.strip() for part in stripped.split(",") if part.strip()]
+                if len(parts) == 4:
+                    r, g, b, cls = parts
+                elif len(parts) == 3:
+                    r, g, b = parts
+                    cls = idx
+                else:
+                    raise ValueError(
+                        "Each line in the colour map file must have 3 (RGB) or 4 (RGB,class) comma-separated values"
+                    )
+                rgb = _validate_rgb((int(r), int(g), int(b)))
+                mapping[rgb] = int(cls)
+    else:
+        entries = [entry.strip() for entry in value.split(";") if entry.strip()]
+        for entry in entries:
+            if "=" in entry:
+                colour_part, class_part = entry.split("=", 1)
+            elif ":" in entry:
+                colour_part, class_part = entry.split(":", 1)
+            else:
+                raise ValueError(
+                    "Inline colour map entries must separate colour and class with '=' or ':'"
+                )
+            rgb_parts = [part.strip() for part in colour_part.split(",") if part.strip()]
+            if len(rgb_parts) != 3:
+                raise ValueError("Colours must be provided as R,G,B")
+            rgb = _validate_rgb((int(rgb_parts[0]), int(rgb_parts[1]), int(rgb_parts[2])))
+            mapping[rgb] = int(class_part.strip())
+
+    if not mapping:
+        raise ValueError("No colours were parsed for the colour map")
+
+    return mapping
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Scribble Annotation Generator CLI",
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    crop_parser = subparsers.add_parser(
+        "crop-field", help="Generate synthetic crop field scribble images."
+    )
+    crop_parser.add_argument(
+        "--output-dir",
+        default="./local/crop_field",
+        help="Directory to write generated crop field images.",
+    )
+    crop_parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=NUM_SAMPLES_TO_GENERATE,
+        help="Number of images to generate.",
+    )
+    crop_parser.add_argument(
+        "--min-rows",
+        type=int,
+        default=4,
+        help="Minimum number of crop rows per sample.",
+    )
+    crop_parser.add_argument(
+        "--max-rows",
+        type=int,
+        default=6,
+        help="Maximum number of crop rows per sample.",
+    )
+    crop_parser.add_argument(
+        "--colour-map",
+        required=True,
+        help=(
+            "Colour map specified inline as 'R,G,B=class;...' or a path to a file "
+            "with one 'R,G,B,class' entry per line."
+        ),
+    )
+
+    train_parser = subparsers.add_parser(
+        "train-nn", help="Train the scribble object generator and run inference."
+    )
+    train_parser.add_argument(
+        "--train-dir",
+        required=True,
+        help="Path to the training dataset directory.",
+    )
+    train_parser.add_argument(
+        "--val-dir",
+        required=True,
+        help="Path to the validation dataset directory.",
+    )
+    train_parser.add_argument(
+        "--checkpoint-dir",
+        default="./local/nn-checkpoints",
+        help="Directory to save model checkpoints.",
+    )
+    train_parser.add_argument(
+        "--inference-dir",
+        default="./local/nn-inference",
+        help="Directory to save inference visualisations.",
+    )
+    train_parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=8,
+        help="Batch size for training.",
+    )
+    train_parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of worker processes for data loading.",
+    )
+    train_parser.add_argument(
+        "--max-epochs",
+        type=int,
+        default=50,
+        help="Maximum number of training epochs.",
+    )
+    train_parser.add_argument(
+        "--num-classes",
+        type=int,
+        default=None,
+        help="Override the number of classes; defaults to the number of unique class IDs in the colour map.",
+    )
+    train_parser.add_argument(
+        "--colour-map",
+        required=True,
+        help=(
+            "Colour map specified inline as 'R,G,B=class;...' or a path to a file "
+            "with one 'R,G,B,class' entry per line."
+        ),
+    )
+
+    return parser
+
+
+def main(argv=None):
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    colour_map = parse_colour_map(args.colour_map)
+
+    if args.command == "crop-field":
+        generate_crop_field_dataset(
+            output_dir=args.output_dir,
+            colour_map=colour_map,
+            num_samples=args.num_samples,
+            min_rows=args.min_rows,
+            max_rows=args.max_rows,
+        )
+    elif args.command == "train-nn":
+        train_and_infer(
+            train_dir=args.train_dir,
+            val_dir=args.val_dir,
+            colour_map=colour_map,
+            checkpoint_dir=args.checkpoint_dir,
+            inference_dir=args.inference_dir,
+            batch_size=args.batch_size,
+            num_workers=args.num_workers,
+            max_epochs=args.max_epochs,
+            num_classes=args.num_classes,
+        )
+    else:
+        parser.error("A subcommand is required.")
+
+
+if __name__ == "__main__":
+    main()

scribble_annotation_generator/crop_field.py

@@ -0,0 +1,366 @@
+import cv2
+import math
+import numpy as np
+import os
+import random
+
+from scribble_annotation_generator.utils import generate_multiclass_scribble
+
+
+NUM_SAMPLES_TO_GENERATE = 200
+ROW_STD = 0.02
+ROW_CURVATURE_MEAN = -0.8
+ROW_CURVATURE_STD = 0.05
+ROW_MIN_LENGTH = 0.1
+ROW_SPARSITY_DISTANCE_MEAN = 0.4
+ROW_SPARSITY_DISTANCE_STD = 0.2
+WEED_MAX_LENGTH = 0.5
+WEED_MIN_LENGTH = 0.001
+WEED_DIRECTIONAL_STD = math.pi / 6
+WEED_CURVATURE_SHIFT_FACTOR = 0.3
+WEED_CURVATURE_SCALE_CONSTANT = 0.2
+WEED_CURVATURE_MIN_STD = 0.4
+
+
+class Point:
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
+
+
+def ccw(A, B, C):
+    return (C.y - A.y) * (B.x - A.x) > (B.y - A.y) * (C.x - A.x)
+
+
+# Return true if line segments AB and CD intersect
+def intersect(A, B, C, D):
+    return ccw(A, C, D) != ccw(B, C, D) and ccw(A, B, C) != ccw(A, B, D)
+
+
+def x_at_y(p1, p2, y):
+    x1, y1 = p1
+    x2, y2 = p2
+
+    t = (y - y1) / (y2 - y1)
+    return x1 + t * (x2 - x1)
+
+
+def split_row(row_object, sparsity):
+    num_splits = np.random.poisson(lam=((1 - sparsity) ** 2) * 5)
+    if num_splits == 0:
+        return [row_object]
+
+    distance_between_splits = np.random.normal(
+        loc=ROW_SPARSITY_DISTANCE_MEAN,
+        scale=ROW_SPARSITY_DISTANCE_STD,
+        size=num_splits,
+    )
+
+    distance_between_splits = list(np.clip(distance_between_splits, 0.05, None))
+
+    # Split the line into num_splits + 1 segments
+    split_ys = np.random.uniform(
+        row_object["start_y"], row_object["end_y"], size=num_splits
+    )
+    split_ys = list(np.sort(split_ys))
+
+    # If split points are too close to each other or to the boundary, remove one
+    i = 1
+    while i < len(split_ys):
+        if (
+            split_ys[i - 1] - (distance_between_splits[i - 1] / 2.0)
+            < row_object["start_y"]
+        ):
+            split_ys.pop(i - 1)
+            distance_between_splits.pop(i - 1)
+        elif split_ys[i] + (distance_between_splits[i] / 2.0) > row_object["end_y"]:
+            split_ys.pop(i)
+            distance_between_splits.pop(i)
+        elif split_ys[i] - split_ys[i - 1] < distance_between_splits[i - 1]:
+            split_ys.pop(i)
+            distance_between_splits.pop(i)
+        else:
+            i += 1
+
+    line_segment_ys = []
+    for i in range(len(split_ys) + 1):
+        if i == 0:
+            segment_start_y = row_object["start_y"]
+        else:
+            segment_start_y = split_ys[i - 1] + (distance_between_splits[i - 1] / 2.0)
+
+        if i == len(split_ys):
+            segment_end_y = row_object["end_y"]
+        else:
+            segment_end_y = split_ys[i] - (distance_between_splits[i] / 2.0)
+
+        line_segment_ys.append((segment_start_y, segment_end_y))
+
+    line_segment_xs = [
+        (
+            x_at_y(
+                (row_object["start_x"], row_object["start_y"]),
+                (row_object["end_x"], row_object["end_y"]),
+                y[0],
+            ),
+            x_at_y(
+                (row_object["start_x"], row_object["start_y"]),
+                (row_object["end_x"], row_object["end_y"]),
+                y[1],
+            ),
+        )
+        for y in line_segment_ys
+    ]
+
+    objects = []
+    for i in range(len(line_segment_ys)):
+        split_row_object = {
+            "start_x": line_segment_xs[i][0],
+            "start_y": line_segment_ys[i][0],
+            "end_x": line_segment_xs[i][1],
+            "end_y": line_segment_ys[i][1],
+            "num_spurs": row_object["num_spurs"],
+            "curvature": row_object["curvature"],
+            "cos_angle": row_object["cos_angle"],
+            "sin_angle": row_object["sin_angle"],
+        }
+        objects.append(split_row_object)
+
+    return objects
+
+
+def generate_row_object(
+    row_starting_x: float,
+    row_class: int,
+    row_sparsity: float = 1.0,
+):
+    row_x0 = np.clip(
+        np.random.normal(loc=row_starting_x, scale=ROW_STD),
+        -1.0,
+        1.0,
+    )
+    row_x1 = np.clip(
+        np.random.normal(loc=row_starting_x, scale=ROW_STD),
+        -1.0,
+        1.0,
+    )
+
+    row_y0 = np.clip(
+        np.random.normal(loc=-1.0, scale=ROW_STD),
+        -1.0,
+        1.0,
+    )
+    row_y1 = np.clip(
+        np.random.normal(loc=1.0, scale=ROW_STD),
+        -1.0,
+        1.0,
+    )
+
+    curvature = np.clip(
+        np.random.normal(
+            loc=ROW_CURVATURE_MEAN,
+            scale=ROW_CURVATURE_STD,
+        ),
+        -1.0,
+        1.0,
+    )
+    num_spurs = 0
+
+    angle = math.atan2(row_y1 - row_y0, row_x1 - row_x0)
+    cos_angle = math.cos(angle)
+    sin_angle = math.sin(angle)
+
+    row_object = {
+        "start_x": row_x0,
+        "start_y": row_y0,
+        "end_x": row_x1,
+        "end_y": row_y1,
+        "num_spurs": num_spurs,
+        "curvature": curvature,
+        "cos_angle": cos_angle,
+        "sin_angle": sin_angle,
+    }
+
+    if row_sparsity < 1.0:
+        objects = split_row(row_object, row_sparsity)
+        classes = [row_class] * len(objects)
+    else:
+        objects = [row_object]
+        classes = [row_class]
+
+    return objects, classes
+
+
+def generate_weed_object():
+    weed_x0 = random.uniform(-1.0, 1.0)
+    weed_y0 = random.uniform(-1.0, 1.0)
+
+    weed_length = random.uniform(WEED_MIN_LENGTH, WEED_MAX_LENGTH)
+    weed_angle = np.random.normal(
+        loc=3 * math.pi / 2,
+        scale=WEED_DIRECTIONAL_STD,
+    )
+
+    weed_x1 = np.clip(weed_x0 + weed_length * math.cos(weed_angle), -1.0, 1.0)
+    weed_y1 = np.clip(weed_y0 + weed_length * math.sin(weed_angle), -1.0, 1.0)
+
+    num_spurs = 0
+    weed_length_factor = (weed_length - WEED_MIN_LENGTH + 1e-6) / (
+        WEED_MAX_LENGTH - WEED_MIN_LENGTH
+    )
+    curvature = np.clip(
+        np.random.normal(
+            loc=((((1 - weed_length_factor) * 2) - 1))
+            * (1 - WEED_CURVATURE_SHIFT_FACTOR)
+            - WEED_CURVATURE_SHIFT_FACTOR,
+            scale=max(
+                (1 - weed_length_factor) * WEED_CURVATURE_SCALE_CONSTANT,
+                WEED_CURVATURE_MIN_STD,
+            ),
+        ),
+        -1.0,
+        1.0,
+    )
+
+    return {
+        "start_x": weed_x0,
+        "start_y": weed_y0,
+        "end_x": weed_x1,
+        "end_y": weed_y1,
+        "num_spurs": num_spurs,
+        "curvature": curvature,
+        "cos_angle": math.cos(weed_angle),
+        "sin_angle": math.sin(weed_angle),
+    }
+
+
+def generate_sample(
+    colour_map: dict[tuple[int, int, int], int],
+    num_rows: int = 5,
+    row_class: int = 1,
+    interspersed: bool = False,
+    interspersed_num_rows: int = 0,
+    interspersed_class: int = 2,
+    row_sparsity: float = 1.0,
+    num_weeds: dict[int, int] = {},
+):
+    objects = []
+    classes = []
+
+    row_offset = 2.0 / (num_rows + 1)
+    initial_row_starting_x = random.uniform(-1.0, -1.0 + row_offset)
+    row_starting_x = initial_row_starting_x
+    for _ in range(num_rows):
+        row_objects, row_classes = generate_row_object(
+            row_starting_x=row_starting_x,
+            row_class=row_class,
+            row_sparsity=row_sparsity,
+        )
+
+        objects.extend(row_objects)
+        classes.extend(row_classes)
+
+        row_starting_x += row_offset
+
+    if interspersed:
+        interspersed_row_starting_x = initial_row_starting_x - (row_offset / 2.0)
+
+        # Ensure interspersed row at index 0 is within bounds
+        if interspersed_row_starting_x < -1.0:
+            interspersed_row_starting_x += row_offset
+
+        # Get maximum number of interspersed rows that fit
+        num_interspersed_row_positions = num_rows
+        half_offset = row_offset / 2.0
+        if initial_row_starting_x - half_offset > -1.0:
+            num_interspersed_row_positions += 1
+        if initial_row_starting_x + (num_rows * row_offset) + half_offset < 1.0:
+            num_interspersed_row_positions += 1
+
+        # Select the starting position for interspersed rows
+        interspersed_row_starting_index = random.randint(
+            0, max(num_interspersed_row_positions - num_interspersed_row_positions, 0)
+        )
+        interspersed_row_starting_x += interspersed_row_starting_index * row_offset
+
+        for _ in range(interspersed_num_rows):
+            row_objects, row_classes = generate_row_object(
+                row_starting_x=interspersed_row_starting_x,
+                row_class=interspersed_class,
+                row_sparsity=row_sparsity,
+            )
+
+            objects.extend(row_objects)
+            classes.extend(row_classes)
+
+            interspersed_row_starting_x += row_offset
+
+    for weed_class, num_weed in num_weeds.items():
+
+        for _ in range(num_weed):
+            intersects = True
+            while intersects:
+                intersects = False
+
+                weed_object = generate_weed_object()
+                weed_start = Point(weed_object["start_x"], weed_object["start_y"])
+                weed_end = Point(weed_object["end_x"], weed_object["end_y"])
+
+                for obj in objects:
+                    obj_start = Point(obj["start_x"], obj["start_y"])
+                    obj_end = Point(obj["end_x"], obj["end_y"])
+                    if intersect(weed_start, weed_end, obj_start, obj_end):
+                        intersects = True
+                        break
+
+            objects.append(weed_object)
+            classes.append(weed_class)
+
+    synthetic = generate_multiclass_scribble(
+        image_shape=(512, 512),
+        objects=objects,
+        classes=classes,
+        colour_map=colour_map,
+    )
+
+    return synthetic
+
+
+def generate_crop_field_dataset(
+    output_dir: str,
+    colour_map: dict,
+    num_samples: int = NUM_SAMPLES_TO_GENERATE,
+    min_rows: int = 4,
+    max_rows: int = 6,
+):
+    os.makedirs(output_dir, exist_ok=True)
+
+    for i in range(num_samples):
+        num_rows = random.randint(min_rows, max_rows)
+        row_class = random.randint(1, 3)
+        interspersed = random.choice([True, False])
+        interspersed_num_rows = random.randint(1, num_rows + 1)
+        interspersed_class = random.choice([c for c in [1, 2, 3] if c != row_class])
+        row_sparsity = random.uniform(0.1, 1.0)
+        num_weeds = {
+            2: random.randint(0, 5),
+            3: random.randint(0, 10),
+            4: random.randint(0, 10),
+        }
+
+        sample = generate_sample(
+            colour_map=colour_map,
+            num_rows=num_rows,
+            row_class=row_class,
+            interspersed=interspersed,
+            interspersed_num_rows=interspersed_num_rows,
+            interspersed_class=interspersed_class,
+            row_sparsity=row_sparsity,
+            num_weeds=num_weeds,
+        )
+
+        if random.random() < 0.5:
+            sample = cv2.flip(sample, 1)
+
+        output_path = os.path.join(output_dir, f"synthetic_{i:04d}.png")
+        cv2.imwrite(output_path, cv2.cvtColor(sample, cv2.COLOR_RGB2BGR))

scribble_annotation_generator/dataset.py

@@ -0,0 +1,96 @@
+import os
+import random
+
+import cv2
+import torch
+import torch.nn.functional as F
+
+from scribble_annotation_generator.utils import (
+    extract_class_masks,
+    extract_object_features,
+    get_objects,
+    rgb_to_indexed,
+)
+
+
+class ScribbleDataset(torch.utils.data.Dataset):
+    def __init__(
+        self, num_classes, data_dir, colour_map=None, max_objects=50, late_shift=False
+    ):
+        self.data_dir = data_dir
+        self.filenames = sorted(os.listdir(data_dir))
+        self.num_classes = num_classes
+        self.colour_map = colour_map
+        self.max_objects = max_objects
+        self.late_shift = late_shift
+
+        if len(self.filenames) == 0:
+            raise ValueError(f"No files found in {data_dir}")
+
+        # Auto-detect format from first image
+        first_image_path = os.path.join(self.data_dir, self.filenames[0])
+        first_img = cv2.imread(first_image_path, cv2.IMREAD_UNCHANGED)
+        if first_img is not None:
+            self.is_rgb = len(first_img.shape) == 3 and first_img.shape[2] >= 3
+        else:
+            raise IOError(f"Could not read {first_image_path}")
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, idx):
+        filepath = os.path.join(self.data_dir, self.filenames[idx])
+
+        if self.is_rgb:
+            if self.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)  # Convert BGR to RGB
+            mask = rgb_to_indexed(mask, self.colour_map)
+        else:
+            mask = cv2.imread(str(filepath), cv2.IMREAD_GRAYSCALE)
+
+        objects = []
+        classes = []
+        class_masks = extract_class_masks(mask)
+
+        for class_id, class_mask in class_masks.items():
+            class_objects = get_objects(class_mask)
+            objects.extend(class_objects)
+            classes.extend([class_id] * len(class_objects))
+
+        objects = [extract_object_features(obj) for obj in objects]
+
+        permutation = list(range(len(objects)))
+        random.shuffle(permutation)
+
+        objects = torch.stack([objects[i] for i in permutation])
+        classes = torch.tensor([classes[i] for i in permutation])
+
+        # Mask everything after a random point
+        if self.late_shift:
+            mask_start = random.randint((len(objects) // 4) * 3, len(objects) - 1)
+        else:
+            mask_start = random.randint(1, len(objects) - 1)
+        mask = torch.ones(len(objects))
+        mask[mask_start:] = 0
+
+        query_cls = classes[mask_start]
+        targets = objects[classes == query_cls, :]
+
+        objects = F.pad(objects, (0, 0, 0, self.max_objects - len(objects)), value=0)
+        classes = F.pad(classes, (0, self.max_objects - len(classes)), value=0)
+        mask = F.pad(mask, (0, self.max_objects - len(mask)), value=0)
+        targets = F.pad(
+            targets, (0, 0, 0, self.max_objects - targets.size(0)), value=1e7
+        )
+
+        return {
+            "objects": objects,
+            "classes": classes,
+            "mask": mask,
+            "query_cls": query_cls,
+            "targets": targets,
+            "counts": torch.bincount(classes, minlength=self.num_classes),
+        }

scribble_annotation_generator/debug.py

@@ -0,0 +1,43 @@
+import cv2
+
+from scribble_annotation_generator.nn.nn import ScribbleDataset
+from scribble_annotation_generator.nn.utils import (
+    generate_multiclass_scribble,
+    unpack_feature_vector,
+)
+
+
+def parameterize_and_unparameterize():
+    colour_map = {
+        (0, 0, 0): 0,
+        (0, 128, 255): 1,
+        (124, 255, 121): 2,
+        (127, 0, 0): 3,
+        (255, 148, 0): 4,
+        (0, 0, 127): 5,
+    }
+    dataset = ScribbleDataset(
+        num_classes=3, data_dir="./local/soybean1", colour_map=colour_map
+    )
+
+    for i in range(len(dataset)):
+        sample = dataset[i]
+
+        objects = sample["objects"]
+        classes = sample["classes"]
+
+        objects = [unpack_feature_vector(obj) for obj in objects.numpy()]
+
+        synthetic = generate_multiclass_scribble(
+            image_shape=(512, 512),
+            objects=objects,
+            classes=classes,
+            colour_map=colour_map,
+        )
+
+        # Save the synthetic scribble
+        output_path = f"./local/nn-out/synthetic_{i:04d}.png"
+
+        # Convert RGB to BGR for saving with OpenCV
+        synthetic_bgr = cv2.cvtColor(synthetic, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(str(output_path), synthetic_bgr)