revert to old version

app.py

@@ -1,77 +1,181 @@
 import gradio as gr
+import onnxruntime as ort
 import numpy as np
 import cv2
-import json
+from huggingface_hub import hf_hub_download
 import os
-from meikiocr import MeikiOCR
+import json
+
+# --- 1. configuration & model loading ---
+# this section runs once when the space starts up.
+
+print("loading models...")
+
+# configuration
+det_model_repo = "rtr46/meiki.text.detect.v0"
+det_model_name = "meiki.text.detect.v0.1.960x544.onnx"
+rec_model_repo = "rtr46/meiki.txt.recognition.v0"
+rec_model_name = "meiki.text.rec.v0.960x32.onnx"
 
-print("Initializing meikiocr...")
+input_det_width = 960
+input_det_height = 544
+input_rec_height = 32
+input_rec_width = 960
+x_overlap_threshold = 0.3
+epsilon = 1e-6
+
+# load models from the hub
 try:
-    ocr = MeikiOCR(provider='CPUExecutionProvider')
-    print("meikiocr initialized successfully.")
+    det_model_path = hf_hub_download(repo_id=det_model_repo, filename=det_model_name)
+    rec_model_path = hf_hub_download(repo_id=rec_model_repo, filename=rec_model_name)
+    
+    # use cpu execution provider for broad compatibility in spaces
+    providers = ['CPUExecutionProvider']
+    det_session = ort.InferenceSession(det_model_path, providers=providers)
+    rec_session = ort.InferenceSession(rec_model_path, providers=providers)
+    
+    print("models loaded successfully.")
 except Exception as e:
