import streamlit as st
import numpy as np
from scipy.ndimage import gaussian_filter

def main():
    st.title("Image Denoising with Diffusion")
    uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png"])

    if uploaded_image:
        image_array = np.array(uploaded_image)
        denoised_image = apply_diffusion(image_array)
        st.image(denoised_image, caption="Denoised Image", use_column_width=True)

if __name__ == "__main__":
    main()
    
def apply_diffusion(image_array):
    # Convert to grayscale if needed
    if len(image_array.shape) == 3:
        image_array = np.mean(image_array, axis=2)

    # Apply Gaussian diffusion
    denoised_image = gaussian_filter(image_array, sigma=1.0)

    return denoised_image