flashstudio 0.0.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of flashstudio might be problematic. Click here for more details.
- flashstudio/__init__.py +5 -0
- flashstudio/app.py +194 -0
- flashstudio/cli.py +18 -0
- flashstudio/components/__init__.py +0 -0
- flashstudio/components/project_manager.py +490 -0
- flashstudio/components/sidebar.py +87 -0
- flashstudio/components/styles.py +346 -0
- flashstudio/components/wizard.py +45 -0
- flashstudio/components/zone_drawer/__init__.py +177 -0
- flashstudio/constants.py +245 -0
- flashstudio/core/__init__.py +3 -0
- flashstudio/core/constants.py +3 -0
- flashstudio/launcher.py +288 -0
- flashstudio/pages/__init__.py +8 -0
- flashstudio/pages/dashboard/__init__.py +3 -0
- flashstudio/pages/dashboard/_common.py +10 -0
- flashstudio/pages/dashboard/overview/__init__.py +3 -0
- flashstudio/pages/dashboard/overview/metrics.py +33 -0
- flashstudio/pages/dashboard/overview/pipeline_status.py +35 -0
- flashstudio/pages/dashboard/overview/project_banner.py +23 -0
- flashstudio/pages/dashboard/overview/quick_actions.py +28 -0
- flashstudio/pages/dashboard/overview/recent_runs.py +99 -0
- flashstudio/pages/dashboard/overview/tab.py +35 -0
- flashstudio/pages/dashboard/page.py +17 -0
- flashstudio/pages/data/__init__.py +3 -0
- flashstudio/pages/data/_common.py +107 -0
- flashstudio/pages/data/download/__init__.py +3 -0
- flashstudio/pages/data/download/tab.py +75 -0
- flashstudio/pages/data/helpers.py +338 -0
- flashstudio/pages/data/page.py +50 -0
- flashstudio/pages/data/preview/__init__.py +3 -0
- flashstudio/pages/data/preview/tab.py +139 -0
- flashstudio/pages/data/upload/__init__.py +3 -0
- flashstudio/pages/data/upload/tab.py +156 -0
- flashstudio/pages/data/verify/__init__.py +3 -0
- flashstudio/pages/data/verify/tab.py +108 -0
- flashstudio/pages/export/__init__.py +3 -0
- flashstudio/pages/export/page.py +54 -0
- flashstudio/pages/export/weights/__init__.py +4 -0
- flashstudio/pages/export/weights/exporter.py +60 -0
- flashstudio/pages/export/weights/settings.py +86 -0
- flashstudio/pages/inference/__init__.py +3 -0
- flashstudio/pages/inference/_common.py +76 -0
- flashstudio/pages/inference/data/__init__.py +3 -0
- flashstudio/pages/inference/data/tab.py +121 -0
- flashstudio/pages/inference/model/__init__.py +3 -0
- flashstudio/pages/inference/model/tab.py +62 -0
- flashstudio/pages/inference/page.py +42 -0
- flashstudio/pages/inference/run/__init__.py +3 -0
- flashstudio/pages/inference/run/detection.py +133 -0
- flashstudio/pages/inference/run/tab.py +161 -0
- flashstudio/pages/inference/run/video.py +195 -0
- flashstudio/pages/inference/solution/__init__.py +3 -0
- flashstudio/pages/inference/solution/tab.py +121 -0
- flashstudio/pages/model/__init__.py +3 -0
- flashstudio/pages/model/_common.py +3 -0
- flashstudio/pages/model/advanced/__init__.py +3 -0
- flashstudio/pages/model/advanced/tab.py +54 -0
- flashstudio/pages/model/architecture/__init__.py +3 -0
- flashstudio/pages/model/architecture/tab.py +90 -0
- flashstudio/pages/model/augment/__init__.py +3 -0
- flashstudio/pages/model/augment/tab.py +18 -0
- flashstudio/pages/model/hyperparams/__init__.py +4 -0
- flashstudio/pages/model/hyperparams/tab.py +89 -0
- flashstudio/pages/model/page.py +30 -0
- flashstudio/pages/model/utils/__init__.py +3 -0
- flashstudio/pages/model/utils/config.py +27 -0
- flashstudio/pages/training/__init__.py +3 -0
- flashstudio/pages/training/_common.py +24 -0
- flashstudio/pages/training/launch/__init__.py +3 -0
- flashstudio/pages/training/launch/controls.py +84 -0
- flashstudio/pages/training/launch/dialogs.py +247 -0
- flashstudio/pages/training/launch/preflight.py +77 -0
- flashstudio/pages/training/launch/runner.py +440 -0
- flashstudio/pages/training/launch/tab.py +134 -0
- flashstudio/pages/training/monitor/__init__.py +3 -0
- flashstudio/pages/training/monitor/checkpoints.py +96 -0
- flashstudio/pages/training/monitor/curves.py +113 -0
- flashstudio/pages/training/monitor/dashboard.py +129 -0
- flashstudio/pages/training/monitor/gt_verification.py +122 -0
- flashstudio/pages/training/monitor/log_viewer.py +43 -0
- flashstudio/pages/training/monitor/parsers.py +253 -0
- flashstudio/pages/training/monitor/run_meta.py +165 -0
- flashstudio/pages/training/monitor/tab.py +142 -0
- flashstudio/pages/training/monitor/visualizations.py +52 -0
- flashstudio/pages/training/page.py +134 -0
- flashstudio/utils/__init__.py +152 -0
- flashstudio/utils/config.py +104 -0
- flashstudio/utils/device.py +67 -0
- flashstudio/utils/filesystem.py +67 -0
- flashstudio/utils/jobs.py +95 -0
- flashstudio/utils/training_hooks.py +266 -0
- flashstudio-0.0.1.dist-info/METADATA +107 -0
- flashstudio-0.0.1.dist-info/RECORD +97 -0
- flashstudio-0.0.1.dist-info/WHEEL +5 -0
- flashstudio-0.0.1.dist-info/entry_points.txt +2 -0
- flashstudio-0.0.1.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,156 @@
|
|
|
1
|
+
"""Upload tab — file upload, class configuration, and format conversion hint."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import streamlit as st
|
|
5
|
+
|
|
6
|
+
from flashstudio.pages.data._common import (
|
|
7
|
+
_HAS_FLASHDET, _auto_detect_classes,
|
|
8
|
+
)
|
|
9
|
+
from flashstudio.pages.data.helpers import _extract_upload, _convert_to_coco
|
|
10
|
+
|
|
11
|
+
try:
|
|
12
|
+
from flashdet.data import detect_dataset_format
|
|
13
|
+
except ImportError:
|
|
14
|
+
detect_dataset_format = None
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def _render_upload():
|
|
18
|
+
col1, col2 = st.columns(2)
|
|
19
|
+
with col1:
|
|
20
|
+
with st.container(border=True):
|
|
21
|
+
st.markdown("#### Train")
|
|
22
|
+
train_upload = st.file_uploader("ZIP/TAR", type=["zip", "tar", "gz"], key="train_data_upload")
|
|
23
|
+
if train_upload:
|
|
24
|
+
st.session_state["dataset_name"] = train_upload.name
|
|
25
|
+
if st.button("Extract", key="extract_train", use_container_width=True, type="primary"):
|
|
26
|
+
_extract_upload(train_upload, "train")
|
|
27
|
+
manual_train = st.text_input(
|
|
28
|
+
"Images dir", placeholder="/path/to/train/",
|
|
29
|
+
value=st.session_state.get("train_img_path", ""),
|
|
30
|
+
key="_train_img_input",
|
|
31
|
+
)
|
|
32
|
+
if manual_train:
|
|
33
|
+
st.session_state["train_img_path"] = manual_train
|
|
34
|
+
|
|
35
|
+
with col2:
|
|
36
|
+
with st.container(border=True):
|
|
37
|
+
st.markdown("#### Val")
|
|
38
|
+
val_upload = st.file_uploader("ZIP/TAR", type=["zip", "tar", "gz"], key="val_data_upload")
|
|
39
|
+
if val_upload:
|
|
40
|
+
if st.button("Extract", key="extract_val", use_container_width=True, type="primary"):
|
|
41
|
+
_extract_upload(val_upload, "val")
|
|
42
|
+
manual_val = st.text_input(
|
|
43
|
+
"Images dir", placeholder="/path/to/valid/",
|
|
44
|
+
value=st.session_state.get("val_img_path", ""),
|
|
45
|
+
key="_val_img_input",
|
|
46
|
+
)
|
|
47
|
+
if manual_val:
|
|
48
|
+
st.session_state["val_img_path"] = manual_val
|
|
49
|
+
|
|
50
|
+
st.selectbox("Format", ["COCO JSON", "Pascal VOC", "YOLO TXT"], key="ann_format")
|
|
51
|
+
|
|
52
|
+
# ── Classes section ──
|
|
53
|
+
_render_class_config()
|
|
54
|
+
|
|
55
|
+
_render_conversion_hint()
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def _render_class_config():
|
|
59
|
+
"""Render class configuration: auto-detect from annotation, load from .txt, or manual edit."""
|
|
60
|
+
from flashstudio.utils import flash
|
|
61
|
+
|
|
62
|
+
with st.container(border=True):
|
|
63
|
+
st.markdown("#### Classes")
|
|
64
|
+
|
|
65
|
+
# Auto-detect button + .txt upload in one row
|
|
66
|
+
cc1, cc2, cc3 = st.columns([2, 2, 1])
|
|
67
|
+
with cc1:
|
|
68
|
+
if st.button("Auto-detect from annotation", key="btn_auto_classes",
|
|
69
|
+
use_container_width=True, type="primary"):
|
|
70
|
+
detected = _auto_detect_classes()
|
|
71
|
+
if detected:
|
|
72
|
+
st.session_state["class_names"] = "\n".join(detected)
|
|
73
|
+
st.session_state["num_classes"] = len(detected)
|
|
74
|
+
flash(f"Detected {len(detected)} classes from annotation", "success")
|
|
75
|
+
st.rerun()
|
|
76
|
+
else:
|
|
77
|
+
flash("No classes found — load a dataset with COCO annotations first", "warning")
|
|
78
|
+
st.rerun()
|
|
79
|
+
|
|
80
|
+
with cc2:
|
|
81
|
+
cls_file = st.file_uploader("Load classes.txt", type=["txt"],
|
|
82
|
+
key="class_txt_upload", label_visibility="collapsed")
|
|
83
|
+
if cls_file:
|
|
84
|
+
content = cls_file.read().decode("utf-8", errors="ignore")
|
|
85
|
+
names = [line.strip() for line in content.strip().split("\n") if line.strip()]
|
|
86
|
+
if names:
|
|
87
|
+
st.session_state["class_names"] = "\n".join(names)
|
|
88
|
+
st.session_state["num_classes"] = len(names)
|
|
89
|
+
flash(f"Loaded {len(names)} classes from file", "success")
|
|
90
|
+
st.rerun()
|
|
91
|
+
|
|
92
|
+
with cc3:
|
|
93
|
+
nc = st.session_state.get("num_classes", 0)
|
|
94
|
+
st.metric("Count", nc if nc else "—")
|
|
95
|
+
|
|
96
|
+
# Auto-detect on first load if classes not set
|
|
97
|
+
current_classes = st.session_state.get("class_names", "")
|
|
98
|
+
if isinstance(current_classes, list):
|
|
99
|
+
current_classes = "\n".join(current_classes)
|
|
100
|
+
if not current_classes.strip() and st.session_state.get("train_img_path"):
|
|
101
|
+
detected = _auto_detect_classes()
|
|
102
|
+
if detected:
|
|
103
|
+
st.session_state["class_names"] = "\n".join(detected)
|
|
104
|
+
st.session_state["num_classes"] = len(detected)
|
|
105
|
+
|
|
106
|
+
# Editable text area
|
|
107
|
+
_cv = st.session_state.get("class_names", "")
|
|
108
|
+
if isinstance(_cv, list):
|
|
109
|
+
_cv = "\n".join(_cv)
|
|
110
|
+
class_text = st.text_area(
|
|
111
|
+
"Class names (one per line)",
|
|
112
|
+
value=_cv,
|
|
113
|
+
height=100,
|
|
114
|
+
key="_class_names_input",
|
|
115
|
+
placeholder="person\ncar\nbicycle\n...",
|
|
116
|
+
)
|
|
117
|
+
if class_text != _cv:
|
|
118
|
+
names = [n.strip() for n in class_text.strip().split("\n") if n.strip()]
|
|
119
|
+
st.session_state["class_names"] = "\n".join(names)
|
|
120
|
+
st.session_state["num_classes"] = len(names)
|
|
121
|
+
|
|
122
|
+
# Download button for current classes
|
|
123
|
+
if current_classes.strip():
|
|
124
|
+
st.download_button(
|
|
125
|
+
"Download classes.txt",
|
|
126
|
+
current_classes.strip() + "\n",
|
|
127
|
+
file_name="classes.txt",
|
|
128
|
+
mime="text/plain",
|
|
129
|
+
use_container_width=True,
|
|
130
|
+
)
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def _render_conversion_hint():
|
|
134
|
+
train_path = st.session_state.get("train_img_path", "")
|
|
135
|
+
if not train_path or not os.path.isdir(train_path):
|
|
136
|
+
return
|
|
137
|
+
ann_format = st.session_state.get("ann_format", "COCO JSON")
|
|
138
|
+
detected = st.session_state.get("detected_format", "")
|
|
139
|
+
needs = "YOLO" in ann_format or "VOC" in ann_format or detected in ("txt", "voc")
|
|
140
|
+
if not needs:
|
|
141
|
+
coco_ann = os.path.join(train_path, "_annotations.coco.json")
|
|
142
|
+
if not os.path.isfile(coco_ann) and _HAS_FLASHDET:
|
|
143
|
+
parent = os.path.dirname(train_path)
|
|
144
|
+
det_fmt = detect_dataset_format(parent)
|
|
145
|
+
if det_fmt in ("txt", "voc"):
|
|
146
|
+
needs = True
|
|
147
|
+
st.session_state["detected_format"] = det_fmt
|
|
148
|
+
if needs:
|
|
149
|
+
c1, c2 = st.columns([1, 4])
|
|
150
|
+
with c1:
|
|
151
|
+
if st.button("Convert to COCO", type="primary", key="manual_convert", use_container_width=True):
|
|
152
|
+
source = os.path.dirname(train_path) if os.path.basename(train_path) in ("images", "train") else train_path
|
|
153
|
+
tfmt = "txt" if "YOLO" in ann_format or detected == "txt" else "voc"
|
|
154
|
+
_convert_to_coco(source, "train", tfmt)
|
|
155
|
+
with c2:
|
|
156
|
+
st.caption(f"Format detected: **{detected or ann_format}** → auto-convert to COCO JSON")
|
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
"""Verify tab — dataset verification and summary statistics."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import streamlit as st
|
|
5
|
+
|
|
6
|
+
from flashstudio.constants import IMG_EXTENSIONS_ALL, MAX_PREVIEW_COLS
|
|
7
|
+
from flashstudio.pages.data._common import _HAS_FLASHDET, _FLASHDET_INSTALL_MSG
|
|
8
|
+
|
|
9
|
+
try:
|
|
10
|
+
from flashdet.data import verify_dataset, summarize_coco_root
|
|
11
|
+
except ImportError:
|
|
12
|
+
pass
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _render_verify():
|
|
16
|
+
train_path = st.session_state.get("train_img_path", "")
|
|
17
|
+
val_path = st.session_state.get("val_img_path", "")
|
|
18
|
+
if not train_path and not val_path:
|
|
19
|
+
st.info("No dataset loaded. Upload or download a dataset first.")
|
|
20
|
+
return
|
|
21
|
+
|
|
22
|
+
dataset_root = st.session_state.get("dataset_output_path", "")
|
|
23
|
+
if not dataset_root and train_path:
|
|
24
|
+
parent = os.path.dirname(os.path.normpath(train_path))
|
|
25
|
+
if os.path.isdir(parent):
|
|
26
|
+
dataset_root = parent
|
|
27
|
+
|
|
28
|
+
# FlashDet verify
|
|
29
|
+
if dataset_root and os.path.isdir(dataset_root) and _HAS_FLASHDET:
|
|
30
|
+
vc1, vc2 = st.columns([1, 5])
|
|
31
|
+
with vc1:
|
|
32
|
+
do_verify = st.button("Verify", type="primary", key="run_verify_btn", use_container_width=True)
|
|
33
|
+
with vc2:
|
|
34
|
+
st.caption(f"Root: `{dataset_root}`")
|
|
35
|
+
if do_verify or st.session_state.get("_last_verify_ok") is not None:
|
|
36
|
+
try:
|
|
37
|
+
ok = verify_dataset(dataset_root)
|
|
38
|
+
st.session_state["_last_verify_ok"] = ok
|
|
39
|
+
st.success("Verification **PASSED**") if ok else st.warning("Verification: issues detected")
|
|
40
|
+
except Exception as e:
|
|
41
|
+
st.caption(str(e)[:60])
|
|
42
|
+
|
|
43
|
+
try:
|
|
44
|
+
s = summarize_coco_root(dataset_root)
|
|
45
|
+
has_summary = s and (s.get("train_images", 0) > 0 or s.get("val_images", 0) > 0)
|
|
46
|
+
if has_summary:
|
|
47
|
+
m1, m2, m3, m4 = st.columns(4)
|
|
48
|
+
with m1:
|
|
49
|
+
st.metric("Train", s.get("train_images", 0))
|
|
50
|
+
with m2:
|
|
51
|
+
st.metric("Val", s.get("val_images", 0))
|
|
52
|
+
with m3:
|
|
53
|
+
st.metric("Annotations", s.get("total_annotations", 0))
|
|
54
|
+
with m4:
|
|
55
|
+
st.metric("Categories", s.get("num_categories", 0))
|
|
56
|
+
if s.get("class_distribution"):
|
|
57
|
+
with st.expander("Class Distribution"):
|
|
58
|
+
top = sorted(s["class_distribution"].items(), key=lambda x: x[1], reverse=True)[:15]
|
|
59
|
+
st.bar_chart({n: c for n, c in top})
|
|
60
|
+
return
|
|
61
|
+
except Exception:
|
|
62
|
+
pass
|
|
63
|
+
|
|
64
|
+
if not _HAS_FLASHDET and dataset_root and os.path.isdir(dataset_root):
|
|
65
|
+
st.info(_FLASHDET_INSTALL_MSG)
|
|
66
|
+
|
|
67
|
+
# Manual count — always show this as fallback
|
|
68
|
+
ti = [f for f in os.listdir(train_path) if f.lower().endswith(IMG_EXTENSIONS_ALL)] if train_path and os.path.isdir(train_path) else []
|
|
69
|
+
vi = [f for f in os.listdir(val_path) if f.lower().endswith(IMG_EXTENSIONS_ALL)] if val_path and os.path.isdir(val_path) else []
|
|
70
|
+
|
|
71
|
+
# Count YOLO labels if present
|
|
72
|
+
train_label_dir = st.session_state.get("train_label_path", "")
|
|
73
|
+
if not train_label_dir and train_path:
|
|
74
|
+
candidate = train_path.replace("/images/", "/labels/").replace("\\images\\", "\\labels\\")
|
|
75
|
+
if os.path.isdir(candidate):
|
|
76
|
+
train_label_dir = candidate
|
|
77
|
+
train_labels = len([f for f in os.listdir(train_label_dir) if f.endswith(".txt")]) if train_label_dir and os.path.isdir(train_label_dir) else 0
|
|
78
|
+
|
|
79
|
+
m1, m2, m3, m4 = st.columns(4)
|
|
80
|
+
with m1:
|
|
81
|
+
st.metric("Train Images", len(ti))
|
|
82
|
+
with m2:
|
|
83
|
+
st.metric("Val Images", len(vi))
|
|
84
|
+
with m3:
|
|
85
|
+
st.metric("Labels", train_labels if train_labels else "—")
|
|
86
|
+
with m4:
|
|
87
|
+
status = "OK" if ti else "Empty"
|
|
88
|
+
if ti and train_labels and train_labels == len(ti):
|
|
89
|
+
status = "Complete"
|
|
90
|
+
elif ti and train_labels and train_labels != len(ti):
|
|
91
|
+
status = f"Partial ({train_labels}/{len(ti)})"
|
|
92
|
+
st.metric("Status", status)
|
|
93
|
+
|
|
94
|
+
if train_path and os.path.isdir(train_path):
|
|
95
|
+
st.caption(f"Train: `{train_path}`")
|
|
96
|
+
if val_path and os.path.isdir(val_path):
|
|
97
|
+
st.caption(f"Val: `{val_path}`")
|
|
98
|
+
|
|
99
|
+
# Sample preview
|
|
100
|
+
if ti:
|
|
101
|
+
with st.expander("Samples"):
|
|
102
|
+
cols = st.columns(min(MAX_PREVIEW_COLS, len(ti)))
|
|
103
|
+
for i, img_name in enumerate(ti[:MAX_PREVIEW_COLS]):
|
|
104
|
+
with cols[i]:
|
|
105
|
+
try:
|
|
106
|
+
st.image(os.path.join(train_path, img_name), caption=img_name[:15], use_container_width=True)
|
|
107
|
+
except Exception:
|
|
108
|
+
pass
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
"""FlashStudio — Export (zero-scroll, split layout)."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import streamlit as st
|
|
5
|
+
from flashstudio.constants import (
|
|
6
|
+
EXPORT_OPSET_DEFAULT,
|
|
7
|
+
CKPT_BEST_INFERENCE, CKPT_BEST, CKPT_LAST_INFERENCE, CKPT_LAST,
|
|
8
|
+
)
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def render_export_page():
|
|
12
|
+
from flashstudio.components.styles import render_page_header
|
|
13
|
+
from flashstudio.utils import show_flashes
|
|
14
|
+
render_page_header("", "Export")
|
|
15
|
+
show_flashes()
|
|
16
|
+
|
|
17
|
+
left, right = st.columns(2)
|
|
18
|
+
|
|
19
|
+
with left:
|
|
20
|
+
with st.container(border=True):
|
|
21
|
+
from flashstudio.pages.export.weights.settings import render_export_settings
|
|
22
|
+
wp, img_sz, export_fmt = render_export_settings()
|
|
23
|
+
|
|
24
|
+
if st.button(f"Export {export_fmt}", use_container_width=True, type="primary"):
|
|
25
|
+
from flashstudio.pages.export.weights.exporter import _run_export
|
|
26
|
+
_run_export(wp, img_sz, st.session_state.get("export_opset", EXPORT_OPSET_DEFAULT),
|
|
27
|
+
st.session_state.get("export_dynamic", True),
|
|
28
|
+
export_fmt.lower())
|
|
29
|
+
|
|
30
|
+
with right:
|
|
31
|
+
with st.container(border=True):
|
|
32
|
+
st.markdown("#### Output")
|
|
33
|
+
exported = st.session_state.get("exported_files", [])
|
|
34
|
+
if not exported:
|
|
35
|
+
st.caption("Auto-saved weights from training:")
|
|
36
|
+
for n, d in [(CKPT_BEST_INFERENCE, "Best, inference-ready"),
|
|
37
|
+
(CKPT_BEST, "Best, full checkpoint"),
|
|
38
|
+
(CKPT_LAST_INFERENCE, "Final, inference-ready"),
|
|
39
|
+
(CKPT_LAST, "Final, full checkpoint")]:
|
|
40
|
+
st.markdown(f'<span style="font-size:0.88rem;color:#4B5563;">• `{n}` — {d}</span>', unsafe_allow_html=True)
|
|
41
|
+
st.caption("Manual:")
|
|
42
|
+
st.code("flashdet export --model model.pth --output model.onnx", language="bash")
|
|
43
|
+
else:
|
|
44
|
+
for exp in exported:
|
|
45
|
+
st.markdown(f"**{exp['format']}** · `{os.path.basename(exp.get('path', ''))}`")
|
|
46
|
+
if exp.get("size"):
|
|
47
|
+
st.caption(exp["size"])
|
|
48
|
+
st.success("OK") if exp.get("success") else st.error("Failed")
|
|
49
|
+
if exp.get("success") and os.path.isfile(exp.get("path", "")):
|
|
50
|
+
with open(exp["path"], "rb") as f:
|
|
51
|
+
st.download_button("Download", f, file_name=os.path.basename(exp["path"]),
|
|
52
|
+
key=f"dl_{exp['format']}", use_container_width=True)
|
|
53
|
+
with st.expander("Post-export"):
|
|
54
|
+
st.code("trtexec --onnx=model.onnx --saveEngine=model.engine --fp16", language="bash")
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
"""FlashStudio — Export engine (model export to ONNX / TorchScript)."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import streamlit as st
|
|
5
|
+
from flashstudio.constants import (
|
|
6
|
+
INFER_NUM_CLASSES, MODEL_SIZE_DEFAULT,
|
|
7
|
+
)
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def _run_export(weights_path, img_size, opset=13, dynamic=True, fmt="onnx"):
|
|
11
|
+
from flashstudio.utils import flash
|
|
12
|
+
|
|
13
|
+
if not weights_path or not os.path.isfile(weights_path):
|
|
14
|
+
flash(f"Weights not found: `{weights_path}`", "error")
|
|
15
|
+
st.rerun()
|
|
16
|
+
return
|
|
17
|
+
|
|
18
|
+
import torch
|
|
19
|
+
try:
|
|
20
|
+
from flashdet import get_config, build_model
|
|
21
|
+
except ImportError:
|
|
22
|
+
flash("FlashDet not installed — run `pip install flashdet`", "error")
|
|
23
|
+
st.rerun()
|
|
24
|
+
return
|
|
25
|
+
|
|
26
|
+
ext = ".onnx" if fmt == "onnx" else ".pt"
|
|
27
|
+
out = weights_path.replace(".pth", ext)
|
|
28
|
+
label = "ONNX" if fmt == "onnx" else "TorchScript"
|
|
29
|
+
with st.spinner(f"Exporting to {label}..."):
|
|
30
|
+
try:
|
|
31
|
+
ckpt = torch.load(weights_path, map_location="cpu", weights_only=False)
|
|
32
|
+
cfg = ckpt.get("config", {})
|
|
33
|
+
nc = cfg.get("num_classes", st.session_state.get("num_classes", INFER_NUM_CLASSES))
|
|
34
|
+
raw_input = cfg.get("input_size", img_size)
|
|
35
|
+
inp_scalar = raw_input[0] if isinstance(raw_input, (tuple, list)) else raw_input
|
|
36
|
+
arch = cfg.get("architecture", "flashdet")
|
|
37
|
+
config = get_config(model_size=cfg.get("model_size", MODEL_SIZE_DEFAULT), input_size=inp_scalar, num_classes=nc)
|
|
38
|
+
config.model.architecture = arch
|
|
39
|
+
model = build_model(config, architecture=arch)
|
|
40
|
+
sd = ckpt.get("model_state_dict", ckpt.get("state_dict", ckpt.get("model", ckpt)))
|
|
41
|
+
model.load_state_dict(sd, strict=False)
|
|
42
|
+
model.eval()
|
|
43
|
+
dummy = torch.randn(1, 3, inp_scalar, inp_scalar)
|
|
44
|
+
|
|
45
|
+
if fmt == "onnx":
|
|
46
|
+
dyn = {"images": {0: "batch"}, "output": {0: "batch"}} if dynamic else None
|
|
47
|
+
torch.onnx.export(model, dummy, out, opset_version=opset,
|
|
48
|
+
input_names=["images"], output_names=["output"],
|
|
49
|
+
dynamic_axes=dyn)
|
|
50
|
+
else:
|
|
51
|
+
traced = torch.jit.trace(model, dummy)
|
|
52
|
+
traced.save(out)
|
|
53
|
+
|
|
54
|
+
sz = f"{os.path.getsize(out) / (1024*1024):.1f}MB" if os.path.isfile(out) else "—"
|
|
55
|
+
st.session_state["exported_files"] = [{"format": label, "path": out, "size": sz, "success": True}]
|
|
56
|
+
flash(f"Export successful: `{os.path.basename(out)}` ({sz})", "success")
|
|
57
|
+
except Exception as e:
|
|
58
|
+
flash(f"Export failed: {e}", "error")
|
|
59
|
+
st.session_state["exported_files"] = [{"format": label, "path": out, "size": "—", "success": False}]
|
|
60
|
+
st.rerun()
|
|
@@ -0,0 +1,86 @@
|
|
|
1
|
+
"""FlashStudio — Export settings UI (weights display & config)."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import streamlit as st
|
|
5
|
+
from flashstudio.constants import (
|
|
6
|
+
DEFAULT_SAVE_DIR, EXPORT_IMG_SIZES, EXPORT_OPSET_MIN,
|
|
7
|
+
EXPORT_OPSET_MAX, EXPORT_OPSET_DEFAULT, EXPORT_FORMATS,
|
|
8
|
+
EXPORT_WEIGHT_MAP, CKPT_BEST, MAX_WEIGHTS_DISPLAY,
|
|
9
|
+
)
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def _get_export_save_dir():
|
|
13
|
+
"""Resolve save_dir: _config_mirror > session_state > DEFAULT_SAVE_DIR."""
|
|
14
|
+
from flashstudio.utils import get_state
|
|
15
|
+
sd = get_state("save_dir")
|
|
16
|
+
return sd if sd else DEFAULT_SAVE_DIR
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def render_export_settings():
|
|
20
|
+
"""Render the weights display and config UI. Returns (weights_path, img_size, export_fmt)."""
|
|
21
|
+
st.markdown("#### Weights")
|
|
22
|
+
save_dir = _get_export_save_dir()
|
|
23
|
+
if os.path.isdir(save_dir):
|
|
24
|
+
wts = []
|
|
25
|
+
for root, _dirs, files in os.walk(save_dir):
|
|
26
|
+
for f in files:
|
|
27
|
+
if f.endswith(".pth"):
|
|
28
|
+
wts.append(os.path.join(root, f))
|
|
29
|
+
if wts:
|
|
30
|
+
cols = st.columns(min(len(wts), MAX_WEIGHTS_DISPLAY))
|
|
31
|
+
for i, wp_disp in enumerate(sorted(wts)[:MAX_WEIGHTS_DISPLAY]):
|
|
32
|
+
sz = os.path.getsize(wp_disp) / (1024 * 1024)
|
|
33
|
+
bn = os.path.basename(wp_disp)
|
|
34
|
+
if "best" in bn and "inference" in bn:
|
|
35
|
+
lbl = "Best Inf"
|
|
36
|
+
elif "best" in bn and "fp16" in bn:
|
|
37
|
+
lbl = "Best FP16"
|
|
38
|
+
elif "best" in bn:
|
|
39
|
+
lbl = "Best"
|
|
40
|
+
elif "last" in bn and "inference" in bn:
|
|
41
|
+
lbl = "Last Inf"
|
|
42
|
+
elif "last" in bn:
|
|
43
|
+
lbl = "Last"
|
|
44
|
+
else:
|
|
45
|
+
lbl = bn[:8]
|
|
46
|
+
with cols[i]:
|
|
47
|
+
st.metric(lbl, f"{sz:.1f}MB")
|
|
48
|
+
else:
|
|
49
|
+
st.caption(f"No weights found in `{save_dir}`")
|
|
50
|
+
else:
|
|
51
|
+
st.caption("No save directory set. Train a model first or set `Save Dir` on Model page.")
|
|
52
|
+
|
|
53
|
+
st.markdown("#### Config")
|
|
54
|
+
src = st.radio("Source", ["Best (inference)", "Best (FP16)", "Last", "Custom"],
|
|
55
|
+
key="weights_source", horizontal=True)
|
|
56
|
+
if src == "Custom":
|
|
57
|
+
st.text_input("Path", placeholder="model.pth", key="export_weights_path")
|
|
58
|
+
|
|
59
|
+
sd = _get_export_save_dir()
|
|
60
|
+
if src != "Custom":
|
|
61
|
+
targets = EXPORT_WEIGHT_MAP.get(src, [CKPT_BEST])
|
|
62
|
+
wp = ""
|
|
63
|
+
if os.path.isdir(sd):
|
|
64
|
+
for root, _dirs, files in os.walk(sd):
|
|
65
|
+
for target in targets:
|
|
66
|
+
if target in files:
|
|
67
|
+
wp = os.path.join(root, target)
|
|
68
|
+
break
|
|
69
|
+
if wp:
|
|
70
|
+
break
|
|
71
|
+
if not wp:
|
|
72
|
+
wp = os.path.join(sd, targets[0])
|
|
73
|
+
else:
|
|
74
|
+
wp = st.session_state.get("export_weights_path", "")
|
|
75
|
+
|
|
76
|
+
oc = st.columns(4)
|
|
77
|
+
with oc[0]:
|
|
78
|
+
export_fmt = st.selectbox("Format", EXPORT_FORMATS, key="export_format")
|
|
79
|
+
with oc[1]:
|
|
80
|
+
img_sz = st.select_slider("Img", EXPORT_IMG_SIZES, value=EXPORT_IMG_SIZES[0], key="export_img_size")
|
|
81
|
+
with oc[2]:
|
|
82
|
+
st.number_input("Opset", EXPORT_OPSET_MIN, EXPORT_OPSET_MAX, EXPORT_OPSET_DEFAULT, key="export_opset")
|
|
83
|
+
with oc[3]:
|
|
84
|
+
st.checkbox("Dynamic", True, key="export_dynamic")
|
|
85
|
+
|
|
86
|
+
return wp, img_sz, export_fmt
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
"""Inference — shared state, constants, and utility functions."""
|
|
2
|
+
|
|
3
|
+
import streamlit as st
|
|
4
|
+
from flashstudio.constants import COCO_CLASSES
|
|
5
|
+
|
|
6
|
+
try:
|
|
7
|
+
from flashstudio.components.zone_drawer import zone_drawer # noqa: F401
|
|
8
|
+
_HAS_ZONE_DRAWER = True
|
|
9
|
+
except ImportError:
|
|
10
|
+
_HAS_ZONE_DRAWER = False
|
|
11
|
+
|
|
12
|
+
try:
|
|
13
|
+
from flashdet import Predictor, FlashTracker # noqa: F401
|
|
14
|
+
_HAS_PREDICTOR = True
|
|
15
|
+
except ImportError:
|
|
16
|
+
_HAS_PREDICTOR = False
|
|
17
|
+
|
|
18
|
+
_SOLUTIONS_AVAILABLE = {}
|
|
19
|
+
_solution_classes = [
|
|
20
|
+
"ObjectCounter", "RegionCounter", "SpeedEstimator", "Heatmap",
|
|
21
|
+
"SecurityAlarm", "QueueManager", "ParkingManager", "TrafficFlow",
|
|
22
|
+
"DwellTimeAnalyzer", "DistanceCalculator", "TrajectoryVisualizer",
|
|
23
|
+
"ObjectBlurrer", "CrowdDensity", "WorkoutMonitor", "ObjectCropper",
|
|
24
|
+
"AnalyticsDashboard",
|
|
25
|
+
]
|
|
26
|
+
for _cls_name in _solution_classes:
|
|
27
|
+
try:
|
|
28
|
+
_mod = __import__("flashdet.solutions", fromlist=[_cls_name])
|
|
29
|
+
_SOLUTIONS_AVAILABLE[_cls_name] = getattr(_mod, _cls_name)
|
|
30
|
+
except (ImportError, AttributeError):
|
|
31
|
+
pass
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
SOLUTIONS = {
|
|
35
|
+
"None (Detection Only)": {"desc": "Standard object detection", "needs_zone": False},
|
|
36
|
+
"Object Counter (Line)": {"desc": "Count objects crossing line", "needs_zone": True, "zone_type": "line"},
|
|
37
|
+
"Region Counter (Polygon)": {"desc": "Count in polygon zones", "needs_zone": True, "zone_type": "polygon"},
|
|
38
|
+
"Speed Estimator": {"desc": "Estimate speed", "needs_zone": True, "zone_type": "line"},
|
|
39
|
+
"Heatmap": {"desc": "Activity heatmap", "needs_zone": False},
|
|
40
|
+
"Security Alarm (Zone)": {"desc": "Alert on zone entry", "needs_zone": True, "zone_type": "polygon"},
|
|
41
|
+
"Trajectory Visualizer": {"desc": "Motion trails", "needs_zone": False},
|
|
42
|
+
"Object Blurrer": {"desc": "Blur detections", "needs_zone": False},
|
|
43
|
+
"Queue Manager": {"desc": "Queue lengths", "needs_zone": True, "zone_type": "polygon"},
|
|
44
|
+
"Crowd Density": {"desc": "Grid density", "needs_zone": False},
|
|
45
|
+
"Parking Manager": {"desc": "Parking occupancy", "needs_zone": True, "zone_type": "polygon"},
|
|
46
|
+
"Traffic Flow": {"desc": "Traffic analysis", "needs_zone": True, "zone_type": "line"},
|
|
47
|
+
"Dwell Time Analyzer": {"desc": "Time in zones", "needs_zone": True, "zone_type": "polygon"},
|
|
48
|
+
"Distance Calculator": {"desc": "Pairwise distances", "needs_zone": False},
|
|
49
|
+
"Workout Monitor": {"desc": "Exercise reps", "needs_zone": False},
|
|
50
|
+
"Object Cropper": {"desc": "Crop detections", "needs_zone": False},
|
|
51
|
+
"Analytics Dashboard": {"desc": "Real-time stats", "needs_zone": False},
|
|
52
|
+
}
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def _get_class_names():
|
|
56
|
+
"""Get class names from session state, falling back to COCO."""
|
|
57
|
+
raw = st.session_state.get("class_names", "")
|
|
58
|
+
if isinstance(raw, list):
|
|
59
|
+
names = [c.strip() for c in raw if c.strip()]
|
|
60
|
+
elif isinstance(raw, str) and raw.strip():
|
|
61
|
+
names = [c.strip() for c in raw.strip().split("\n") if c.strip()]
|
|
62
|
+
else:
|
|
63
|
+
names = []
|
|
64
|
+
return names if names else COCO_CLASSES
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def _get_device_options():
|
|
68
|
+
devices = ["cpu"]
|
|
69
|
+
try:
|
|
70
|
+
import torch
|
|
71
|
+
if torch.cuda.is_available():
|
|
72
|
+
for i in range(torch.cuda.device_count()):
|
|
73
|
+
devices.append(f"cuda:{i} ({torch.cuda.get_device_name(i)})")
|
|
74
|
+
except ImportError:
|
|
75
|
+
pass
|
|
76
|
+
return devices
|