flashstudio 0.0.1__py3-none-any.whl

flashstudio/pages/inference/data/tab.py

@@ -0,0 +1,121 @@
+"""Inference — Data tab."""
+
+import os
+import tempfile
+import streamlit as st
+from PIL import Image
+from flashstudio.constants import (
+    INFER_MAX_FRAMES, INFER_STREAM_DURATION,
+    INFER_FRAME_SKIP, MAX_PREVIEW_IMAGES,
+)
+
+
+def _tab_data():
+    input_type = st.radio("Input", ["Images", "Video", "RTSP"], horizontal=True, key="infer_input_type")
+
+    if input_type == "Images":
+        uploaded = st.file_uploader("Images", type=["jpg", "jpeg", "png", "webp", "bmp"],
+                                    accept_multiple_files=True, key="infer_images")
+        if uploaded:
+            cols = st.columns(min(len(uploaded), MAX_PREVIEW_IMAGES))
+            for i, f in enumerate(uploaded[:MAX_PREVIEW_IMAGES]):
+                with cols[i]:
+                    st.image(f, caption=f.name[:10], use_container_width=True)
+            if len(uploaded) > MAX_PREVIEW_IMAGES:
+                st.caption(f"+{len(uploaded) - MAX_PREVIEW_IMAGES} more")
+
+    elif input_type == "Video":
+        vc1, vc2 = st.columns([3, 1])
+        with vc1:
+            st.file_uploader("Video", type=["mp4", "avi", "mov", "mkv"], key="infer_video")
+            vid = st.session_state.get("infer_video")
+            if vid:
+                st.caption(f"{vid.name} ({vid.size / 1e6:.1f}MB)")
+                meta = _get_video_metadata(vid)
+                if meta:
+                    mc = st.columns(4)
+                    with mc[0]: st.metric("Res", f"{meta['width']}x{meta['height']}")
+                    with mc[1]: st.metric("FPS", f"{meta['fps']:.0f}")
+                    with mc[2]: st.metric("Frames", meta["total_frames"])
+                    with mc[3]: st.metric("Dur", f"{meta['total_frames'] / max(meta['fps'], 1):.0f}s")
+        with vc2:
+            st.number_input("Max frames (0=all)", 0, 100000, INFER_MAX_FRAMES, key="infer_max_frames")
+            st.select_slider("Skip N", [1, 2, 3, 5, 10], value=INFER_FRAME_SKIP, key="infer_frame_skip")
+            st.checkbox("Save output", True, key="infer_save_video")
+
+    else:
+        sc1, sc2 = st.columns([4, 1])
+        with sc1:
+            st.text_input("RTSP URL", placeholder="rtsp://192.168.1.100:554/stream", key="infer_stream_url")
+        with sc2:
+            st.number_input("Duration (s)", 0, 3600, INFER_STREAM_DURATION, key="infer_stream_duration")
+        if st.session_state.get("infer_stream_url") and st.button("Test", key="test_stream"):
+            _test_stream()
+
+
+def _get_video_metadata(video_file):
+    tmp_path = None
+    try:
+        import cv2
+        tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+        tmp_path = tmp.name
+        tmp.write(video_file.read()); tmp.flush(); tmp.close()
+        video_file.seek(0)
+        cap = cv2.VideoCapture(tmp_path)
+        if cap.isOpened():
+            meta = {"width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
+                    "fps": cap.get(cv2.CAP_PROP_FPS) or 30, "total_frames": int(cap.get(cv2.CAP_PROP_FRAME_COUNT))}
+            cap.release()
+            return meta
+        cap.release()
+    except Exception:
+        pass
+    finally:
+        if tmp_path and os.path.exists(tmp_path):
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+    return None
+
+
+def _get_first_frame():
+    video = st.session_state.get("infer_video")
+    if video:
+        try:
+            import cv2
+            tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+            tmp.write(video.read()); tmp.flush(); tmp.close()
+            video.seek(0)
+            cap = cv2.VideoCapture(tmp.name)
+            ret, frame = cap.read(); cap.release()
+            os.unlink(tmp.name)
+            if ret:
+                return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        except Exception:
+            pass
+    images = st.session_state.get("infer_images")
+    if images:
+        try:
+            img = Image.open(images[0]).convert("RGB"); images[0].seek(0)
+            return img
+        except Exception:
+            pass
+    return None
+
+
+def _test_stream():
+    try:
+        import cv2
+        cap = cv2.VideoCapture(st.session_state["infer_stream_url"])
+        if cap.isOpened():
+            ret, frame = cap.read(); cap.release()
+            if ret:
+                st.success("Connected")
+                st.image(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)), use_container_width=True)
+            else:
+                st.error("No frame")
+        else:
+            st.error("Failed to connect")
+    except Exception as e:
+        st.error(str(e)[:60])

