flashstudio 0.0.1__py3-none-any.whl

flashstudio/components/project_manager.py

@@ -0,0 +1,490 @@
+"""FlashStudio — Project Management System.
+
+Provides project isolation: each project has its own dataset, model config,
+training runs, exports, and inference results. Users can create, switch,
+duplicate, and delete projects.
+"""
+
+import os
+import json
+import shutil
+from datetime import datetime
+from typing import Optional
+
+import streamlit as st
+
+PROJECTS_ROOT = os.path.join(os.path.expanduser("~"), ".flashstudio", "projects")
+PROJECTS_INDEX = os.path.join(os.path.expanduser("~"), ".flashstudio", "projects.json")
+
+
+def _ensure_dirs():
+    os.makedirs(PROJECTS_ROOT, exist_ok=True)
+
+
+def _load_index() -> dict:
+    """Load the global project index."""
+    _ensure_dirs()
+    if os.path.isfile(PROJECTS_INDEX):
+        try:
+            with open(PROJECTS_INDEX, "r") as f:
+                data = json.load(f)
+            if isinstance(data, dict) and "projects" in data:
+                return data
+        except (json.JSONDecodeError, OSError):
+            pass
+    return {"projects": [], "active": None}
+
+
+def _save_index(index: dict):
+    _ensure_dirs()
+    with open(PROJECTS_INDEX, "w") as f:
+        json.dump(index, f, indent=2)
+
+
+def get_project_dir(project_id: str) -> str:
+    return os.path.join(PROJECTS_ROOT, project_id)
+
+
+def list_projects() -> list:
+    """Return all projects with metadata."""
+    index = _load_index()
+    projects = []
+    for p in index.get("projects", []):
+        proj_dir = get_project_dir(p["id"])
+        p["exists"] = os.path.isdir(proj_dir)
+        if p["exists"]:
+            p["size"] = _dir_size_str(proj_dir)
+            runs_dir = os.path.join(proj_dir, "runs")
+            p["num_runs"] = len(os.listdir(runs_dir)) if os.path.isdir(runs_dir) else 0
+        projects.append(p)
+    return projects
+
+
+def get_active_project() -> Optional[dict]:
+    """Get the currently active project, or None."""
+    index = _load_index()
+    active_id = index.get("active")
+    if not active_id:
+        return None
+    for p in index.get("projects", []):
+        if p["id"] == active_id:
+            return p
+    return None
+
+
+def create_project(name: str, description: str = "") -> dict:
+    """Create a new project and set it as active."""
+    index = _load_index()
+    project_id = _slugify(name) + "_" + datetime.now().strftime("%Y%m%d_%H%M%S")
+    proj_dir = get_project_dir(project_id)
+
+    os.makedirs(proj_dir, exist_ok=True)
+    os.makedirs(os.path.join(proj_dir, "data"), exist_ok=True)
+    os.makedirs(os.path.join(proj_dir, "runs"), exist_ok=True)
+    os.makedirs(os.path.join(proj_dir, "exports"), exist_ok=True)
+    os.makedirs(os.path.join(proj_dir, "inference_results"), exist_ok=True)
+
+    project = {
+        "id": project_id,
+        "name": name,
+        "description": description,
+        "created": datetime.now().isoformat(),
+        "last_modified": datetime.now().isoformat(),
+    }
+
+    # Save project config
+    with open(os.path.join(proj_dir, "project.json"), "w") as f:
+        json.dump(project, f, indent=2)
+
+    index["projects"].append(project)
+    index["active"] = project_id
+    _save_index(index)
+
+    return project
+
+
+def switch_project(project_id: str):
+    """Switch active project and load its state."""
+    _clear_session_for_new_project()
+    index = _load_index()
+    index["active"] = project_id
+    _save_index(index)
+    load_project_state(project_id)
+    st.session_state["_project_state_loaded"] = True
+
+
+def delete_project(project_id: str):
+    """Delete a project and all its data."""
+    index = _load_index()
+    index["projects"] = [p for p in index["projects"] if p["id"] != project_id]
+    if index.get("active") == project_id:
+        index["active"] = index["projects"][0]["id"] if index["projects"] else None
+    _save_index(index)
+
+    proj_dir = get_project_dir(project_id)
+    if os.path.isdir(proj_dir):
+        shutil.rmtree(proj_dir)
+
+
+def delete_all_projects():
+    """Delete all projects and reset everything."""
+    index = _load_index()
+    for p in index.get("projects", []):
+        proj_dir = get_project_dir(p["id"])
+        if os.path.isdir(proj_dir):
+            shutil.rmtree(proj_dir)
+    _save_index({"projects": [], "active": None})
+
+
+def save_project_state():
+    """Save current session state to the active project."""
+    index = _load_index()
+    active_id = index.get("active")
+    if not active_id:
+        return
+
+    proj_dir = get_project_dir(active_id)
+    if not os.path.isdir(proj_dir):
+        return
+
+    # Save relevant session state keys
+    state_keys = [
+        "dataset_name", "dataset_classes", "dataset_id", "dataset_output_path",
+        "train_img_path", "val_img_path", "train_ann_path", "val_ann_path",
+        "ann_format", "class_names", "upload_num_classes", "num_classes",
+        "model_arch", "arch_family", "pico_backbone",
+        "epochs", "batch_size", "lr", "img_size",
+        "optimizer", "weight_decay", "val_interval", "lr_final_ratio",
+        "finetune_strategy", "pretrain_option", "custom_weights",
+        "lora_variant", "lora_rank", "lora_alpha", "lora_dropout", "lora_targets",
+        "qlora", "qlora_dtype", "warmup_epochs", "patience", "num_workers",
+        "amp", "grad_clip", "multiscale",
+        "aug_mosaic", "aug_mixup", "aug_copypaste",
+        "activation_checkpointing", "activation_offloading", "optimizer_in_bwd",
+        "use_8bit_optimizer", "compile_model", "ddp", "grad_accum",
+        "chunked_loss", "chunk_size", "class_file",
+        "save_dir", "save_best", "run_name", "resume_training", "resume_path",
+        "training_status",
+        "infer_model_arch", "infer_conf", "infer_nms", "infer_img_size",
+        "selected_solution",
+    ]
+
+    state = {}
+    for key in state_keys:
+        if key in st.session_state:
+            val = st.session_state[key]
+            if isinstance(val, (str, int, float, bool, list, dict, type(None))):
+                state[key] = val
+
+    state_file = os.path.join(proj_dir, "session_state.json")
+    with open(state_file, "w") as f:
+        json.dump(state, f, indent=2)
+
+    # Update last_modified
+    for p in index["projects"]:
+        if p["id"] == active_id:
+            p["last_modified"] = datetime.now().isoformat()
+            break
+    _save_index(index)
+
+
+def load_project_state(project_id: str):
+    """Load a project's saved state into session."""
+    proj_dir = get_project_dir(project_id)
+    state_file = os.path.join(proj_dir, "session_state.json")
+
+    if not os.path.isfile(state_file):
+        return
+
+    with open(state_file, "r") as f:
+        state = json.load(f)
+
+    for key, val in state.items():
+        st.session_state[key] = val
+
+
+def duplicate_project(project_id: str, new_name: str) -> dict:
+    """Duplicate an existing project with a new name."""
+    index = _load_index()
+    prev_active = index.get("active")
+    src_dir = get_project_dir(project_id)
+    new_project = create_project(new_name)
+    dst_dir = get_project_dir(new_project["id"])
+
+    # Copy data and config (not runs — those are large)
+    for subdir in ("data",):
+        src = os.path.join(src_dir, subdir)
+        dst = os.path.join(dst_dir, subdir)
+        if os.path.isdir(src):
+            shutil.rmtree(dst, ignore_errors=True)
+            shutil.copytree(src, dst)
+
+    # Copy session state
+    src_state = os.path.join(src_dir, "session_state.json")
+    if os.path.isfile(src_state):
+        shutil.copy2(src_state, os.path.join(dst_dir, "session_state.json"))
+
+    # Restore previously active project so duplication doesn't switch context
+    if prev_active and prev_active != new_project["id"]:
+        idx = _load_index()
+        idx["active"] = prev_active
+        _save_index(idx)
+
+    return new_project
+
+
+def get_project_stats(project_id: str) -> dict:
+    """Get statistics for a project."""
+    import re
+    import csv
+    import glob
+    proj_dir = get_project_dir(project_id)
+    stats = {"runs": 0, "best_map": None, "total_size": "0 KB", "has_export": False}
+
+    # Check both project/runs/ and workspace/
+    from flashstudio.constants import DEFAULT_SAVE_DIR
+    runs_dirs = [os.path.join(proj_dir, "runs")]
+    if os.path.isdir(DEFAULT_SAVE_DIR):
+        runs_dirs.append(DEFAULT_SAVE_DIR)
+
+    best_map = 0
+    total_runs = 0
+    for runs_dir in runs_dirs:
+        if not os.path.isdir(runs_dir):
+            continue
+        runs = [d for d in os.listdir(runs_dir) if os.path.isdir(os.path.join(runs_dir, d))]
+        total_runs += len(runs)
+
+        for run in runs:
+            rd = os.path.join(runs_dir, run)
+            from flashstudio.constants import TRAINING_LOG_CSV
+            csv_path = os.path.join(rd, TRAINING_LOG_CSV)
+            if os.path.isfile(csv_path):
+                try:
+                    with open(csv_path, "r") as f:
+                        rows = list(csv.DictReader(f))
+                    vals = [float(r.get("mAP@0.5", 0)) for r in rows if r.get("mAP@0.5", "").strip()]
+                    if vals and max(vals) > best_map:
+                        best_map = max(vals)
+                except Exception:
+                    pass
+            if best_map == 0:
+                from flashstudio.constants import TRAINING_LOG_GLOB
+                logs = glob.glob(os.path.join(rd, TRAINING_LOG_GLOB))
+                for lf in logs:
+                    try:
+                        with open(lf, "r", encoding="utf-8", errors="replace") as f:
+                            for ln in f.readlines()[-10:]:
+                                bm = re.search(r"Best mAP@0\.5:\s*([\d.]+)", ln)
+                                if bm and float(bm.group(1)) > best_map:
+                                    best_map = float(bm.group(1))
+                    except OSError:
+                        pass
+
+    stats["runs"] = total_runs
+    if best_map > 0:
+        stats["best_map"] = best_map
+
+    exports_dir = os.path.join(proj_dir, "exports")
+    if os.path.isdir(exports_dir) and os.listdir(exports_dir):
+        stats["has_export"] = True
+    if not stats["has_export"]:
+        for runs_dir in runs_dirs:
+            if os.path.isdir(runs_dir):
+                for d in os.listdir(runs_dir):
+                    rd = os.path.join(runs_dir, d)
+                    if os.path.isdir(rd) and any(f.endswith(".onnx") for f in os.listdir(rd)):
+                        stats["has_export"] = True
+                        break
+
+    stats["total_size"] = _dir_size_str(proj_dir)
+    return stats
+
+
+def render_project_selector():
+    """Render the project management UI in the sidebar or top of page."""
+    index = _load_index()
+    active_id = index.get("active")
+    projects = index.get("projects", [])
+
+    if not projects:
+        return False  # No projects exist — show creation UI
+
+    # Quick selector
+    active_idx = 0
+    for i, p in enumerate(projects):
+        if p["id"] == active_id:
+            active_idx = i
+            break
+
+    selected_idx = st.selectbox(
+        "Active Project",
+        range(len(projects)),
+        index=active_idx,
+        format_func=lambda i: projects[i]['name'],
+        key="project_selector",
+    )
+
+    if projects[selected_idx]["id"] != active_id:
+        save_project_state()
+        switch_project(projects[selected_idx]["id"])
+        st.rerun()
+
+    return True
+
+
+def render_project_manager_page():
+    """Full project management page — create, switch, delete, manage projects."""
+    st.markdown(
+        "<div style='text-align:center; padding:0.4rem 0;'>"
+        "<b style='font-size:1.2rem;'>Project Manager</b>"
+        "</div>",
+        unsafe_allow_html=True,
+    )
+
+    index = _load_index()
+    projects = index.get("projects", [])
+    active_id = index.get("active")
+
+    # ─── Create New Project ───
+    with st.container(border=True):
+        st.markdown("#### Create New Project")
+        col_name, col_desc, col_btn = st.columns([2, 3, 1])
+        with col_name:
+            new_name = st.text_input("Project Name", placeholder="e.g. PPE Detection v2",
+                                     key="new_project_name")
+        with col_desc:
+            new_desc = st.text_input("Description (optional)",
+                                     placeholder="Safety equipment detection for construction sites",
+                                     key="new_project_desc")
+        with col_btn:
+            st.markdown("<div style='height:28px'></div>", unsafe_allow_html=True)
+            if st.button("Create", type="primary", key="create_project_btn",
+                         use_container_width=True, disabled=not new_name):
+                create_project(new_name, new_desc)
+                st.success(f"Project **{new_name}** created!")
+                _clear_session_for_new_project()
+                st.rerun()
+
+    # ─── Project List ───
+    if not projects:
+        st.info("No projects yet. Create your first project above to get started.")
+        return
+
+    st.markdown(f"#### Your Projects ({len(projects)})")
+
+    for p in sorted(projects, key=lambda x: x.get("last_modified", ""), reverse=True):
+        is_active = p["id"] == active_id
+        stats = get_project_stats(p["id"])
+
+        with st.container(border=True):
+            col_info, col_stats, col_actions = st.columns([3, 2, 2])
+
+            with col_info:
+                active_badge = " **[ACTIVE]**" if is_active else ""
+                st.markdown(f"**{p['name']}**{active_badge}")
+                if p.get("description"):
+                    st.caption(p["description"])
+                created = p.get("created", "")[:10]
+                modified = p.get("last_modified", "")[:10]
+                st.caption(f"Created: {created} · Modified: {modified}")
+
+            with col_stats:
+                sc1, sc2, sc3 = st.columns(3)
+                with sc1:
+                    st.metric("Runs", stats["runs"])
+                with sc2:
+                    map_str = f"{stats['best_map']:.3f}" if stats["best_map"] else "—"
+                    st.metric("Best mAP", map_str)
+                with sc3:
+                    st.metric("Size", stats["total_size"])
+
+            with col_actions:
+                ac1, ac2, ac3 = st.columns(3)
+                with ac1:
+                    if not is_active:
+                        if st.button("Open", key=f"switch_{p['id']}", help="Switch to this project"):
+                            save_project_state()
+                            switch_project(p["id"])
+                            st.rerun()
+                with ac2:
+                    if st.button("Copy", key=f"dup_{p['id']}", help="Duplicate project"):
+                        dup = duplicate_project(p["id"], f"{p['name']} (copy)")
+                        st.success(f"Duplicated as '{dup['name']}'")
+                        st.rerun()
+                with ac3:
+                    if st.button("Del", key=f"del_{p['id']}", help="Delete project"):
+                        st.session_state[f"confirm_delete_{p['id']}"] = True
+
+            # Delete confirmation
+            if st.session_state.get(f"confirm_delete_{p['id']}"):
+                st.error(f"Delete project **{p['name']}** and ALL its data? This cannot be undone.")
+                dc1, dc2, dc3 = st.columns([1, 1, 3])
+                with dc1:
+                    if st.button("Yes, Delete", key=f"confirm_del_{p['id']}", type="primary"):
+                        delete_project(p["id"])
+                        st.session_state.pop(f"confirm_delete_{p['id']}", None)
+                        st.rerun()
+                with dc2:
+                    if st.button("Cancel", key=f"cancel_del_{p['id']}"):
+                        st.session_state.pop(f"confirm_delete_{p['id']}", None)
+                        st.rerun()
+
+    # ─── Danger Zone ───
+    st.divider()
+    with st.expander("Danger Zone", expanded=False):
+        st.error("These actions are irreversible!")
+        col_d1, col_d2 = st.columns(2)
+        with col_d1:
+            if st.button("Delete ALL Projects", key="delete_all_projects"):
+                st.session_state["confirm_delete_all"] = True
+
+        if st.session_state.get("confirm_delete_all"):
+            st.warning(f"This will delete **{len(projects)} projects** and all training data!")
+            if st.button("CONFIRM: Delete Everything", type="primary", key="confirm_delete_all_btn"):
+                delete_all_projects()
+                _clear_session_for_new_project()
+                st.session_state.pop("confirm_delete_all", None)
+                st.rerun()
+            if st.button("Cancel", key="cancel_delete_all"):
+                st.session_state.pop("confirm_delete_all", None)
+                st.rerun()
+
+
+def _clear_session_for_new_project():
+    """Clear session state keys when switching to a new/blank project."""
+    keys_to_clear = [
+        "dataset_name", "dataset_classes", "dataset_id", "dataset_output_path",
+        "train_img_path", "val_img_path", "train_ann_path", "val_ann_path",
+        "training_active", "training_status", "training_pid",
+        "infer_img_results", "video_results", "exported_files",
+        "run_name", "detected_format",
+    ]
+    for key in keys_to_clear:
+        st.session_state.pop(key, None)
+    st.session_state["current_step"] = 0
+
+
+def _slugify(text: str) -> str:
+    """Convert text to a safe directory name."""
+    import re
+    text = text.lower().strip()
+    text = re.sub(r'[^\w\s-]', '', text)
+    text = re.sub(r'[\s-]+', '_', text)
+    return text[:30]
+
+
+def _dir_size_str(path: str) -> str:
+    """Get human-readable directory size."""
+    total = 0
+    for dirpath, _dirnames, filenames in os.walk(path):
+        for f in filenames:
+            try:
+                total += os.path.getsize(os.path.join(dirpath, f))
+            except OSError:
+                pass
+    if total > 1024 * 1024 * 1024:
+        return f"{total / (1024**3):.1f} GB"
+    elif total > 1024 * 1024:
+        return f"{total / (1024**2):.0f} MB"
+    return f"{total / 1024:.0f} KB"

