Spaces:
Running
Running
| 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 calculate_accuracy(prediction, target): | |
| """Calculate character-by-character accuracy.""" | |
| if not prediction or not target: | |
| return "0%" | |
| correct_chars = 0 | |
| min_len = min(len(prediction), len(target)) | |
| for i in range(min_len): | |
| if prediction[i] == target[i]: | |
| correct_chars += 1 | |
| if min_len == 0: | |
| return "0%" | |
| accuracy = (correct_chars / min_len) * 100 | |
| return f"{accuracy:.1f}%" | |
| 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) | |
| # Enable and turn Solve green now that an image exists | |
| solve_btn_state = gr.update(interactive=True, variant="primary") | |
| return img, text, filepath, solve_btn_state | |
| def ui_solve(path_hint: str, ground_truth: str): | |
| 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) | |
| # Calculate accuracy | |
| accuracy = calculate_accuracy(pred, ground_truth) | |
| return accuracy, pred | |
| return "0%", "No image generated yet. Click Generate CAPTCHA first." | |
| with gr.Blocks(title="CAPTCHA OCR (checkpoint)") as demo: | |
| gr.Markdown("## CAPTCHA OCR ") | |
| with gr.Row(): | |
| # Left column: Generate button + Solve button stacked vertically | |
| with gr.Column(scale=1): | |
| gen_btn = gr.Button("Generate CAPTCHA", variant="primary") | |
| solve_btn = gr.Button("Solve", interactive=False, variant="secondary") | |
| # Right column: Ground Truth | |
| gt_out = gr.Textbox(label="Ground Truth", interactive=False, text_align="center") | |
| with gr.Row(): | |
| img_out = gr.Image(label="Generated CAPTCHA", type="pil") | |
| path_box = gr.Textbox(label="Internal Path", interactive=False, visible=False) | |
| # Prediction row split into two columns | |
| with gr.Row(): | |
| accuracy_out = gr.Textbox(label="Character Accuracy", interactive=False, text_align="center") | |
| pred_out = gr.Textbox(label="Prediction", interactive=False, text_align="center") | |
| # Generate: outputs image, ground truth, path, and enables Solve (green) | |
| gen_btn.click( | |
| fn=ui_generate, | |
| outputs=[img_out, gt_out, path_box, solve_btn], | |
| ) | |
| # Solve: only uses the internal path (no upload option anymore) | |
| solve_btn.click( | |
| fn=ui_solve, | |
| inputs=[path_box, gt_out], | |
| outputs=[accuracy_out, pred_out], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |