Updated app.py

app.py

@@ -67,7 +67,9 @@ def download_image_for_gradio(url: str, timeout: int = 20) -> Image.Image | None
     """Downloads an image from a URL for Gradio, returns PIL Image or raises gr.Error."""
     print(f"Attempting to download image from: {url}")
     if not url or not url.startswith(('http://', 'https://')):
-        raise gr.Error("Please enter a valid HTTP or HTTPS image URL.")
+        # Don't raise error here, let the caller decide if URL is optional
+        print("Invalid or empty URL provided.")
+        return None # Indicate failure without raising error immediately
 
     try:
         headers = {'User-Agent': 'Gradio-Image-Processor/1.1'}
@@ -108,19 +110,27 @@ def download_image_for_gradio(url: str, timeout: int = 20) -> Image.Image | None
 
 # --- Processing Function (Handles the ML part with progress) ---
 def run_processing(input_pil_image: Image.Image, progress=gr.Progress(track_tqdm=True)):
-    """Processes the already downloaded PIL image and returns the result."""
+    """Processes the input PIL image (from upload or download) and returns the result."""
     if MODEL is None:
         # Include the more specific error message if loading failed
         error_msg = f"Model is not loaded. Cannot process image. Load error: {MODEL_LOAD_ERROR}" if MODEL_LOAD_ERROR else "Model is not loaded. Cannot process image."
         raise gr.Error(error_msg)
     if input_pil_image is None:
-        return None, None
+        # This case might happen if the input component feeding this function is cleared
+        # before processing starts, or if the previous step failed silently.
+        print("run_processing called with None input.")
+        return None, None # Return None for both outputs
 
     start_time = time.time()
     print("Starting image processing...")
     progress(0, desc="Preparing for processing...")
 
     try:
