refactor

app.py

@@ -1,11 +1,15 @@
+import argparse
 import os
 from functools import partial
+
 import gradio as gr
 from PIL import Image
 from huggingface_hub import hf_hub_download
+
 from die_model import UNetDIEModel
 from utils import resize_image, make_image_square, cast_pil_image_to_torch_tensor_with_4_channel_dim, remove_square_padding
 
+
 def die_inference(image_raw, num_of_die_iterations, die_model, device):
     """
     Applies the DIE model for document enhancement on a provided image.
@@ -31,51 +35,80 @@ def die_inference(image_raw, num_of_die_iterations, die_model, device):
         resize_back_to_original=True
     )
 
-description = """
-Welcome to the Document Image Enhancement (DIE) model demo on Hugging Face!
-
-This application showcases a specialized AI model by the Artificial Intelligence group at the Alfréd Rényi Institute of Mathematics, aimed at enhancing and restoring archival document images. This model removes domain-specific noise, preserving clarity and improving OCR accuracy, particularly for aged and historical documents.
-
-Contact: [email protected]
-
-"""
-
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("## Document Image Enhancement (DIE) Model")
 
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown(description)
-        with gr.Column():
-            # Displaying the QR code directly as an image in Gradio
-            gr.Image(value=Image.open("logo/qr-code.png"), label="QR Code")
-
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Upload Degraded Document Image")
-            num_iterations = gr.Dropdown([1, 2, 3], label="Number of DIE Iterations", value=1)
-            run_button = gr.Button("Enhance Image")
-
-        with gr.Column():
-            output_image = gr.Image(type="pil", label="Enhanced Document Image")
+def main():
+    """
+    Main function to set up and run the Gradio demo.
+    """
+    args = parse_arguments()
 
-    # Load model
+    # Set up model
     die_token = os.getenv("DIE_TOKEN")
     model_path = hf_hub_download(
         repo_id="gabar92/die",
-        filename="2024_08_09_model_epoch_89.pt",
+        filename=args.die_model_path,
         use_auth_token=die_token
     )
     die_model = UNetDIEModel(args=model_path)
-    device = "cpu"  # or "cuda" based on your setup
+    device = args.device
+
+    # Prepare example images
+    example_image_list = [
+        [Image.open(os.path.join(args.example_image_path, image_path))]
+        for image_path in os.listdir(args.example_image_path)
+    ]
+
+    description = """
+    Welcome to the Document Image Enhancement (DIE) model demo on Hugging Face!
+
+    This application showcases a specialized AI model by the Artificial Intelligence group at the Alfréd Rényi Institute of Mathematics, aimed at enhancing and restoring archival document images. This model removes domain-specific noise, preserving clarity and improving OCR accuracy, particularly for aged and historical documents.
 
-    # Partial function for inference
+    Contact: [email protected]
+
+    """
+
+    # Partial function for inference with model and device arguments
     partial_die_inference = partial(die_inference, die_model=die_model, device=device)
 
-    # Define button behavior
-    run_button.click(partial_die_inference, [input_image, num_iterations], output_image)
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("## Document Image Enhancement (DIE) Model")
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown(description)
+            with gr.Column():
+                # Display QR code as an image in Gradio
+                gr.Image(value=Image.open("path/to/qr-code.png"), label="QR Code")
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(type="pil", label="Upload Degraded Document Image")
+                num_iterations = gr.Dropdown([1, 2, 3], label="Number of DIE Iterations", value=1)
+                run_button = gr.Button("Enhance Image")
+
+            with gr.Column():
+                output_image = gr.Image(type="pil", label="Enhanced Document Image")
+
+        # Button trigger for inference
+        run_button.click(partial_die_inference, [input_image, num_iterations], output_image)
+
+    demo.launch()
+
+
+def parse_arguments():
+    """
+    Parses command-line arguments.
+    :return: argument namespace
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--die_model_path", default="2024_08_09_model_epoch_89.pt", help="Path to the DIE model checkpoint")
+    parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"], help="Device to run the model on")
+    parser.add_argument("--example_image_path", default="example_images", help="Path to directory with example images")
+    return parser.parse_args()
+    
 
-demo.launch()
+if __name__ == "__main__":
+    main()