import streamlit as st
from PIL import Image, ImageEnhance
import numpy as np
import cv2
from rembg import remove
from super_image import EdsrModel, ImageLoader
import io

# Fungsi untuk Menghapus Latar Belakang
def remove_background(image):
    image_byte = io.BytesIO()
    image.save(image_byte, format="PNG")
    image_byte.seek(0)
    processed_image_bytes = remove(image_byte.read())
    return Image.open(io.BytesIO(processed_image_bytes))

# Fungsi untuk Meningkatkan Kualitas Gambar
def enhance_image(image, brightness, contrast, sharpness):
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(brightness)
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(contrast)
    enhancer = ImageEnhance.Sharpness(image)
    image = enhancer.enhance(sharpness)
    return image

# Fungsi untuk Memperbesar Skala Gambar
def upscale_image(image):
    st.info("Proses upscale sedang berjalan, ini mungkin memakan waktu beberapa saat...")
    model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=4)
    inputs = ImageLoader.load_image(image)
    preds = model(inputs)
    # Konversi tensor kembali ke gambar PIL
    # Pastikan tensor berada di CPU dan dalam format yang benar
    processed_image = ImageLoader.save_image(preds, 'temp_upscaled.png', single_batch=True)
    return Image.open('temp_upscaled.png')


# Konfigurasi Halaman Streamlit
st.set_page_config(page_title="Aplikasi Pengolah Gambar Otomatis", layout="wide")

st.title("Aplikasi Pengolah Gambar Otomatis 🖼️")
st.write("Unggah gambar Anda dan pilih proses yang ingin dilakukan.")

# Sidebar untuk unggah dan kontrol
with st.sidebar:
    st.header("Unggah Gambar")
    uploaded_file = st.file_uploader("Pilih sebuah gambar...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        st.header("Pilih Proses")
        process_type = st.selectbox(
            "Pilih jenis proses:",
            ["Tampilkan Asli", "Hapus Latar Belakang", "Peningkatan Gambar (Enhance)", "Perbesar Skala (Upscale)"]
        )

        if process_type == "Peningkatan Gambar (Enhance)":
            st.subheader("Pengaturan Peningkatan")
            brightness = st.slider("Kecerahan", 0.1, 3.0, 1.0)
            contrast = st.slider("Kontras", 0.1, 3.0, 1.0)
            sharpness = st.slider("Ketajaman", 0.1, 3.0, 1.0)

# Tampilan Utama
if uploaded_file is not None:
    original_image = Image.open(uploaded_file).convert("RGB")

    col1, col2 = st.columns(2)

    with col1:
        st.subheader("Gambar Asli")
        st.image(original_image, use_column_width=True)

    with col2:
        st.subheader("Hasil Proses")
        if process_type == "Tampilkan Asli":
            st.image(original_image, use_column_width=True)
        elif process_type == "Hapus Latar Belakang":
            with st.spinner('Menghapus latar belakang...'):
                processed_image = remove_background(original_image)
                st.image(processed_image, use_column_width=True)
        elif process_type == "Peningkatan Gambar (Enhance)":
            with st.spinner('Meningkatkan kualitas gambar...'):
                processed_image = enhance_image(original_image, brightness, contrast, sharpness)
                st.image(processed_image, use_column_width=True)
        elif process_type == "Perbesar Skala (Upscale)":
            with st.spinner('Memperbesar skala gambar...'):
                processed_image = upscale_image(original_image)
                st.image(processed_image, use_column_width=True)

        # Tombol Unduh
        if 'processed_image' in locals():
            buf = io.BytesIO()
            processed_image.save(buf, format="PNG")
            byte_im = buf.getvalue()
            st.download_button(
                label="Unduh Gambar Hasil",
                data=byte_im,
                file_name="hasil_proses.png",
                mime="image/png"
            )
else:
    st.info("Silakan unggah gambar terlebih dahulu melalui sidebar.")