Update app.py

app.py

@@ -12,6 +12,7 @@ from transformers import pipeline
 import numpy as np
 import json
 import torch
+import logging
 
 # Load Stable Diffusion Model
 def load_stable_diffusion_model():
@@ -354,10 +355,15 @@ def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_s
     
     return original_images, processed_images, output_zip_path, processing_time
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+
 def gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers):
+    """Handles input files and processes them accordingly."""
     progress = gr.Progress()
     watermark_path = watermark.name if watermark else None
     
+    # Check the type of input_files and call the process_images function
     if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
         return process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
     elif isinstance(input_files, list):
@@ -366,31 +372,41 @@ def gradio_interface(input_files, bg_method, watermark, canvas_size, output_form
         return process_images(input_files.name, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
 
 def show_color_picker(bg_choice):
-    if bg_choice == 'custom':
-        return gr.update(visible=True)
-    return gr.update(visible=False)
+    """Shows the color picker if 'custom' background is selected."""
+    return gr.update(visible=(bg_choice == 'custom'))
 
 def update_compare(evt: gr.SelectData):
+    """Updates the displayed images and their ratios when a processed image is selected."""
     if isinstance(evt.value, dict) and 'caption' in evt.value:
-        input_path = evt.value['caption']
+        input_path = evt.value['caption'].split("Input: ")[-1]
         output_path = evt.value['image']['path']
-        input_path = input_path.split("Input: ")[-1]
         
+        # Open the original and processed images
         original_img = Image.open(input_path)
         processed_img = Image.open(output_path)
         
+        # Calculate the aspect ratios
         original_ratio = f"{original_img.width}x{original_img.height}"
         processed_ratio = f"{processed_img.width}x{processed_img.height}"
         
-        return gr.update(value=input_path), gr.update(value=output_path), gr.update(value=original_ratio), gr.update(value=processed_ratio)
+        return (gr.update(value=input_path), 
+                gr.update(value=output_path), 
+                gr.update(value=original_ratio), 
+                gr.update(value=processed_ratio)
+               )
     else:
-        print("No caption found in selection")
-        return gr.update(value=None), gr.update(value=None), gr.update(value=None), gr.update(value=None)
+        logging.warning("No caption found in selection")
+        return (gr.update(value=None),) * 4
 
 def process(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers):
-    _, processed_images, zip_path, time_taken = gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers)
-    processed_images_with_captions = [(img, f"Input: {caption}") for img, caption in processed_images]
-    return processed_images_with_captions, zip_path, f"{time_taken:.2f} seconds"
+    """Processes the images and returns the results."""
+    try:
+        _, processed_images, zip_path, time_taken = gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers)
+        processed_images_with_captions = [(img, f"Input: {caption}") for img, caption in processed_images]
+        return processed_images_with_captions, zip_path, f"{time_taken:.2f} seconds"
+    except Exception as e:
+        logging.error(f"Error in processing images: {e}")
+        return [], None, "Error during processing"
 
 with gr.Blocks(theme="NoCrypt/miku@1.2.2") as iface:
     gr.Markdown("# Image Background Removal and Resizing with Optional Watermark")