import streamlit as st
import os
from pathlib import Path
from PIL import Image

# Define the main dataset directory
DATASET_PATH = "Test dat"  # Update to your actual folder path

# Function to get subfolders (digit categories 0-9)
def get_subfolders(path):
    return sorted([f.name for f in Path(path).iterdir() if f.is_dir()])

# Function to get images from a selected subfolder
def get_images_from_folder(folder_path):
    return sorted([f for f in Path(folder_path).glob("*.png")])  # Change "*.png" if using JPG

# Streamlit UI
st.title("MNIST Image Viewer")

# Select a digit category (0-9)
subfolders = get_subfolders(DATASET_PATH)
selected_folder = st.selectbox("Select a digit:", subfolders)

# Get images from selected digit folder
image_files = get_images_from_folder(Path(DATASET_PATH) / selected_folder)

# Show images
if image_files:
    st.write(f"Showing images from **digit {selected_folder}**")
    selected_image = st.selectbox("Choose an image:", image_files)
    
    # Display the selected image
    image = Image.open(selected_image)
    st.image(image, caption=f"Digit: {selected_folder}", use_container_width=True)
else:
    st.warning("No images found in this folder.")