app.py

# import gradio as gr
# from ultralytics import YOLO
# from PIL import Image, ImageDraw, ImageFont
# import torch
# import logging
# import os
# from datetime import datetime

# # # ── Quiet startup ───────────────────────────────────────────────────────
# # os.environ['HF_HUB_DISABLE_PROGRESS_BARS'] = '1'
# # logging.getLogger('ultralytics').setLevel(logging.WARNING)

# # logging.basicConfig(
# #     level=logging.INFO,
# #     format='%(asctime)s | %(level)-5s | %(message)s'
# # )
# # logger = logging.getLogger(__name__)

# os.environ['HF_HUB_DISABLE_PROGRESS_BARS'] = '1'
# logging.getLogger('ultralytics').setLevel(logging.WARNING)

# # FIXED logging format: use levelname, not level
# logging.basicConfig(
#     level=logging.INFO,
#     format='%(asctime)s | %(levelname)-5s | %(message)s',   # ← changed level → levelname
#     datefmt='%Y-%m-%d %H:%M:%S'
# )
# logger = logging.getLogger(__name__)

# logger.info("Initializing region detector...")

# device = "cuda" if torch.cuda.is_available() else "cpu"
# logger.info(f"Device: {device}")

# # ── Load YOLO ───────────────────────────────────────────────────────────
# try:
#     region_pt = 'regions.pt'
#     if not os.path.exists(region_pt):
#         for f in os.listdir('.'):
#             name = f.lower()
#             if name.endswith('.pt') and 'region' in name:
#                 region_pt = f
#                 break

#     if not os.path.exists(region_pt):
#         raise FileNotFoundError("No regions.pt (or similar *.pt) found in current directory")

#     logger.info(f"Loading model: {region_pt}")
#     model = YOLO(region_pt)
#     logger.info("Region detector loaded")

# except Exception as e:
#     logger.error(f"Model loading failed → {e}", exc_info=True)
#     raise


# def visualize_regions(
#     image,
#     conf_thresh: float = 0.25,
#     min_size: int = 60,
#     padding: int = 0,
#     show_labels: bool = True,
#     save_debug_crops: bool = False,
#     imgsz: int = 1024,
# ):
#     start = datetime.now().strftime("%H:%M:%S")
#     logs = [f"[{start}] Processing started"]

#     if image is None:
#         logs.append("No image uploaded")
#         return None, "\n".join(logs)

#     # Load & convert
#     if isinstance(image, str):
#         img = Image.open(image).convert("RGB")
#     else:
#         img = image.convert("RGB")

#     w, h = img.size
#     logs.append(f"Image size: {w} × {h}")

#     debug_img = img.copy()
#     draw = ImageDraw.Draw(debug_img)

#     try:
#         # Font for drawing labels (fallback to default)
#         try:
#             font = ImageFont.truetype("arial.ttf", 18)
#         except:
#             font = ImageFont.load_default()

#         # ── Run detection ───────────────────────────────────────────────
#         results = model(
#             img,
#             conf=conf_thresh,
#             imgsz=imgsz,
#             verbose=False
#         )[0]

#         boxes = results.boxes
#         logs.append(f"Detected {len(boxes)} region candidate(s)")

#         kept = 0

#         # Sort top → bottom
#         if len(boxes) > 0:
#             ys = boxes.xyxy[:, 1].cpu().numpy()
#             order = ys.argsort()

#             for idx in order:
#                 box = boxes[idx]
#                 conf = float(box.conf)
#                 if conf < conf_thresh:
#                     continue

#                 x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
#                 bw, bh = x2 - x1, y2 - y1

#                 if bw < min_size or bh < min_size:
#                     continue

#                 # Optional padding (mostly for crop saving)
#                 px1 = max(0, x1 - padding)
#                 py1 = max(0, y1 - padding)
#                 px2 = min(w, x2 + padding)
#                 py2 = min(h, y2 + padding)

#                 # Draw box
#                 draw.rectangle((x1, y1, x2, y2), outline="lime", width=3)

