flashstudio 0.0.1__py3-none-any.whl

flashstudio/pages/data/helpers.py

@@ -0,0 +1,338 @@
+"""Ingest helper functions — extraction, conversion, download runners."""
+
+import os
+import streamlit as st
+
+from flashstudio.constants import IMG_EXTENSIONS, DEFAULT_DATA_DIR, INFER_NUM_CLASSES
+from flashstudio.pages.data._common import (
+    _HAS_FLASHDET, _FLASHDET_INSTALL_MSG,
+    _auto_detect_classes, _check_network, _dataset_already_downloaded,
+)
+
+try:
+    from flashdet import download_dataset
+    from flashdet.data import convert_dataset, detect_dataset_format
+except ImportError:
+    pass
+
+
+def _extract_upload(uploaded_file, split):
+    import zipfile
+    import tarfile
+    import tempfile
+    import shutil
+    name = uploaded_file.name
+    out = os.path.join(DEFAULT_DATA_DIR, "uploaded", split)
+    os.makedirs(out, exist_ok=True)
+    try:
+        if name.endswith(".zip"):
+            tmp = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
+            tmp.write(uploaded_file.read()); tmp.flush(); tmp.close()
+            uploaded_file.seek(0)
+            if not zipfile.is_zipfile(tmp.name):
+                os.unlink(tmp.name)
+                st.error(f"❌ **Invalid ZIP file:** `{name}` is not a valid zip archive. Please upload a correct file.")
+                return
+            with zipfile.ZipFile(tmp.name, "r") as zf:
+                members = zf.namelist()
+                total = len(members)
+                if total == 0:
+                    os.unlink(tmp.name)
+                    st.error(f"❌ **Empty archive:** `{name}` contains no files.")
+                    return
+                progress = st.progress(0, text=f"Extracting {name}  —  0 / {total} files")
+                img_count = 0
+                for i, member in enumerate(members, 1):
+                    zf.extract(member, out)
+                    if member.lower().endswith(IMG_EXTENSIONS):
+                        img_count += 1
+                    progress.progress(i / total, text=f"Extracting {name}  —  {i} / {total} files  ·  {img_count} images")
+                progress.progress(1.0, text=f"✓ Extracted {total} files  ·  {img_count} images")
+            os.unlink(tmp.name)
+        elif name.endswith((".tar", ".tar.gz", ".tgz")) or (name.endswith(".gz") and ".tar." in name):
+            tmp = tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False)
+            tmp.write(uploaded_file.read()); tmp.flush(); tmp.close()
+            uploaded_file.seek(0)
+            try:
+                with tarfile.open(tmp.name, "r:*") as tf:
+                    members = tf.getmembers()
+                    total = len(members)
+                    if total == 0:
+                        os.unlink(tmp.name)
+                        st.error(f"❌ **Empty archive:** `{name}` contains no files.")
+                        return
+                    progress = st.progress(0, text=f"Extracting {name}  —  0 / {total} files")
+                    img_count = 0
+                    for i, member in enumerate(members, 1):
+                        tf.extract(member, out)
+                        if member.name.lower().endswith(IMG_EXTENSIONS):
+                            img_count += 1
+                        progress.progress(i / total, text=f"Extracting {name}  —  {i} / {total} files  ·  {img_count} images")
+                    progress.progress(1.0, text=f"✓ Extracted {total} files  ·  {img_count} images")
+            except tarfile.TarError:
+                os.unlink(tmp.name)
+                st.error(f"❌ **Invalid TAR file:** `{name}` is not a valid tar archive. Please upload a correct file.")
+                return
+            os.unlink(tmp.name)
+        else:
+            st.error(f"Unsupported: {name}")
+            return
+
+        # --- Post-extraction: validate, detect format, find images ---
+        if not os.listdir(out):
+            shutil.rmtree(out, ignore_errors=True)
+            st.error(f"❌ **Empty archive:** `{name}` contains no files. Please upload a valid dataset zip.")
+            return
+
+        root = _find_root(out)
+        st.session_state["dataset_output_path"] = root
+
+        fmt = "unknown"
+        if _HAS_FLASHDET:
+            fmt = detect_dataset_format(root)
+        st.session_state["detected_format"] = fmt
+
+        if fmt in ("txt", "voc"):
+            _convert_to_coco(root, split, fmt)
+            return
+
+        img_dir = _find_images(root)
+        if img_dir:
+            st.session_state[f"{split}_img_path"] = img_dir
+            ann = _find_ann(root)
+            if ann:
+                st.session_state[f"{split}_ann_path"] = ann
+            label_dir = _find_labels(root)
+            if label_dir:
+                st.session_state[f"{split}_label_path"] = label_dir
+            st.session_state["dataset_name"] = name.rsplit(".", 1)[0]
+
+            parent = os.path.dirname(img_dir)
+            for sibling, key in [("valid", "val"), ("val", "val"), ("test", "test"), ("train", "train")]:
+                sib_path = os.path.join(parent, sibling)
+                if os.path.isdir(sib_path) and key != split:
+                    sk = "val_img_path" if key in ("val",) else f"{key}_img_path"
+                    if not st.session_state.get(sk):
+                        st.session_state[sk] = sib_path
+
+            detected_cls = _auto_detect_classes()
+            _existing_cls = st.session_state.get("class_names", "")
+            if isinstance(_existing_cls, list):
+                _existing_cls = "\n".join(_existing_cls)
+            if detected_cls and not _existing_cls.strip():
+                st.session_state["class_names"] = "\n".join(detected_cls)
+                st.session_state["num_classes"] = len(detected_cls)
+
+            from flashstudio.utils import flash
+            flash(f"Extracted `{name}` to {img_dir}", "success")
+            st.rerun()
+        else:
+            st.session_state[f"{split}_img_path"] = root
+            st.session_state["dataset_name"] = name.rsplit(".", 1)[0]
+            from flashstudio.utils import flash
+            flash(f"Extracted `{name}` (no images dir found, using root)", "warning")
+            st.rerun()
+    except Exception as e:
+        st.error(f"Extract failed: {e}")
+
+
+def _find_root(base):
+    entries = os.listdir(base)
+    if len(entries) == 1:
+        s = os.path.join(base, entries[0])
+        if os.path.isdir(s):
+            return s
+    for e in entries:
+        full = os.path.join(base, e)
+        if os.path.isdir(full) and _HAS_FLASHDET and detect_dataset_format(full) != "unknown":
+            return full
+    return base
+
+
+def _convert_to_coco(src, split, src_fmt):
+    if not _HAS_FLASHDET:
+        st.error(_FLASHDET_INSTALL_MSG)
+        return
+    out = os.path.join(DEFAULT_DATA_DIR, "uploaded", f"{split}_coco")
+    cls_names = None
+    raw = st.session_state.get("class_names", "")
+    if isinstance(raw, list):
+        raw = "\n".join(raw)
+    if raw.strip():
+        cls_names = [c.strip() for c in raw.strip().replace(",", "\n").split("\n") if c.strip()]
+    with st.spinner("Converting to COCO..."):
+        try:
+            stats = convert_dataset(source_dir=src, output_dir=out, target_format="coco", class_names=cls_names)
+            if stats.get("status") == "already_in_target_format":
+                out = src
+            st.session_state["dataset_output_path"] = out
+            st.session_state["ann_format"] = "COCO JSON"
+            t = os.path.join(out, "train")
+            v = os.path.join(out, "valid")
+            if not os.path.isdir(v):
+                v = os.path.join(out, "val")
+            if os.path.isdir(t):
+                st.session_state[f"{split}_img_path"] = t
+            if os.path.isdir(v):
+                st.session_state["val_img_path"] = v
+            from flashstudio.utils import flash
+            flash("Converted to COCO format", "success")
+            st.rerun()
+        except Exception as e:
+            st.error(f"Conversion failed: {e}")
+
+
+def _find_images(base):
+    best = None
+    best_count = 0
+    for root, dirs, files in os.walk(base):
+        img_count = len([f for f in files if f.lower().endswith(IMG_EXTENSIONS)])
+        if img_count > best_count:
+            best = root
+            best_count = img_count
+        for d in dirs:
+            if d.lower() in ("images", "imgs"):
+                c = os.path.join(root, d)
+                for sub_root, _, sub_files in os.walk(c):
+                    sc = len([f for f in sub_files if f.lower().endswith(IMG_EXTENSIONS)])
+                    if sc > best_count:
+                        best = sub_root
+                        best_count = sc
+    return best if best_count >= 1 else None
+
+
+def _find_ann(base):
+    for root, _, files in os.walk(base):
+        for f in files:
+            if f.endswith(".json") and ("annotation" in f.lower() or "coco" in f.lower()):
+                return os.path.join(root, f)
+    return None
+
+
+def _find_labels(base):
+    """Find YOLO-format label directory (folder named 'labels' with .txt files)."""
+    for root, dirs, files in os.walk(base):
+        for d in dirs:
+            if d.lower() == "labels":
+                lbl_dir = os.path.join(root, d)
+                for sub_root, _, sub_files in os.walk(lbl_dir):
+                    if any(f.endswith(".txt") for f in sub_files):
+                        return sub_root
+    return None
+
+
+def _use_existing(path, ds_info):
+    from flashstudio.utils import flash
+    st.session_state["dataset_output_path"] = path
+    name = ds_info.get("name", os.path.basename(path))
+    st.session_state["dataset_name"] = name
+    st.session_state["dataset_classes"] = ds_info.get("classes", ds_info.get("cls", INFER_NUM_CLASSES))
+    if ds_info.get("id"):
+        st.session_state["dataset_id"] = ds_info["id"]
+    t = os.path.join(path, "train")
+    v = os.path.join(path, "valid")
+    if not os.path.isdir(v):
+        v = os.path.join(path, "val")
+    if os.path.isdir(t):
+        st.session_state["train_img_path"] = t
+    if os.path.isdir(v):
+        st.session_state["val_img_path"] = v
+
+    detected_cls = _auto_detect_classes()
+    if detected_cls:
+        st.session_state["class_names"] = "\n".join(detected_cls)
+        st.session_state["num_classes"] = len(detected_cls)
+
+    flash(f"Dataset `{name}` loaded", "success")
+    st.rerun()
+
+
+def _run_flashdet_download(dataset_id):
+    if not _HAS_FLASHDET:
+        st.error(_FLASHDET_INSTALL_MSG)
+        return
+    out = st.session_state.get("download_output_dir") or os.path.join(DEFAULT_DATA_DIR, dataset_id)
+    existing = _dataset_already_downloaded(dataset_id, out)
+    if existing:
+        st.session_state["dataset_output_path"] = existing
+        for d, k in [("train", "train_img_path"), ("valid", "val_img_path")]:
+            p = os.path.join(existing, d)
+            if os.path.isdir(p):
+                st.session_state[k] = p
+        detected_cls = _auto_detect_classes()
+        if detected_cls:
+            st.session_state["class_names"] = "\n".join(detected_cls)
+            st.session_state["num_classes"] = len(detected_cls)
+        st.rerun()
+        return
+    if not _check_network():
+        st.error("Network unavailable.")
+        return
+    with st.spinner(f"Downloading {dataset_id}..."):
+        try:
+            result = download_dataset(dataset_id=dataset_id, output_dir=out)
+            st.session_state["dataset_output_path"] = result
+            for d, k in [("train", "train_img_path"), ("valid", "val_img_path")]:
+                p = os.path.join(result, d)
+                if os.path.isdir(p):
+                    st.session_state[k] = p
+            detected_cls = _auto_detect_classes()
+            if detected_cls:
+                st.session_state["class_names"] = "\n".join(detected_cls)
+                st.session_state["num_classes"] = len(detected_cls)
+            from flashstudio.utils import flash
+            flash(f"Downloaded `{dataset_id}` successfully", "success")
+            st.rerun()
+        except Exception as e:
+            st.error(f"Download failed: {e}")
+
+
+def _run_external_download(ds):
+    import urllib.request
+    import urllib.error
+    import zipfile
+    import tempfile
+    url, name = ds.get("url", ""), ds.get("name", "dataset")
+    if not url or "VisDrone" in url:
+        st.info(f"Manual: [{url}]({url})")
+        return
+    out = os.path.join(DEFAULT_DATA_DIR, name.lower().replace(" ", "_").replace("(", "").replace(")", ""))
+    if os.path.isdir(out) and os.listdir(out):
+        st.session_state.update({"dataset_output_path": out, "dataset_name": name})
+        st.rerun()
+        return
+    with st.spinner(f"Downloading {name}..."):
+        tmp_path = None
+        try:
+            tmp = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
+            tmp_path = tmp.name; tmp.close()
+            urllib.request.urlretrieve(url, tmp_path)
+            if os.path.getsize(tmp_path) < 1000 or not zipfile.is_zipfile(tmp_path):
+                st.error(f"Invalid download. Manual: [{url}]({url})")
+                return
+            os.makedirs(out, exist_ok=True)
+            with zipfile.ZipFile(tmp_path, "r") as zf:
+                members = zf.namelist()
+                total = len(members)
+                progress = st.progress(0, text=f"Extracting {name}  —  0 / {total}")
+                img_count = 0
+                for i, member in enumerate(members, 1):
+                    zf.extract(member, out)
+                    if member.lower().endswith(IMG_EXTENSIONS):
+                        img_count += 1
+                    progress.progress(i / total, text=f"Extracting {name}  —  {i} / {total} files  ·  {img_count} images")
+                progress.progress(1.0, text=f"Extracted {total} files  ·  {img_count} images")
+            st.session_state.update({"dataset_output_path": out, "dataset_name": name})
+            for d, k in [("train", "train_img_path"), ("valid", "val_img_path")]:
+                p = os.path.join(out, d)
+                if os.path.isdir(p):
+                    st.session_state[k] = p
+            st.rerun()
+        except Exception as e:
+            st.error(f"Error: {e}")
+        finally:
+            if tmp_path and os.path.exists(tmp_path):
+                try:
+                    os.unlink(tmp_path)
+                except OSError:
+                    pass

