flashstudio 0.1.0__py3-none-any.whl

flashstudio/pages/dashboard.py

@@ -0,0 +1,168 @@
+"""FlashStudio — Dashboard Page."""
+
+import os
+import json
+import glob as glob_module
+
+import streamlit as st
+from flashstudio.components.styles import render_page_header
+from flashstudio.utils.device import get_gpu_info, get_colab_runtime_type
+
+
+def render_dashboard():
+    """Render the main dashboard with overview and quick-start."""
+    render_page_header("⚡", "FlashStudio Dashboard",
+                       "End-to-end object detection: Data → Model → Train → Export → Inference")
+
+    # Top metrics row
+    col1, col2, col3, col4 = st.columns(4)
+    gpu = get_gpu_info()
+    with col1:
+        st.metric("🖥️ GPU", gpu["name"] if gpu["available"] else "CPU")
+    with col2:
+        st.metric("🌐 Runtime", get_colab_runtime_type())
+    with col3:
+        mem = f"{gpu['memory_total']:.1f} GB" if gpu["available"] else "—"
+        st.metric("💾 VRAM", mem)
+    with col4:
+        st.metric("📦 FlashDet", _get_flashdet_version())
+
+    st.markdown("<div style='height:1rem;'></div>", unsafe_allow_html=True)
+
+    # Main content
+    col_left, col_right = st.columns([2, 1])
+
+    with col_left:
+        with st.container(border=True):
+            st.markdown("### 🚀 Quick Start Guide")
+            st.markdown("""
+| Step | Action | Description |
+|:----:|--------|-------------|
+| 1 | **📦 Data** | Upload your dataset or download an open-source one |
+| 2 | **🧠 Model** | Choose architecture & configure training parameters |
+| 3 | **🏋️ Train** | Start training and monitor real-time progress |
+| 4 | **📤 Export** | Convert to ONNX for deployment |
+| 5 | **🔍 Inference** | Test your model on images, video, or RTSP |
+            """)
+            st.caption("Click **Next** below or use the sidebar to jump to any step.")
+
+    with col_right:
+        with st.container(border=True):
+            st.markdown("### 📊 Session Info")
+            dataset = st.session_state.get("dataset_name", "Not selected")
+            model = st.session_state.get("model_arch", "FlashDet-Pico")
+            status = st.session_state.get("training_status", "Not started")
+
+            st.markdown(f"**Dataset:** {dataset}")
+            st.markdown(f"**Model:** {model}")
+            st.markdown(f"**Training:** {status}")
+
+            if st.session_state.get("best_map"):
+                st.metric("Best mAP", f"{st.session_state['best_map']:.4f}")
+            if st.session_state.get("exported_model"):
+                st.markdown(f"**Exported:** {st.session_state['exported_model']}")
+
+    st.markdown("<div style='height:1rem;'></div>", unsafe_allow_html=True)
+
+    # Recent training runs
+    _render_recent_runs()
+
+    # Supported models table (correct data)
+    with st.container(border=True):
+        st.markdown("### 🔗 Supported Architectures")
+        models_data = [
+            ["FlashDet-Pico", "LiteBackbone (0.5x)", "~298K", "Ultra-fast", "Edge / MCU"],
+            ["FlashDet-Nano", "FlashBackbone (stem=32)", "~790K", "Very fast", "Embedded / IoT"],
+            ["FlashDet-Small", "FlashBackbone (stem=48)", "~1.8M", "Fast", "General purpose"],
+            ["FlashDet-Medium", "FlashBackbone (stem=64)", "~3.6M", "Balanced", "High accuracy"],
+            ["FlashDet-Large", "FlashBackbone (stem=80)", "~5.8M", "Accurate", "High accuracy"],
+            ["FlashDet-X", "FlashBackbone (stem=96)", "~9.0M", "Max accuracy", "Server"],
+            ["YOLOv8", "YOLOv8Backbone", "Varies", "Fast", "General YOLO"],
+            ["YOLOv9", "YOLOv9Backbone", "Varies", "Fast", "PGI-based"],
+            ["YOLOv10", "YOLOv10Backbone", "Varies", "Fast", "PSA-enhanced"],
+            ["YOLOv11", "YOLOv11Backbone", "Varies", "Fast", "C2PSA-based"],
+            ["YOLOX", "YOLOXBackbone", "Varies", "Fast", "Anchor-free"],
+        ]
+        st.dataframe(
+            {
+                "Model": [m[0] for m in models_data],
+                "Backbone": [m[1] for m in models_data],
+                "Params": [m[2] for m in models_data],
+                "Speed": [m[3] for m in models_data],
+                "Best For": [m[4] for m in models_data],
+            },
+            use_container_width=True,
+            hide_index=True,
+        )
+        st.caption("All architectures use fixed backbone/neck/head — only FlashDet-Pico allows backbone choice (lite vs pico_v2).")
+
+
+def _render_recent_runs():
+    """Show recent training runs from workspace."""
+    workspace_candidates = [
+        os.path.join(os.getcwd(), "workspace"),
+        os.path.join(os.getcwd(), "..", "FlashDet", "workspace"),
+        os.path.join(os.path.dirname(__file__), "..", "..", "..", "FlashDet", "workspace"),
+    ]
+
+    workspace = None
+    for c in workspace_candidates:
+        c = os.path.abspath(c)
+        if os.path.isdir(c):
+            workspace = c
+            break
+
+    if not workspace:
+        return
+
+    runs = sorted(
+        [d for d in os.listdir(workspace) if os.path.isdir(os.path.join(workspace, d))],
+        key=lambda d: os.path.getmtime(os.path.join(workspace, d)),
+        reverse=True,
+    )[:5]
+
+    if not runs:
+        return
+
+    with st.container(border=True):
+        st.markdown("### 🕐 Recent Training Runs")
+
+        for run_name in runs:
+            run_dir = os.path.join(workspace, run_name)
+            # Try to get results summary
+            results_path = os.path.join(run_dir, "results.json")
+            log_files = glob_module.glob(os.path.join(run_dir, "train_*.log"))
+
+            col_name, col_info, col_status = st.columns([2, 3, 1])
+            with col_name:
+                st.markdown(f"**{run_name}**")
+            with col_info:
+                if os.path.isfile(results_path):
+                    with open(results_path) as f:
+                        results = json.load(f)
+                    mAP = results.get("best_mAP50", 0)
+                    epochs = results.get("epochs_trained", "?")
+                    st.caption(f"mAP: {mAP:.4f} · {epochs} epochs")
+                elif log_files:
+                    st.caption("Training in progress / incomplete")
+                else:
+                    st.caption("No data")
+            with col_status:
+                has_best = os.path.isfile(os.path.join(run_dir, "checkpoint_best.pth"))
+                if has_best:
+                    st.caption("✅ Done")
+                elif log_files:
+                    st.caption("🔄 Active")
+                else:
+                    st.caption("—")
+
+
+def _get_flashdet_version() -> str:
+    try:
+        import importlib.util
+        if importlib.util.find_spec("flashdet"):
+            import flashdet
+            return getattr(flashdet, "__version__", "installed")
+    except Exception:
+        pass
+    return "not installed"