-    ocr = None
-    print(f"Error initializing meikiocr: {e}")
-    # Display a persistent error in the Gradio interface if model loading fails.
-    raise gr.Error(f"Failed to load OCR models. Please check the space logs for details. Error: {e}")
+    det_session, rec_session = None, None
+    print(f"error loading models: {e}")
+    raise gr.Error(f"failed to load models. please check space logs. error: {e}")
+
+# --- 2. ocr pipeline helper functions ---
+# (these functions remain unchanged)
+
+def preprocess_for_detection(image):
+    h_orig, w_orig, _ = image.shape
+    resized = cv2.resize(image, (input_det_width, input_det_height), interpolation=cv2.INTER_LINEAR)
+    input_tensor = resized.astype(np.float32) / 255.0
+    input_tensor = np.transpose(input_tensor, (2, 0, 1))
+    input_tensor = np.expand_dims(input_tensor, axis=0)
+    scale_x = w_orig / input_det_width
+    scale_y = h_orig / input_det_height
+    return input_tensor, scale_x, scale_y
+
+def postprocess_detection_results(raw_outputs, scale_x, scale_y, conf_threshold):
+    _, boxes, scores = raw_outputs
+    boxes, scores = boxes[0], scores[0]
+    text_boxes = []
+    for box, score in zip(boxes, scores):
+        if score < conf_threshold: continue
+        x1, y1, x2, y2 = box
+        x1_orig, y1_orig = int(x1 * scale_x), int(y1 * scale_y)
+        x2_orig, y2_orig = int(x2 * scale_x), int(y2 * scale_y)
+        text_boxes.append({'bbox': [x1_orig, y1_orig, x2_orig, y2_orig]})
+    text_boxes.sort(key=lambda tb: tb['bbox'][1])
+    return text_boxes
+
+def preprocess_for_recognition(image, text_boxes):
+    tensors, valid_indices, crop_metadata = [], [], []
+    for i, tb in enumerate(text_boxes):
+        x1, y1, x2, y2 = tb['bbox']
+        width, height = x2 - x1, y2 - y1
+        if width < height or width == 0 or height == 0: continue
+        crop = image[y1:y2, x1:x2]
+        h, w, _ = crop.shape
+        new_h, new_w = input_rec_height, int(round(w * (input_rec_height / h)))
+        if new_w > input_rec_width:
+            scale = input_rec_width / new_w
+            new_w, new_h = input_rec_width, int(round(new_h * scale))
+        resized = cv2.resize(crop, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+        pad_w, pad_h = input_rec_width - new_w, input_rec_height - new_h
+        padded = np.pad(resized, ((0, pad_h), (0, pad_w), (0, 0)), constant_values=0)
+        tensor = (padded.astype(np.float32) / 255.0)
+        tensor = np.transpose(tensor, (2, 0, 1))
+        tensors.append(tensor)
+        valid_indices.append(i)
+        crop_metadata.append({'orig_bbox': [x1, y1, x2, y2], 'effective_w': new_w})
+    if not tensors: return None, [], []
+    return np.stack(tensors, axis=0), valid_indices, crop_metadata
+
+def postprocess_recognition_results(raw_rec_outputs, valid_indices, crop_metadata, rec_conf_threshold, num_total_boxes):
+    labels_batch, boxes_batch, scores_batch = raw_rec_outputs
+    full_results = [{'text': '', 'chars': []} for _ in range(num_total_boxes)]
+    for i, (labels, boxes, scores) in enumerate(zip(labels_batch, boxes_batch, scores_batch)):
+        meta = crop_metadata[i]
+        gx1, gy1, gx2, gy2 = meta['orig_bbox']
+        crop_w, crop_h = gx2 - gx1, gy2 - gy1
+        effective_w = meta['effective_w']
+        candidates = []
+        for lbl, box, scr in zip(labels, boxes, scores):
+            if scr < rec_conf_threshold: continue
+            char = chr(lbl)
+            rx1, ry1, rx2, ry2 = box
+            rx1, rx2 = min(rx1, effective_w), min(rx2, effective_w)
+            cx1, cx2 = (rx1 / effective_w) * crop_w, (rx2 / effective_w) * crop_w
+            cy1, cy2 = (ry1 / input_rec_height) * crop_h, (ry2 / input_rec_height) * crop_h
+            gx1_char, gy1_char = gx1 + int(cx1), gy1 + int(cy1)
+            gx2_char, gy2_char = gx1 + int(cx2), gy1 + int(cy2)
+            candidates.append({'char': char, 'bbox': [gx1_char, gy1_char, gx2_char, gy2_char], 'x_interval': (gx1_char, gx2_char), 'conf': float(scr)})
+        candidates.sort(key=lambda c: c['conf'], reverse=True)
+        accepted = []
+        for cand in candidates:
+            x1_c, x2_c = cand['x_interval']
+            width_c = x2_c - x1_c + epsilon
+            is_overlap = any((max(0, min(x2_c, x2_a) - max(x1_c, x1_a)) / width_c) > x_overlap_threshold for x1_a, x2_a in (acc['x_interval'] for acc in accepted))
+            if not is_overlap: accepted.append(cand)
+        accepted.sort(key=lambda c: c['x_interval'][0])
+        text = ''.join(c['char'] for c in accepted)
+        final_chars = [{'char': c['char'], 'bbox': c['bbox'], 'conf': c['conf']} for c in accepted]
+        full_results[valid_indices[i]] = {'text': text, 'chars': final_chars}
+    return full_results
+
+# --- 3. main gradio processing function ---
 
 def run_ocr_pipeline(input_image, det_threshold, rec_threshold):
-    """
-    Takes a user-uploaded image and confidence thresholds, runs the OCR pipeline,
-    and returns the results formatted for the Gradio interface.
-    """
     if input_image is None:
-        raise gr.Error("Please upload an image to process.")
+        raise gr.Error("please upload an image to process.")
+
+    det_input, sx, sy = preprocess_for_detection(input_image)
+    det_raw = det_session.run(None, {det_session.get_inputs()[0].name: det_input, det_session.get_inputs()[1].name: np.array([[input_det_width, input_det_height]], dtype=np.int64)})
+    text_boxes = postprocess_detection_results(det_raw, sx, sy, det_threshold)
 
-    results = ocr.run_ocr(input_image, det_threshold=det_threshold, rec_threshold=rec_threshold)
+    if not text_boxes:
+        return input_image, "no text detected. try lowering the 'detection confidence' slider.", ""
 
-    if not results:
-        return input_image, "No text detected. Try lowering the 'Detection Confidence' slider.", ""
+    rec_batch, valid_indices, crop_metadata = preprocess_for_recognition(input_image, text_boxes)
+    rec_raw = rec_session.run(None, {"images": rec_batch, "orig_target_sizes": np.array([[input_rec_width, input_rec_height]], dtype=np.int64)})
+    results = postprocess_recognition_results(rec_raw, valid_indices, crop_metadata, rec_threshold, len(text_boxes))
 
-    # Prepare the outputs for Gradio
     output_image = input_image.copy()
-    full_text_lines = []
-    
-    # Draw bounding boxes and collect text
-    for line_result in results:
-        if line_result['text']:
-            full_text_lines.append(line_result['text'])
-        # Draw a green rectangle for each recognized character
-        for char_info in line_result['chars']:
+    full_text = []
+    for res in results:
+        if res['text']: full_text.append(res['text'])
+        for char_info in res['chars']:
             x1, y1, x2, y2 = char_info['bbox']
             cv2.rectangle(output_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
     
-    # Format the full text and JSON output
-    recognized_text = "\n".join(full_text_lines)
     json_output = json.dumps(results, indent=2, ensure_ascii=False)
     
-    return output_image, recognized_text, json_output
+    return output_image, "\n".join(full_text), json_output
+
+# --- 4. gradio interface definition ---
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# meikiocr: Japanese Video Game OCR")
+    gr.Markdown("# meikiocr: japanese video game ocr")
     gr.Markdown(
-        "Upload a screenshot from a Japanese video game to see the high-accuracy OCR in action. "
-        "The pipeline first detects text lines, then recognizes the characters in each line. "
-        "Adjust the confidence sliders if text is missed or incorrectly detected."
+        "upload a screenshot from a japanese video game to see the high-accuracy ocr in action. "
+        "the pipeline first detects text lines, then recognizes the characters in each line. "
+        "adjust the confidence sliders if text is missed or incorrectly detected."
     )
     
     with gr.Row():
         with gr.Column(scale=1):
-            input_image = gr.Image(type="numpy", label="Upload Image")
-            det_threshold = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="Detection Confidence")
-            rec_threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.05, label="Recognition Confidence")
-            run_button = gr.Button("Run OCR", variant="primary")
+            input_image = gr.Image(type="numpy", label="upload image")
+            det_threshold = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="detection confidence")
+            rec_threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.05, label="recognition confidence")
+            run_button = gr.Button("run ocr", variant="primary")
             
         with gr.Column(scale=2):
