import streamlit as st from pathlib import Path from io import BytesIO from PIL import Image, UnidentifiedImageError import matplotlib.pyplot as plt from imgshape.shape import get_shape from imgshape.analyze import analyze_type from imgshape.recommender import recommend_preprocessing from imgshape.augmentations import AugmentationRecommender from imgshape.report import generate_markdown_report, generate_html_report from imgshape.viz import plot_shape_distribution from imgshape.torchloader import to_torch_transform # Page config st.set_page_config(page_title="imgshape v2.1.3", layout="wide") st.title("🖼️ imgshape — Smart Dataset Assistant (v2.1.3)") st.markdown( "Upload an image or provide a dataset folder to analyze, " "recommend preprocessing, generate reports, and even get PyTorch transforms." ) # Sidebar for inputs st.sidebar.header("📂 Input") uploaded_file = st.sidebar.file_uploader( "Upload an image", type=["jpg", "jpeg", "png", "bmp", "tiff"] ) dataset_path = st.sidebar.text_input("Dataset folder path", "assets/sample_images") tabs = st.tabs(["📐 Shape", "🔍 Analyze", "🧠 Recommend", "📄 Report", "🔗 TorchLoader"]) # ------------------------- helpers ------------------------- def cache_uploaded_bytes(): """ Read uploaded_file once and cache raw bytes in session_state['uploaded_bytes']. Returns bytes or None. """ if uploaded_file is None: return None if "uploaded_bytes" not in st.session_state: try: st.session_state["uploaded_bytes"] = uploaded_file.read() except Exception as e: st.session_state["uploaded_bytes"] = None st.error(f"Error reading upload: {e}") return None return st.session_state["uploaded_bytes"] def load_uploaded_image_from_bytes(bytes_data): """ Build fresh BytesIO and PIL.Image from raw bytes. Returns (PIL.Image, BytesIO) or (None, None) on error. """ if not bytes_data: return None, None try: buf = BytesIO(bytes_data) pil_img = Image.open(BytesIO(bytes_data)).convert("RGB") return pil_img, buf except UnidentifiedImageError: return None, None except Exception as e: st.error(f"Unexpected error opening image: {e}") return None, None # ------------------------- SHAPE TAB ------------------------- with tabs[0]: st.subheader("📐 Shape Detection") bytes_data = cache_uploaded_bytes() if bytes_data: pil_img, buf = load_uploaded_image_from_bytes(bytes_data) if pil_img is None: st.error("Uploaded file is not a valid image. Please upload a PNG/JPEG/etc.") else: st.image(pil_img, caption="Uploaded Image", use_column_width=True) try: # get_shape accepts a PIL.Image or path depending on implementation shape = get_shape(pil_img) st.json({"shape": shape}) except Exception as e: st.error(f"Error in shape detection: {e}") else: st.info("Upload an image to see its shape.") # ------------------------- ANALYZE TAB ------------------------- with tabs[1]: st.subheader("🔍 Image Analysis") bytes_data = cache_uploaded_bytes() if bytes_data: pil_img, buf = load_uploaded_image_from_bytes(bytes_data) if pil_img is None: st.error("Uploaded file is not a valid image. Please upload a PNG/JPEG/etc.") else: buf.seek(0) try: analysis = analyze_type(buf) st.json(analysis) except Exception as e: st.error(f"Error in analysis: {e}") else: st.info("Upload an image to analyze.") st.subheader("📊 Dataset Visualization") if st.button("Plot Shape Distribution"): try: fig = plt.figure() plot_shape_distribution(dataset_path, save=False) st.pyplot(fig) except Exception as e: st.error(f"Error plotting dataset: {e}") # ------------------------- RECOMMEND TAB ------------------------- with tabs[2]: st.subheader("🧠 Preprocessing + Augmentation Recommendations") bytes_data = cache_uploaded_bytes() if bytes_data: pil_img, buf = load_uploaded_image_from_bytes(bytes_data) if pil_img is None: st.error("Uploaded file is not a valid image. Please upload a PNG/JPEG/etc.") else: try: rec = recommend_preprocessing(pil_img) # ✅ pass PIL.Image st.json({"preprocessing": rec}) except Exception as e: st.error(f"Error in preprocessing recommendation: {e}") # Augmentation plan try: ar = AugmentationRecommender(seed=42) buf.seek(0) analysis = analyze_type(buf) # still safe with BytesIO plan = ar.recommend_for_dataset( {"entropy_mean": analysis.get("entropy", 5.0), "image_count": 1} ) st.json({ "augmentation_plan": { "order": plan.recommended_order, "augmentations": [a.__dict__ for a in plan.augmentations] } }) except Exception as e: st.error(f"Error in augmentation plan: {e}") else: st.info("Upload an image to get recommendations.") # ------------------------- REPORT TAB ------------------------- with tabs[3]: st.subheader("📄 Dataset Report") if st.button("Generate Markdown + HTML Report"): try: stats = {"image_count": 1, "source_dir": dataset_path} rec = {} bytes_data = cache_uploaded_bytes() if bytes_data: # use bytes for preprocessing recommendation _, buf = load_uploaded_image_from_bytes(bytes_data) if buf is not None: buf.seek(0) rec = recommend_preprocessing(buf) ar = AugmentationRecommender(seed=42) plan = ar.recommend_for_dataset({"entropy_mean": 5.0, "image_count": 10}) md_path = Path("report.md") html_path = Path("report.html") generate_markdown_report( md_path, stats, {}, rec, {"augmentations": [a.__dict__ for a in plan.augmentations]} ) generate_html_report(md_path, html_path) st.success("Reports generated!") st.download_button("⬇️ Download Markdown", md_path.read_text(), file_name="report.md") st.download_button("⬇️ Download HTML", html_path.read_text(), file_name="report.html") except Exception as e: st.error(f"Error generating report: {e}") # ------------------------- TORCHLOADER TAB ------------------------- with tabs[4]: st.subheader("🔗 TorchLoader Export") bytes_data = cache_uploaded_bytes() if bytes_data: pil_img, buf = load_uploaded_image_from_bytes(bytes_data) if pil_img is None: st.error("Uploaded file is not a valid image. Please upload a PNG/JPEG/etc.") else: try: rec = recommend_preprocessing(pil_img) # ✅ use PIL.Image snippet_or_transform = to_torch_transform({}, rec) if isinstance(snippet_or_transform, str): st.code(snippet_or_transform, language="python") else: st.success("✅ torchvision.transforms.Compose object created") st.write(snippet_or_transform) except Exception as e: st.error(f"Error building Torch transform: {e}") else: st.info("Upload an image to export Torch transforms.") # ------------------------- FOOTER ------------------------- st.markdown("---") st.markdown( """
Connect with me