asterocr 1.0.0__py3-none-any.whl

asterocr/studio.py

@@ -0,0 +1,371 @@
+"""
+AsterOCR Standalone Training Studio & Neural Flow Visualizer
+Independent server running on its own dedicated port (default: 7860).
+Completely decoupled from any other application or document extraction project.
+"""
+
+import os
+import sys
+import webbrowser
+import threading
+import base64
+from io import BytesIO
+
+# Ensure parent directory is in sys.path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+if sys.platform == "win32":
+    try:
+        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
+        sys.stderr.reconfigure(encoding='utf-8', errors='replace')
+    except Exception:
+        pass
+
+import time
+import torch
+from flask import Flask, render_template, request, jsonify, send_file, Response
+from asterocr.trainer.trainer import global_trainer
+from asterocr.synthesizer.generator import default_generator
+
+template_dir = os.path.join(os.path.dirname(__file__), "web", "templates")
+app = Flask(__name__, template_folder=template_dir)
+
+@app.route("/")
+def index():
+    return render_template("index.html")
+
+@app.route("/api/asterocr/status")
+def api_status():
+    return jsonify(global_trainer.get_state())
+
+session_start_chunk = None
+
+@app.route("/api/asterocr/start", methods=["POST"])
+def api_start():
+    global session_start_chunk
+    session_start_chunk = None # Reset relative chunk mapping for new session
+    cfg = request.json or {}
+    success, msg = global_trainer.start(cfg)
+    return jsonify({"status": "ok" if success else "error", "message": msg})
+
+@app.route("/api/asterocr/stop", methods=["POST"])
+def api_stop():
+    global_trainer.stop()
+    return jsonify({"status": "ok", "message": "Stop requested."})
+
+@app.route("/api/asterocr/reset", methods=["POST"])
+def api_reset():
+    global_trainer.reset_history()
+    return jsonify({"status": "ok", "message": "Đã làm sạch toàn bộ biểu đồ và lịch sử về 0."})
+
+@app.route("/api/asterocr/preview_synthetic", methods=["POST"])
+def api_preview_synthetic():
+    req_data = request.json or {}
+    text = req_data.get("text", "CỘNG HÒA XÃ HỘI CHỦ NGHĨA VIỆT NAM")
+    images_b64 = []
+    
+    for style in ["clean", "scan", "aged"]:
+        try:
+            im = default_generator.render_text(text, target_height=32)
+            buf = BytesIO()
+            im.save(buf, format="PNG")
+            images_b64.append("data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("utf-8"))
+        except Exception as e:
+            print("Preview synth error:", e)
+            
+    return jsonify({"status": "ok", "images": images_b64})
+
+# =========================================================================
+# CLOUD GPU STREAMING PIPELINE ENDPOINTS
+# =========================================================================
+from asterocr.kaggle_bridge.chunker import default_chunker
+from asterocr.kaggle_bridge.tunnel import default_tunnel
+import json
+
+@app.route("/api/tunnel/info", methods=["GET"])
+def api_tunnel_info():
+    cfg_path = os.path.join(os.path.dirname(__file__), "tunnel_config.json")
+    custom_u = ""
+    if os.path.exists(cfg_path):
+        try:
+            with open(cfg_path, "r", encoding="utf-8") as f:
+                custom_u = json.load(f).get("custom_url", "")
+                if "trycloudflare.com" in custom_u:
+                    custom_u = ""
+        except Exception:
+            pass
+    active_u = default_tunnel.get_url() or ""
+    return jsonify({
+        "active_url": active_u,
+        "custom_url": custom_u,
+        "is_alive": default_tunnel.is_alive()
+    })
+
+@app.route("/api/tunnel/config", methods=["POST"])
+def api_tunnel_config():
+    data = request.json or {}
+    custom_u = data.get("custom_url", "").strip()
+    if "trycloudflare.com" in custom_u:
+        return jsonify({
+            "status": "error",
+            "message": "URL trycloudflare.com là URL tạm thời ngẫu nhiên của Cloudflare, không thể lưu làm tên miền cố định. Vui lòng để trống ô này để hệ thống tự động sinh và kết nối tự động (Zero-Touch)!"
+        })
+    cfg_path = os.path.join(os.path.dirname(__file__), "tunnel_config.json")
+    try:
+        existing_cfg = {}
+        if os.path.exists(cfg_path):
+            try:
+                with open(cfg_path, "r", encoding="utf-8") as f:
+                    existing_cfg = json.load(f)
+            except Exception:
+                pass
+        existing_cfg["custom_url"] = custom_u
+        if "tunnel_token" in data and data["tunnel_token"]:
+            existing_cfg["tunnel_token"] = data["tunnel_token"]
+        with open(cfg_path, "w", encoding="utf-8") as f:
+            json.dump(existing_cfg, f, indent=2)
+        default_tunnel.tunnel_url = custom_u if custom_u else None
+        return jsonify({"status": "ok", "custom_url": custom_u, "message": "Đã lưu cấu hình Tunnel." if custom_u else "Đã xóa tên miền tùy chỉnh. Hệ thống chuyển sang chế độ Tự Động (Zero-Touch)."})
+    except Exception as e:
+        return jsonify({"status": "error", "message": str(e)})
+
+@app.route("/api/tunnel/start", methods=["POST"])
+def api_tunnel_start():
+    try:
+        url = default_tunnel.start(timeout=35)
+        return jsonify({"status": "ok", "url": url, "message": "Cloudflare Tunnel đã kích hoạt thành công."})
+    except Exception as e:
+        return jsonify({"status": "error", "message": str(e)})
+
+@app.route("/api/tunnel/stop", methods=["POST"])
+def api_tunnel_stop():
+    try:
+        default_tunnel.stop()
+        return jsonify({"status": "ok", "message": "Đã dừng Cloudflare Tunnel."})
+    except Exception as e:
+        return jsonify({"status": "error", "message": str(e)})
+
+@app.route("/api/colab/worker.py")
+def api_colab_worker():
+    worker_p = os.path.join(os.path.dirname(__file__), "kaggle_bridge", "kaggle_worker.py")
+    if os.path.exists(worker_p):
+        with open(worker_p, "r", encoding="utf-8") as f:
+            code = f.read()
+            
+        import re, json
+        # Đọc cấu hình tạm do _train_colab_worker xuất ra
+        config_path = os.path.join(os.path.dirname(__file__), "..", "scratch", "colab_run_config.json")
+        if os.path.exists(config_path):
+            with open(config_path, "r", encoding="utf-8") as f:
+                cfg = json.load(f)
+            
+            is_tpu = (cfg.get("device") == "colab_tpu")
+            batch_sz = cfg.get("batch_size", 256)
+            lr = cfg.get("lr", 0.0005)
+            epochs = cfg.get("epochs", 5)
+            
+            # Lấy tunnel url
+            tunnel_url = "https://asterai.sglocalhub.com"
+            from asterocr.kaggle_bridge.tunnel import default_tunnel
+            if default_tunnel.tunnel_url:
+                tunnel_url = default_tunnel.tunnel_url
+
+            code = re.sub(r'IS_TPU_DEVICE\s*=\s*(True|False)', f'IS_TPU_DEVICE = {is_tpu}', code)
+            code = re.sub(r'TUNNEL_URL\s*=\s*["\'][^"\']+["\']', f'TUNNEL_URL = "{tunnel_url}"', code)
+            code = re.sub(r'BATCH_SIZE\s*=\s*\d+', f'BATCH_SIZE = {batch_sz}', code)
+            code = re.sub(r'LEARNING_RATE\s*=\s*[\d\.]+', f'LEARNING_RATE = {lr}', code)
+            code = re.sub(r'MAX_EPOCHS\s*=\s*\d+', f'MAX_EPOCHS = {epochs}', code)
+
+        return Response(code, mimetype="text/plain; charset=utf-8")
+    return ("Worker not found", 404)
+
+@app.route("/api/kaggle/stream/ping")
+def api_kaggle_ping():
+    return jsonify({"status": "ok", "service": "AsterOCR Local Streamer", "time": time.time()})
+
+@app.route("/api/kaggle/stream/heartbeat", methods=["POST"])
+def api_kaggle_heartbeat():
+    data = request.json or {}
+    global_trainer.record_cloud_heartbeat(data)
+    return jsonify({"status": "ok", "time": time.time()})
+
+@app.route("/api/kaggle/stream/control")
+def api_kaggle_control():
+    return jsonify({"should_stop": bool(global_trainer.stop_requested)})
+
+@app.route("/api/kaggle/stream/init_weights")
+def api_kaggle_init_weights():
+    global_trainer.record_cloud_heartbeat({"event": "init_weights"})
+    best_p = os.path.join(global_trainer.output_dir, "asternet_best.pth")
+    last_p = os.path.join(global_trainer.output_dir, "asternet_last.pth")
+    
+    # Intelligently select latest / highest step checkpoint
+    chosen = None
+    if os.path.exists(last_p) and os.path.exists(best_p):
+        try:
+            c_last = torch.load(last_p, map_location="cpu", weights_only=False)
+            c_best = torch.load(best_p, map_location="cpu", weights_only=False)
+            s_last = c_last.get("step", 0) or 0
+            s_best = c_best.get("step", 0) or 0
+            chosen = last_p if s_last >= s_best else best_p
+        except Exception:
+            chosen = last_p
+    elif os.path.exists(last_p):
+        chosen = last_p
+    elif os.path.exists(best_p):
+        chosen = best_p
+
+    if chosen and os.path.exists(chosen):
+        return send_file(chosen, mimetype="application/octet-stream")
+    return ("", 204)
+
+@app.route("/api/kaggle/stream/get_chunk")
+def api_kaggle_get_chunk():
+    global session_start_chunk
+    global_trainer.record_cloud_heartbeat({"event": "get_chunk"})
+    chunk_idx = int(request.args.get("chunk_idx", 0))
+    epochs = int(global_trainer.state.get("total_epochs", 5))
+    total_chunks = default_chunker.get_total_chunks()
+
+    # Dynamic session chunk calibration:
+    # When worker connects, it may restore an old checkpoint's chunk_idx (e.g. 10,078).
+    # We calibrate the relative chunk to start from 0 for the currently active dataset!
+    if session_start_chunk is None or chunk_idx < session_start_chunk:
+        session_start_chunk = chunk_idx
+
+    relative_chunk = max(0, chunk_idx - session_start_chunk)
+
+    buf = default_chunker.get_chunk_bytes(relative_chunk, max_epochs=epochs)
+    if buf is None:
+        return ("", 204)
+        
+    cur_epoch = min(epochs, (relative_chunk // total_chunks) + 1)
+    
+    resp = send_file(buf, mimetype="application/octet-stream")
+    resp.headers["X-Current-Epoch"] = str(cur_epoch)
+    resp.headers["X-Total-Epochs"] = str(epochs)
+    resp.headers["X-Total-Chunks"] = str(total_chunks)
+    return resp
+
+@app.route("/api/kaggle/stream/telemetry", methods=["POST"])
+def api_kaggle_telemetry():
+    data = request.json or {}
+    global_trainer.update_remote_telemetry(data)
+    return jsonify({"status": "ok"})
+
+@app.route("/api/kaggle/stream/save_checkpoint", methods=["POST"])
+def api_kaggle_save_checkpoint():
+    file = request.files.get("checkpoint")
+    if file:
+        os.makedirs(global_trainer.output_dir, exist_ok=True)
+        last_p = os.path.join(global_trainer.output_dir, "asternet_last.pth")
+        best_p = os.path.join(global_trainer.output_dir, "asternet_best.pth")
+        temp_p = os.path.join(global_trainer.output_dir, "temp_incoming.pth")
+        
+        file.save(temp_p)
+        try:
+            test_ckpt = torch.load(temp_p, map_location="cpu", weights_only=False)
+            sd = test_ckpt.get("state_dict", test_ckpt.get("model", test_ckpt))
+            has_nan = any(torch.isnan(v).any() or torch.isinf(v).any() for v in sd.values() if isinstance(v, torch.Tensor))
+            if has_nan:
+                print(f"[!] CẢNH BÁO: Checkpoint từ Cloud chứa giá trị NaN/Inf! Đã từ chối lưu để bảo vệ mô hình an toàn!")
+                if os.path.exists(temp_p):
+                    os.remove(temp_p)
+                return jsonify({"status": "rejected", "message": "Checkpoint contains NaN"}), 400
+        except Exception as e:
+            print(f"[!] Lỗi kiểm tra tính toàn vẹn của checkpoint: {e}")
+            if os.path.exists(temp_p):
+                os.remove(temp_p)
+            return jsonify({"status": "error"}), 400
+
+        # Trọng số sạch: chuyển vào last_p
+        import shutil
+        shutil.move(temp_p, last_p)
+
+        step = int(request.form.get("step", 0))
+        chunk_idx = int(request.form.get("chunk_idx", 0))
+        try:
+            loss = float(request.form.get("loss", 999.0))
+            best = float(request.form.get("best_loss", 999.0))
+            with global_trainer.lock:
+                cur_best = float(global_trainer.ledger.get("best_loss", 999.0))
+                # Protect against corrupted/dummy zero loss (< 0.01 is invalid in CTC)
+                if cur_best < 0.01:
+                    cur_best = 999.0
+                    global_trainer.ledger["best_loss"] = 999.0
+                    global_trainer.state["best_loss"] = 999.0
+
+                # Only update best model file if loss strictly improved, valid (> 0.01) and reasonable (< 2.0)
+                # Chặn không cho giá trị 0.1500 cố định giả tạo ghi đè kỷ lục nếu không thật
+                if (0.01 < loss < cur_best and loss < 2.0) or not os.path.exists(best_p):
+                    shutil.copyfile(last_p, best_p)
+                    global_trainer.ledger["best_loss"] = round(loss, 4)
+                    global_trainer.state["best_loss"] = round(loss, 4)
+                    print(f"[+] Kỷ lục Loss mới đạt được: {loss:.4f}! Đã lưu vào {best_p}")
+                
+                # Maintain strictly monotonic step and chunk accounting
+                if step > global_trainer.ledger.get("last_step", 0):
+                    global_trainer.ledger["last_step"] = step
+                if chunk_idx > global_trainer.ledger.get("last_chunk_idx", 0):
+                    global_trainer.ledger["last_chunk_idx"] = chunk_idx
+                global_trainer.ledger["last_updated"] = time.strftime("%Y-%m-%dT%H:%M:%S")
+                
+                # Sync total_samples from cumulative state if available
+                ct_samp = global_trainer.state.get("cumulative_training", {}).get("total_samples", 0)
+                if ct_samp > global_trainer.ledger.get("total_samples", 0):
+                    global_trainer.ledger["total_samples"] = ct_samp
+                
+                global_trainer._save_ledger(global_trainer.ledger)
+        except Exception as e:
+            print(f"[!] Lỗi khi đồng bộ checkpoint metadata vào ledger: {e}")
+        print(f"[+] Checkpoint từ Cloud GPU (Step {step}, Chunk {chunk_idx}) đã được đồng bộ an toàn về: {last_p}")
+    return jsonify({"status": "ok"})
+
+@app.route("/api/kaggle/status", methods=["GET"])
+def api_kaggle_status():
+    import subprocess
+    slug = request.args.get("slug", "creyt2012/asterocr-cloud-gpu-worker-v2")
+    env = os.environ.copy()
+    tok_p = os.path.expanduser("~/.kaggle/access_token")
+    if os.path.exists(tok_p):
+        try:
+            with open(tok_p, "r", encoding="utf-8") as f:
+                env["KAGGLE_API_TOKEN"] = f.read().strip()
+        except Exception:
+            pass
+    try:
+        res = subprocess.run([sys.executable, "-m", "kaggle.cli", "kernels", "status", slug], env=env, capture_output=True, text=True, timeout=15)
+        out = (res.stdout + " " + res.stderr).strip()
+        return jsonify({"status": "ok", "output": out})
+    except Exception as e:
+        return jsonify({"status": "error", "message": str(e)})
+
+@app.route("/api/kaggle/launch_finetune", methods=["POST"])
+def api_kaggle_launch_finetune():
+    data = request.json or {}
+    epochs = int(data.get("epochs", 5))
+    batch_size = int(data.get("batch_size", 96))
+    lr = float(data.get("lr", 0.00005))
+    dataset_size = int(data.get("dataset_size", 1558564))
+    
+    cfg = {
+        "dataset_mode": "pure_real_scans",
+        "dataset_size": dataset_size,
+        "epochs": epochs,
+        "batch_size": batch_size,
+        "lr": lr,
+        "device": "kaggle_t4"
+    }
+    success, msg = global_trainer.start(cfg)
+    return jsonify({"status": "ok" if success else "error", "message": msg})
+
+def open_browser(port):
+    import time
+    time.sleep(1.2)
+    webbrowser.open(f"http://127.0.0.1:{port}")
+
+if __name__ == "__main__":
+    PORT = 7860
+    print(f"[*] Khởi động AsterOCR Studio độc lập tại: http://127.0.0.1:{PORT}")
+    threading.Thread(target=open_browser, args=(PORT,), daemon=True).start()
+    app.run(host="127.0.0.1", port=PORT, debug=False, threaded=True)