flashseg 1.0.0__py3-none-any.whl

flashseg/__init__.py

@@ -0,0 +1,20 @@
+"""FlashSeg - Ultra-lightweight real-time image segmentation."""
+
+__version__ = "1.0.0"
+
+from flashseg.cfg.config import get_config
+from flashseg.engine.trainer import Trainer
+from flashseg.engine.predictor import Predictor
+from flashseg.engine.exporter import Exporter
+from flashseg.engine.validator import Validator
+from flashseg.models.build import build_model
+
+__all__ = [
+    "__version__",
+    "get_config",
+    "build_model",
+    "Trainer",
+    "Predictor",
+    "Exporter",
+    "Validator",
+]

flashseg/analytics/__init__.py

@@ -0,0 +1,6 @@
+"""Analytics and benchmarking tools."""
+
+from flashseg.analytics.benchmark import Benchmark
+from flashseg.analytics.profiler import Profiler
+
+__all__ = ["Benchmark", "Profiler"]

flashseg/analytics/benchmark.py

@@ -0,0 +1,65 @@
+"""Model benchmarking for segmentation."""
+
+import time
+import logging
+
+import torch
+
+from flashseg.cfg.config import get_config
+from flashseg.models.build import build_model
+
+logger = logging.getLogger(__name__)
+
+
+class Benchmark:
+    """Benchmark FlashSeg model speed and efficiency."""
+
+    def __init__(self, model_path: str = None, model_size: str = "m", input_size: int = 512, num_classes: int = 21, device: str = "cuda"):
+        self.device = torch.device(device if torch.cuda.is_available() or device == "cpu" else "cpu")
+        config = get_config(model_size=model_size, input_size=input_size, num_classes=num_classes)
+        self.model = build_model(config).to(self.device)
+
+        if model_path:
+            self.model.load_state_dict(torch.load(model_path, map_location=self.device))
+
+        self.model.eval()
+        self.input_size = input_size
+
+    def run(self, warmup: int = 10, iterations: int = 100) -> dict:
+        """Run benchmark and return timing results."""
+        dummy = torch.randn(1, 3, self.input_size, self.input_size).to(self.device)
+
+        # Warmup
+        with torch.no_grad():
+            for _ in range(warmup):
+                self.model(dummy)
+
+        if self.device.type == "cuda":
+            torch.cuda.synchronize()
+
+        # Benchmark
+        times = []
+        with torch.no_grad():
+            for _ in range(iterations):
+                start = time.perf_counter()
+                self.model(dummy)
+                if self.device.type == "cuda":
+                    torch.cuda.synchronize()
+                times.append(time.perf_counter() - start)
+
+        avg_ms = sum(times) / len(times) * 1000
+        fps = 1000.0 / avg_ms
+        params = sum(p.numel() for p in self.model.parameters())
+
+        results = {
+            "latency_ms": round(avg_ms, 2),
+            "fps": round(fps, 1),
+            "params": params,
+            "params_m": round(params / 1e6, 2),
+            "size_mb": round(params * 4 / 1024 / 1024, 2),
+            "device": str(self.device),
+            "input_size": self.input_size,
+        }
+
+        logger.info(f"Benchmark: {fps:.1f} FPS, {avg_ms:.2f}ms, {params / 1e6:.2f}M params")
+        return results

flashseg/analytics/profiler.py

@@ -0,0 +1,42 @@
+"""Layer-by-layer profiling."""
+
+import logging
+
+import torch
+
+from flashseg.cfg.config import get_config
+from flashseg.models.build import build_model
+
+logger = logging.getLogger(__name__)
+
+
+class Profiler:
+    """Profile FlashSeg model layer-by-layer."""
+
+    def __init__(self, model_path: str = None, model_size: str = "m", input_size: int = 512, num_classes: int = 21):
+        config = get_config(model_size=model_size, input_size=input_size, num_classes=num_classes)
+        self.model = build_model(config)
+
+        if model_path:
+            self.model.load_state_dict(torch.load(model_path, map_location="cpu"))
+
+        self.model.eval()
+        self.input_size = input_size
+
+    def run(self) -> dict:
+        """Profile model and print per-module statistics."""
+        dummy = torch.randn(1, 3, self.input_size, self.input_size)
+
+        results = {}
+        for name, module in self.model.named_children():
+            params = sum(p.numel() for p in module.parameters())
+            results[name] = {
+                "params": params,
+                "params_m": round(params / 1e6, 3),
+            }
+            print(f"  {name:20s} | {params:>10,} params | {params / 1e6:.3f}M")
+
+        total = sum(p.numel() for p in self.model.parameters())
+        print(f"  {'TOTAL':20s} | {total:>10,} params | {total / 1e6:.3f}M")
+        results["total"] = {"params": total, "params_m": round(total / 1e6, 3)}
+        return results