flashstudio/pages/data.py

@@ -0,0 +1,272 @@
+"""FlashStudio — Data Upload / Download Page."""
+
+import streamlit as st
+
+FLASHDET_NATIVE_DATASETS = [
+    {
+        "name": "Sample (demo)",
+        "description": "Built-in demo dataset — fastest way to test training pipeline",
+        "images": 80,
+        "classes": 10,
+        "size": "~5 MB",
+        "format": "COCO JSON",
+        "command": "flashdet download sample",
+        "native": True,
+    },
+    {
+        "name": "COCO 2017",
+        "description": "Full COCO 2017 — 118K train + 5K val images, 80 object classes",
+        "images": 123287,
+        "classes": 80,
+        "size": "~20 GB",
+        "format": "COCO JSON",
+        "command": "flashdet download coco2017",
+        "native": True,
+    },
+    {
+        "name": "COCO 2017 (val only)",
+        "description": "COCO 2017 validation set only — for quick evaluation",
+        "images": 5000,
+        "classes": 80,
+        "size": "~1 GB",
+        "format": "COCO JSON",
+        "command": "flashdet download coco2017-val",
+        "native": True,
+    },
+    {
+        "name": "VOC 2007",
+        "description": "Pascal VOC 2007 — classic detection benchmark, 20 classes",
+        "images": 9963,
+        "classes": 20,
+        "size": "~450 MB",
+        "format": "VOC XML → COCO JSON",
+        "command": "flashdet download voc2007",
+        "native": True,
+    },
+    {
+        "name": "VOC 2012",
+        "description": "Pascal VOC 2012 — extended VOC with more images",
+        "images": 17125,
+        "classes": 20,
+        "size": "~2 GB",
+        "format": "VOC XML → COCO JSON",
+        "command": "flashdet download voc2012",
+        "native": True,
+    },
+]
+
+EXTERNAL_DATASETS = [
+    {
+        "name": "COCO128 (tiny subset)",
+        "description": "128 COCO images — quick pipeline testing",
+        "images": 128,
+        "classes": 80,
+        "size": "7 MB",
+        "format": "COCO JSON",
+        "command": "wget -q https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128.zip && unzip -q coco128.zip -d datasets/",
+        "native": False,
+    },
+    {
+        "name": "VisDrone (drone aerial)",
+        "description": "Drone aerial images — pedestrians, vehicles from above",
+        "images": 6471,
+        "classes": 10,
+        "size": "~2 GB",
+        "format": "Custom → COCO JSON",
+        "command": "# See: https://github.com/VisDrone/VisDrone-Dataset",
+        "native": False,
+    },
+]
+
+
+def render_data_page():
+    """Render data upload and open-source dataset download page."""
+    from flashstudio.components.styles import render_page_header
+    render_page_header("📦", "Data Setup",
+                       "Upload your own dataset or download a ready-to-use open-source dataset.")
+
+    data_mode = st.radio(
+        "Choose data source",
+        ["📤 Upload Your Dataset", "📥 Download Open-Source Dataset"],
+        horizontal=True,
+        key="data_mode",
+    )
+
+    if data_mode == "📤 Upload Your Dataset":
+        _render_upload_tab()
+    else:
+        _render_download_tab()
+
+    st.divider()
+    _render_dataset_status()
+
+
+def _render_upload_tab():
+    """Upload custom dataset."""
+    st.markdown("### Upload Your Dataset")
+    st.info("Supported formats: **COCO JSON**, **Pascal VOC XML**, **YOLO TXT**. Upload as ZIP/TAR.")
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        with st.container(border=True):
+            st.markdown("**Training Data**")
+            train_upload = st.file_uploader(
+                "Upload training dataset (zip/tar)",
+                type=["zip", "tar", "gz"],
+                key="train_data_upload",
+            )
+            st.text_input(
+                "Or enter path to annotations",
+                placeholder="/content/datasets/train/annotations.json",
+                key="train_ann_path",
+            )
+            st.text_input(
+                "Images directory",
+                placeholder="/content/datasets/train/images/",
+                key="train_img_path",
+            )
+
+    with col2:
+        with st.container(border=True):
+            st.markdown("**Validation Data**")
+            val_upload = st.file_uploader(
+                "Upload validation dataset (zip/tar)",
+                type=["zip", "tar", "gz"],
+                key="val_data_upload",
+            )
+            st.text_input(
+                "Or enter path to annotations",
+                placeholder="/content/datasets/val/annotations.json",
+                key="val_ann_path",
+            )
+            st.text_input(
+                "Images directory",
+                placeholder="/content/datasets/val/images/",
+                key="val_img_path",
+            )
+
+    with st.container(border=True):
+        st.markdown("**Dataset Format**")
+        col_fmt, col_cls = st.columns(2)
+        with col_fmt:
+            st.selectbox(
+                "Annotation Format",
+                ["COCO JSON (recommended)", "Pascal VOC XML", "YOLO TXT"],
+                key="ann_format",
+            )
+        with col_cls:
+            st.number_input("Number of Classes", min_value=1, max_value=1000, value=80, key="upload_num_classes")
+        st.text_area(
+            "Class Names (one per line, optional)",
+            placeholder="person\ncar\nbicycle\n...",
+            height=100,
+            key="class_names",
+        )
+
+    if train_upload:
+        st.session_state["dataset_name"] = train_upload.name
+        st.success(f"✅ Training data uploaded: {train_upload.name} ({train_upload.size / 1e6:.1f} MB)")
+
+    if val_upload:
+        st.success(f"✅ Validation data uploaded: {val_upload.name} ({val_upload.size / 1e6:.1f} MB)")
+
+
+def _render_download_tab():
+    """Download datasets — FlashDet native first, then external."""
+    st.markdown("### FlashDet Native Datasets")
+    st.caption("These use `flashdet download` — auto-handled format conversion and split.")
+
+    for i, ds in enumerate(FLASHDET_NATIVE_DATASETS):
+        with st.container(border=True):
+            col_info, col_stats, col_action = st.columns([3, 2, 1])
+
+            with col_info:
+                st.markdown(f"**{ds['name']}**")
+                st.caption(ds["description"])
+
+            with col_stats:
+                st.markdown(
+                    f"📸 {ds['images']} images · 🏷️ {ds['classes']} classes · "
+                    f"💾 {ds['size']} · 📋 {ds['format']}"
+                )
+
+            with col_action:
+                if st.button("⬇️ Download", key=f"native_dl_{i}", use_container_width=True):
+                    st.session_state["dataset_name"] = ds["name"]
+                    st.session_state["dataset_command"] = ds["command"]
+                    st.session_state["dataset_classes"] = ds["classes"]
+                    st.rerun()
+
+    st.divider()
+    st.markdown("### External Datasets")
+    st.caption("Third-party datasets — may need manual format conversion.")
+
+    for i, ds in enumerate(EXTERNAL_DATASETS):
+        with st.container(border=True):
+            col_info, col_stats, col_action = st.columns([3, 2, 1])
+
+            with col_info:
+                st.markdown(f"**{ds['name']}**")
+                st.caption(ds["description"])
+
+            with col_stats:
+                st.markdown(
+                    f"📸 {ds['images']} images · 🏷️ {ds['classes']} classes · "
+                    f"💾 {ds['size']} · 📋 {ds['format']}"
+                )
+
+            with col_action:
+                if st.button("⬇️ Download", key=f"ext_dl_{i}", use_container_width=True):
+                    st.session_state["dataset_name"] = ds["name"]
+                    st.session_state["dataset_command"] = ds["command"]
+                    st.session_state["dataset_classes"] = ds["classes"]
+                    st.rerun()
+
+    if "dataset_command" in st.session_state and st.session_state.get("dataset_command"):
+        st.divider()
+        st.markdown("### 📋 Download Command")
+        st.code(st.session_state["dataset_command"], language="bash")
+        st.caption("Run this in a terminal or Colab cell.")
+
+        if st.button("🚀 Run Download", type="primary"):
+            _auto_download(st.session_state["dataset_command"])
+
+
+def _render_dataset_status():
+    """Show current dataset selection status."""
+    st.markdown("### ✅ Current Dataset")
+    dataset = st.session_state.get("dataset_name", None)
+
+    if dataset:
+        col1, col2, col3 = st.columns(3)
+        with col1:
+            st.metric("Dataset", dataset)
+        with col2:
+            num_cls = st.session_state.get("dataset_classes", st.session_state.get("upload_num_classes", 80))
+            st.metric("Classes", num_cls)
+        with col3:
+            fmt = st.session_state.get("ann_format", "COCO JSON")
+            st.metric("Format", fmt.split("(")[0].strip() if "(" in fmt else fmt)
+        st.success("Dataset ready! Click **Next** to choose your model.")
+    else:
+        st.warning("No dataset selected yet. Upload or download a dataset above.")
+
+
+def _auto_download(command: str):
+    """Execute download command (works in Colab)."""
+    import subprocess
+    with st.spinner("Downloading dataset..."):
+        try:
+            result = subprocess.run(
+                command, shell=True, capture_output=True, text=True, timeout=300
+            )
+            if result.returncode == 0:
+                st.success("✅ Dataset downloaded successfully!")
+                st.session_state["train_ann_path"] = "/content/datasets/"
+            else:
+                st.error(f"Download failed: {result.stderr[:200]}")
+        except subprocess.TimeoutExpired:
+            st.error("Download timed out (5 min limit). Try manually.")
+        except Exception as e:
+            st.error(f"Error: {e}")

