Save each result (image + braille + english) to braille-reader-results dataset

app.py

@@ -1,14 +1,18 @@
 """Braille Reader — Upload a braille image, get English text."""
 
 import json
+import os
 import tempfile
+import uuid
+from datetime import datetime
+from pathlib import Path
 
 import cv2
 import gradio as gr
 import numpy as np
 import spaces
 import torch
-from huggingface_hub import hf_hub_download
+from huggingface_hub import CommitScheduler, hf_hub_download
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from ultralytics import YOLO
 
@@ -16,6 +20,7 @@ from ultralytics import YOLO
 
 YOLO_REPO = "prasanthmj/yolov8-braille"
 BYT5_REPO = "prasanthmj/braille-byt5-v3"
+DATASET_REPO = "prasanthmj/braille-reader-results"
 
 print("Loading models...")
 
@@ -33,6 +38,47 @@ byt5_model.eval()
 
 print("Models loaded (CPU). GPU allocated per-request via ZeroGPU.")
 
+# --- Result saving via CommitScheduler ---
+
+RESULTS_DIR = Path("./results")
+RESULTS_DIR.mkdir(exist_ok=True)
+(RESULTS_DIR / "images").mkdir(exist_ok=True)
+
+scheduler = CommitScheduler(
+    repo_id=DATASET_REPO,
+    repo_type="dataset",
+    folder_path=RESULTS_DIR,
+    every=5,  # push every 5 minutes
+    token=os.environ.get("HF_TOKEN"),
+)
+
+
+def save_result(image: np.ndarray, braille_text: str, english_text: str,
+                total_cells: int, num_lines: int, avg_conf: float):
+    """Save image and result to the dataset (batched by CommitScheduler)."""
+    entry_id = datetime.utcnow().strftime("%Y%m%d_%H%M%S") + "_" + uuid.uuid4().hex[:6]
+
+    # Save image
+    image_filename = f"images/{entry_id}.jpg"
+    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    cv2.imwrite(str(RESULTS_DIR / image_filename), image_bgr)
+
+    # Append to JSONL
+    record = {
+        "id": entry_id,
+        "image": image_filename,
+        "braille_unicode": braille_text,
+        "english": english_text,
+        "cells": total_cells,
+        "lines": num_lines,
+        "avg_confidence": round(avg_conf, 4),
+        "timestamp": datetime.utcnow().isoformat(),
+    }
+
+    with scheduler.lock:
+        with open(RESULTS_DIR / "results.jsonl", "a") as f:
+            f.write(json.dumps(record) + "\n")
+
 # --- CLAHE Preprocessing ---
 
 def preprocess_clahe(image_path: str) -> str:
@@ -125,7 +171,7 @@ def transcribe(image) -> str:
 
     # Stats
     total_cells = sum(len(line) for line in lines)
-    avg_conf = np.mean([cell["confidence"] for line in lines for cell in line])
+    avg_conf = float(np.mean([cell["confidence"] for line in lines for cell in line]))
 
     # Stage 2: Interpret each line with ByT5 on GPU
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -151,6 +197,9 @@ def transcribe(image) -> str:
     braille_text = "\n".join(braille_lines)
     english_text = "\n".join(english_lines)
 
+    # Save to dataset
+    save_result(image, braille_text, english_text, total_cells, len(lines), avg_conf)
+
     output = f"{english_text}\n\n"
     output += f"--- Details ---\n"
     output += f"Cells detected: {total_cells}\n"