Update app.py

app.py

@@ -1,180 +1,57 @@
+import os
 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 HfApi, hf_hub_download
 from PIL import Image
-from huggingface_hub import hf_hub_download
-import requests
-from io import BytesIO
 
-# Custom CSS
-def set_css(style):
-    st.markdown(f"<style>{style}</style>", unsafe_allow_html=True)
-
-# Combined dark mode styles
-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;
-    }
-"""
-
-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 Art</span></div>', unsafe_allow_html=True)
-st.markdown('<div class="custom-text">Convert Photos to Art using CycleGAN</div>', unsafe_allow_html=True)
-
-# Define CycleGAN model architecture
-class CycleGAN(tf.keras.Model):
-    def __init__(self):
-        super(CycleGAN, self).__init__()
-        self.generatorP = self.build_generator()  # Generates Photo from Style
-        self.generatorS = self.build_generator()  # Generates Style from Photo
-
-    def build_generator(self):
-        def ConvBlock(filters, kernel, strides, x):
-            x = tf.keras.layers.Conv2D(filters, kernel, strides, 'same', kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02), use_bias=False)(x)
-            if filters == 3:
-                return tf.keras.layers.Activation('tanh')(x)
-            x = InstanceNormalization(gamma_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02))(x)
-            return tf.keras.layers.ReLU()(x)
-
-        def ResBlock(filters, inputs):
-            x = tf.keras.layers.Conv2D(filters, 3, 1, padding='same', kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02), use_bias=False)(inputs)
-            x = InstanceNormalization(gamma_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02))(x)
-            x = tf.keras.layers.ReLU()(x)
-            x = tf.keras.layers.Conv2D(filters, 3, 1, padding='same', kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02), use_bias=False)(x)
-            x = InstanceNormalization(gamma_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02))(x)
-            return tf.keras.layers.Add()([x, inputs])
-
-        def TransBlock(filters, x):
-            x = tf.keras.layers.Conv2DTranspose(filters, 3, 2, 'same', kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02), use_bias=False)(x)
-            x = InstanceNormalization(gamma_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02))(x)
-            return tf.keras.layers.ReLU()(x)
-
-        input = tf.keras.Input(shape=(256, 256, 3))
-        x = ConvBlock(64, 7, 1, input)
-        x = ConvBlock(128, 3, 2, x)
-        x = ConvBlock(256, 3, 2, x)
-        for _ in range(9):
-            x = ResBlock(256, x)
-        x = TransBlock(128, x)
-        x = TransBlock(64, x)
-        out = ConvBlock(3, 7, 1, x)
-        return tf.keras.Model(inputs=input, outputs=out)
-
-# Load the model weights from Hugging Face
-try:
-    repo_id = "Hammad712/CycleGAN-Model"
-    generatorS_path = hf_hub_download(repo_id=repo_id, filename="generatorS.h5")
-    generatorP_path = hf_hub_download(repo_id=repo_id, filename="generatorP.h5")
-
-    model = CycleGAN()
-    model.generatorS.load_weights(generatorS_path)
-    model.generatorP.load_weights(generatorP_path)
-    st.success("Model loaded successfully!")
-except Exception as e:
-    st.error(f"An error occurred while loading the model: {e}")
-    st.stop()
-
-def load_and_preprocess_image(image_path_or_url):
-    if isinstance(image_path_or_url, str) and image_path_or_url.startswith(('http://', 'https://')):
-        response = requests.get(image_path_or_url)
-        img = Image.open(BytesIO(response.content)).convert("RGB")
-    else:
-        img = Image.open(image_path_or_url).convert("RGB")
-    img = img.resize((256, 256))
-    img = np.array(img) / 127.5 - 1  # Normalize to [-1, 1]
+# Hugging Face credentials
+api = HfApi()
+
+# Set your Hugging Face username and model repository name
+username = "Hammad712"
+repo_name = "CycleGAN-Model"
+repo_id = f"{username}/{repo_name}"
+
+# Download model files from Hugging Face
+local_dir = "CycleGAN"  # Changed to a relative path
+os.makedirs(local_dir, exist_ok=True)
+for file in api.list_repo_files(repo_id=repo_id, repo_type="model"):
+    hf_hub_download(repo_id=repo_id, filename=file, local_dir=local_dir, token=token)
+
+# Load the model
+custom_objects = {'InstanceNormalization': tf.keras.layers.Layer}  # Adjust custom objects as needed
+loaded_model = tf.keras.models.load_model(local_dir, custom_objects=custom_objects)
+
+# Helper functions
+def load_and_preprocess_image(image):
+    img = image.resize((256, 256))
+    img = np.array(img)
+    img = (img - 127.5) / 127.5  # Normalize to [-1, 1]
     img = np.expand_dims(img, axis=0)  # Add batch dimension
     return 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)
+def infer_image(model, image):
+    preprocessed_img = load_and_preprocess_image(image)
+    generated_img = model(preprocessed_img, training=False)
+    generated_img = np.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 generated_img
 
-    # Postprocess and return the generated image
-    return postprocess_and_display_image(generated_img)
+# Streamlit UI
+st.title("CycleGAN Inference App")
 
-# Input for image URL or path
-with st.expander("Input Options", expanded=True):
-    image_path_or_url = st.text_input("Enter image URL", "", key="image_url", placeholder="Enter image URL", help="Enter the URL of the image to convert")
-    uploaded_file = st.file_uploader("Or upload an image", type=["jpg", "jpeg", "png", "webp"], key="upload_file", help="Upload an image file to convert")
+uploaded_file = st.file_uploader("Choose an image...", type="jpg")
 
 if uploaded_file is not None:
-    image_path_or_url = uploaded_file
-
-# Run inference button
-if st.button("Convert"):
-    if image_path_or_url:
-        with st.spinner('Processing...'):
-            try:
-                generated_image = perform_inference(model, image_path_or_url)
-                original_image = load_and_preprocess_image(image_path_or_url)
-
-                # Display original and generated images side by side
-                st.markdown("### Result")
-                col1, col2 = st.columns(2)
-
-                with col1:
-                    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='Generated Art Image', use_column_width=True)
-
-                # Provide a download button for the generated image
-                img_byte_arr = BytesIO()
-                generated_image.save(img_byte_arr, format='JPEG')
-                img_byte_arr = img_byte_arr.getvalue()
-
-                st.download_button(
-                    label="Download Art Image",
-                    data=img_byte_arr,
-                    file_name="art_image.jpg",
-                    mime="image/jpeg"
-                )
-
-                st.success("Image processed successfully!")
-
-            except Exception as e:
-                st.error(f"An error occurred: {e}")
-    else:
-        st.error("Please enter a valid image path or URL.")
+    # Load and display the uploaded image
+    image = Image.open(uploaded_file)
+    st.image(image, caption='Uploaded Image', use_column_width=True)
+    
+    if st.button("Run Inference"):
+        # Perform inference
+        generated_image = infer_image(loaded_model, image)
+        
+        # Display the result
+        st.image(generated_image, caption='Generated Image', use_column_width=True)