Update app.py

app.py

@@ -7,6 +7,8 @@ from skimage.color import rgb2lab, lab2rgb
 import numpy as np
 import matplotlib.pyplot as plt
 from io import BytesIO
+import requests
+from io import BytesIO
 
 # Download the model from Hugging Face Hub
 repo_id = "Hammad712/GAN-Colorization-Model"
@@ -31,8 +33,8 @@ G_net.load_state_dict(torch.load(model_path, map_location=device))
 G_net.eval()
 
 # Preprocessing function
-def preprocess_image(img_path):
-    img = Image.open(img_path).convert("RGB")
+def preprocess_image(img):
+    img = img.convert("RGB")
     img = transforms.Resize((256, 256), Image.BICUBIC)(img)
     img = np.array(img)
     img_to_lab = rgb2lab(img).astype("float32")
@@ -41,8 +43,8 @@ def preprocess_image(img_path):
     return L.unsqueeze(0).to(device)
 
 # Inference function
-def colorize_image(img_path, model):
-    L = preprocess_image(img_path)
+def colorize_image(img, model):
+    L = preprocess_image(img)
     with torch.no_grad():
         ab = model(L)
     L = (L + 1.) * 50.
@@ -67,8 +69,7 @@ combined_css = """
     .title {
         font-size: 3rem;
         font-weight: bold;
-        display: flex;
-        align-items: center;
+        display: flex; align-items: center;
         justify-content: center;
     }
     .colorful-text {
@@ -101,19 +102,31 @@ 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">Image</span> <span class="black-white-text">Colorization</span></div>', unsafe_allow_html=True)
+st.markdown('<div class="title"><span class="black-white-text">Image</span> <span class="colorful-text">Colorization</span></div>', unsafe_allow_html=True)
 st.markdown('<div class="custom-text">Convert black and white images to color using AI</div>', unsafe_allow_html=True)
 
 # Input for image URL or file upload
 with st.expander("Input Options", expanded=True):
     uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png", "webp"], key="upload_file", help="Upload an image file to convert")
+    url_input = st.text_input("Or enter an image URL", key="url_input", help="Enter the URL of an image to convert")
 
 # Run inference button
 if st.button("Colorize"):
+    img = None
+
     if uploaded_file is not None:
+        img = Image.open(uploaded_file)
+    elif url_input:
+        try:
+            response = requests.get(url_input)
+            img = Image.open(BytesIO(response.content))
+        except Exception as e:
+            st.error(f"Error fetching the image from URL: {e}")
+
+    if img is not None:
         with st.spinner('Processing...'):
             try:
-                colorized_images = colorize_image(uploaded_file, G_net)
+                colorized_images = colorize_image(img, G_net)
                 colorized_image = colorized_images[0]
 
                 # Display original and colorized images side by side
@@ -121,7 +134,7 @@ if st.button("Colorize"):
                 col1, col2 = st.columns(2)
 
                 with col1:
-                    st.image(uploaded_file, caption='Original Image', use_column_width=True)
+                    st.image(img, caption='Original Image', use_column_width=True)
                 with col2:
                     st.image(colorized_image, caption='Colorized Image', use_column_width=True)
 
@@ -141,6 +154,5 @@ if st.button("Colorize"):
 
             except Exception as e:
                 st.error(f"An error occurred: {e}")
-                logging.error("Error during inference", exc_info=True)
     else:
-        st.error("Please upload an image file.")
+        st.error("Please upload an image file or provide a valid URL.")