Update app.py

app.py

@@ -39,9 +39,53 @@ def infer_image(model, image):
     generated_img = (generated_img * 127.5 + 127.5).numpy().astype(np.uint8)  # De-normalize to [0, 255]
     return generated_img
 
-# Streamlit UI
-st.title("CycleGAN Inference App")
+# Custom CSS
+combined_css = """
+    .main, .sidebar .sidebar-content { background-color: #1c1c1c; color: #f0f2f6; }
+    .block-container { padding: 1rem 2rem; background-color: #333; border-radius: 10px; box-shadow: 0px 4px 10px rgba(0, 0, 0, 0.5); }
+    .stButton>button, .stDownloadButton>button { background: linear-gradient(135deg, #ff7e5f, #feb47b); color: white; border: none; padding: 10px 24px; text-align: center; text-decoration: none; display: inline-block; font-size: 16px; margin: 4px 2px; cursor: pointer; border-radius: 5px; }
+    .stSpinner { color: #4CAF50; }
+    .title {
+        font-size: 3rem;
+        font-weight: bold;
+        display: flex;
+        align-items: center;
+        justify-content: center;
+    }
+    .colorful-text {
+        background: -webkit-linear-gradient(135deg, #ff7e5f, #feb47b);
+        -webkit-background-clip: text;
+        -webkit-text-fill-color: transparent;
+    }
+    .black-white-text {
+        color: black;
+    }
+    .small-input .stTextInput>div>input {
+        height: 2rem;
+        font-size: 0.9rem;
+    }
+    .small-file-uploader .stFileUploader>div>div {
+        height: 2rem;
+        font-size: 0.9rem;
+    }
+    .custom-text {
+        font-size: 1.2rem;
+        color: #feb47b;
+        text-align: center;
+        margin-top: -20px;
+        margin-bottom: 20px;
+    }
+"""
+
+# Streamlit application
+st.set_page_config(layout="wide")
+
+st.markdown(f"<style>{combined_css}</style>", unsafe_allow_html=True)
+
+st.markdown('<div class="title"><span class="colorful-text">CycleGAN</span> <span class="black-white-text">Inference App</span></div>', unsafe_allow_html=True)
+st.markdown('<div class="custom-text">Convert your images using CycleGAN</div>', unsafe_allow_html=True)
 
+# Streamlit UI
 uploaded_file = st.file_uploader("Choose an image...", type="jpg")
 
 if uploaded_file is not None:
@@ -51,7 +95,22 @@ if uploaded_file is not None:
     
     if st.button("Run Inference"):
         # Perform inference
-        generated_image = infer_image(loaded_model, image)
+        with st.spinner('Processing...'):
+            generated_image = infer_image(loaded_model, image)
         
         # Display the result
         st.image(generated_image, caption='Generated Image', use_column_width=True)
+        
+        # Provide a download button for the generated image
+        img_byte_arr = BytesIO()
+        Image.fromarray(generated_image).save(img_byte_arr, format='JPEG')
+        img_byte_arr = img_byte_arr.getvalue()
+
+        st.download_button(
+            label="Download Generated Image",
+            data=img_byte_arr,
+            file_name="generated_image.jpg",
+            mime="image/jpeg"
+        )
+
+        st.success("Image processed successfully!")