Update app.py

app.py

@@ -1,28 +1,13 @@
 import os
 import shutil
 import tempfile
+import zipfile
+import rarfile
+import gradio as gr
 import cv2
 import numpy as np
-import gradio as gr
-import zipfile
-import patoolib
-import threading
 from paddleocr import PaddleOCR
 from PIL import Image
-import io
-
-def delayed_cleanup(path, delay=30):
-    def cleanup():
-        import time
-        time.sleep(delay)
-        try:
-            if os.path.isdir(path):
-                shutil.rmtree(path)
-            elif os.path.exists(path):
-                os.remove(path)
-        except Exception as e:
-            print(f"Cleanup failed: {e}")
-    threading.Thread(target=cleanup).start()
 
 ocr = PaddleOCR(use_angle_cls=True, lang='en', det_model_dir='models/det', rec_model_dir='models/rec', cls_model_dir='models/cls')
 
@@ -75,7 +60,9 @@ def remove_text_dynamic_fill(img_path, output_path):
     image = cv2.imread(img_path)
     if image is None:
         return
-    if len(image.shape) == 2 or image.shape[2] == 1:
+    if len(image.shape) == 2:
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+    elif image.shape[2] == 1:
         image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
     else:
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
@@ -88,57 +75,50 @@ def remove_text_dynamic_fill(img_path, output_path):
         x_min, y_min, x_max, y_max = bbox
         height = y_max - y_min
         padding = 2 if height <= 30 else 4 if height <= 60 else 6
+
         x_min_p = max(0, x_min - padding)
         y_min_p = max(0, y_min - padding)
         x_max_p = min(image.shape[1]-1, x_max + padding)
         y_max_p = min(image.shape[0]-1, y_max + padding)
+
         sample_crop = image[y_min_p:y_max_p, x_min_p:x_max_p]
         avg_color = np.mean(sample_crop.reshape(-1, 3), axis=0)
         fill_color = classify_background_color(avg_color)
         if fill_color is None:
             fill_color = sample_border_color(image, (x_min, y_min, x_max, y_max))
+
         cv2.rectangle(image, (x_min_p, y_min_p), (x_max_p, y_max_p), fill_color, -1)
 
     image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
     cv2.imwrite(output_path, image)
 
-def extract_comic_archive(archive_path, extract_to):
-    if archive_path.endswith(".cbz"):
-        with zipfile.ZipFile(archive_path, 'r') as zip_ref:
-            zip_ref.extractall(extract_to)
-    elif archive_path.endswith(".cbr"):
-        patoolib.extract_archive(archive_path, outdir=extract_to)
-
-def process_cbz_cbr(input_files):
-    temp_output = tempfile.mkdtemp()
-
-    for file in input_files:
-        filename = os.path.basename(file)
-        output_path = os.path.join(temp_output, filename)
-        remove_text_dynamic_fill(file, output_path)
+def process_cbz_cbr(files):
+    final_output = tempfile.mkdtemp()
+
+    for file_path in files:
+        if file_path.endswith(".cbz"):
+            with zipfile.ZipFile(file_path, 'r') as archive:
+                extract_dir = tempfile.mkdtemp()
+                archive.extractall(extract_dir)
+        elif file_path.endswith(".cbr"):
+            with rarfile.RarFile(file_path) as archive:
+                extract_dir = tempfile.mkdtemp()
+                archive.extractall(extract_dir)
+        else:
+            continue
 
-    zip_bytes = io.BytesIO()
-    with zipfile.ZipFile(zip_bytes, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
-        for root, _, files in os.walk(temp_output):
+        for root, _, imgs in os.walk(extract_dir):
+            for img in imgs:
+                if img.lower().endswith(('.jpg', '.jpeg', '.png')):
+                    input_path = os.path.join(root, img)
+                    output_path = os.path.join(final_output, os.path.basename(img))
+                    remove_text_dynamic_fill(input_path, output_path)
+
+    # Create output zip
+    zip_buffer = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
+    with zipfile.ZipFile(zip_buffer.name, 'w', zipfile.ZIP_DEFLATED) as zf:
+        for root, _, files in os.walk(final_output):
             for file in files:
-                file_path = os.path.join(root, file)
-                arcname = os.path.relpath(file_path, temp_output)
-                zf.write(file_path, arcname)
-
-    zip_bytes.seek(0)
-    zip_path = os.path.join(tempfile.gettempdir(), "cleaned_output.zip")
-    with open(zip_path, "wb") as f:
-        f.write(zip_bytes.read())
-
-    return zip_path
-
-
-demo = gr.Interface(
-    fn=process_cbz_cbr,
-    inputs=gr.File(file_types=[".cbz", ".cbr"], label="Upload Comic Archive (.cbz or .cbr)"),
-    outputs=gr.File(label="Download Cleaned Zip"),
-    title="Comic Text Cleaner (.cbz/.cbr)",
-    description="Upload a .cbz or .cbr file and get a zip of cleaned comic images (text removed using PaddleOCR)."
-)
+                zf.write(os.path.join(root, file), arcname=file)
 
-demo.launch()
+    return zip_buffer.name