Spaces:
Running
Running
File size: 2,509 Bytes
322be7d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import os
import random
import gradio as gr
from PIL import Image
import torch
# Import your inference module
import inference as inf
from src.generateCaptcha import generate_captcha
from src.config import cfg # sizes, charset, dirs
# Device and one-time model load
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL = inf.load_model("checkpoints/best_model.pth").to(DEVICE).eval()
# Ensure results dir exists
os.makedirs(cfg.RESULT_DIR, exist_ok=True)
def random_text():
L = random.randint(cfg.CAPTCHA_LEN_LOWER_LIMIT, cfg.CAPTCHA_LEN_UPPER_LIMIT)
return "".join(random.choices(cfg.chars, k=L))
def ui_generate():
text = random_text()
filename = f"{text}_{random.randint(1000,9999)}.png"
# Use generateCaptcha.py directly
img = generate_captcha(text, width=cfg.W_max, height=cfg.H)
# Save to results directory
filepath = os.path.join(cfg.RESULT_DIR, filename)
img.save(filepath)
return img, text, filepath
def ui_solve(img: Image.Image, path_hint: str):
# Prefer uploaded image
if img is not None:
tmp_path = os.path.join(cfg.RESULT_DIR, f"upload_{random.randint(1000,9999)}.png")
img.save(tmp_path)
tensor = inf.preprocess_image(tmp_path, (cfg.W_max, cfg.H))
pred = inf.predict_captcha(MODEL, tensor, DEVICE)
return pred
# Otherwise, solve the last generated image
if path_hint and os.path.exists(path_hint):
tensor = inf.preprocess_image(path_hint, (cfg.W_max, cfg.H))
pred = inf.predict_captcha(MODEL, tensor, DEVICE)
return pred
return "No image provided. Generate or upload first."
with gr.Blocks(title="CAPTCHA OCR (checkpoint)") as demo:
gr.Markdown("## CAPTCHA OCR demo")
with gr.Row():
gen_btn = gr.Button("Generate CAPTCHA", variant="primary")
gt_out = gr.Textbox(label="Ground Truth", interactive=False)
with gr.Row():
img_out = gr.Image(label="Generated CAPTCHA", type="pil")
path_box = gr.Textbox(label="Internal Path", interactive=False, visible=False)
gen_btn.click(fn=ui_generate, outputs=[img_out, gt_out, path_box])
gr.Markdown("### Solve")
with gr.Row():
img_in = gr.Image(label="Upload CAPTCHA (optional)", type="pil")
solve_btn = gr.Button("Solve")
pred_out = gr.Textbox(label="Prediction", interactive=False)
solve_btn.click(fn=ui_solve, inputs=[img_in, path_box], outputs=[pred_out])
if __name__ == "__main__":
demo.launch() |