flashseg/cfg/__init__.py

@@ -0,0 +1,3 @@
+from flashseg.cfg.config import get_config, load_yaml_config
+
+__all__ = ["get_config", "load_yaml_config"]

flashseg/cfg/config.py

@@ -0,0 +1,113 @@
+"""Configuration management for FlashSeg."""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+
+
+MODEL_SIZE_MAP = {
+    "n": {"width_mult": 0.25, "depth_mult": 0.33},
+    "s": {"width_mult": 0.50, "depth_mult": 0.33},
+    "m": {"width_mult": 0.75, "depth_mult": 0.67},
+    "l": {"width_mult": 1.00, "depth_mult": 1.00},
+}
+
+
+@dataclass
+class Config:
+    """FlashSeg configuration."""
+
+    # Model
+    model_size: str = "m"
+    num_classes: int = 21
+    input_size: int = 512
+    width_mult: float = 0.75
+    depth_mult: float = 0.67
+    backbone: str = "shufflenetv2"
+    neck: str = "fpn"
+    head: str = "seg_head"
+
+    # Training
+    epochs: int = 100
+    batch_size: int = 16
+    lr: float = 0.01
+    momentum: float = 0.9
+    weight_decay: float = 5e-4
+    warmup_epochs: int = 5
+    scheduler: str = "cosine"
+    amp: bool = False
+    multi_gpu: bool = False
+
+    # Data
+    train_images: str = ""
+    train_masks: str = ""
+    val_images: str = ""
+    val_masks: str = ""
+    num_workers: int = 4
+    augment: bool = True
+
+    # LoRA
+    use_lora: bool = False
+    lora_rank: int = 8
+    lora_alpha: int = 16
+    lora_variant: str = "standard"
+
+    # Knowledge Distillation
+    use_kd: bool = False
+    teacher_checkpoint: str = ""
+    teacher_size: str = "l"
+    kd_temperature: float = 4.0
+    kd_alpha: float = 0.5
+
+    # Pretrained
+    pretrained: bool = True
+
+    # Paths
+    save_dir: str = "workspace"
+    device: str = "cuda"
+
+    # Extra
+    extra: Dict[str, Any] = field(default_factory=dict)
+
+
+def get_config(
+    model_size: str = "m",
+    input_size: int = 512,
+    num_classes: int = 21,
+    **overrides,
+) -> Config:
+    """Create a config with sensible defaults for the given model size."""
+    size_params = MODEL_SIZE_MAP.get(model_size, MODEL_SIZE_MAP["m"])
+
+    config = Config(
+        model_size=model_size,
+        input_size=input_size,
+        num_classes=num_classes,
+        width_mult=size_params["width_mult"],
+        depth_mult=size_params["depth_mult"],
+    )
+
+    for key, value in overrides.items():
+        if hasattr(config, key):
+            setattr(config, key, value)
+
+    return config
+
+
+def load_yaml_config(path: str) -> Config:
+    """Load configuration from a YAML file."""
+    with open(path, "r") as f:
+        data = yaml.safe_load(f)
+
+    model_size = data.pop("model_size", "m")
+    input_size = data.pop("input_size", 512)
+    num_classes = data.pop("num_classes", 21)
+
+    return get_config(
+        model_size=model_size,
+        input_size=input_size,
+        num_classes=num_classes,
+        **data,
+    )

flashseg/cli.py

