Update app.py

app.py

@@ -1,144 +1,136 @@
-# app.py
-import os
-import shutil
-import tempfile
-import cv2
-import numpy as np
-import gradio as gr
+# app.py  –  2025‑06 update
+import os, cv2, time, psutil, shutil, tempfile, zipfile, numpy as np, gradio as gr
 from paddleocr import PaddleOCR
-import psutil
-import time
-#import spaces
 
-#ocr = PaddleOCR(use_angle_cls=True, lang='en')
-#ocr = PaddleOCR(use_angle_cls=True, lang='en', det_model_dir='models/det', rec_model_dir='models/rec')
-ocr = PaddleOCR(use_angle_cls=True, lang='en')
+# ------------------------------------------------------------------------
+# 1. PaddleOCR – latest lightweight v5‑mobile with angle‑cls disabled
+# ------------------------------------------------------------------------
+ocr = PaddleOCR(
+    lang="en",
+    det_model_dir="ppocr_v5_det",
+    rec_model_dir="ppocr_v5_rec",
+    use_angle_cls=False,            # comic bubbles are already horizontal
+    show_log=False
+)
 
-def classify_background_color(avg_color, white_thresh=230, black_thresh=50, yellow_thresh=100):
-    r, g, b = avg_color
-    if r >= white_thresh and g >= white_thresh and b >= white_thresh:
+# ------------------------------------------------------------------------
+# 2. Utility helpers
+# ------------------------------------------------------------------------
+def wait_for_cpu(th=90, interval=3, timeout=30):
+    """Pause if CPU is saturated (helps on free‑tier Spaces)."""
+    start = time.time()
+    while psutil.cpu_percent(interval=1) > th:
+        time.sleep(interval)
+        if time.time() - start > timeout:
+            break
+
+def classify_bg(avg, w=230, b=50, y=100):
+    r, g, b_ = avg
+    if r >= w and g >= w and b_ >= w:          # white
         return (255, 255, 255)
-    if r <= black_thresh and g <= black_thresh and b <= black_thresh:
+    if r <= b and g <= b and b_ <= b:          # black
         return (0, 0, 0)
-    if r >= yellow_thresh and g >= yellow_thresh and b < yellow_thresh:
+    if r >= y and g >= y and b_ < y:           # yellowish narration box
         return (255, 255, 0)
     return None
 
-def sample_border_color(image, box, padding=2):
-    h, w = image.shape[:2]
-    x_min, y_min, x_max, y_max = box
-    x_min = max(0, x_min - padding)
-    x_max = min(w-1, x_max + padding)
-    y_min = max(0, y_min - padding)
-    y_max = min(h-1, y_max + padding)
-
-    top = image[y_min:y_min+padding, x_min:x_max]
-    bottom = image[y_max-padding:y_max, x_min:x_max]
-    left = image[y_min:y_max, x_min:x_min+padding]
-    right = image[y_min:y_max, x_max-padding:x_max]
-
-    border_pixels = np.vstack((top.reshape(-1, 3), bottom.reshape(-1, 3),
-                               left.reshape(-1, 3), right.reshape(-1, 3)))
-    if border_pixels.size == 0:
-        return (255, 255, 255)
-    median_color = np.median(border_pixels, axis=0)
-    return tuple(map(int, median_color))
-
-  
-def detect_text_boxes(image):
-    results = ocr.ocr(image, cls=True)
-    boxes = []
-    if results and results[0]:
-        for line in results[0]:
-            box, (text, confidence) = line
-            if text.strip():
-                x_min = int(min(pt[0] for pt in box))
-                x_max = int(max(pt[0] for pt in box))
-                y_min = int(min(pt[1] for pt in box))
-                y_max = int(max(pt[1] for pt in box))
-                boxes.append(((x_min, y_min, x_max, y_max), text, confidence))
-    else:
-        print("No text detected in the image.")
-    return boxes
-
-def remove_text_dynamic_fill(img_path, output_path):
-    image = cv2.imread(img_path)
-    if image is None:
+def sample_border(img, box, pad=2):
+    h, w = img.shape[:2]
+    x1, y1, x2, y2 = box
+    x1, x2 = max(0, x1 - pad), min(w - 1, x2 + pad)
+    y1, y2 = max(0, y1 - pad), min(h - 1, y2 + pad)
+    border = np.concatenate([
+        img[y1:y1+pad, x1:x2], img[y2-pad:y2, x1:x2],
+        img[y1:y2, x1:x1+pad], img[y1:y2, x2-pad:x2]
+    ], axis=0)
+    return tuple(np.median(border.reshape(-1, 3), axis=0).astype(int))
+
+# ------------------------------------------------------------------------
+# 3. Bubble‑mask (simple heuristic: very‑light regions enclosed)
+# ------------------------------------------------------------------------
+def make_bubble_mask(rgb):
+    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
+    # threshold near‑white & narration‑yellow
+    _, white = cv2.threshold(gray, 230, 255, cv2.THRESH_BINARY)
+    # small morph closing to join dotted edges
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
+    mask = cv2.morphologyEx(white, cv2.MORPH_CLOSE, kernel, iterations=2)
+    return mask // 255   # 0/1
+
+# ------------------------------------------------------------------------
+# 4. Detect & clean
+# ------------------------------------------------------------------------
+def remove_text_in_bubbles(img_path, dst_path):
+    bgr = cv2.imread(img_path)
+    if bgr is None:  # skip unreadable
         return
