Update app.py

app.py

@@ -1,127 +1,178 @@
 import streamlit as st
-from fastai.vision import open_image, load_learner, show_image
-import PIL.Image
+from fastai.vision import open_image, load_learner
 from PIL import Image
-from io import BytesIO
 import requests
+import os
+import logging
 import torch
 import torch.nn as nn
-import os
-import tempfile
-import shutil
+from io import BytesIO
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Define the FeatureLoss class
 class FeatureLoss(nn.Module):
     def __init__(self, m_feat, layer_ids, layer_wgts):
         super().__init__()
         self.m_feat = m_feat
         self.loss_features = [self.m_feat[i] for i in layer_ids]
-        self.hooks = hook_outputs(self.loss_features, detach=False)
+        self.hooks = [module.register_forward_hook(self.hook_fn) for module in self.loss_features]
         self.wgts = layer_wgts
-        self.metric_names = ['pixel',] + [f'feat_{i}' for i in range(len(layer_ids))] + [f'gram_{i}' for i in range(len(layer_ids))]
+        self.metric_names = ['pixel'] + [f'feat_{i}' for i in range(len(layer_ids))] + [f'gram_{i}' for i in range(len(layer_ids))]
 
-    def make_features(self, x, clone=False):
-        self.m_feat(x)
-        return [(o.clone() if clone else o) for o in self.hooks.stored]
+    def hook_fn(self, module, input, output):
+        self.stored = output.detach().clone()
 
     def forward(self, input, target):