-            output_image = gr.Image(type="numpy", label="OCR Result")
-            output_text = gr.Textbox(label="Recognized Text", lines=5)
-            output_json = gr.Code(label="JSON Output", language="json", lines=5)
+            output_image = gr.Image(type="numpy", label="ocr result")
+            output_text = gr.Textbox(label="recognized text", lines=5)
+            output_json = gr.Code(label="json output", language="json", lines=5)
 
-    # The function for handling examples is also simplified.
     def process_example(img):
+        # examples are pre-loaded as numpy by gradio, so we can pass them directly
         return run_ocr_pipeline(img, 0.5, 0.1)
 
-    # Load the example image if it exists
     example_image_path = os.path.join(os.path.dirname(__file__), "example.jpg")
     if os.path.exists(example_image_path):
         gr.Examples(
@@ -82,7 +186,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             cache_examples=True
         )
             
-    # Connect the button click to the main processing function
     run_button.click(
         fn=run_ocr_pipeline,
         inputs=[input_image, det_threshold, rec_threshold],
@@ -92,10 +195,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
         ---
-        ### Official GitHub Repository
-        The full source code, documentation, and local command-line script for `meikiocr` are available on GitHub.
+        ### official github repository
+        the full source code, documentation, and local command-line script for `meikiocr` are available on github.
         **[github.com/rtr46/meikiocr](https://github.com/rtr46/meikiocr)**
         """
     )
 
+# --- 5. launch the app ---
 demo.launch()