flashstudio/components/sidebar.py

@@ -0,0 +1,87 @@
+"""FlashStudio sidebar — clean icon-based navigation."""
+
+import os
+import streamlit as st
+from flashstudio.utils.device import get_gpu_info
+from flashstudio.components.project_manager import (
+    get_active_project, save_project_state,
+)
+from flashstudio.constants import (
+    COLOR_SUCCESS, COLOR_WARNING, COLOR_ERROR,
+    COLOR_TEXT_PRIMARY, COLOR_TEXT_BODY, COLOR_TEXT_SECONDARY,
+    SIDEBAR_LABEL_TRUNCATE,
+)
+
+NAV = [
+    ("1", "Dashboard"),
+    ("2", "Data"),
+    ("3", "Model"),
+    ("4", "Training"),
+    ("5", "Export"),
+    ("6", "Inference"),
+]
+
+
+def _get_logo_b64():
+    """Load logo as base64 for inline HTML rendering."""
+    import base64
+    logo_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "..", "assets", "logo.png")
+    if not os.path.isfile(logo_path):
+        logo_path = os.path.join(os.getcwd(), "assets", "logo.png")
+    if os.path.isfile(logo_path):
+        with open(logo_path, "rb") as f:
+            return base64.b64encode(f.read()).decode()
+    return None
+
+
+def render_sidebar():
+    with st.sidebar:
+        logo_b64 = _get_logo_b64()
+        if logo_b64:
+            st.markdown(
+                f'<div style="display:flex;flex-direction:column;align-items:center;gap:0.3rem;padding:0.6rem 0.2rem 0.5rem;">'
+                f'<img src="data:image/png;base64,{logo_b64}" style="width:56px;height:56px;">'
+                f'<b style="font-size:1.1rem;color:{COLOR_TEXT_PRIMARY};">FlashStudio</b>'
+                f'</div>',
+                unsafe_allow_html=True,
+            )
+        else:
+            st.markdown(
+                f'<div style="display:flex;flex-direction:column;align-items:center;gap:0.2rem;padding:0.5rem 0 0.3rem;">'
+                f'<b style="font-size:1.5rem;">FS</b>'
+                f'<b style="font-size:1rem;color:{COLOR_TEXT_PRIMARY};">FlashStudio</b>'
+                f'</div>',
+                unsafe_allow_html=True,
+            )
+
+        project = get_active_project()
+        if project:
+            if st.button(project['name'][:SIDEBAR_LABEL_TRUNCATE], key="sb_project", use_container_width=True):
+                save_project_state()
+                st.session_state["show_project_manager"] = True
+                st.rerun()
+
+        current = st.session_state.get("current_step", 0)
+        for i, (icon, label) in enumerate(NAV):
+            btn_type = "primary" if i == current else "secondary"
+            if st.button(f"{icon}  {label}", key=f"sb_{i}", type=btn_type, use_container_width=True):
+                st.session_state["current_step"] = i
+                st.rerun()
+
+        st.divider()
+
+        gpu = get_gpu_info()
+        from flashstudio.constants import GPU_NAME_TRUNCATE, DATASET_NAME_TRUNCATE
+        train_status = st.session_state.get("training_status", "Idle")
+        status_color = COLOR_SUCCESS if train_status == "Complete" else (
+            COLOR_WARNING if train_status == "Running" else (
+            COLOR_ERROR if train_status == "Failed" else COLOR_TEXT_SECONDARY))
+        st.markdown(
+            f'<div style="font-size:0.82rem;color:{COLOR_TEXT_BODY};padding:0.1rem 0.5rem;line-height:1.9;">'
+            f'GPU: {gpu["name"][:GPU_NAME_TRUNCATE] if gpu["available"] else "CPU"}<br>'
+            f'Data: {(st.session_state.get("dataset_name") or "—")[:DATASET_NAME_TRUNCATE]}<br>'
+            f'Model: {st.session_state.get("model_arch", "Pico").replace("FlashDet-", "")}<br>'
+            f'Status: <span style="color:{status_color};font-weight:600;">{train_status[:12]}</span>'
+            f'</div>',
+            unsafe_allow_html=True,
+        )