Spaces:
Sleeping
Sleeping
| import os | |
| import io | |
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| from peft import PeftModel | |
| from transformers import TrOCRProcessor, VisionEncoderDecoderModel | |
| from cnn_model import CharacterClassifier | |
| from preprocessing import preprocess_for_ocr | |
| # --- ROBUST GLOBAL PATCH FOR GRADIO 4.x --- | |
| import gradio_client.utils | |
| def robust_get_type(schema): | |
| if isinstance(schema, bool): return "Any" | |
| if not isinstance(schema, dict): return "Any" | |
| if "const" in schema: return "Any" | |
| return original_get_type(schema) | |
| if hasattr(gradio_client.utils, "get_type"): | |
| original_get_type = gradio_client.utils.get_type | |
| gradio_client.utils.get_type = robust_get_type | |
| # ------------------------------------------ | |
| # --- CONFIGURATION --- | |
| BASE_MODEL_ID = "paudelanil/trocr-devanagari-2" | |
| ADAPTER_ID = "manishw10/devgen-trocr-devanagari-lora" | |
| CNN_MODEL_PATH = "devanagari-cnn-classifier.pt" | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # --- ENGINE INITIALIZATION --- | |
| print("System: Initializing Full Combined Suite...") | |
| processor = TrOCRProcessor.from_pretrained(BASE_MODEL_ID) | |
| base_model = VisionEncoderDecoderModel.from_pretrained(BASE_MODEL_ID) | |
| base_model.config.decoder_start_token_id = processor.tokenizer.cls_token_id | |
| base_model.config.pad_token_id = processor.tokenizer.pad_token_id | |
| base_model.config.eos_token_id = processor.tokenizer.sep_token_id | |
| base_model.config.vocab_size = base_model.config.decoder.vocab_size | |
| peft_model = PeftModel.from_pretrained(base_model, ADAPTER_ID) | |
| try: | |
| model = peft_model.merge_and_unload() | |
| except Exception: | |
| model = peft_model | |
| model.to(device); model.eval() | |
| cnn_engine = CharacterClassifier(model_path=CNN_MODEL_PATH, device=device) | |
| # --- ORIGINAL ROUTING LOGIC --- | |
| def _flood_fill(binary, visited, start_y, start_x, h, w): | |
| stack = [(start_y, start_x)] | |
| size = 0 | |
| while stack: | |
| y, x = stack.pop() | |
| if y<0 or y>=h or x<0 or x>=w or visited[y,x] or not binary[y,x]: continue | |
| visited[y,x] = True; size += 1 | |
| stack.extend([(y+1,x),(y-1,x),(y,x+1),(y,x-1)]) | |
| return size | |
| def count_blobs(binary): | |
| h, w = binary.shape; visited = np.zeros_like(binary, dtype=bool); count = 0 | |
| for y in range(h): | |
| for x in range(w): | |
| if binary[y,x] and not visited[y,x]: | |
| size = _flood_fill(binary, visited, y, x, h, w) | |
| if size >= max(binary.size * 0.001, 10): count += 1 | |
| return count | |
| def original_classify_input(image): | |
| gray = image.convert("L"); arr = np.array(gray) | |
| threshold = min(arr.mean() * 0.75, 200) | |
| binary = (arr < threshold).astype(np.uint8) | |
| rows, cols = np.any(binary, axis=1), np.any(binary, axis=0) | |
| if not rows.any() or not cols.any(): return "character", 1.0, 1 | |
| y0, x0 = np.where(rows)[0][0], np.where(cols)[0][0] | |
| y1, x1 = np.where(rows)[0][-1], np.where(cols)[0][-1] | |
| w, h = x1-x0+1, y1-y0+1 | |
| ar, bc = w/h, count_blobs(binary) | |
| is_char = True | |
| if ar > 2.5: is_char = False | |
| elif ar > 1.8 and bc >= 3: is_char = False | |
| elif bc >= 4: is_char = False | |
| elif ar < 1.3 and bc <= 2: is_char = True | |
| elif bc == 1 and ar < 1.5: is_char = True | |
| elif ar < 1.75 and bc <= 2: is_char = True | |
| elif ar > 1.6: is_char = False | |
| return ("character" if is_char else "word"), ar, bc | |
| def get_confidence_html(confidence): | |
| color = "#10b981" if confidence > 0.9 else "#f59e0b" if confidence > 0.7 else "#ef4444" | |
| return f"""<div style="display: flex; flex-direction: column; align-items: center; background: rgba(0,0,0,0.2); border-radius: 20px; padding: 15px;"> | |
| <svg width="100" height="100" viewBox="0 0 100 100"> | |
| <circle cx="50" cy="50" r="45" fill="none" stroke="rgba(255,255,255,0.1)" stroke-width="8" /> | |
| <circle cx="50" cy="50" r="45" fill="none" stroke="{color}" stroke-width="8" stroke-dasharray="282.7" stroke-dashoffset="{282.7 * (1 - confidence)}" stroke-linecap="round" style="transition: stroke-dashoffset 1s;" /> | |
| <text x="50" y="55" font-family="Outfit" font-size="20" font-weight="bold" fill="{color}" text-anchor="middle">{int(confidence * 100)}%</text> | |
| </svg> | |
| </div>""" | |
| # --- PREDICT --- | |
| def predict(image, manual_mode): | |
| if image is None: return None, None, "Upload image.", "", "" | |
| buf = io.BytesIO(); image.save(buf, format="PNG") | |
| pre_pil = preprocess_for_ocr(buf.getvalue()) | |
| if manual_mode == "Automatic": | |
| mode, ar, bc = original_classify_input(pre_pil) | |
| status = f"**System**: {mode.upper()} detected (AR: {ar:.2f}, Blobs: {bc})" | |
| else: | |
| mode = manual_mode.lower(); status = f"**Manual Mode**: {mode.upper()}" | |
| try: | |
| if mode == "character" and cnn_engine.available: | |
| res = cnn_engine.predict(pre_pil) | |
| return pre_pil, res["text"], status, "CNN Classifier", get_confidence_html(res["confidence"]) | |
| else: | |
| pixel_values = processor(pre_pil, return_tensors="pt").pixel_values.to(device) | |
| with torch.no_grad(): | |
| out = model.generate(pixel_values, num_beams=4, max_length=128, early_stopping=True, return_dict_in_generate=True, output_scores=True, decoder_start_token_id=model.config.decoder_start_token_id) | |
| scores = torch.exp(model.compute_transition_scores(out.sequences, out.scores, normalize_logits=True)[0]) | |
| txt = processor.batch_decode(out.sequences, skip_special_tokens=True)[0] | |
| return pre_pil, txt, status, "TrOCR + LoRA", get_confidence_html(float(scores.mean().item())) | |
| except Exception as e: | |
| return pre_pil, f"Error: {str(e)}", "Failed", "None", "" | |
| # --- PREMIUM UI --- | |
| CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Outfit:wght@400;600&family=Inter:wght@400;500&display=swap'); | |
| .gradio-container { background: #0f172a !important; color: white !important; font-family: 'Inter', sans-serif !important; } | |
| .premium-card { background: rgba(30, 41, 59, 0.7) !important; backdrop-filter: blur(12px); border: 1px solid rgba(255,255,255,0.1); border-radius: 24px; padding: 2rem; box-shadow: 0 25px 50px -12px rgba(0,0,0,0.5); } | |
| .result-box { font-size: 3rem !important; font-weight: 600; text-align: center; color: #818cf8; background: transparent !important; border: none !important; } | |
| """ | |
| with gr.Blocks(css=CSS, theme=gr.themes.Default()) as demo: | |
| with gr.Column(elem_classes="premium-card"): | |
| gr.Markdown("# 🕉️ DevGen OCR") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| img_in = gr.Image(type="pil", label="Input Handwriting") | |
| mode_ctrl = gr.Radio(["Automatic", "Word", "Character"], value="Automatic", label="Logic Mode") | |
| sub_btn = gr.Button("Recognize", variant="primary") | |
| with gr.Column(scale=1): | |
| conf_html = gr.HTML() | |
| text_out = gr.Textbox(label="Result", elem_classes="result-box", interactive=False, show_label=False) | |
| status_md = gr.Markdown("Engine ready.") | |
| engine_txt = gr.Textbox(label="Active Model", interactive=False) | |
| with gr.Column(): | |
| gr.Markdown("### 🛠️ Visual Debug: What the Model Sees") | |
| img_proc = gr.Image(type="pil", label="Preprocessed Input", interactive=False, show_label=False) | |
| sub_btn.click(predict, [img_in, mode_ctrl], [img_proc, text_out, status_md, engine_txt, conf_html]) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |