Update app.py

app.py

@@ -1,36 +1,110 @@
 import streamlit as st
 from PIL import Image
 import torch
-
 from model import ModelColorization
-
 from utils import process_gs_image, inverse_transform_cs
 
-# create the model
+# Custom CSS for styling
+st.markdown(
+    """
+    <style>
+    .main {
+        background-color: #f9f9f9;
+    }
+    .title {
+        color: #2a3f5f;
+        font-size: 2.5em;
+        text-align: center;
+        margin-bottom: 0.5em;
+    }
+    .subheader {
+        color: #5a5a5a;
+        font-size: 1.1em;
+        text-align: center;
+        margin-bottom: 2em;
+    }
+    .upload-box {
+        background-color: #ffffff;
+        border-radius: 10px;
+        padding: 2em;
+        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+        margin-bottom: 2em;
+    }
+    .result-box {
+        background-color: #ffffff;
+        border-radius: 10px;
+        padding: 2em;
+        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+        margin-top: 2em;
+    }
+    .stButton>button {
+        background-color: #4CAF50;
+        color: white;
+        border-radius: 5px;
+        padding: 0.5em 1em;
+        font-size: 1em;
+        width: 100%;
+    }
+    .stButton>button:hover {
+        background-color: #45a049;
+    }
+    .group-banner {
+        text-align: center;
+        font-size: 0.9em;
+        color: #777777;
+        margin-top: 2em;
+    }
+    </style>
+    """,
+    unsafe_allow_html=True
+)
+
+# Load model
 model = ModelColorization().from_pretrained("sebastiansarasti/AutoEncoderImageColorization")
 
-# create the streamlit app
-st.title("Image Colorization App")
-st.write("This is an app to colorize black and white images.")
+# App header
+st.markdown('<p class="title">🎨 Neural Image Colorizer</p>', unsafe_allow_html=True)
+st.markdown('<p class="subheader">Bring black & white photos to life with AI</p>', unsafe_allow_html=True)
 
-# create a botton to upload the image
-uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+# Upload section
+with st.container():
+    st.markdown("### 📤 Upload Your Image")
+    uploaded_file = st.file_uploader(
+        "Choose a black & white photo...", 
+        type=["jpg", "jpeg", "png"],
+        label_visibility="collapsed"
+    )
 
-# check if the image is uploaded
+# Processing section
 if uploaded_file is not None:
-    # display the image
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image.", use_container_width=True)
+    with st.container():
+        col1, col2 = st.columns(2)
+        with col1:
+            st.markdown("### ⬆ Original")
+            original_img = Image.open(uploaded_file)
+            st.image(original_img, use_column_width=True)
+        
+        with col2:
+            st.markdown("### 🎨 Colorized")
+            if st.button("✨ Colorize Image", type="primary"):
+                with st.spinner("Colorizing your image..."):
+                    # Process image
+                    image, original_size = process_gs_image(original_img)
+                    
+                    # Run model
+                    model.eval()
+                    with torch.no_grad():
+                        result = model(image)
+                    
+                    # Get colorized image
+                    colorized_image = inverse_transform_cs(result.squeeze(0), original_size)
+                    
+                    # Display result
+                    st.image(colorized_image, use_column_width=True)
+                    st.success("Colorization complete!")
 
-    # create a button to colorize the image
-    if st.button("Colorize"):
-        # process the grayscale image
-        image, original_size = process_gs_image(image)
-        # run the model
-        model.eval()
-        with torch.no_grad():
-            result = model(image)
-        # colorize the image
-        colorized_image = inverse_transform_cs(result.squeeze(0), original_size)
-        # display the colorized image
-        st.image(colorized_image, caption="Colorized Image.", use_container_width=True)
+# Footer
+st.markdown(
+    '<p class="group-banner">Developed with ❤️ by Group 9 |  Computer Vision Project</p>',
+    unsafe_allow_html=True
+)