import streamlit as st
import numpy as np
from PIL import Image
import os
import cv2
import random
import numpy as np
from glob import glob
from PIL import Image, ImageOps
import matplotlib.pyplot as plt

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers


def plot_results(images, titles, figure_size=(12, 12)):
    fig = plt.figure(figsize=figure_size)
    for i in range(len(images)):
        fig.add_subplot(1, len(images), i + 1).set_title(titles[i])
        _ = plt.imshow(images[i])
        plt.axis("off")
    plt.show()


def infer(original_image):
    image = keras.utils.img_to_array(original_image)
    image = image.astype("float32") / 255.0
    image = np.expand_dims(image, axis=0)
    output = model.predict(image)
    output_image = output[0] * 255.0
    output_image = output_image.clip(0, 255)
    output_image = output_image.reshape(
        (np.shape(output_image)[0], np.shape(output_image)[1], 3)
    )
    output_image = Image.fromarray(np.uint8(output_image))
    original_image = Image.fromarray(np.uint8(original_image))
    return output_image

# Mock model for image prediction (replace this with your actual model)
def predict_image(img):
   
original_image = Image.open(uploaded_image)
enhanced_image = infer(original_image)
plot_results(
  [original_image, enhanced_image],
  ["Original", "MIRNet Enhanced"],
    (20, 12),
    )

# Streamlit UI
st.title("Image Prediction App")

# Upload image window
uploaded_image = st.file_uploader("Choose an image...", type="jpg")

if uploaded_image is not None:
    # Display uploaded image
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image.", use_column_width=True)
    
    # Button to run prediction
    if st.button("Predict"):
        # Get the prediction
        prediction = predict_image(image)

        # Display the prediction
        st.write(f"Prediction: {prediction}")
        st.image(image, caption=f"Predicted as {prediction}.", use_column_width=True)