flashstudio/pages/inference/model/__init__.py

@@ -0,0 +1,3 @@
+"""Inference — Model tab sub-package."""
+
+from flashstudio.pages.inference.model.tab import _tab_model  # noqa: F401

flashstudio/pages/inference/model/tab.py

@@ -0,0 +1,62 @@
+"""Inference — Model tab."""
+
+import os
+import streamlit as st
+from flashstudio.constants import (
+    INFER_CONF_THRESHOLD, INFER_NMS_THRESHOLD, INFER_IMG_SIZE,
+    INFER_NUM_CLASSES, DEFAULT_SAVE_DIR, BEST_WEIGHT_PRIORITY,
+    CKPT_BEST_INFERENCE, FLASHDET_MODELS,
+)
+from flashstudio.pages.inference._common import _get_class_names, _get_device_options
+
+
+def _tab_model():
+    col1, col2 = st.columns(2)
+    with col1:
+        with st.container(border=True):
+            st.markdown("#### Architecture & Params")
+            st.selectbox("Model", list(FLASHDET_MODELS.keys()),
+                         key="infer_model_arch")
+            pc1, pc2 = st.columns(2)
+            with pc1:
+                st.slider("Confidence", 0.0, 1.0, INFER_CONF_THRESHOLD, 0.05, key="infer_conf")
+                st.number_input("Img Size", 320, 1920, INFER_IMG_SIZE, 32, key="infer_img_size")
+            with pc2:
+                st.slider("NMS IoU", 0.0, 1.0, INFER_NMS_THRESHOLD, 0.05, key="infer_nms")
+                st.selectbox("Device", _get_device_options(), key="infer_device")
+            dataset_nc = st.session_state.get("num_classes", 0)
+            default_nc = dataset_nc if dataset_nc > 0 else INFER_NUM_CLASSES
+            st.number_input("Classes", 1, 1000, default_nc, key="infer_num_classes")
+            filter_classes = _get_class_names()
+            st.multiselect("Filter (empty=all)", filter_classes, default=[], key="infer_class_filter")
+
+    with col2:
+        with st.container(border=True):
+            st.markdown("#### Weights")
+            ws = st.radio("Source", ["Upload", "Path", "Training output"], key="weight_source", horizontal=True)
+            if ws == "Upload":
+                up = st.file_uploader("File", type=["pt", "pth", "onnx", "engine"], key="infer_weights_file")
+                if up:
+                    st.success(f"{up.name} ({up.size / 1e6:.1f}MB)")
+            elif ws == "Path":
+                st.text_input("Path", placeholder="/path/best.pth", key="infer_weights_path")
+            else:
+                sd = st.session_state.get("save_dir", DEFAULT_SAVE_DIR)
+                best_path = ""
+                for candidate in BEST_WEIGHT_PRIORITY:
+                    if os.path.isdir(sd):
+                        for root, _d, files in os.walk(sd):
+                            if candidate in files:
+                                best_path = os.path.join(root, candidate)
+                                break
+                    if best_path:
+                        break
+                if not best_path:
+                    best_path = os.path.join(sd, CKPT_BEST_INFERENCE)
+                st.session_state["infer_weights_path"] = best_path
+                exists = os.path.isfile(best_path)
+                if exists:
+                    sz = os.path.getsize(best_path) / (1024 * 1024)
+                    st.success(f"`{os.path.basename(best_path)}` ({sz:.1f}MB)")
+                else:
+                    st.caption(f"Not found: `{best_path}`")