@@ -0,0 +1,216 @@
+"""FlashSeg CLI."""
+
+import argparse
+import logging
+import sys
+
+from flashseg import __version__
+
+
+def main():
+    """FlashSeg CLI entry point."""
+    parser = argparse.ArgumentParser(
+        prog="flashseg",
+        description="FlashSeg — Ultra-lightweight real-time image segmentation",
+    )
+    subparsers = parser.add_subparsers(dest="command")
+
+    # Train
+    train_parser = subparsers.add_parser("train", help="Train a segmentation model")
+    train_parser.add_argument("--model-size", default="m", choices=["n", "s", "m", "l"])
+    train_parser.add_argument("--train-images", required=True)
+    train_parser.add_argument("--train-masks", required=True)
+    train_parser.add_argument("--val-images", required=True)
+    train_parser.add_argument("--val-masks", required=True)
+    train_parser.add_argument("--num-classes", type=int, default=21)
+    train_parser.add_argument("--input-size", type=int, default=512)
+    train_parser.add_argument("--epochs", type=int, default=100)
+    train_parser.add_argument("--batch-size", type=int, default=16)
+    train_parser.add_argument("--lr", type=float, default=0.01)
+    train_parser.add_argument("--device", default="cuda")
+    train_parser.add_argument("--save-dir", default="workspace")
+    train_parser.add_argument("--amp", action="store_true")
+    train_parser.add_argument("--lora", action="store_true")
+    train_parser.add_argument("--config", type=str, help="YAML config file")
+
+    # Predict
+    pred_parser = subparsers.add_parser("predict", help="Run segmentation inference")
+    pred_parser.add_argument("--model", required=True)
+    pred_parser.add_argument("--source", required=True)
+    pred_parser.add_argument("--model-size", default="m")
+    pred_parser.add_argument("--num-classes", type=int, default=21)
+    pred_parser.add_argument("--input-size", type=int, default=512)
+    pred_parser.add_argument("--device", default="cuda")
+    pred_parser.add_argument("--save-dir", default="output")
+
+    # Validate
+    val_parser = subparsers.add_parser("val", help="Validate model")
+    val_parser.add_argument("--model", required=True)
+    val_parser.add_argument("--val-images", required=True)
+    val_parser.add_argument("--val-masks", required=True)
+    val_parser.add_argument("--model-size", default="m")
+    val_parser.add_argument("--num-classes", type=int, default=21)
+    val_parser.add_argument("--input-size", type=int, default=512)
+    val_parser.add_argument("--device", default="cuda")
+
+    # Export
+    export_parser = subparsers.add_parser("export", help="Export to ONNX")
+    export_parser.add_argument("--model", required=True)
+    export_parser.add_argument("--output", default="model.onnx")
+    export_parser.add_argument("--model-size", default="m")
+    export_parser.add_argument("--num-classes", type=int, default=21)
+    export_parser.add_argument("--input-size", type=int, default=512)
+    export_parser.add_argument("--simplify", action="store_true")
+
+    # Utility commands
+    subparsers.add_parser("version", help="Print version")
+    subparsers.add_parser("check", help="Run health check")
+    subparsers.add_parser("settings", help="Show system info")
+
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+    if args.command == "version":
+        print(f"flashseg {__version__}")
+
+    elif args.command == "check":
+        _run_check()
+
+    elif args.command == "settings":
+        _show_settings()
+
+    elif args.command == "train":
+        from flashseg.engine.trainer import Trainer
+        trainer = Trainer(
+            model_size=args.model_size,
+            train_images=args.train_images,
+            train_masks=args.train_masks,
+            val_images=args.val_images,
+            val_masks=args.val_masks,
+            num_classes=args.num_classes,
+            input_size=args.input_size,
+            epochs=args.epochs,
+            batch_size=args.batch_size,
+            lr=args.lr,
+            device=args.device,
+            save_dir=args.save_dir,
+            amp=args.amp,
+            use_lora=args.lora,
+            config_path=args.config,
+        )
+        trainer.train()
+
+    elif args.command == "predict":
+        from flashseg.engine.predictor import Predictor
+        predictor = Predictor(
+            model_path=args.model,
+            model_size=args.model_size,
+            num_classes=args.num_classes,
+            input_size=args.input_size,
+            device=args.device,
+        )
+        predictor.predict_directory(args.source, save_dir=args.save_dir)
+
+    elif args.command == "val":
+        from flashseg.engine.validator import Validator
+        validator = Validator(
+            model_path=args.model,
+            val_images=args.val_images,
+            val_masks=args.val_masks,
+            model_size=args.model_size,
+            num_classes=args.num_classes,
+            input_size=args.input_size,
+            device=args.device,
+        )
+        validator.validate()
+
+    elif args.command == "export":
+        from flashseg.engine.exporter import Exporter
+        exporter = Exporter(
+            model_path=args.model,
+            model_size=args.model_size,
+            num_classes=args.num_classes,
+            input_size=args.input_size,
+        )
+        exporter.export(output=args.output, simplify=args.simplify)
+
+    else:
+        parser.print_help()
+
+
+def _run_check():
+    """Run health check."""
+    print("FlashSeg Health Check")
+    print("=" * 40)
+    checks = []
+
+    try:
+        import torch
+        checks.append(("PyTorch", f"{torch.__version__}"))
+        checks.append(("CUDA available", str(torch.cuda.is_available())))
+        if torch.cuda.is_available():
+            checks.append(("GPU", torch.cuda.get_device_name(0)))
+    except ImportError:
+        checks.append(("PyTorch", "NOT INSTALLED"))
+
+    try:
+        import cv2
+        checks.append(("OpenCV", cv2.__version__))
+    except ImportError:
+        checks.append(("OpenCV", "NOT INSTALLED"))
+
+    try:
+        import flashseg
+        checks.append(("FlashSeg", flashseg.__version__))
+    except Exception as e:
+        checks.append(("FlashSeg", f"ERROR: {e}"))
+
+    try:
+        from flashseg.models.build import build_model
+        from flashseg.cfg.config import get_config
+        config = get_config(model_size="m", input_size=512, num_classes=21)
+        model = build_model(config)
+        params = sum(p.numel() for p in model.parameters())
+        checks.append(("Model build", f"OK ({params:,} params)"))
+    except Exception as e:
+        checks.append(("Model build", f"FAILED: {e}"))
+
+    for name, status in checks:
+        print(f"  {name:20s}: {status}")
+
+    print("=" * 40)
+    print("All checks passed!" if all("NOT INSTALLED" not in s and "FAILED" not in s for _, s in checks) else "Some checks failed.")
+
+
+def _show_settings():
+    """Show system settings."""
+    import platform
+    print("FlashSeg System Info")
+    print("=" * 40)
+    print(f"  Python:    {platform.python_version()}")
+    print(f"  Platform:  {platform.platform()}")
+
+    try:
+        import torch
+        print(f"  PyTorch:   {torch.__version__}")
+        print(f"  CUDA:      {torch.version.cuda or 'N/A'}")
+        if torch.cuda.is_available():
+            print(f"  GPU:       {torch.cuda.get_device_name(0)}")
+            mem = torch.cuda.get_device_properties(0).total_mem / 1024**3
+            print(f"  GPU RAM:   {mem:.1f} GB")
+        else:
+            print("  GPU:       Not available")
+    except ImportError:
+        print("  PyTorch:   Not installed")
+
+    try:
+        import flashseg
+        print(f"  FlashSeg:  {flashseg.__version__}")
+    except ImportError:
+        pass
+    print("=" * 40)
+
+
+if __name__ == "__main__":
+    main()