#                 if show_labels:
#                     label = f"conf {conf:.2f}  {bw}×{bh}"
#                     tw, th = draw.textbbox((0,0), label, font=font)[2:]
#                     draw.rectangle(
#                         (x1, y1 - th - 4, x1 + tw + 8, y1),
#                         fill=(0, 180, 0, 160)
#                     )
#                     draw.text((x1 + 4, y1 - th - 2), label, fill="white", font=font)

#                 kept += 1

#                 # Optional: save individual crops
#                 if save_debug_crops:
#                     os.makedirs("debug_regions", exist_ok=True)
#                     crop = img.crop((px1, py1, px2, py2))
#                     fname = f"debug_regions/r{kept:02d}_conf{conf:.2f}_{bw}x{bh}.png"
#                     crop.save(fname)
#                     logs.append(f"Saved crop → {fname}")

#         if kept == 0:
#             msg = f"No regions kept after filters (conf ≥ {conf_thresh}, size ≥ {min_size}px)"
#             logs.append(msg)
#         else:
#             logs.append(f"Visualized {kept} region(s)")

#         logs.append("Finished.")

#         return debug_img, "\n".join(logs)

#     except Exception as e:
#         logs.append(f"Error during inference: {str(e)}")
#         logger.exception("Inference failed")
#         return debug_img, "\n".join(logs)



# # ── Gradio Interface ────────────────────────────────────────────────────
# demo = gr.Interface(
#     fn=visualize_regions,
#     inputs=[
#         gr.Image(type="pil", label="Upload image (handwritten document)"),
#         gr.Slider(0.10, 0.60, step=0.02, value=0.25, label="Confidence threshold"),
#         gr.Slider(30,  300,  step=10,  value=60,   label="Minimum region width/height (px)"),
#         gr.Slider(0,   40,   step=4,   value=0,    label="Padding around box (for crops only)"),
#         gr.Checkbox(label="Draw confidence + size labels on boxes", value=True),
#         gr.Checkbox(label="Save individual region crops to debug_regions/", value=False),
#         gr.Slider(640, 1280, step=64, value=1024, label="Inference image size (imgsz)"),
#     ],
#     outputs=[
#         gr.Image(label="Detected text regions (green boxes)"),
#         gr.Textbox(label="Log / debug info", lines=14),
#     ],
#     title="Region Detector Debug View",
#     description=(
#         "Only shows what the region YOLO model sees.\n\n"
#         "• Green boxes = detected text regions\n"
#         "• Tune confidence and min size until boxes look reasonable\n"
#         "• Use logs to see exact confidences and sizes\n"
#         "• Save crops if you want to manually check what is being detected"
#     ),
#     # theme=gr.themes.Soft(),          # ← comment out or remove (moved to launch)
#     # allow_flagging="never",          # ← remove this line completely
# )

# if __name__ == "__main__":
#     logger.info("Launching debug interface...")
#     demo.launch()













# import gradio as gr
# from ultralytics import YOLO
# from transformers import TrOCRProcessor, VisionEncoderDecoderModel
# from PIL import Image, ImageDraw
# import torch
# import logging
# import os
# import warnings
# import time
# from datetime import datetime

# # ── Suppress noisy logs ──────────────────────────────────────────────────────
# os.environ['HF_HUB_DISABLE_PROGRESS_BARS'] = '1'
# warnings.filterwarnings('ignore')
# logging.getLogger('transformers').setLevel(logging.ERROR)
# logging.getLogger('ultralytics').setLevel(logging.WARNING)

# # Clean logging
# logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(levelname)-5s | %(message)s')
# logger = logging.getLogger(__name__)

# logger.info("Initializing models...")
# device = "cuda" if torch.cuda.is_available() else "cpu"
# logger.info(f"Device: {device}")

