Update app.py

app.py

@@ -110,16 +110,14 @@ def redact_image(pdf_image_path, redaction_score_threshold):
     return redacted_image_path
 
 # Convert the set of redacted images to a pdf
-def stich_images_to_pdf(redacted_image_files, input_pdf_name):
+def stich_images_to_pdf(redacted_image_files, input_pdf_path):
 
     # Sort the redacted images
     redacted_image_files.sort()
 
     # Convert the redacted images to a single PDF
     print("Converting redacted images to PDF...")
-    redacted_pdf_folder = "/tmp/gradio/redacted"
-    os.makedirs(redacted_pdf_folder, exist_ok=True )
-    redacted_pdf_path = f'{redacted_pdf_folder}/{input_pdf_name}_redacted.pdf'
+    redacted_pdf_path = input_pdf_path.replace('.pdf', '_redacted.pdf')
 
     doc = pymupdf.open()
     for redacted_image_file in redacted_image_files:
@@ -165,10 +163,6 @@ def predict(input_pdf_path, sensitivity):
     # Convert sensitivity to threshold
     redaction_score_threshold = (100-sensitivity)/100
 
-    # Get file name
-    print("Getting filename")
-    input_pdf_name = input_pdf_path.split('.')[-2]
-
     # Convert the PDF to images
     print("Converting pdf to images")
     pdf_images_dir = convert_to_images(input_pdf_path)
@@ -189,7 +183,7 @@ def predict(input_pdf_path, sensitivity):
 
     # Convert the redacted images to a single PDF
     print("Stitching images to pdf")
-    redacted_pdf_path = stich_images_to_pdf(redacted_image_files, input_pdf_name)
+    redacted_pdf_path = stich_images_to_pdf(redacted_image_files, input_pdf_path)
 
     print("Cleaning up")
     cleanup(redacted_image_files, pdf_images, pdf_images_dir, input_pdf_path)