Skip to main content

A comprehensive computer vision framework for training, evaluation, and visualization with support for multiple architectures and GradCAM analysis

Project description

VisFlow

PyPI version Python Version License: MIT Downloads

A comprehensive computer vision framework for training, evaluation, and visualization

English | 中文


🚀 Features

  • 🎯 Easy Training: Simple YAML configuration for model training
  • 🔥 GradCAM Visualization: Built-in support for model interpretability
  • 🏗️ Multiple Architectures: Support for 50+ pre-trained models from torchvision
  • 🎨 Extensible: Easy to add custom models with the registration system
  • ⚡ CLI & Programmatic: Use via command line or Python API
  • 📊 Rich Logging: Beautiful terminal output with progress tracking

📦 Installation

pip install visflow

🎯 Quick Start

Training a Model

  1. Create a configuration file (config.yml):
model:
  architecture: resnet18
  pretrained: true
  num_classes: 2

training:
  device: cuda
  batch_size: 32
  epochs: 10
  learning_rate: 0.001
  optimizer: adam

data:
  train_dir: ./data/train
  val_dir: ./data/val
  test_dir: ./data/test

output:
  output_dir: ./output
  experiment_name: my-experiment
  1. Train via CLI:
visflow train --config config.yml
  1. Or train via Python API:
from visflow.resources.configs import TrainConfig
from visflow.pipelines.train import TrainPipeline

pipeline = TrainPipeline(TrainConfig.from_yaml('config.yml'))
pipeline()

GradCAM Visualization

visflow gradcam \
    --ckpt-path model.pth \
    --image-path image.jpg \
    --output-dir ./output \
    --target-layer layer4 \
    --colormap jet

🏗️ Supported Architectures

Visflow supports 50+ architectures from torchvision:

  • ResNet: resnet18, resnet34, resnet50, resnet101, resnet152, resnext50_32x4d, etc.
  • EfficientNet: efficientnet_b0 through efficientnet_b7, efficientnet_v2_s/m/l
  • Vision Transformers: vit_b_16, vit_b_32, vit_l_16, swin_t, swin_s, swin_b
  • MobileNet: mobilenet_v2, mobilenet_v3_small, mobilenet_v3_large
  • DenseNet: densenet121, densenet169, densenet201, densenet161
  • And many more: VGG, ConvNeXt, RegNet, MaxViT, etc.

🎨 Custom Models

Easily add your own models:

from visflow.resources.models import BaseClassifier, register_model

@register_model('my_custom_model')
class MyCustomModel(BaseClassifier):
    def __init__(self, num_classes: int):
        super().__init__(num_classes=num_classes)
        # Your model implementation
        
    def forward(self, x):
        # Forward pass implementation
        pass
        
    def last_conv(self):
        # Return the last convolutional layer for GradCAM
        return self.conv_layer

📖 CLI Reference

Training Command

visflow train [OPTIONS]

Options:
  -c, --config PATH     Path to training configuration file [required]
  -v, --verbose         Enable verbose logging
  --help                Show this message and exit

GradCAM Command

visflow gradcam [OPTIONS]

Options:
  -k, --ckpt-path PATH      Path to model checkpoint [required]
  -i, --image-path PATH     Path to input image [required]
  -o, --output-dir PATH     Output directory [default: ./output]
  -l, --target-layer TEXT   Target layer name
  -t, --target-class TEXT   Target class (index or name)
  -c, --colormap TEXT       Colormap [default: jet]
  --heatmap-only            Save only heatmap
  --eigen-smooth            Apply eigen smoothing
  --aug-smooth              Apply augmented smoothing
  -d, --device TEXT         Device (cpu/cuda)
  -v, --verbose             Enable verbose logging

📋 Configuration Reference

Complete configuration example
logging:
  backend: native  # Options: native, loguru
  loglevel: info   # Options: debug, info, warning, error, critical

seed: 42

model:
  architecture: resnet18
  pretrained: true
  num_classes: 2
  weights_path: ~  # Optional custom weights

training:
  device: cuda
  shuffle: true
  batch_size: 32
  weighted_sampling: false
  drop_last: false
  epochs: 10
  learning_rate: 0.001
  momentum: 0.9
  weight_decay: 0.0001
  optimizer: adam  # Options: sgd, adam, adamw
  lr_scheduler: ~  # Options: step, cosine, plateau
  early_stopping: true
  early_stopping_patience: 5
  label_smoothing: 0.0

testing:
  batch_size: 32

data:
  train_dir: ./data/train
  val_dir: ./data/val
  test_dir: ./data/test
  num_workers: 4
  pin_memory: false

resize:
  size: 224
  interpolation: bicubic
  antialias: true

normalization:
  enabled: true
  mean: [0.485, 0.456, 0.406]
  std: [0.229, 0.224, 0.225]

augmentation:
  horizontal_flip:
    enabled: true
    p: 0.5
  rotation:
    enabled: false
    degrees: 30
  color_jitter:
    enabled: false
    brightness: 0.2
    contrast: 0.2
    saturation: 0.2
    hue: 0.1
  # ... more augmentation options

output:
  output_dir: ./output
  experiment_name: vision-research
  checkpoint_frequency: 10

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

visflow-1.0.0.tar.gz (47.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

visflow-1.0.0-py3-none-any.whl (61.2 kB view details)

Uploaded Python 3

File details

Details for the file visflow-1.0.0.tar.gz.

File metadata

  • Download URL: visflow-1.0.0.tar.gz
  • Upload date:
  • Size: 47.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/2.1.4 CPython/3.11.13 Linux/6.11.0-1018-azure

File hashes

Hashes for visflow-1.0.0.tar.gz
Algorithm Hash digest
SHA256 913e0941c803884d405c93a4dc1b847d2fea8d645fffc064f04b3b621efcfde5
MD5 c37eb24d6ceff7a151c549979f33aff3
BLAKE2b-256 8d2af4eb4802a2cd89729f949953d50e7d348caa6259230cdcd439e000378b24

See more details on using hashes here.

File details

Details for the file visflow-1.0.0-py3-none-any.whl.

File metadata

  • Download URL: visflow-1.0.0-py3-none-any.whl
  • Upload date:
  • Size: 61.2 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/2.1.4 CPython/3.11.13 Linux/6.11.0-1018-azure

File hashes

Hashes for visflow-1.0.0-py3-none-any.whl
Algorithm Hash digest
SHA256 be3df01001d51c865d67e5aa5b57c1def587d8cd6c8d0c2ead980cf95e5a6f3d
MD5 4616037fd550e18ab5312b42e3be5783
BLAKE2b-256 ebd10868bb1fbd2f6c0358c9a1277a6b2bb78f03f0eaf05b39e3354fb9c58933

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page