# def load_with_retry(cls, name, token=None, retries=4, delay=6):
#     for attempt in range(1, retries + 1):
#         try:
#             logger.info(f"Loading {name} (attempt {attempt}/{retries})")
#             if "Processor" in str(cls):
#                 return cls.from_pretrained(name, token=token)
#             return cls.from_pretrained(name, token=token).to(device)
#         except Exception as e:
#             logger.warning(f"Load failed: {e}")
#             if attempt < retries:
#                 time.sleep(delay)
#     raise RuntimeError(f"Failed to load {name} after {retries} attempts")


# try:
#     # Locate local YOLO line detection weights
#     line_pt = 'lines.pt'

#     if not os.path.exists(line_pt):
#         for f in os.listdir('.'):
#             name = f.lower()
#             if 'line' in name and name.endswith('.pt'):
#                 line_pt = f
#                 break

#     if not os.path.exists(line_pt):
#         raise FileNotFoundError("Could not find lines.pt (or similar *.pt file containing 'line' in name)")

#     logger.info("Loading YOLO line model...")
#     line_model = YOLO(line_pt)
#     logger.info("YOLO line model loaded")

#     hf_token = os.getenv("HF_TOKEN")
#     processor = load_with_retry(TrOCRProcessor, "microsoft/trocr-base-handwritten", hf_token)
#     trocr     = load_with_retry(VisionEncoderDecoderModel, "microsoft/trocr-base-handwritten", hf_token)
#     logger.info("TrOCR loaded → ready")

# except Exception as e:
#     logger.error(f"Model loading failed: {e}", exc_info=True)
#     raise


# def run_ocr(crop: Image.Image) -> str:
#     if crop.width < 20 or crop.height < 12:
#         return ""
#     pixels = processor(images=crop, return_tensors="pt").pixel_values.to(device)
#     ids = trocr.generate(pixels, max_new_tokens=128)
#     return processor.batch_decode(ids, skip_special_tokens=True)[0].strip()


# def process_document(
#     image,
#     enable_debug_crops: bool = False,
#     line_imgsz: int = 768,
#     conf_thresh: float = 0.25,
# ):
#     start_ts = datetime.now().strftime("%H:%M:%S")
#     logs = []

#     def log(msg: str, level: str = "INFO"):
#         line = f"[{start_ts}] {level:5} {msg}"
#         logs.append(line)
#         if level == "ERROR":
#             logger.error(msg)
#         else:
#             logger.info(msg)

#     log("Start processing")

#     if image is None:
#         log("No image uploaded", "ERROR")
#         return None, "Upload an image", "\n".join(logs)

#     try:
#         # ── Prepare ─────────────────────────────────────────────────────────────
#         if not isinstance(image, Image.Image):
#             img = Image.open(image).convert("RGB")
#         else:
#             img = image.convert("RGB")

#         debug_img = img.copy()
#         draw = ImageDraw.Draw(debug_img)
#         w, h = img.size
#         log(f"Input image: {w} × {h} px")

#         debug_dir = "debug_crops"
#         if enable_debug_crops:
#             os.makedirs(debug_dir, exist_ok=True)
#             log(f"Debug crops will be saved to {debug_dir}/")

#         extracted = []

#         # ── Line detection on full image ────────────────────────────────────────
#         # Adaptive size based on image dimensions
#         max_dim = max(w, h)
#         if max_dim > 2200:
#             used_sz = 1280
#         elif max_dim > 1400:
#             used_sz = 1024
#         elif max_dim < 600:
#             used_sz = 640
#         else:
#             used_sz = line_imgsz

#         log(f"Running line detection (imgsz={used_sz}, conf≥{conf_thresh}) …")

#         res = line_model(img, conf=conf_thresh, imgsz=used_sz, verbose=False)[0]
#         boxes = res.boxes

#         log(f"→ Detected {len(boxes)} line candidate(s)")

#         if len(boxes) == 0:
#             msg = "No text lines detected"
#             log(msg, "WARNING")
#             return debug_img, msg, "\n".join(logs)

#         # Sort top → bottom
#         ys = boxes.xyxy[:, 1].cpu().numpy()   # y_min
#         order = ys.argsort()

