Update app.py

app.py

@@ -5,131 +5,82 @@ from io import BytesIO
 from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import StreamingResponse
 from tensorflow.keras.models import load_model
-from tensorflow.keras.layers import Conv2DTranspose  # Import Conv2DTranspose
 import numpy as np
 import io
 import warnings
 
+# Suppress warnings
+warnings.filterwarnings('ignore')
+
 # Set Streamlit page configuration
 st.set_page_config(
     page_title="Sketch to Image using GAN",
     layout="centered",
     page_icon="🖌️",
-    initial_sidebar_state="expanded",
-)
-
-# Custom CSS for styling
-st.markdown(
-    """
-    <style>
-    body {
-        background-color: #2a2a2a;  /* Cool dark background */
-        color:#ffffff;  /* White text */
-        font-family: 'Courier New', monospace; /* Cool font */
-    }
-    h1, h2 {
-        color: #ff6347; /* Tomato color for titles */
-        font-weight: bold; /* Bold text */
-    }
-    .stButton>button {
-        color: #ffffff;
-        background-color: #ff6347; /* Tomato color for buttons */
-        border-radius: 10px;
-        border: 2px solid #ffffff;
-        font-weight: bold; /* Bold text */
-    }
-    </style>
-    """,
-    unsafe_allow_html=True,
 )
 
-# Title with colors and emojis
+# Title and description
 st.markdown("<h1 style='text-align: center; color: #ff6347;'>Sketch to Image using GAN 🖌️</h1>", unsafe_allow_html=True)
+st.markdown("<h2 style='text-align: center; color: #ff6347;'>Upload your sketch to generate an image!</h2>", unsafe_allow_html=True)
 
-# Description Section
-st.markdown("<h2 style='text-align: center; color: #ff6347;'>Empowering Multiple Fields with GANs 🌐</h2>", unsafe_allow_html=True)
-
-# Display Logo Image
-logo_image = Image.open("home1.jpeg")  # Make sure to replace with the actual path
-st.image(logo_image, width=300)
-
-st.write("The application of Generative Adversarial Networks (GANs) in the Sketch to Image project extends beyond creative endeavors, finding significant utility in various fields. The ability to transform sketches into vibrant and detailed images has far-reaching implications, especially in sectors such as law enforcement, forensic science, and more.")
-
-# Upload Section
-st.markdown("<h2 style='text-align: center; color: #ff6347;'>Upload Your Sketch 📤</h2>", unsafe_allow_html=True)
-uploaded_file = st.file_uploader("Choose an image... 📤", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-    # Display uploaded image in the center
-    col1, col2, col3 = st.columns([1, 2, 1])
-    with col2:
-        st.image(uploaded_file, caption="Uploaded Image 🖼️", width=300)
-
-    # Button to generate the image
-    if st.button('Generate 🚀'):
-        with st.spinner('Wait for it... Generating your image 🎨'):
-            try:
-                # Prepare file for sending
-                files = {"file": uploaded_file.getvalue()}
-                
-                # Send POST request to FastAPI server
-                response = requests.post("http://127.0.0.1:8000/generate-image/", files=files)
-
-                if response.status_code == 200:
-                    # Convert response content to an image
-                    generated_image = Image.open(BytesIO(response.content))
-                    
-                    # Display generated image in the center
-                    col1, col2, col3 = st.columns([1, 2, 1])
-                    with col2:
-                        st.image(generated_image, caption="Generated Image ✨", width=300)
-                else:
-                    st.error("Error in image generation 😢")
-            except requests.ConnectionError:
-                st.error("Unable to connect to the FastAPI server. Please make sure it is running.")
-
-# FastAPI Section
-warnings.filterwarnings('ignore')
-
-# Ensure that Conv2DTranspose is included when loading the model
-custom_objects = {"Conv2DTranspose": Conv2DTranspose}
+# Upload file widget
+uploaded_file = st.file_uploader("Upload a sketch (jpg, jpeg, png):", type=["jpg", "jpeg", "png"])
 
-# Load generator model (you may need to update the model path)
+# Model Loading
 try:
-    generator_model = load_model('model.h5', custom_objects=custom_objects)
-    st.success("Model loaded successfully! 🎉")
+    generator_model = load_model('model.h5')  # Update this path to your actual model file
+    st.success("Model loaded successfully!")
 except Exception as e:
-    st.error(f"Failed to load model: {str(e)}")
+    st.error(f"Error loading the model: {str(e)}")
 
-# FastAPI application
-app = FastAPI()
-
-@app.post("/generate-image/")
-async def generate_image(file: UploadFile = File(...)):
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert('RGB')
+# Image processing function
+def process_and_generate_image(image_data):
+    image = Image.open(io.BytesIO(image_data)).convert('RGB')
     image = image.resize((256, 256))
 
     # Preprocess image
     image_array = np.array(image)
-    image_array = (image_array - 127.5) / 127.5  # Normalize between -1 and 1
+    image_array = (image_array - 127.5) / 127.5  # Normalize to [-1, 1]
     image_array = np.expand_dims(image_array, axis=0)
 
     # Generate fake image
     fake_image = generator_model.predict(image_array)
     fake_image = (fake_image + 1) / 2.0  # Rescale to [0, 1]
     fake_image = np.squeeze(fake_image)
+    fake_image = (fake_image * 255).astype(np.uint8)
+
+    return Image.fromarray(fake_image)
 
-    fake_image = (fake_image * 255).astype(np.uint8)  # Convert to uint8
-    fake_image = Image.fromarray(fake_image)
+# Display uploaded image and handle generation
+if uploaded_file is not None:
+    st.image(uploaded_file, caption="Uploaded Sketch", width=300)
 
-    # Prepare image for streaming
+    if st.button("Generate Image"):
+        with st.spinner('Generating...'):
+            try:
+                # Generate the image
+                generated_image = process_and_generate_image(uploaded_file.getvalue())
+                
+                # Display the generated image
+                st.image(generated_image, caption="Generated Image", width=300)
+            except Exception as e:
+                st.error(f"Error generating image: {str(e)}")
+
+# FastAPI app for backend
+app = FastAPI()
+
+@app.post("/generate-image/")
+async def generate_image(file: UploadFile = File(...)):
+    contents = await file.read()
+    generated_image = process_and_generate_image(contents)
+    
     img_io = io.BytesIO()
-    fake_image.save(img_io, 'JPEG', quality=70)
+    generated_image.save(img_io, 'JPEG')
     img_io.seek(0)
 
-    return StreamingResponse(img_io, media_type='image/jpeg')
+    return StreamingResponse(img_io, media_type="image/jpeg")
 
+# Running FastAPI app if script is executed directly
 if __name__ == '__main__':
     import uvicorn
-    uvicorn.run(app, host='127.0.0.1', port=8000)
+    uvicorn.run(app, host="127.0.0.1", port=8000)