alchemydetect 0.1.0__py3-none-any.whl

alchemydetect/app.py

@@ -0,0 +1,23 @@
+"""Entry point for the AlchemyDetect application."""
+
+import multiprocessing
+import sys
+
+from PyQt6.QtWidgets import QApplication
+
+from alchemydetect.gui.main_window import MainWindow
+
+
+def main():
+    app = QApplication(sys.argv)
+    app.setApplicationName("AlchemyDetect")
+
+    window = MainWindow()
+    window.show()
+
+    sys.exit(app.exec())
+
+
+if __name__ == "__main__":
+    multiprocessing.freeze_support()
+    main()

alchemydetect/core/config_builder.py

@@ -0,0 +1,90 @@
+"""Build Detectron2 config from user-specified parameters."""
+
+import torch
+from detectron2 import model_zoo
+from detectron2.config import get_cfg
+
+from .dataset_utils import get_num_classes, register_coco_dataset
+from .model_catalog import get_config_path
+
+
+def build_cfg(
+    model_name,
+    train_images_dir,
+    train_json,
+    output_dir,
+    lr=0.0025,
+    max_iter=1000,
+    batch_size=2,
+    val_images_dir=None,
+    val_json=None,
+    resume=False,
+    weights_path=None,
+):
+    """Build a Detectron2 CfgNode from user selections.
+
+    Args:
+        model_name: Key from MODEL_ZOO (e.g. "Faster R-CNN (R50-FPN)")
+        train_images_dir: Path to training images directory
+        train_json: Path to COCO JSON annotations for training
+        output_dir: Directory to save checkpoints and logs
+        lr: Base learning rate
+        max_iter: Maximum training iterations
+        batch_size: Images per batch
+        val_images_dir: Optional path to validation images directory
+        val_json: Optional path to validation COCO JSON
+        resume: Whether to resume from last checkpoint in output_dir
+        weights_path: Optional path to custom weights (.pth). If None, uses model zoo pretrained.
+
+    Returns:
+        CfgNode ready for training.
+    """
+    config_path = get_config_path(model_name)
+    num_classes = get_num_classes(train_json)
+
+    # Register datasets
+    train_dataset_name = "alchemy_train"
+    register_coco_dataset(train_dataset_name, train_json, train_images_dir)
+
+    val_dataset_name = None
+    if val_json and val_images_dir:
+        val_dataset_name = "alchemy_val"
+        register_coco_dataset(val_dataset_name, val_json, val_images_dir)
+
+    cfg = get_cfg()
+    cfg.merge_from_file(model_zoo.get_config_file(config_path))
+
+    # Datasets
+    cfg.DATASETS.TRAIN = (train_dataset_name,)
+    cfg.DATASETS.TEST = (val_dataset_name,) if val_dataset_name else ()
+
+    # Dataloader
+    cfg.DATALOADER.NUM_WORKERS = 2
+
+    # Model weights
+    if weights_path:
+        cfg.MODEL.WEIGHTS = weights_path
+    else:
+        cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(config_path)
+
+    # Solver
+    cfg.SOLVER.IMS_PER_BATCH = batch_size
+    cfg.SOLVER.BASE_LR = lr
+    cfg.SOLVER.MAX_ITER = max_iter
+    cfg.SOLVER.STEPS = []  # No LR decay for simplicity
+    cfg.SOLVER.CHECKPOINT_PERIOD = max(max_iter // 5, 100)
+
+    # Number of classes
+    if hasattr(cfg.MODEL, "ROI_HEADS"):
+        cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes
+    if hasattr(cfg.MODEL, "RETINANET"):
+        cfg.MODEL.RETINANET.NUM_CLASSES = num_classes
+
+    # Device
+    cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Output
+    cfg.OUTPUT_DIR = output_dir
+
+    cfg.freeze()
+    return cfg

alchemydetect/core/dataset_utils.py

@@ -0,0 +1,62 @@
+"""COCO dataset registration helpers for Detectron2."""
+
+import json
+from pathlib import Path
+
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from detectron2.data.datasets import register_coco_instances
+
+
+def register_coco_dataset(name, json_path, image_root):
+    """Register a COCO-format dataset, skipping if already registered."""
+    if name in DatasetCatalog.list():
+        DatasetCatalog.remove(name)
+        MetadataCatalog.remove(name)
+    register_coco_instances(name, {}, json_path, image_root)
+
+
+def validate_coco_json(json_path, image_root):
+    """Validate a COCO JSON file. Returns (is_valid, error_message)."""
+    json_path = Path(json_path)
+    image_root = Path(image_root)
+
+    if not json_path.exists():
+        return False, f"Annotation file not found: {json_path}"
+
+    if not image_root.exists():
+        return False, f"Image directory not found: {image_root}"
+
+    try:
+        with open(json_path, "r") as f:
+            data = json.load(f)
+    except json.JSONDecodeError as e:
+        return False, f"Invalid JSON: {e}"
+
+    for key in ("images", "annotations", "categories"):
+        if key not in data:
+            return False, f"Missing required key '{key}' in COCO JSON"
+
+    if len(data["categories"]) == 0:
+        return False, "No categories found in COCO JSON"
+
+    if len(data["images"]) == 0:
+        return False, "No images found in COCO JSON"
+
+    # Check that at least some image files exist
+    found = 0
+    for img_info in data["images"][:10]:
+        img_file = image_root / img_info["file_name"]
+        if img_file.exists():
+            found += 1
+
+    if found == 0:
+        return False, "None of the first 10 image files were found in the image directory"
+
+    return True, f"Valid COCO dataset: {len(data['images'])} images, {len(data['categories'])} categories"
+
+
+def get_num_classes(json_path):
+    """Return the number of categories in a COCO JSON file."""
+    with open(json_path, "r") as f:
+        data = json.load(f)
+    return len(data["categories"])

alchemydetect/core/inferencer.py

@@ -0,0 +1,67 @@
+"""Inference wrapper around Detectron2's DefaultPredictor."""
+
+import cv2
+from detectron2.config import get_cfg
+from detectron2.data import MetadataCatalog
+from detectron2.engine import DefaultPredictor
+from detectron2.utils.visualizer import ColorMode, Visualizer
+
+
+def load_predictor(config_yaml_path, weights_path, threshold=0.5):
+    """Load a Detectron2 predictor from a saved config + weights pair.
+
+    Args:
+        config_yaml_path: Path to the saved config.yaml file.
+        weights_path: Path to the .pth weights file.
+        threshold: Confidence threshold for predictions.
+
+    Returns:
+        (DefaultPredictor, CfgNode)
+    """
+    cfg = get_cfg()
+    cfg.merge_from_file(config_yaml_path)
+    cfg.MODEL.WEIGHTS = weights_path
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = threshold
+    if hasattr(cfg.MODEL, "RETINANET"):
+        cfg.MODEL.RETINANET.SCORE_THRESH_TEST = threshold
+    cfg.freeze()
+
+    predictor = DefaultPredictor(cfg)
+    return predictor, cfg
+
+
+def run_inference_single(predictor, image_path):
+    """Run inference on a single image.
+
+    Args:
+        predictor: A DefaultPredictor instance.
+        image_path: Path to the image file.
+
+    Returns:
+        (original_image_bgr, instances) where instances is the prediction output.
+    """
+    img = cv2.imread(str(image_path))
+    if img is None:
+        raise ValueError(f"Could not read image: {image_path}")
+    outputs = predictor(img)
+    return img, outputs["instances"].to("cpu")
+
+
+def visualize_predictions(image_bgr, instances, metadata=None):
+    """Draw predictions on an image using Detectron2's Visualizer.
+
+    Args:
+        image_bgr: Original image in BGR format (numpy array).
+        instances: Detectron2 Instances object (on CPU).
+        metadata: Optional MetadataCatalog metadata for class names.
+
+    Returns:
+        Annotated image as RGB numpy array.
+    """
+    image_rgb = image_bgr[:, :, ::-1]
+    if metadata is None:
+        metadata = MetadataCatalog.get("__empty")
+
+    v = Visualizer(image_rgb, metadata=metadata, scale=1.0, instance_mode=ColorMode.IMAGE)
+    out = v.draw_instance_predictions(instances)
+    return out.get_image()

alchemydetect/core/model_catalog.py

@@ -0,0 +1,56 @@
+"""Available Detectron2 model zoo entries."""
+
+# Maps user-friendly name -> (model_zoo config path, task type)
+MODEL_ZOO = {
+    # Object Detection
+    "Faster R-CNN (R50-FPN)": {
+        "config": "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml",
+        "task": "detection",
+    },
+    "Faster R-CNN (R101-FPN)": {
+        "config": "COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml",
+        "task": "detection",
+    },
+    "RetinaNet (R50-FPN)": {
+        "config": "COCO-Detection/retinanet_R_50_FPN_3x.yaml",
+        "task": "detection",
+    },
+    "RetinaNet (R101-FPN)": {
+        "config": "COCO-Detection/retinanet_R_101_FPN_3x.yaml",
+        "task": "detection",
+    },
+    # Instance Segmentation
+    "Mask R-CNN (R50-FPN)": {
+        "config": "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml",
+        "task": "instance_segmentation",
+    },
+    "Mask R-CNN (R101-FPN)": {
+        "config": "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml",
+        "task": "instance_segmentation",
+    },
+}
+
+
+def get_model_names():
+    """Return list of all available model names."""
+    return list(MODEL_ZOO.keys())
+
+
+def get_detection_models():
+    """Return names of detection-only models."""
+    return [k for k, v in MODEL_ZOO.items() if v["task"] == "detection"]
+
+
+def get_segmentation_models():
+    """Return names of instance segmentation models."""
+    return [k for k, v in MODEL_ZOO.items() if v["task"] == "instance_segmentation"]
+
+
+def get_config_path(model_name):
+    """Return the model zoo config path for a given model name."""
+    return MODEL_ZOO[model_name]["config"]
+
+
+def get_task(model_name):
+    """Return the task type ('detection' or 'instance_segmentation')."""
+    return MODEL_ZOO[model_name]["task"]

alchemydetect/core/trainer.py

@@ -0,0 +1,67 @@
+"""Custom Detectron2 trainer with metric emission for the GUI."""
+
+import logging
+import sys
+
+from detectron2.engine import DefaultTrainer, HookBase
+
+
+class MetricEmitterHook(HookBase):
+    """Hook that pushes training metrics to a multiprocessing.Queue."""
+
+    def __init__(self, queue, stop_event, period=20):
+        """
+        Args:
+            queue: multiprocessing.Queue to push metric dicts into.
+            stop_event: multiprocessing.Event; when set, training stops.
+            period: Emit metrics every N iterations.
+        """
+        self._queue = queue
+        self._stop_event = stop_event
+        self._period = period
+
+    def after_step(self):
+        # Check for stop request
+        if self._stop_event.is_set():
+            self._queue.put({"type": "log", "msg": "Training stopped by user."})
+            sys.exit(0)
+
+        iter_num = self.trainer.iter
+        if (iter_num + 1) % self._period == 0:
+            storage = self.trainer.storage
+            metrics = {}
+            for k, v in storage.latest().items():
+                if isinstance(v, tuple):
+                    metrics[k] = v[0]
+            metrics["iter"] = iter_num + 1
+            metrics["type"] = "metrics"
+            self._queue.put(metrics)
+
+
+class QueueLogHandler(logging.Handler):
+    """Logging handler that sends log records to a multiprocessing.Queue."""
+
+    def __init__(self, queue):
+        super().__init__()
+        self._queue = queue
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            self._queue.put({"type": "log", "msg": msg})
+        except Exception:
+            pass
+
+
+class AlchemyTrainer(DefaultTrainer):
+    """Detectron2 trainer that emits metrics to a queue for GUI consumption."""
+
+    def __init__(self, cfg, metric_queue, stop_event):
+        self._metric_queue = metric_queue
+        self._stop_event = stop_event
+        super().__init__(cfg)
+
+    def build_hooks(self):
+        hooks = super().build_hooks()
+        hooks.append(MetricEmitterHook(self._metric_queue, self._stop_event))
+        return hooks

alchemydetect/gui/dialogs.py

@@ -0,0 +1,81 @@
+"""File dialogs for model save/load and dataset selection."""
+
+import shutil
+from pathlib import Path
+
+from PyQt6.QtWidgets import QFileDialog, QMessageBox
+
+
+def browse_directory(parent, title="Select Directory", start_dir=""):
+    """Open a directory picker dialog. Returns path string or empty string."""
+    path = QFileDialog.getExistingDirectory(parent, title, start_dir)
+    return path
+
+
+def browse_file(parent, title="Select File", start_dir="", filter_str="All Files (*)"):
+    """Open a file picker dialog. Returns path string or empty string."""
+    path, _ = QFileDialog.getOpenFileName(parent, title, start_dir, filter_str)
+    return path
+
+
+def save_model_dialog(parent, output_dir):
+    """Save trained model (.pth + config.yaml) to a user-chosen directory.
+
+    Args:
+        parent: Parent widget.
+        output_dir: The training output directory containing model_final.pth and config.yaml.
+
+    Returns:
+        Path to the saved directory, or None if cancelled.
+    """
+    output_path = Path(output_dir)
+    weights_file = output_path / "model_final.pth"
+    config_file = output_path / "config.yaml"
+
+    if not weights_file.exists():
+        QMessageBox.warning(parent, "Save Model", "No model_final.pth found. Train a model first.")
+        return None
+
+    dest_dir = QFileDialog.getExistingDirectory(parent, "Save Model To Directory")
+    if not dest_dir:
+        return None
+
+    dest = Path(dest_dir)
+    try:
+        shutil.copy2(weights_file, dest / "model_final.pth")
+        if config_file.exists():
+            shutil.copy2(config_file, dest / "config.yaml")
+        QMessageBox.information(parent, "Save Model", f"Model saved to:\n{dest}")
+        return str(dest)
+    except Exception as e:
+        QMessageBox.critical(parent, "Save Model Error", str(e))
+        return None
+
+
+def load_model_dialog(parent):
+    """Open dialogs to select a model weights file and its config.
+
+    Returns:
+        (config_yaml_path, weights_path) or (None, None) if cancelled.
+    """
+    weights_path, _ = QFileDialog.getOpenFileName(
+        parent, "Select Model Weights", "", "PyTorch Weights (*.pth);;All Files (*)"
+    )
+    if not weights_path:
+        return None, None
+
+    # Try to auto-find config.yaml in the same directory
+    weights_dir = Path(weights_path).parent
+    auto_config = weights_dir / "config.yaml"
+
+    if auto_config.exists():
+        return str(auto_config), weights_path
+
+    # Ask user to select config manually
+    config_path, _ = QFileDialog.getOpenFileName(
+        parent, "Select Config YAML", str(weights_dir), "YAML Files (*.yaml *.yml);;All Files (*)"
+    )
+    if not config_path:
+        return None, None
+
+    return config_path, weights_path

alchemydetect/gui/image_viewer.py

@@ -0,0 +1,61 @@
+"""Image viewer widget that displays images with detection overlays."""
+
+from PyQt6.QtCore import Qt
+from PyQt6.QtGui import QImage, QPixmap
+from PyQt6.QtWidgets import QLabel, QScrollArea, QVBoxLayout, QWidget
+
+
+class ImageViewer(QWidget):
+    """Scrollable image display widget."""
+
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self._label = QLabel()
+        self._label.setAlignment(Qt.AlignmentFlag.AlignCenter)
+
+        scroll = QScrollArea()
+        scroll.setWidget(self._label)
+        scroll.setWidgetResizable(True)
+
+        layout = QVBoxLayout(self)
+        layout.setContentsMargins(0, 0, 0, 0)
+        layout.addWidget(scroll)
+
+        self._current_pixmap = None
+
+    def set_image_rgb(self, image_rgb):
+        """Display an RGB numpy array image.
+
+        Args:
+            image_rgb: numpy array of shape (H, W, 3) with dtype uint8.
+        """
+        if image_rgb is None:
+            self._label.clear()
+            self._current_pixmap = None
+            return
+
+        h, w, ch = image_rgb.shape
+        bytes_per_line = ch * w
+        qimg = QImage(image_rgb.data, w, h, bytes_per_line, QImage.Format.Format_RGB888)
+        self._current_pixmap = QPixmap.fromImage(qimg)
+        self._update_display()
+
+    def _update_display(self):
+        """Scale pixmap to fit the label while keeping aspect ratio."""
+        if self._current_pixmap is None:
+            return
+        scaled = self._current_pixmap.scaled(
+            self._label.size(),
+            Qt.AspectRatioMode.KeepAspectRatio,
+            Qt.TransformationMode.SmoothTransformation,
+        )
+        self._label.setPixmap(scaled)
+
+    def resizeEvent(self, event):
+        super().resizeEvent(event)
+        self._update_display()
+
+    def clear_image(self):
+        """Clear the displayed image."""
+        self._label.clear()
+        self._current_pixmap = None