Update app.py

app.py

@@ -5,9 +5,7 @@ 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.utils import get_custom_objects
-
-
+from tensorflow.keras.layers import Conv2DTranspose
 import numpy as np
 import io
 import warnings
@@ -41,7 +39,7 @@ st.markdown(
         font-weight: bold; /* Bold text */
     }
     </style>
-""",
+    """,
     unsafe_allow_html=True,
 )
 
@@ -51,38 +49,37 @@ st.markdown("<h1 style='text-align: center; color: #ff6347;'>Sketch to Image usi
 # Description Section
 st.markdown("<h2 style='text-align: center; color: #ff6347;'>Empowering Multiple Fields with GANs 🌐</h2>", unsafe_allow_html=True)
 
-# Logo Image
-logo_image = Image.open("home1.jpeg")
+# 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 Pic Section
+# 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 the uploaded image in the center
+    # 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 with emoji
+    # Button to generate the image
     if st.button('Generate 🚀'):
-        # Display a message while generating the image
         with st.spinner('Wait for it... Generating your image 🎨'):
             try:
-                # Prepare the file for sending
+                # 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 the response content to an image
+                    # Convert response content to an image
                     generated_image = Image.open(BytesIO(response.content))
-
-                    # Display the generated image in the center
+                    
+                    # 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)
@@ -93,9 +90,18 @@ if uploaded_file is not None:
 
 # FastAPI Section
 warnings.filterwarnings('ignore')
-#generator_model = load_model('model.h5')  # Update this with your generator model's path
+
+# Ensure that Conv2DTranspose is included when loading the model
 custom_objects = {"Conv2DTranspose": Conv2DTranspose}
-generator_model = load_model('model.h5', custom_objects=custom_objects)
+
+# Load generator model (you may need to update the model path)
+try:
+    generator_model = load_model('model.h5', custom_objects=custom_objects)
+    st.success("Model loaded successfully! 🎉")
+except Exception as e:
+    st.error(f"Failed to load model: {str(e)}")
+
+# FastAPI application
 app = FastAPI()
 
 @app.post("/generate-image/")
@@ -104,17 +110,20 @@ async def generate_image(file: UploadFile = File(...)):
     image = Image.open(io.BytesIO(contents)).convert('RGB')
     image = image.resize((256, 256))
 
+    # Preprocess image
     image_array = np.array(image)
-    image_array = (image_array - 127.5) / 127.5
+    image_array = (image_array - 127.5) / 127.5  # Normalize between -1 and 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
+    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)
+    fake_image = (fake_image * 255).astype(np.uint8)  # Convert to uint8
     fake_image = Image.fromarray(fake_image)
 
+    # Prepare image for streaming
     img_io = io.BytesIO()
     fake_image.save(img_io, 'JPEG', quality=70)
     img_io.seek(0)