+        # Ensure image is RGB before processing
+        if input_pil_image.mode != 'RGB':
+            print(f"Converting input image from {input_pil_image.mode} to RGB.")
+            input_pil_image = input_pil_image.convert('RGB')
+
         output_pil_image = process_image_from_data(
             input_pil_image=input_pil_image,
             model=MODEL,
@@ -150,26 +160,54 @@ def run_processing(input_pil_image: Image.Image, progress=gr.Progress(track_tqdm
 
     return output_pil_image, str(save_path)
 
-# --- Wrapper Function for Button Click ---
-def download_and_clear(url):
-    """Calls download, returns (image, None) on success, raises gr.Error on failure."""
-    print("Download and clear triggered.") # Debug print
-    try:
-        img = download_image_for_gradio(url)
-        # Return the downloaded image for the original slot, and None for the processed slot
-        return img, None
-    except gr.Error as e:
-        # If download fails, raise the Gradio error again
-        # Gradio will handle clearing outputs and displaying the error message
-        print(f"Download failed: {e}") # Debug print
-        raise e
-    except Exception as e:
-        # Catch other unexpected errors during the wrapper logic
-        print(f"Unexpected error in download_and_clear: {e}") # Debug print
-        import traceback
-        traceback.print_exc()
-        # Raise a Gradio error so the user sees something meaningful
-        raise gr.Error(f"An unexpected error occurred preparing the image: {e}")
+# --- NEW Wrapper Function for Button Click (Handles URL or Upload) ---
+def handle_input(url, uploaded_image):
+    """
+    Determines the input source (URL or Upload), prepares the image,
+    and returns it for the 'original_image' display.
+    Clears the other input method and the processed image output.
+    Raises gr.Error if input is invalid or download fails.
+    """
+    print("Handle input triggered.") # Debug print
+    input_image = None
+    final_url = url # Keep original url unless cleared
+    final_uploaded = None # Clear upload by default
+
+    if uploaded_image is not None:
+        print("Processing uploaded image.")
+        input_image = uploaded_image
+        final_url = "" # Clear URL field if upload is used
+
+    elif url:
+        print("Processing URL.")
+        try:
+            input_image = download_image_for_gradio(url)
+            if input_image is None:
+                # download_image_for_gradio returns None for non-http/https URLs
+                 raise gr.Error("Invalid URL provided. Please enter a valid HTTP or HTTPS image URL.")
+            # Keep the url in the box if download succeeds
+        except gr.Error as e:
+            # If download fails, raise the Gradio error again
+            print(f"Download failed: {e}") # Debug print
+            raise e
+        except Exception as e:
+            # Catch other unexpected errors during the wrapper logic
+            print(f"Unexpected error in handle_input (download branch): {e}") # Debug print
+            import traceback
+            traceback.print_exc()
+            # Raise a Gradio error so the user sees something meaningful
+            raise gr.Error(f"An unexpected error occurred preparing the image from URL: {e}")
+    else:
+        # Neither URL nor upload provided
+        raise gr.Error("Please provide an image URL or upload an image file.")
+
+    # At this point, input_image should be a valid PIL image or an error was raised.
+    # Return:
+    # 1. The input PIL image for 'original_image'
+    # 2. None to clear 'processed_image'
+    # 3. The final URL value (empty if upload was used, original if URL was used)
+    # 4. None to clear the file upload input
+    return input_image, None, final_url, final_uploaded
 
 
 def show_image(index):
@@ -192,68 +230,98 @@ with gr.Blocks() as demo:
     gr.Markdown(f"""
         # Watermark Remover
 
-        **Disclaimer: This project was made to showcase my Deep Learning skills with no intention to cause harm to any business or infringe on any IP and it will be decommissioned as soon as I get a decent job offer. 
+        **Disclaimer: This project was made to showcase my Deep Learning skills with no intention to cause harm to any business or infringe on any IP and it will be decommissioned as soon as I get a decent job offer.
         If you still have any issue, please reach out to me at [[email protected]](mailto:[email protected]).**
     """)
 
-    gr.Markdown(f"""
-        This is a demo of a DL model which takes in URL of an image with watermark and gives out the image with its watermark removed.
-        The image is broken into overlapping patches of 256x256 pixels with a stride of 64 before feeding them to the model, 
-        the model infers on each patch separately and then they are stitched together to form the whole image again to give the final output.
-    """)
-
-    if MODEL_LOAD_ERROR:
-        gr.Markdown(f"<font color='red'>**Model Loading Error:** {MODEL_LOAD_ERROR}</font>")
-
     with gr.Row():
         with gr.Column(scale=1):
+            gr.Markdown(f"""
+                This is a demo of a DL model which takes in an image with a watermark (either via URL or direct upload) and gives out the image with its watermark removed.
+                The image is broken into overlapping patches of {PATCH_KERNEL_SIZE}x{PATCH_KERNEL_SIZE} pixels with a stride of {PATCH_STRIDE} before feeding them to the model,
+                the model infers on each patch separately and then they are stitched together to form the whole image again to give the final output.
+            """)
+
+            if MODEL_LOAD_ERROR:
+                gr.Markdown(f"<font color='red'>**Model Loading Error:** {MODEL_LOAD_ERROR}</font>")
+
+            # Input Method 1: URL
             image_url = gr.Textbox(
                 label="Image URL",
-                placeholder="Paste image URL (https://www.shutterstock.com/shutterstock/photos/...)",
-                info="Enter the direct link to a publicly accessible image file (jpg, png, etc.)."
+                placeholder="Paste image URL (e.g., https://www.shutterstock.com/shutterstock/photos/.../image.jpg)",
+                info="Enter the direct link to a publicly accessible image file."
             )
+
+            # Input Method 2: File Upload
+            file_upload_input = gr.Image(
+                label="Or Upload Image File",
+                type="pil", # Keep as PIL for direct use
+                sources=["upload"], # Specify only upload source
+                height=130
+            )
+
             submit_button = gr.Button("Process Image", variant="primary")
-            gr.Markdown("---") # Separator
+            gr.Markdown("---")  # Separator
             gr.Examples(
-                 examples=[
-                    ["https://www.shutterstock.com/shutterstock/photos/2429728793/display_1500/stock-photo-monkey-funny-open-mouth-and-showing-teeth-crazy-nature-angry-short-hair-brown-grand-bassin-2429728793.jpg"],
-                    ["https://www.shutterstock.com/shutterstock/photos/1905929728/display_1500/stock-photo-skeptic-surprised-cat-thinking-dont-know-what-to-do-big-eyes-closeup-tabby-cat-look-side-dont-1905929728.jpg"],
-                    ["https://www.shutterstock.com/shutterstock/photos/2501926843/display_1500/stock-photo-brunette-woman-laying-on-couch-cuddling-light-brown-dog-and-brown-tabby-cat-happy-2501926843.jpg"]
-                 ],
-                 inputs=image_url # Link examples to the URL textbox
+                examples=[
+                    [
+                        "https://www.shutterstock.com/shutterstock/photos/2429728793/display_1500/stock-photo-monkey-funny-open-mouth-and-showing-teeth-crazy-nature-angry-short-hair-brown-grand-bassin-2429728793.jpg", None], # Need None for the upload input
+                    [
+                        "https://www.shutterstock.com/shutterstock/photos/1905929728/display_1500/stock-photo-skeptic-surprised-cat-thinking-dont-know-what-to-do-big-eyes-closeup-tabby-cat-look-side-dont-1905929728.jpg", None],
+                    [
+                        "https://www.shutterstock.com/shutterstock/photos/2501926843/display_1500/stock-photo-brunette-woman-laying-on-couch-cuddling-light-brown-dog-and-brown-tabby-cat-happy-2501926843.jpg", None]
+                ],
+                inputs=[image_url, file_upload_input] # Link examples to both inputs
             )
 
         with gr.Column(scale=2):
+            gr.Markdown(f"""### Some Samples""")
+            img_display = gr.Image(type="pil")
+            with gr.Row():
+                prev_btn = gr.Button("<")
+                next_btn = gr.Button(">")
+
+            prev_btn.click(fn=prev_img, inputs=index_state, outputs=[img_display, index_state])
+            next_btn.click(fn=next_img, inputs=index_state, outputs=[img_display, index_state])
+
+            # Load the first image
+            demo.load(fn=show_image, inputs=index_state, outputs=[img_display, index_state])
+
             with gr.Row():
+                # Make original_image interactive=False as it's now purely a display
                 original_image = gr.Image(label="Original Input Image", type="pil", interactive=False)
                 processed_image = gr.Image(label="Processed Output Image", type="pil", interactive=False)
             download_button = gr.DownloadButton("Download Processed Image", visible=True, variant="secondary")
-            gr.Markdown("---") # Separator
-            gr.Markdown(f"""
-                **A bit about the model:** In this project, I have trained a GAN network,
-                with the Generator being inspired from Pix2Pix and Pix2PixHD architectures and the Discriminator is very similar to PatchGAN in Pix2Pix and CycleGAN. 
-                For the loss, I have also added Perceptual Loss using VGG like in Pix2PixHD and SRGAN papers apart from the L1 and BCE loss.
-            """)
 
-            gr.Markdown(f"""If you liked this project, you can find my CV [here]() or reach me out at [[email protected]](mailto:[email protected]).""")
+    gr.Markdown("---")  # Separator
+    gr.Markdown(f"""
+                    **A bit about the model:** In this project, I have trained a GAN network,
+                    with the Generator being inspired from Pix2Pix and Pix2PixHD architectures and the Discriminator is very similar to PatchGAN in Pix2Pix.
+                    For the loss, I have also added Perceptual Loss using VGG like in Pix2PixHD and SRGAN papers apart from the L1 and BCE loss.
+                """)
 
+    gr.Markdown(
+        f"""If you liked this project, you can find my CV [here]() or reach me out at [[email protected]](mailto:[email protected]).""")
 
     # --- Event Handling Logic ---
 
     # 1. When Button is clicked:
-    #    - Input: URL from textbox
-    #    - Action: Call 'download_and_clear' wrapper function
-    #    - Output: Update 'original_image' with downloaded image, clear 'processed_image' by returning None for it.
+    #    - Input: URL from textbox AND image from file upload
+    #    - Action: Call 'handle_input' wrapper function to decide which input to use,
+    #              download if needed, and handle errors.
+    #    - Output: Update 'original_image' with the chosen/downloaded image,
+    #              clear 'processed_image', clear the URL field if upload was used,
+    #              clear the file upload component.
     submit_button.click(
-        fn=download_and_clear,  # <--- CORRECTED: Use the wrapper function
-        inputs=image_url,
-        outputs=[original_image, processed_image] # Target both outputs
+        fn=handle_input,
+        inputs=[image_url, file_upload_input],
+        outputs=[original_image, processed_image, image_url, file_upload_input] # Target outputs
     )
 
-    # 2. When 'original_image' component *changes* (i.e., after successful download):
+    # 2. When 'original_image' component *changes* (i.e., after successful input handling):
     #    - Input: The PIL image data from 'original_image'
     #    - Action: Call 'run_processing' (this function has the progress bar)
-    #    - Output: Update the 'processed_image' component.
+    #    - Output: Update the 'processed_image' component and the saved_image_path state.
     original_image.change(
         fn=run_processing,
         inputs=original_image,
@@ -261,24 +329,26 @@ with gr.Blocks() as demo:
         # concurrency_limit=1 # Optional: Prevent multiple simultaneous processing runs if needed
     )
 
+    # 3. When 'processed_image' changes (after processing finishes or is cleared):
+    #    - Action: Update the visibility/style of the download button.
     processed_image.change(
         fn=lambda img: gr.update(variant="primary" if img is not None else "secondary"),
         inputs=processed_image,
         outputs=download_button
     )
 
+    # 4. When Download Button is clicked:
+    #    - Action: Provide the file path stored in saved_image_path state to the download component.
     download_button.click(
         fn=lambda path: path if path else None,
         inputs=saved_image_path,
-        outputs=download_button
+        outputs=download_button # The output of the click event for DownloadButton is the file path itself
     )
 
 
-
 # --- Launch the Application ---
 if __name__ == "__main__":
     print("Launching Gradio Blocks interface...")
     # Set queue=True for better handling under load, especially with long-running processing
-    # demo.queue() # Consider adding this
-    demo.launch(share=False, server_name="0.0.0.0", show_api=False)
-
+    demo.queue() # Use queue for better user experience with processing time
+    demo.launch(share=False, server_name="0.0.0.0", show_api=False)