flashstudio/pages/data/page.py

@@ -0,0 +1,50 @@
+"""FlashStudio — Data Page (ultra-compact, no-scroll)."""
+
+import os
+import streamlit as st
+
+from flashstudio.constants import IMG_EXTENSIONS
+from flashstudio.pages.data.upload.tab import _render_upload
+from flashstudio.pages.data.download.tab import _render_download
+from flashstudio.pages.data.preview.tab import _render_preview
+from flashstudio.pages.data.verify.tab import _render_verify
+
+
+def render_data_page():
+    from flashstudio.components.styles import render_page_header
+    from flashstudio.utils import show_flashes
+    render_page_header("", "Data")
+    show_flashes()
+
+    # Status banner
+    dataset = st.session_state.get("dataset_name")
+    if dataset:
+        train_path = st.session_state.get("train_img_path", "")
+        val_path = st.session_state.get("val_img_path", "")
+        tc = len([f for f in os.listdir(train_path) if f.lower().endswith(IMG_EXTENSIONS)]) if train_path and os.path.isdir(train_path) else 0
+        vc = len([f for f in os.listdir(val_path) if f.lower().endswith(IMG_EXTENSIONS)]) if val_path and os.path.isdir(val_path) else 0
+        dc1, dc2 = st.columns([8, 1])
+        with dc1:
+            st.success(f"**{dataset}** — {tc} train / {vc} val images")
+        with dc2:
+            if st.button("Clear", key="clear_dataset_btn", use_container_width=True):
+                from flashstudio.utils import flash
+                for k in ["dataset_name", "train_img_path", "val_img_path",
+                           "train_label_path", "val_label_path",
+                           "dataset_output_path", "dataset_id", "dataset_classes",
+                           "class_names", "num_classes",
+                           "detected_format", "_last_verify_ok"]:
+                    st.session_state.pop(k, None)
+                flash("Dataset cleared", "info")
+                st.rerun()
+
+    tab_upload, tab_download, tab_preview, tab_verify = st.tabs(["Upload", "Download", "Preview", "Verify"])
+
+    with tab_upload:
+        _render_upload()
+    with tab_download:
+        _render_download()
+    with tab_preview:
+        _render_preview()
+    with tab_verify:
+        _render_verify()