flashstudio/pages/inference/page.py

@@ -0,0 +1,42 @@
+"""FlashStudio — Inference Page (ultra-compact, no-scroll)."""
+
+import streamlit as st
+from flashstudio.components.styles import render_page_header
+from flashstudio.pages.inference._common import SOLUTIONS
+
+
+def render_inference_page():
+    from flashstudio.utils import show_flashes
+    from flashstudio.pages.inference.model.tab import _tab_model
+    from flashstudio.pages.inference.data.tab import _tab_data
+    from flashstudio.pages.inference.solution.tab import _tab_solution
+    from flashstudio.pages.inference.run.tab import _tab_run
+
+    render_page_header("", "Inference")
+    show_flashes()
+
+    has_model = bool(st.session_state.get("infer_weights_file") or st.session_state.get("infer_weights_path"))
+    has_data = bool(st.session_state.get("infer_images") or st.session_state.get("infer_video") or st.session_state.get("infer_stream_url"))
+    sol = SOLUTIONS.get(st.session_state.get("selected_solution", "None (Detection Only)"), {})
+    has_zone = not sol.get("needs_zone") or bool(st.session_state.get("zone_line_points") or st.session_state.get("zone_polygons"))
+    def dot(ok):
+        return '<span style="color:#22C55E;">&#9679;</span>' if ok else '<span style="color:#EF4444;">&#9679;</span>'
+    st.markdown(
+        f'<div class="ds-card-stats" style="padding:0.2rem 0;">'
+        f'<span>{dot(has_model)} Model</span>'
+        f'<span>{dot(has_data)} Data</span>'
+        f'<span>{dot(has_zone)} Zone</span>'
+        f'<span style="margin-left:auto;color:#9CA3AF;">Ready: {sum([has_model, has_data, has_zone])}/3</span>'
+        f'</div>',
+        unsafe_allow_html=True,
+    )
+
+    tab_model, tab_data, tab_solution, tab_run = st.tabs(["Model", "Data", "Solution", "Run"])
+    with tab_model:
+        _tab_model()
+    with tab_data:
+        _tab_data()
+    with tab_solution:
+        _tab_solution()
+    with tab_run:
+        _tab_run()

flashstudio/pages/inference/run/__init__.py

@@ -0,0 +1,3 @@
+"""Inference — Run tab sub-package."""
+
+from flashstudio.pages.inference.run.tab import _tab_run  # noqa: F401

flashstudio/pages/inference/run/detection.py

