Update app.py

app.py

@@ -1,18 +1,13 @@
-import os
-os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
-
 import streamlit as st
 import tensorflow as tf
 from tensorflow_addons.layers import InstanceNormalization
+import numpy as np
+import matplotlib.pyplot as plt
+from huggingface_hub import hf_hub_download
 from PIL import Image
 import requests
 from io import BytesIO
-import logging
-from huggingface_hub import HfApi, hf_hub_download
-import numpy as np
-
-# Setup logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+import os
 
 # Custom CSS
 def set_css(style):
@@ -52,31 +47,28 @@ combined_css = """
         color: #feb47b;
         text-align: center;
         margin-top: -20px;
-        margin-bottom: 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">Photo</span> <span class="black-white-text">to Van Gogh</span></div>', unsafe_allow_html=True)
-st.markdown('<div class="custom-text">Convert photos to Van Gogh style using AI</div>', unsafe_allow_html=True)
+st.markdown('<div class="title"><span class="colorful-text">Photo</span> <span class="black-white-text">to Art</span></div>', unsafe_allow_html=True)
+st.markdown('<div class="custom-text">Convert Photos to Art using CycleGAN</div>', unsafe_allow_html=True)
 
-# Download and load the model
-MODEL_REPO = "Hammad712/CycleGAN-Model"
-MODEL_DIR = "/content/CycleGAN"
-if not os.path.exists(MODEL_DIR):
-    os.makedirs(MODEL_DIR)
+# Define your Hugging Face repository details
+username = "Hammad712"  # Replace with your Hugging Face username
+repo_name = "CycleGAN-Model"
+repo_id = f"{username}/{repo_name}"
+model_filename = "CycleGAN.h5"
 
-# Load the model if not already loaded
-if 'model' not in st.session_state:
-    api = HfApi()
-    for file in api.list_repo_files(repo_id=MODEL_REPO, repo_type="model"):
-        hf_hub_download(repo_id=MODEL_REPO, filename=file, local_dir=MODEL_DIR)
-    st.session_state['model'] = tf.keras.models.load_model(MODEL_DIR, custom_objects={'InstanceNormalization': InstanceNormalization})
+# Download the model file from the repository
+model_path = hf_hub_download(repo_id=repo_id, filename=model_filename)
 
-model = st.session_state['model']
+# Load the model
+model = tf.keras.models.load_model(model_path, custom_objects={"InstanceNormalization": InstanceNormalization})
 
 def load_and_preprocess_image(image_path_or_url):
     if isinstance(image_path_or_url, str) and image_path_or_url.startswith(('http://', 'https://')):
@@ -85,16 +77,24 @@ def load_and_preprocess_image(image_path_or_url):
     else:
         img = Image.open(image_path_or_url).convert("RGB")
     img = img.resize((256, 256))
-    img = (np.array(img) - 127.5) / 127.5  # Normalize to [-1, 1]
-    img = tf.expand_dims(img, axis=0)  # Add batch dimension
+    img = np.array(img) / 127.5 - 1  # Normalize to [-1, 1]
+    img = np.expand_dims(img, axis=0)  # Add batch dimension
     return img
 
-def infer_image(model, image_path_or_url):
-    preprocessed_img = load_and_preprocess_image(image_path_or_url)
-    generated_img = model(preprocessed_img, training=False)
-    generated_img = tf.squeeze(generated_img, axis=0)  # Remove batch dimension
-    generated_img = (generated_img * 127.5 + 127.5).numpy().astype(np.uint8)  # De-normalize to [0, 255]
-    return Image.fromarray(generated_img)
+def postprocess_and_display_image(img_tensor):
+    img = tf.squeeze(img_tensor, axis=0)  # Remove batch dimension
+    img = (img * 127.5 + 127.5).numpy().astype(np.uint8)  # De-normalize to [0, 255]
+    return Image.fromarray(img)
+
+def perform_inference(model, image_path_or_url):
+    # Load and preprocess the image
+    input_img = load_and_preprocess_image(image_path_or_url)
+
+    # Perform inference
+    generated_img = model.generatorS(input_img, training=False)
+
+    # Postprocess and return the generated image
+    return postprocess_and_display_image(generated_img)
 
 # Input for image URL or path
 with st.expander("Input Options", expanded=True):
@@ -109,19 +109,17 @@ if st.button("Convert"):
     if image_path_or_url:
         with st.spinner('Processing...'):
             try:
-                generated_image = infer_image(model, image_path_or_url)
+                generated_image = perform_inference(model, image_path_or_url)
                 original_image = load_and_preprocess_image(image_path_or_url)
-                original_image = (tf.squeeze(original_image, axis=0) * 127.5 + 127.5).numpy().astype(np.uint8)  # De-normalize to [0, 255]
-                original_image = Image.fromarray(original_image)
 
                 # Display original and generated images side by side
                 st.markdown("### Result")
                 col1, col2 = st.columns(2)
 
                 with col1:
-                    st.image(original_image, caption='Original Image', use_column_width=True)
+                    st.image(np.array(original_image[0] * 127.5 + 127.5, dtype=np.uint8), caption='Original Image', use_column_width=True)
                 with col2:
-                    st.image(generated_image, caption='Van Gogh Styled Image', use_column_width=True)
+                    st.image(generated_image, caption='Generated Art Image', use_column_width=True)
 
                 # Provide a download button for the generated image
                 img_byte_arr = BytesIO()
@@ -129,9 +127,9 @@ if st.button("Convert"):
                 img_byte_arr = img_byte_arr.getvalue()
 
                 st.download_button(
-                    label="Download Styled Image",
+                    label="Download Art Image",
                     data=img_byte_arr,
-                    file_name="vangogh_styled_image.jpg",
+                    file_name="art_image.jpg",
                     mime="image/jpeg"
                 )
 
@@ -139,6 +137,5 @@ if st.button("Convert"):
 
             except Exception as e:
                 st.error(f"An error occurred: {e}")
-                logging.error("Error during inference", exc_info=True)
     else:
         st.error("Please enter a valid image path or URL.")