flashstudio/pages/export.py

@@ -0,0 +1,212 @@
+"""FlashStudio — Model Export Page (matches actual FlashDet capabilities)."""
+
+import os
+import streamlit as st
+
+
+def render_export_page():
+    """Render model export/conversion page."""
+    from flashstudio.components.styles import render_page_header
+    render_page_header("📤", "Export Model",
+                       "Convert your trained model for production deployment.")
+
+    _render_training_summary()
+    st.divider()
+
+    col_config, col_output = st.columns([1, 1])
+
+    with col_config:
+        _render_export_config()
+
+    with col_output:
+        _render_export_output()
+
+
+def _render_training_summary():
+    """Show what weights are already available from training."""
+    save_dir = st.session_state.get("save_dir", "")
+
+    if save_dir and os.path.isdir(save_dir):
+        st.markdown("#### Available Weights from Training")
+        weights_files = [f for f in os.listdir(save_dir) if f.endswith(".pth")]
+        if weights_files:
+            cols = st.columns(min(len(weights_files), 4))
+            for i, wf in enumerate(sorted(weights_files)[:4]):
+                fpath = os.path.join(save_dir, wf)
+                size_mb = os.path.getsize(fpath) / (1024 * 1024)
+                with cols[i]:
+                    label = "Best" if "best" in wf else "Last" if "last" in wf else "Other"
+                    precision = "FP16" if "fp16" in wf else "FP32"
+                    st.metric(f"{label} ({precision})", f"{size_mb:.1f} MB")
+            st.caption(f"Directory: `{save_dir}`")
+        else:
+            st.info("No weights found in save directory.")
+    else:
+        st.info("Complete training first, or provide a weights path below.")
+
+
+def _render_export_config():
+    """Export configuration — only options FlashDet actually supports."""
+    st.markdown("### ⚙️ Export Settings")
+
+    with st.container(border=True):
+        st.markdown("**Model Weights**")
+        weights_source = st.radio(
+            "Source",
+            ["Best model (inference weights)", "Best model (FP16)",
+             "Last model (inference weights)", "Custom path"],
+            key="weights_source",
+        )
+        if weights_source == "Custom path":
+            st.text_input("Weights path (.pth)", placeholder="/path/to/model.pth",
+                          key="export_weights_path")
+
+        save_dir = st.session_state.get("save_dir", "/content/flashstudio_runs")
+        weights_map = {
+            "Best model (inference weights)": "model_best_inference.pth",
+            "Best model (FP16)": "model_best_fp16.pth",
+            "Last model (inference weights)": "model_last_inference.pth",
+        }
+        if weights_source in weights_map:
+            weights_path = os.path.join(save_dir, weights_map[weights_source])
+        else:
+            weights_path = st.session_state.get("export_weights_path", "")
+
+        st.caption(f"📁 `{weights_path}`")
+
+    with st.container(border=True):
+        st.markdown("**Export Format**")
+        st.markdown("FlashDet supports **ONNX** export with dynamic batch size (opset 13).")
+
+        st.checkbox("ONNX (recommended — universal format)", value=True, key="export_onnx",
+                    disabled=True)
+        st.caption("Works with: ONNX Runtime, OpenVINO, TensorRT (via onnx2trt), TFLite (via onnx2tf)")
+
+        st.divider()
+        st.caption("**Not natively supported** (convert from ONNX using external tools):")
+        st.markdown("""
+        - TensorRT: `trtexec --onnx=model.onnx --saveEngine=model.engine --fp16`
+        - OpenVINO: `mo --input_model model.onnx --output_dir openvino/`
+        - CoreML: `python -m coremltools.converters.onnx model.onnx`
+        - NCNN: `onnx2ncnn model.onnx model.param model.bin`
+        """)
+
+    with st.container(border=True):
+        st.markdown("**ONNX Export Options**")
+        img_size = st.select_slider(
+            "Input Size", [320, 416, 640], value=320, key="export_img_size",
+            help="Should match your training input size"
+        )
+        st.checkbox("Dynamic batch size", value=True, key="export_dynamic",
+                    help="Allows variable batch size at inference (default in FlashDet)")
+        st.number_input("Opset Version", 11, 18, 13, key="export_opset",
+                        help="FlashDet uses opset 13 by default")
+
+    if st.button("🚀 Export to ONNX", use_container_width=True, type="primary"):
+        _run_export(weights_path, img_size)
+
+
+def _render_export_output():
+    """Show export results and commands."""
+    st.markdown("### 📦 Export Output")
+
+    exported = st.session_state.get("exported_files", [])
+
+    if not exported:
+        st.markdown("#### Pre-Generated Weights")
+        st.info(
+            "FlashDet **automatically saves** FP16 inference weights during training:\n\n"
+            "- `model_best_inference.pth` — best mAP, FP32, o2m head stripped\n"
+            "- `model_best_fp16.pth` — best mAP, FP16 (half memory)\n"
+            "- `model_last_inference.pth` — final epoch, FP32\n"
+            "- `model_last_fp16.pth` — final epoch, FP16\n\n"
+            "These are ready to use with `FlashDetPredictor` — no export step needed for PyTorch inference."
+        )
+
+        st.divider()
+        st.markdown("#### ONNX Export Command (Manual)")
+        weights = st.session_state.get("export_weights_path", "model_best_inference.pth")
+        img_size = st.session_state.get("export_img_size", 320)
+        st.code(
+            f"flashdet export --model {weights} --output model.onnx\n"
+            f"# Output: model.onnx (opset 13, dynamic batch)",
+            language="bash"
+        )
+        return
+
+    for exp in exported:
+        with st.container(border=True):
+            col_info, col_action = st.columns([3, 1])
+            with col_info:
+                st.markdown(f"**{exp['format']}** — `{exp['path']}`")
+                if exp.get("size"):
+                    st.caption(f"Size: {exp['size']}")
+                if exp.get("success"):
+                    st.success("Export successful")
+                else:
+                    st.error("Export failed")
+            with col_action:
+                if exp.get("success") and os.path.isfile(exp.get("path", "")):
+                    with open(exp["path"], "rb") as f:
+                        st.download_button("📥 Download", f, file_name=os.path.basename(exp["path"]),
+                                           key=f"dl_{exp['format']}")
+
+    st.divider()
+    st.success("✅ Export complete! Click **Next** to test inference.")
+
+
+def _run_export(weights_path: str, img_size: int):
+    """Run ONNX export."""
+    if not weights_path:
+        st.error("Please select or provide weights path.")
+        return
+
+    if not os.path.isfile(weights_path):
+        st.warning(f"Weights file not found: `{weights_path}`. Attempting export anyway...")
+
+    import importlib.util
+    if importlib.util.find_spec("flashdet") is not None:
+        _flashdet_export(weights_path, img_size)
+    else:
+        st.warning("FlashDet not installed. Run export manually:")
+        st.code(f"flashdet export --model {weights_path} --output model.onnx", language="bash")
+        st.session_state["exported_files"] = [{
+            "format": "ONNX", "path": weights_path.replace(".pth", ".onnx"),
+            "size": "—", "success": False
+        }]
+        st.rerun()
+
+
+def _flashdet_export(weights_path: str, img_size: int):
+    """Actual FlashDet ONNX export."""
+    import subprocess
+    import sys
+
+    output_path = weights_path.replace(".pth", ".onnx")
+
+    cmd = [
+        sys.executable, "-m", "flashdet.cli", "export",
+        "--model", weights_path,
+        "--output", output_path,
+    ]
+
+    with st.spinner("Exporting to ONNX..."):
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
+
+    if result.returncode == 0:
+        size_str = "—"
+        if os.path.isfile(output_path):
+            size_mb = os.path.getsize(output_path) / (1024 * 1024)
+            size_str = f"{size_mb:.1f} MB"
+
+        st.session_state["exported_files"] = [{
+            "format": "ONNX", "path": output_path, "size": size_str, "success": True
+        }]
+        st.toast("ONNX export complete!", icon="✅")
+    else:
+        st.error(f"Export failed:\n{result.stderr[:500]}")
+        st.session_state["exported_files"] = [{
+            "format": "ONNX", "path": output_path, "size": "—", "success": False
+        }]
+
+    st.rerun()