flashstudio/pages/data/preview/__init__.py

@@ -0,0 +1,3 @@
+"""Preview tab sub-package."""
+
+from flashstudio.pages.data.preview.tab import _render_preview  # noqa: F401

flashstudio/pages/data/preview/tab.py

@@ -0,0 +1,139 @@
+"""Preview tab — image display with COCO/YOLO bbox overlay."""
+
+import os
+import json
+import streamlit as st
+
+from flashstudio.constants import IMG_EXTENSIONS, BBOX_COLORS_RGB, MAX_PREVIEW_COLS
+
+
+def _render_preview():
+    train_path = st.session_state.get("train_img_path", "")
+    val_path = st.session_state.get("val_img_path", "")
+    if not train_path or not os.path.isdir(train_path):
+        st.info("No dataset loaded yet.")
+        return
+
+    c1, c2, c3 = st.columns([2, 1, 1])
+    with c1:
+        split = st.radio("Split", ["Train", "Val"], horizontal=True, key="prev_split")
+    preview_path = train_path if split == "Train" else val_path
+    if not preview_path or not os.path.isdir(preview_path):
+        st.warning("Path not set.")
+        return
+
+    all_imgs = sorted([f for f in os.listdir(preview_path) if f.lower().endswith(IMG_EXTENSIONS)])
+    if not all_imgs:
+        st.warning("No images.")
+        return
+
+    with c2:
+        n = st.slider("Count", 1, min(12, len(all_imgs)), min(4, len(all_imgs)), key="prev_n")
+    with c3:
+        bbox = st.checkbox("Boxes", value=True, key="prev_bb")
+
+    if "prev_start" not in st.session_state:
+        st.session_state["prev_start"] = 0
+    if st.session_state["prev_start"] >= len(all_imgs):
+        st.session_state["prev_start"] = 0
+    start = st.session_state["prev_start"]
+
+    nav_prev, nav_info, nav_next = st.columns([1, 3, 1])
+    with nav_prev:
+        if st.button("Prev", disabled=(start <= 0), key="prev_btn"):
+            st.session_state["prev_start"] = max(0, start - n)
+            st.rerun()
+    with nav_info:
+        end_idx = min(start + n, len(all_imgs))
+        st.markdown(f'<p style="text-align:center;margin:0;padding:0.3rem 0;color:#6B7280;">'
+                    f'{start + 1}–{end_idx} of {len(all_imgs)}</p>', unsafe_allow_html=True)
+    with nav_next:
+        if st.button("Next", disabled=(start + n >= len(all_imgs)), key="next_btn"):
+            st.session_state["prev_start"] = min(start + n, max(len(all_imgs) - n, 0))
+            st.rerun()
+
+    # Load COCO annotations if available
+    ann_file = os.path.join(preview_path, "_annotations.coco.json")
+    ann_data = None
+    if os.path.isfile(ann_file):
+        try:
+            with open(ann_file) as f:
+                ann_data = json.load(f)
+        except Exception:
+            pass
+
+    img_id_map, ann_by_img, cat_map = {}, {}, {}
+    if ann_data:
+        for c in ann_data.get("categories", []):
+            cat_map[c["id"]] = c["name"]
+        for im in ann_data.get("images", []):
+            img_id_map[im["file_name"]] = im["id"]
+        for a in ann_data.get("annotations", []):
+            ann_by_img.setdefault(a["image_id"], []).append(a)
+
+    # YOLO label directory — check sibling 'labels' folder or session state
+    yolo_label_dir = st.session_state.get(f"{split.lower()}_label_path", "")
+    if not yolo_label_dir or not os.path.isdir(yolo_label_dir):
+        candidate = preview_path.replace("/images/", "/labels/").replace("\\images\\", "\\labels\\")
+        if os.path.isdir(candidate):
+            yolo_label_dir = candidate
+        else:
+            yolo_label_dir = ""
+
+    from PIL import Image, ImageDraw
+    selected = all_imgs[start:start + n]
+
+    total_boxes = 0
+    for row in range(0, len(selected), MAX_PREVIEW_COLS):
+        cols = st.columns(MAX_PREVIEW_COLS)
+        for j, col in enumerate(cols):
+            idx = row + j
+            if idx >= len(selected):
+                break
+            img_name = selected[idx]
+            with col:
+                try:
+                    img = Image.open(os.path.join(preview_path, img_name)).convert("RGB")
+                    w_img, h_img = img.size
+                    if bbox:
+                        draw = ImageDraw.Draw(img)
+                        drawn = False
+                        if ann_data and img_name in img_id_map:
+                            for a in ann_by_img.get(img_id_map[img_name], []):
+                                b = a.get("bbox", [])
+                                if len(b) == 4:
+                                    x, y, w, h = b
+                                    color = BBOX_COLORS_RGB[a.get("category_id", 0) % len(BBOX_COLORS_RGB)]
+                                    draw.rectangle([x, y, x+w, y+h], outline=color, width=2)
+                                    total_boxes += 1
+                                    drawn = True
+                        if not drawn and yolo_label_dir:
+                            lbl_name = os.path.splitext(img_name)[0] + ".txt"
+                            lbl_path = os.path.join(yolo_label_dir, lbl_name)
+                            if os.path.isfile(lbl_path):
+                                with open(lbl_path) as lf:
+                                    for line in lf:
+                                        parts = line.strip().split()
+                                        if len(parts) >= 5:
+                                            cls_id = int(parts[0])
+                                            cx, cy, bw, bh = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4])
+                                            x1 = int((cx - bw / 2) * w_img)
+                                            y1 = int((cy - bh / 2) * h_img)
+                                            x2 = int((cx + bw / 2) * w_img)
+                                            y2 = int((cy + bh / 2) * h_img)
+                                            color = BBOX_COLORS_RGB[cls_id % len(BBOX_COLORS_RGB)]
+                                            draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
+                                            total_boxes += 1
+                    st.image(img, caption=f"{img_name[:15]} {w_img}x{h_img}", use_container_width=True)
+                except Exception as e:
+                    st.error(str(e)[:40])
+
+    m1, m2, m3 = st.columns(3)
+    with m1:
+        st.metric("Images", len(all_imgs))
+    with m2:
+        ann_count = len(ann_data.get("annotations", [])) if ann_data else total_boxes
+        st.metric("Annotations", ann_count)
+    with m3:
+        cat_count = len(cat_map) if cat_map else "—"
+        st.metric("Categories", cat_count)

flashstudio/pages/data/upload/__init__.py

@@ -0,0 +1,3 @@
+"""Upload tab sub-package."""
+
+from flashstudio.pages.data.upload.tab import _render_upload  # noqa: F401