Update app.py

app.py

@@ -2,16 +2,15 @@ import spacy
 import os
 import shutil
 import gradio as gr
+import subprocess
 from booknlp.booknlp import BookNLP
 from spacy.cli import download
 
 # Ensure Spacy model is downloaded
 def ensure_spacy_model():
     try:
-        # Try to load the model to see if it's already installed
         spacy.load("en_core_web_sm")
     except OSError:
-        # If not installed, download the model
         download("en_core_web_sm")
 
 # Initialize Spacy model
@@ -24,12 +23,23 @@ model_params = {
 }
 booknlp = BookNLP("en", model_params)
 
+# Define function to convert ebook to txt using Calibre
+def convert_to_txt(input_file):
+    output_txt = os.path.splitext(input_file)[0] + ".txt"
+    if not os.path.exists(output_txt):
+        subprocess.run(["ebook-convert", input_file, output_txt], check=True)
+    return output_txt
+
 # Define function to process file
 def process_book(file):
     input_file = file.name
     output_dir = "output_dir/booknlp_output/"
     book_id = os.path.splitext(os.path.basename(input_file))[0]
     
+    # Check if the file is already a .txt file
+    if not input_file.endswith(".txt"):
+        input_file = convert_to_txt(input_file)
+
     # Create output directory if it doesn't exist
     if os.path.exists(output_dir):
         shutil.rmtree(output_dir)
@@ -45,15 +55,26 @@ def process_book(file):
 
 # Gradio Interface
 def gradio_interface():
-    file_input = gr.File(file_types=['.txt'], label="Upload a .txt file")
+    # Define supported file formats
+    supported_formats = [
+        '.azw', '.azw3', '.azw4', '.cbz', '.cbr', '.cb7', '.cbc', '.chm', 
+        '.djvu', '.docx', '.epub', '.fb2', '.fbz', '.html', '.htmlz', '.lit', 
+        '.lrf', '.mobi', '.odt', '.pdf', '.prc', '.pdb', '.pml', '.rb', 
+        '.rtf', '.snb', '.tcr', '.txt', '.txtz'
+    ]
+
+    file_input = gr.File(file_types=supported_formats, label="Upload an ebook file (.azw, .epub, .pdf, .txt, etc.)")
     file_output = gr.File(label="Download the output files")
     
+    # Show supported formats in the description
+    description = f"Upload any of the supported formats: {', '.join(supported_formats)}. If a .txt file is uploaded, it will directly be processed by BookNLP. Otherwise, it will be converted to .txt using Calibre first."
+
     gr.Interface(
         fn=process_book, 
         inputs=file_input, 
         outputs=file_output, 
-        title="BookNLP Processor",
-        description="Upload a .txt book and process it using BookNLP. Download the processed output files."
+        title="BookNLP Processor with Ebook Support",
+        description=description
     ).launch()
 
 if __name__ == "__main__":