flashseg/data/__init__.py

@@ -0,0 +1,4 @@
+from flashseg.data.dataset import SegmentationDataset
+from flashseg.data.transforms import get_train_transforms, get_val_transforms
+
+__all__ = ["SegmentationDataset", "get_train_transforms", "get_val_transforms"]

flashseg/data/dataset.py

@@ -0,0 +1,93 @@
+"""Segmentation dataset classes."""
+
+import logging
+from pathlib import Path
+from typing import Callable, List, Optional, Tuple
+
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+logger = logging.getLogger(__name__)
+
+
+class SegmentationDataset(Dataset):
+    """Dataset for semantic segmentation with image-mask pairs."""
+
+    SUPPORTED_FORMATS = (".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff")
+
+    def __init__(
+        self,
+        images_dir: str,
+        masks_dir: str,
+        input_size: int = 512,
+        num_classes: int = 21,
+        transform: Optional[Callable] = None,
+        augment: bool = False,
+    ):
+        self.images_dir = Path(images_dir)
+        self.masks_dir = Path(masks_dir)
+        self.input_size = input_size
+        self.num_classes = num_classes
+        self.transform = transform
+        self.augment = augment
+
+        self.image_files = sorted(
+            [f for f in self.images_dir.iterdir() if f.suffix.lower() in self.SUPPORTED_FORMATS]
+        )
+        self.mask_files = sorted(
+            [f for f in self.masks_dir.iterdir() if f.suffix.lower() in self.SUPPORTED_FORMATS]
+        )
+
+        assert len(self.image_files) == len(self.mask_files), (
+            f"Mismatch: {len(self.image_files)} images vs {len(self.mask_files)} masks"
+        )
+
+        logger.info(f"Loaded {len(self.image_files)} image-mask pairs")
+
+    def __len__(self) -> int:
+        return len(self.image_files)
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        image = cv2.imread(str(self.image_files[idx]))
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        mask = cv2.imread(str(self.mask_files[idx]), cv2.IMREAD_GRAYSCALE)
+
+        image = cv2.resize(image, (self.input_size, self.input_size))
+        mask = cv2.resize(mask, (self.input_size, self.input_size), interpolation=cv2.INTER_NEAREST)
+
+        if self.augment:
+            image, mask = self._augment(image, mask)
+
+        if self.transform:
+            image = self.transform(image)
+        else:
+            image = image.astype(np.float32) / 255.0
+            image = np.transpose(image, (2, 0, 1))
+            image = torch.from_numpy(image)
+
+        mask = torch.from_numpy(mask.astype(np.int64))
+        return image, mask
+
+    def _augment(self, image: np.ndarray, mask: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """Apply random augmentations."""
+        if np.random.random() > 0.5:
+            image = np.fliplr(image).copy()
+            mask = np.fliplr(mask).copy()
+
+        if np.random.random() > 0.5:
+            image = np.flipud(image).copy()
+            mask = np.flipud(mask).copy()
+
+        if np.random.random() > 0.5:
+            k = np.random.randint(1, 4)
+            image = np.rot90(image, k).copy()
+            mask = np.rot90(mask, k).copy()
+
+        # Color jitter (image only)
+        if np.random.random() > 0.5:
+            factor = np.random.uniform(0.8, 1.2)
+            image = np.clip(image * factor, 0, 255).astype(np.uint8)
+
+        return image, mask

flashseg/data/transforms.py

@@ -0,0 +1,34 @@
+"""Data transforms for segmentation."""
+
+from typing import Tuple
+
+import numpy as np
+import torch
+
+
+def get_train_transforms(input_size: int = 512):
+    """Get training transforms."""
+
+    def transform(image: np.ndarray) -> torch.Tensor:
+        image = image.astype(np.float32) / 255.0
+        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+        image = (image - mean) / std
+        image = np.transpose(image, (2, 0, 1))
+        return torch.from_numpy(image.copy())
+
+    return transform
+
+
+def get_val_transforms(input_size: int = 512):
+    """Get validation transforms."""
+
+    def transform(image: np.ndarray) -> torch.Tensor:
+        image = image.astype(np.float32) / 255.0
+        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+        image = (image - mean) / std
+        image = np.transpose(image, (2, 0, 1))
+        return torch.from_numpy(image.copy())
+
+    return transform

flashseg/engine/__init__.py

@@ -0,0 +1,6 @@
+from flashseg.engine.trainer import Trainer
+from flashseg.engine.predictor import Predictor
+from flashseg.engine.exporter import Exporter
+from flashseg.engine.validator import Validator
+
+__all__ = ["Trainer", "Predictor", "Exporter", "Validator"]

flashseg/engine/exporter.py

@@ -0,0 +1,68 @@
+"""FlashSeg ONNX exporter."""
+
+import logging
+from pathlib import Path
+
+import torch
+
+from flashseg.cfg.config import get_config
+from flashseg.models.build import build_model
+
+logger = logging.getLogger(__name__)
+
+
+class Exporter:
+    """Export FlashSeg models to ONNX format."""
+
+    def __init__(
+        self,
+        model_path: str,
+        model_size: str = "m",
+        num_classes: int = 21,
+        input_size: int = 512,
+    ):
+        self.model_path = model_path
+        self.model_size = model_size
+        self.num_classes = num_classes
+        self.input_size = input_size
+
+    def export(self, output: str = "model.onnx", simplify: bool = True, opset: int = 11) -> str:
+        """Export model to ONNX."""
+        return self.export_onnx(output, simplify, opset)
+
+    def export_onnx(self, output: str = "model.onnx", simplify: bool = True, opset: int = 11) -> str:
+        """Export to ONNX format."""
+        config = get_config(model_size=self.model_size, input_size=self.input_size, num_classes=self.num_classes)
+        model = build_model(config)
+        model.load_state_dict(torch.load(self.model_path, map_location="cpu"))
+        model.eval()
+
+        dummy_input = torch.randn(1, 3, self.input_size, self.input_size)
+
+        torch.onnx.export(
+            model,
+            dummy_input,
+            output,
+            opset_version=opset,
+            input_names=["images"],
+            output_names=["output"],
+            dynamic_axes={"images": {0: "batch"}, "output": {0: "batch"}},
+        )
+        logger.info(f"Exported ONNX model to {output}")
+
+        if simplify:
+            try:
+                import onnx
+                from onnxsim import simplify as onnx_simplify
+
+                model_onnx = onnx.load(output)
+                model_simple, check = onnx_simplify(model_onnx)
+                if check:
+                    onnx.save(model_simple, output)
+                    logger.info("ONNX model simplified")
+            except ImportError:
+                logger.warning("onnxsim not installed, skipping simplification")
+
+        file_size = Path(output).stat().st_size / 1024 / 1024
+        logger.info(f"ONNX model size: {file_size:.2f} MB")
+        return output