@@ -0,0 +1,133 @@
+"""Inference — Detection utilities."""
+
+import os
+import json
+import tempfile
+import streamlit as st
+import numpy as np
+from PIL import ImageDraw, ImageFont
+from flashstudio.constants import (
+    COCO_CLASSES, BBOX_COLORS_HEX, BBOX_COLORS_RGB,
+    INFER_CONF_THRESHOLD, INFER_NMS_THRESHOLD, INFER_IMG_SIZE,
+    INFER_DEFAULT_RESOLUTION, FONT_PATH_LINUX, FONT_SIZE_DEFAULT,
+)
+from flashstudio.pages.inference._common import _get_class_names, _HAS_PREDICTOR
+
+try:
+    from flashdet import Predictor
+except ImportError:
+    pass
+
+
+def _detect(image):
+    weights = st.session_state.get("infer_weights_path", "")
+    if not weights:
+        cached = st.session_state.get("_infer_weights_tmp_path", "")
+        if cached and os.path.isfile(cached):
+            weights = cached
+        else:
+            up = st.session_state.get("infer_weights_file")
+            if up:
+                tmp = tempfile.NamedTemporaryFile(suffix=".pth", delete=False)
+                tmp.write(up.read()); tmp.flush(); up.seek(0)
+                weights = tmp.name
+                st.session_state["_infer_weights_tmp_path"] = weights
+    if weights:
+        try:
+            return _detect_real(image, weights)
+        except Exception as e:
+            from flashstudio.utils import flash
+            flash(f"Model error: {e} — showing demo detections", "warning")
+    return _detect_demo(image)
+
+
+def _detect_real(image, weights_path):
+    if not _HAS_PREDICTOR:
+        raise ImportError("Predictor unavailable")
+    custom_cls = _get_class_names()
+    predictor = Predictor(model_path=weights_path, device=st.session_state.get("infer_device", "cpu").split(" ")[0],
+                          conf_thresh=st.session_state.get("infer_conf", INFER_CONF_THRESHOLD),
+                          nms_thresh=st.session_state.get("infer_nms", INFER_NMS_THRESHOLD),
+                          input_size=st.session_state.get("infer_img_size", INFER_IMG_SIZE),
+                          class_names=custom_cls if custom_cls != COCO_CLASSES else None)
+    raw = predictor(np.array(image))
+    class_filter = st.session_state.get("infer_class_filter", [])
+    label_list = custom_cls or COCO_CLASSES
+    dets = []
+    for bbox, score, cls_id in raw:
+        name = label_list[int(cls_id)] if int(cls_id) < len(label_list) else f"class_{cls_id}"
+        if class_filter and name not in class_filter:
+            continue
+        dets.append([name, f"{float(score):.2f}", str([int(x) for x in (bbox.tolist() if hasattr(bbox, 'tolist') else list(bbox))])])
+    return dets
+
+
+def _detect_demo(image):
+    w, h = image.size
+    rng = np.random.default_rng(hash(image.tobytes()[:100]) % (2 ** 31))
+    conf_thr = st.session_state.get("infer_conf", INFER_CONF_THRESHOLD)
+    class_filter = st.session_state.get("infer_class_filter", [])
+    dets = []
+    for _ in range(rng.integers(3, 8)):
+        cid = rng.integers(0, 10)
+        name = COCO_CLASSES[cid]
+        if class_filter and name not in class_filter:
+            continue
+        conf = rng.uniform(0.3, 0.98)
+        if conf < conf_thr:
+            continue
+        x1 = int(rng.integers(0, max(w - 100, 1)))
+        y1 = int(rng.integers(0, max(h - 100, 1)))
+        dets.append([name, f"{conf:.2f}", str([x1, y1, min(x1 + int(rng.integers(60, 200)), w), min(y1 + int(rng.integers(60, 200)), h)])])
+    return sorted(dets, key=lambda x: float(x[1]), reverse=True)
+
+
+def _draw_boxes_cv2(frame, raw_results):
+    """Draw bounding boxes on a BGR cv2 frame from Predictor output [(bbox, score, cls_id), ...]."""
+    import cv2
+    label_list = _get_class_names() or COCO_CLASSES
+    for bbox, score, cls_id in raw_results:
+        x1, y1, x2, y2 = [int(v) for v in (bbox.tolist() if hasattr(bbox, 'tolist') else list(bbox))]
+        cid = int(cls_id)
+        name = label_list[cid] if cid < len(label_list) else f"cls_{cid}"
+        label = f"{name} {float(score):.2f}"
+        cv2.rectangle(frame, (x1, y1), (x2, y2), BBOX_COLORS_RGB[cid % len(BBOX_COLORS_RGB)], 2)
+        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+        cv2.rectangle(frame, (x1, y1 - th - 6), (x1 + tw + 4, y1), BBOX_COLORS_RGB[cid % len(BBOX_COLORS_RGB)], -1)
+        cv2.putText(frame, label, (x1 + 2, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+    return frame
+
+
+def _draw_boxes(image, dets):
+    draw = ImageDraw.Draw(image)
+    try:
+        font = ImageFont.truetype(FONT_PATH_LINUX, FONT_SIZE_DEFAULT)
+    except (OSError, IOError):
+        font = ImageFont.load_default()
+    for i, (cls, conf, bbox_str) in enumerate(dets):
+        bbox = json.loads(bbox_str.replace("'", '"'))
+        color = BBOX_COLORS_HEX[i % len(BBOX_COLORS_HEX)]
+        draw.rectangle(bbox, outline=color, width=3)
+        label = f"{cls} {conf}"
+        tb = draw.textbbox((0, 0), label, font=font)
+        tw, th = tb[2] - tb[0], tb[3] - tb[1]
+        ly = max(bbox[1] - th - 6, 0)
+        draw.rectangle([bbox[0], ly, bbox[0] + tw + 8, ly + th + 6], fill=color)
+        draw.text((bbox[0] + 4, ly + 2), label, fill="white", font=font)
+    return image
+
+
+def _to_coco_format(results):
+    coco = {"images": [], "annotations": [],
+            "categories": [{"id": i, "name": n} for i, n in enumerate(COCO_CLASSES)]}
+    aid = 1
+    for iid, r in enumerate(results, 1):
+        coco["images"].append({"id": iid, "file_name": r["name"], "width": r.get("width", INFER_DEFAULT_RESOLUTION[0]), "height": r.get("height", INFER_DEFAULT_RESOLUTION[1])})
+        for cls, conf, bbox_str in r["dets"]:
+            bbox = json.loads(bbox_str.replace("'", '"'))
+            coco["annotations"].append({"id": aid, "image_id": iid,
+                                         "category_id": COCO_CLASSES.index(cls) if cls in COCO_CLASSES else 0,
+                                         "bbox": [bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]],
+                                         "area": (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), "score": float(conf), "iscrowd": 0})
+            aid += 1
+    return coco

flashstudio/pages/inference/run/tab.py

@@ -0,0 +1,161 @@
+"""Inference — Run tab."""
+
+import io
+import time
+import json
+import streamlit as st
+from PIL import Image
+from flashstudio.constants import DEFAULT_MODEL_ARCH
+from flashstudio.pages.inference.run.detection import (
+    _detect, _draw_boxes, _to_coco_format,
+)
+from flashstudio.pages.inference.run.video import _run_video
+
+
+def _tab_run():
+    m1, m2, m3, m4 = st.columns(4)
+    with m1:
+        st.metric("Model", st.session_state.get("infer_model_arch", DEFAULT_MODEL_ARCH).replace("FlashDet-", ""))
+    with m2:
+        st.metric("Input", st.session_state.get("infer_input_type", "Images"))
+    with m3:
+        sol_name = st.session_state.get("selected_solution", "Detection")
+        sol_display = "Detect" if "Detection Only" in sol_name else sol_name[:12]
+        st.metric("Solution", sol_display)
+    with m4:
+        st.metric("Device", st.session_state.get("infer_device", "cpu").split(" ")[0])
+
+    input_type = st.session_state.get("infer_input_type", "Images")
+    has_data = st.session_state.get("infer_images") or st.session_state.get("infer_video") or st.session_state.get("infer_stream_url")
+    if not has_data:
+        st.warning("No data. Upload in Data tab.")
+        return
+
+    rc1, rc2 = st.columns([1, 3])
+    with rc1:
+        run = st.button("Run Inference", type="primary", key="run_infer_btn", use_container_width=True)
+    with rc2:
+        st.checkbox("Per-frame details", False, key="show_per_frame")
+
+    if run:
+        if input_type == "Images":
+            _run_images()
+        elif input_type == "Video":
+            _run_video()
+        else:
+            st.info("Stream — coming soon")
+
+    if input_type == "Images":
+        _show_image_results()
+    else:
+        _show_video_results()
+
+
+def _run_images():
+    images = st.session_state.get("infer_images", [])
+    if not images:
+        return
+    results = []
+    total_time = 0
+    progress = st.progress(0)
+    for i, f in enumerate(images):
+        img = Image.open(f).convert("RGB")
+        t0 = time.perf_counter()
+        dets = _detect(img)
+        elapsed = time.perf_counter() - t0
+        total_time += elapsed
+        results.append({"name": f.name, "dets": dets, "annotated": _draw_boxes(img.copy(), dets),
+                         "inference_time_ms": elapsed * 1000, "width": img.width, "height": img.height})
+        progress.progress((i + 1) / len(images))
+    progress.empty()
+    st.session_state["infer_img_results"] = results
+    st.session_state["infer_total_time"] = total_time
+    st.rerun()
+
+
+def _show_image_results():
+    if "infer_img_results" not in st.session_state:
+        return
+    results = st.session_state["infer_img_results"]
+    total_dets = sum(len(r["dets"]) for r in results)
+    total_time = st.session_state.get("infer_total_time", 1)
+
+    c1, c2, c3, c4 = st.columns(4)
+    with c1:
+        st.metric("FPS", f"{len(results) / max(total_time, 0.001):.1f}")
+    with c2:
+        st.metric("Latency", f"{total_time / max(len(results), 1) * 1000:.1f}ms")
+    with c3:
+        st.metric("Detections", total_dets)
+    with c4:
+        st.metric("Images", len(results))
+
+    for r in results:
+        ci, ct = st.columns([1, 1])
+        with ci:
+            st.image(r["annotated"], caption=f"{r['name'][:15]} ({r['inference_time_ms']:.0f}ms)", use_container_width=True)
+        with ct:
+            if r["dets"]:
+                import pandas as pd
+                st.dataframe(pd.DataFrame(r["dets"], columns=["Class", "Conf", "BBox"]),
+                             use_container_width=True, hide_index=True, height=180)
+
+    with st.expander("Export"):
+        ec1, ec2, ec3 = st.columns(3)
+        with ec1:
+            import pandas as pd
+            all_d = [[r["name"]] + d for r in results for d in r["dets"]]
+            if all_d:
+                st.download_button("CSV", pd.DataFrame(all_d, columns=["File", "Class", "Conf", "BBox"]).to_csv(index=False),
+                                   "detections.csv", "text/csv", use_container_width=True)
+        with ec2:
+            st.download_button("JSON", json.dumps(_to_coco_format(results), indent=2),
+                               "coco.json", "application/json", use_container_width=True)
+        with ec3:
+            if results:
+                buf = io.BytesIO()
+                results[0]["annotated"].save(buf, format="PNG")
+                st.download_button("Image", buf.getvalue(), "annotated.png", "image/png", use_container_width=True)
+
+
+def _show_video_results():
+    if "video_results" not in st.session_state:
+        return
+    res = st.session_state["video_results"]
+    c1, c2, c3, c4 = st.columns(4)
+    with c1:
+        st.metric("Frames", res.get("frames", 0))
+    with c2:
+        st.metric("Detections", res.get("detections", 0))
+    with c3:
+        st.metric("FPS", f"{res.get('processing_fps', 0):.1f}")
+    with c4:
+        st.metric("Time", f"{res.get('elapsed_time', 0):.1f}s")
+
+    sol = res.get("solution_output", {})
+    in_count = sol.get("in_count") or sol.get("in", 0)
+    out_count = sol.get("out_count") or sol.get("out", 0)
+    if in_count or out_count:
+        sc1, sc2, sc3 = st.columns(3)
+        with sc1: st.metric("In", in_count)
+        with sc2: st.metric("Out", out_count)
+        with sc3: st.metric("Net", in_count - out_count)
+    region_data = sol.get("region_counts") or {k: v for k, v in sol.items()
+                                                if k not in ("in", "out", "in_count", "out_count", "total", "speeds") and isinstance(v, (int, float))}
+    if region_data:
+        cols = st.columns(min(len(region_data), 6))
+        for i, (r, c) in enumerate(list(region_data.items())[:6]):
+            with cols[i]:
+                st.metric(r, c)
+    if "speeds" in sol:
+        import pandas as pd
+        st.dataframe(pd.DataFrame(sol["speeds"], columns=["ID", "km/h"]), hide_index=True, height=150)
+
+    if res.get("frames_preview"):
+        cols = st.columns(min(len(res["frames_preview"]), 4))
+        for i, frame in enumerate(res["frames_preview"][:4]):
+            with cols[i]:
+                st.image(frame, use_container_width=True)
+
+    if res.get("output_path"):
+        st.caption(f"Output: `{res['output_path']}`")