+    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+    bubble_mask = make_bubble_mask(rgb)
 
-    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)
-
-    boxes = detect_text_boxes(image)
+    results = ocr.ocr(rgb, cls=False)
+    if not results or not results[0]:
+        cv2.imwrite(dst_path, bgr)
+        return
 
-    for (bbox, text, confidence) in boxes:
-        if confidence < 0.4 or not text.strip():
+    for line in results[0]:
+        box = line[0]
+        text, conf = line[1]
+        if conf < 0.4 or not text.strip():
             continue
 
-        x_min, y_min, x_max, y_max = bbox
-        height = y_max - y_min
-
-        if height <= 30:
-            padding = 2
-        elif height <= 60:
-            padding = 4
-        else:
-            padding = 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)
-
-#@spaces.GPU  
-def process_folder(input_files):
-    temp_output = tempfile.mkdtemp()
-    
-
-    wait_for_cpu();    
-    for file in input_files:
-        filename = os.path.basename(file.name)
-        output_path = os.path.join(temp_output, filename)
-        remove_text_dynamic_fill(file.name, output_path)
-
-    zip_path = shutil.make_archive(temp_output, 'zip', temp_output)
+        # bounding box
+        xs, ys = zip(*box)
+        x1, x2 = int(min(xs)), int(max(xs))
+        y1, y2 = int(min(ys)), int(max(ys))
+
+        # skip if box centre is outside bubble mask
+        cx, cy = int((x1+x2)/2), int((y1+y2)/2)
+        if bubble_mask[cy, cx] == 0:
+            continue  # text is not inside a bubble
+
+        # dynamic padding by height
+        h_box = y2 - y1
+        pad = 2 if h_box <= 30 else 4 if h_box <= 60 else 6
+        x1p, y1p = max(0, x1-pad), max(0, y1-pad)
+        x2p, y2p = min(rgb.shape[1]-1, x2+pad), min(rgb.shape[0]-1, y2+pad)
+
+        # choose fill colour
+        crop = rgb[y1p:y2p, x1p:x2p]
+        fill = classify_bg(np.mean(crop.reshape(-1,3), axis=0))
+        if fill is None:
+            fill = sample_border(rgb, (x1, y1, x2, y2))
+
+        cv2.rectangle(bgr, (x1p, y1p), (x2p, y2p), fill, thickness=-1)
+
+    cv2.imwrite(dst_path, bgr)
+
+# ------------------------------------------------------------------------
+# 5. Gradio batch wrapper
+# ------------------------------------------------------------------------
+def process_folder(files):
+    wait_for_cpu()
+    out_dir = tempfile.mkdtemp()
+    for f in files:
+        fname = os.path.basename(f)
+        remove_text_in_bubbles(f, os.path.join(out_dir, fname))
+    zip_path = shutil.make_archive(out_dir, 'zip', out_dir)
     return zip_path
 
-
-
-def wait_for_cpu(threshold=90, interval=3, timeout=30):
-    start = time.time()
-    while psutil.cpu_percent(interval=1) > threshold:
-        print("High CPU usage detected, waiting...")
-        time.sleep(interval)
-        if time.time() - start > timeout:
-            print("Timed out waiting for CPU to cool down.")
-            break
-
-
+# ------------------------------------------------------------------------
+# 6. Gradio UI
+# ------------------------------------------------------------------------
 demo = gr.Interface(
     fn=process_folder,
-    inputs=gr.File(file_types=[".jpg", ".jpeg", ".png"], file_count="multiple", label="Upload Comic Images"),
-    outputs=gr.File(label="Download Cleaned Zip"),
-    concurrency_limit=1,
-    title="Comic Text Cleaner",
-    description="Upload comic images and get a zip of cleaned versions (text removed). Uses PaddleOCR for detection."
+    inputs=gr.File(file_types=[".jpg", ".jpeg", ".png"],
+                   label="Upload comic page images",
+                   file_count="multiple"),
+    outputs=gr.File(label="Download cleaned .zip"),
+    title="Comic Bubble Text Cleaner – PP‑OCRv5",
+    description=("Removes speech/thought/narration bubble text only, "
+                 "leaving outside FX or captions untouched. "
+                 "Powered by PaddleOCR PP‑OCRv5‑mobile."),
+    concurrency_limit=1
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()