#         for j, idx in enumerate(order, 1):
#             conf = float(boxes.conf[idx])
#             x1, y1, x2, y2 = map(round, boxes.xyxy[idx].cpu().tolist())

#             lw, lh = x2 - x1, y2 - y1
#             log(f"  Line {j}/{len(boxes)}  conf={conf:.3f}  {x1},{y1} → {x2},{y2}  ({lw}×{lh})")

#             # Skip very small detections
#             if lw < 60 or lh < 20:
#                 log(f"    → skipped (too small)")
#                 continue

#             draw.rectangle((x1, y1, x2, y2), outline="red", width=3)

#             line_crop = img.crop((x1, y1, x2, y2))

#             if enable_debug_crops:
#                 fname = f"{debug_dir}/line_{j:02d}_conf{conf:.2f}.png"
#                 line_crop.save(fname)

#             text = run_ocr(line_crop)
#             log(f"    OCR → '{text}'")

#             if text.strip():
#                 extracted.append(text)

#         # ── Finalize ────────────────────────────────────────────────────────────
#         if not extracted:
#             msg = "No readable text found after OCR"
#             log(msg, "WARNING")
#             return debug_img, msg, "\n".join(logs)

#         log(f"Success — extracted {len(extracted)} line(s)")
#         if enable_debug_crops:
#             log(f"Debug crops saved to {debug_dir}/")

#         return debug_img, "\n".join(extracted), "\n".join(logs)

#     except Exception as e:
#         log(f"Processing failed: {e}", "ERROR")
#         logger.exception("Traceback:")
#         return debug_img, f"Error: {str(e)}", "\n".join(logs)


# demo = gr.Interface(
#     fn=process_document,
#     inputs=[
#         gr.Image(type="pil", label="Handwritten document"),
#         gr.Checkbox(label="Save debug crops", value=False),
#         gr.Slider(512, 1280, step=64, value=768, label="Line detection size (imgsz)"),
#         gr.Slider(0.15, 0.5, step=0.05, value=0.25, label="Confidence threshold"),
#     ],
#     outputs=[
#         gr.Image(label="Debug (red = detected text lines)"),
#         gr.Textbox(label="Extracted Text", lines=10),
#         gr.Textbox(label="Detailed Logs (copy if alignment is wrong)", lines=16),
#     ],
#     title="Handwritten Line Detection + TrOCR",
#     description=(
#         "Red boxes = text lines detected by YOLO → sent to TrOCR for recognition\n\n"
#         "Use **Detailed Logs** to check coordinates, sizes & confidence values if results look off."
#     ),
#     theme=gr.themes.Soft(),
#     flagging_mode="never",
# )

# if __name__ == "__main__":
#     logger.info("Launching interface…")
#     demo.launch()



















# app.py - FIXED VERSION with empty crop protection
import gradio as gr
from ultralytics import YOLO
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image
import torch
import numpy as np

# Load models
region_model = YOLO("regions.pt")
line_model = YOLO("lines.pt")

processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def get_crop(image: Image.Image, result, idx: int, padding: int = 15):
    img_np = np.array(image)  # shape: (H_full, W_full, 3)

    if result.masks is not None:
        # Get the ORIGINAL bounding box (before any upsampling)
        box = result.boxes.xyxy[idx].cpu().numpy().astype(int)  # [x1, y1, x2, y2]
        x1, y1, x2, y2 = box

        # Get the mask – but make sure we use the mask at ORIGINAL size
        # In many cases masks.data[idx] is already at input resolution → we crop it directly
        mask = result.masks.data[idx].cpu().numpy()  # shape likely (H_full, W_full)
        mask_bool = mask > 0.5

        # Crop both image and mask using the **same box coordinates**
        crop_img = img_np[y1:y2, x1:x2]          # shape ~ (h_box, w_box, 3)
        crop_mask = mask_bool[y1:y2, x1:x2]      # shape ~ (h_box, w_box)

        if crop_img.size == 0 or crop_mask.size == 0:
            return None

        # Now apply **padding** around the cropped region
        h, w = crop_img.shape[:2]
        pad_top    = min(padding, y1)
        pad_bottom = min(padding, img_np.shape[0] - y2)
        pad_left   = min(padding, x1)
        pad_right  = min(padding, img_np.shape[1] - x2)

        # Padded coordinates in full image
        y_start = y1 - pad_top
        y_end   = y2 + pad_bottom
        x_start = x1 - pad_left
        x_end   = x2 + pad_right

        # Extract padded crops
        padded_img  = img_np[y_start:y_end, x_start:x_end]
        padded_mask = mask_bool[y_start:y_end, x_start:x_end]

        # Set background (outside mask) to white
        padded_img[~padded_mask] = 255

        return Image.fromarray(padded_img)

    else:
        # Bounding box fallback (no mask)
        xyxy = result.boxes.xyxy[idx].cpu().numpy().astype(int)
        x1, y1, x2, y2 = xyxy

        x1 = max(0, x1 - padding)
        y1 = max(0, y1 - padding)
        x2 = min(image.width, x2 + padding)
        y2 = min(image.height, y2 + padding)

        if x2 <= x1 or y2 <= y1:
            return None

        return image.crop((x1, y1, x2, y2))
        
def process_image(image: Image.Image):
    if image is None:
        return "Please upload an image."

    results = region_model(image)
    region_result = results[0]

    if region_result.boxes is None or len(region_result.boxes) == 0:
        return "No text regions detected."

    regions_with_pos = []
    for i in range(len(region_result.boxes)):
        y1 = region_result.boxes.xyxy[i][1].item()
        crop = get_crop(image, region_result, i, padding=20)
        if crop and crop.size[0] > 0 and crop.size[1] > 0:
            regions_with_pos.append((y1, crop))

    if not regions_with_pos:
        return "No valid text regions after cropping."

    regions_with_pos.sort(key=lambda x: x[0])

    full_text_parts = []

    for region_idx, (_, region_crop) in enumerate(regions_with_pos):
        line_results = line_model(region_crop)
        line_result = line_results[0]

        if line_result.boxes is None or len(line_result.boxes) == 0:
            continue

        lines_with_pos = []
        for j in range(len(line_result.boxes)):
            rel_y1 = line_result.boxes.xyxy[j][1].item()
            rel_x1 = line_result.boxes.xyxy[j][0].item()
            line_crop = get_crop(region_crop, line_result, j, padding=15)

            if line_crop is None or line_crop.size[0] < 10 or line_crop.size[1] < 8:
                # Skip tiny/invalid crops to prevent TrOCR crash
                # print(f"Skipped tiny line {j} in region {region_idx}")
                continue

            try:
                pixel_values = processor(line_crop, return_tensors="pt").pixel_values.to(device)
                generated_ids = model.generate(pixel_values)
                text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
                if text:  # only add non-empty
                    lines_with_pos.append((rel_y1, rel_x1, text))
            except Exception as e:
                # Catch any remaining processing errors
                # print(f"TrOCR error on line {j}: {e}")
                continue

        lines_with_pos.sort(key=lambda x: (x[0], x[1]))
        region_text = "\n".join([item[2] for item in lines_with_pos if item[2]])
        if region_text:
            full_text_parts.append(region_text)

    if not full_text_parts:
        return "No readable text recognized (possibly due to small/tiny lines or model limitations). Try a clearer document or larger padding."

    return "\n\n".join(full_text_parts)

# Gradio interface
demo = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil", label="Upload handwritten document"),
    outputs=gr.Textbox(label="Recognized Text"),
    title="Handwritten Text Recognition (YOLO + TrOCR)",
    description="Local models: regions.pt / lines.pt + microsoft/trocr-base-handwritten. Mask-based cropping + safeguards against empty crops.",
    flagging_mode="never"
)

if __name__ == "__main__":
    demo.launch()