app.py

import streamlit as st
import cv2
import numpy as np
from colorize import colorize

st.set_page_config(page_title="Image Colorizer", layout="wide")

st.title("🎨 Image Colorization")
st.caption("Compare Local vs Pretrained Models")

uploaded_file = st.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])

if uploaded_file:

    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    image = cv2.imdecode(file_bytes, 1)

    st.image(
        cv2.cvtColor(image, cv2.COLOR_BGR2RGB), caption="Uploaded Image", width=200
    )

    mode_option = st.selectbox(
        "Choose Mode", ["Use Both (Compare)", "Local Only", "Pretrained Only"]
    )

    mode_map = {
        "Use Both (Compare)": "both",
        "Local Only": "local",
        "Pretrained Only": "pretrained",
    }

    mode = mode_map[mode_option]

    if st.button("🚀 Colorize"):

        with st.spinner("Processing... (first run downloads model)"):

            result = colorize(image, mode=mode)

        st.markdown("---")

        if mode == "both":

            col1, col2, col3 = st.columns(3)

            with col1:
                st.subheader("Original")
                st.image(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), width="stretch")

            with col2:
                st.subheader("Local")
                st.image(
                    cv2.cvtColor(result["local"], cv2.COLOR_BGR2RGB), width="stretch"
                )

            with col3:
                st.subheader("Pretrained")
                st.image(
                    cv2.cvtColor(result["pretrained"], cv2.COLOR_BGR2RGB),
                    width="stretch",
                )

        else:

            col1, col2 = st.columns(2)

            with col1:
                st.subheader("Original")
                st.image(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), width="stretch")

            with col2:
                st.subheader("Colorized")
                st.image(cv2.cvtColor(result, cv2.COLOR_BGR2RGB), width="stretch")

        final_img = result["pretrained"] if mode == "both" else result

        _, buffer = cv2.imencode(".png", final_img)

        st.download_button(
            "⬇️ Download", buffer.tobytes(), file_name="colorized.png", mime="image/png"
        )