-        out_feat = self.make_features(target, clone=True)
-        in_feat = self.make_features(input)
-        self.feat_losses = [base_loss(input, target)]
-        self.feat_losses += [base_loss(f_in, f_out) * w for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
-        self.feat_losses += [base_loss(gram_matrix(f_in), gram_matrix(f_out)) * w**2 * 5e3 for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
+        self.m_feat(target)
+        out_feat = [self.stored.clone()]
+        self.m_feat(input)
+        in_feat = [self.stored]
+        self.feat_losses = [torch.nn.functional.mse_loss(input, target)]
+        self.feat_losses += [torch.nn.functional.mse_loss(f_in, f_out) * w for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
+        self.feat_losses += [torch.nn.functional.mse_loss(self.gram_matrix(f_in), self.gram_matrix(f_out)) * w**2 * 5e3 for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
         self.metrics = dict(zip(self.metric_names, self.feat_losses))
         return sum(self.feat_losses)
 
-    def __del__(self): self.hooks.remove()
+    @staticmethod
+    def gram_matrix(input):
+        b, c, h, w = input.size()
+        features = input.view(b, c, h * w)
+        G = torch.bmm(features, features.transpose(1, 2))
+        return G.div(c * h * w)
+
+def fetch_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")
+    return img
 
-def add_margin(pil_img, top, right, bottom, left, color):
-    width, height = pil_img.size
-    new_width = width + right + left
-    new_height = height + top + bottom
-    result = Image.new(pil_img.mode, (new_width, new_height), color)
-    result.paste(pil_img, (left, top))
-    return result
+def inference(image_path_or_url, learn):
+    img = fetch_image(image_path_or_url)
+    img_with_margin = Image.new('RGB', (img.width + 500, img.height + 500), (255, 255, 255))
+    img_with_margin.paste(img, (250, 250))
+    temp_image_path = "temp_image.jpg"
+    img_with_margin.save(temp_image_path, quality=95)
+    img_fastai = open_image(temp_image_path)
+    _, img_hr, _ = learn.predict(img_fastai)
+    return tensor_to_pil(img_hr)
 
 def tensor_to_pil(tensor):
-    """
-    Convert a tensor to a PIL Image.
-    """
     tensor = tensor.cpu().clamp(0, 1)
     array = tensor.numpy().transpose(1, 2, 0)
     return Image.fromarray((array * 255).astype('uint8'))
 
-def inference(image_path_or_url, learn):
-    """
-    Perform inference on an image from a local path or a URL.
-
-    Parameters:
-    image_path_or_url (str): Path to the local image or URL of the image.
-    learn (Learner): The trained model.
-
-    Returns:
-    PIL.Image.Image: The high-resolution image generated by the model.
-    """
-    if image_path_or_url.startswith('http://') or image_path_or_url.startswith('https://'):
-        response = requests.get(image_path_or_url)
-        img = PIL.Image.open(BytesIO(response.content)).convert("RGB")
-    else:
-        img = PIL.Image.open(image_path_or_url).convert("RGB")
-
-    im_new = add_margin(img, 250, 250, 250, 250, (255, 255, 255))
-    im_new.save("test.jpg", quality=95)
-    img = open_image("test.jpg")
-    p, img_hr, b = learn.predict(img)
-    return tensor_to_pil(img_hr)
+def load_model(model_url, model_file_path):
+    if not os.path.exists(model_file_path):
+        with st.spinner('Downloading model...'):
+            response = requests.get(model_url)
+            with open(model_file_path, 'wb') as f:
+                f.write(response.content)
+        st.success('Model downloaded successfully!')
+    learn = load_learner(os.path.dirname(model_file_path), model_file_path)
+    return learn
+
+# 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;
+    }
+"""
 
 # Streamlit application
-st.title("Image to Sketch")
-st.write("Convert any image to its sketch version using a Pix2Pix GAN Model.")
-
-# Download the model file from the Hugging Face repository
-model_url = "https://huggingface.co/Hammad712/image2sketch/resolve/main/image2sketch.pkl"
-model_file_path = 'image2sketch.pkl'
+st.set_page_config(layout="wide")
 
-if not os.path.exists(model_file_path):
-    with st.spinner('Downloading model...'):
-        response = requests.get(model_url)
-        with open(model_file_path, 'wb') as f:
-            f.write(response.content)
-    st.success('Model downloaded successfully!')
+st.markdown(f"<style>{combined_css}</style>", unsafe_allow_html=True)
 
-# Create a temporary directory for the model
-with tempfile.TemporaryDirectory() as tmpdirname:
-    shutil.move(model_file_path, os.path.join(tmpdirname, 'export.pkl'))
-    learn = load_learner(tmpdirname)
+st.markdown('<div class="title"><span class="colorful-text">Image</span> <span class="black-white-text">to Sketch</span></div>', unsafe_allow_html=True)
+st.markdown('<div class="custom-text">Jana\'s embroidery studio. Convert Photo\'s to Drawings using AI</div>', unsafe_allow_html=True)
 
-# Sidebar to display previous generated images
-st.sidebar.title("Previous Results")
-if 'generated_images' not in st.session_state:
-    st.session_state['generated_images'] = []
+# Download and load the model
+MODEL_URL = "https://huggingface.co/Hammad712/image2sketch/resolve/main/image2sketch.pkl"
+MODEL_FILE_PATH = 'image2sketch.pkl'
 
-for img in st.session_state['generated_images']:
-    st.sidebar.image(img, use_column_width=True)
+if 'learn' not in st.session_state:
+    st.session_state['learn'] = load_model(MODEL_URL, MODEL_FILE_PATH)
+learn = st.session_state['learn']
 
 # Input for image URL or path
-image_path_or_url = st.text_input("Enter image URL", "")
+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")
+
+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...'):
-            high_res_image = inference(image_path_or_url, learn)
-            original_image = PIL.Image.open(BytesIO(requests.get(image_path_or_url).content)) if image_path_or_url.startswith('http') else PIL.Image.open(image_path_or_url)
-
-            # Display original and high-res images side by side
-            col1, col2 = st.columns(2)
-
-            with col1:
-                st.image(original_image, caption='Original Image', use_column_width=True)
-            with col2:
-                st.image(high_res_image, caption='Sketch Image', use_column_width=True)
-
-            # Save the generated image to session state for sidebar display
-            st.session_state['generated_images'].append(high_res_image)
-
+            try:
+                high_res_image = inference(image_path_or_url, learn)
+                original_image = fetch_image(image_path_or_url)
+
+                # Display original and high-res 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)
+                with col2:
+                    st.image(high_res_image, caption='Sketch Image', use_column_width=True)
+
+                # Provide a download button for the generated image
+                img_byte_arr = BytesIO()
+                high_res_image.save(img_byte_arr, format='JPEG')
+                img_byte_arr = img_byte_arr.getvalue()
+
+                st.download_button(
+                    label="Download Sketch Image",
+                    data=img_byte_arr,
+                    file_name="sketch_image.jpg",
+                    mime="image/jpeg"
+                )
+
+                st.success("Image processed successfully!")
+
+            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.")