Spaces:
Sleeping
Sleeping
File size: 4,718 Bytes
f7316f8 | 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """
Author: Juan Pablo Triana Martinez
Gradio app — LinkNet Binary Text Segmentation (DocLayNet).
Entry-point for HuggingFace Spaces.
"""
import gradio as gr
import numpy as np
import torch
from PIL import Image
from torchvision import transforms
from model import create_binary_model
# ── Constants ──────────────────────────────────────────────────────────────────
MEAN = [0.9329, 0.9343, 0.9341]
STD = [0.1651, 0.1592, 0.1623]
IMG_SIZE = 512
THRESHOLD = 0.5
MODEL_PATH = "linknet_binary_doclaynet_20_percent_seed_7_ce_0.25_dice_1.0_.pth"
# ── Transforms ─────────────────────────────────────────────────────────────────
inference_transform = transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ToTensor(),
transforms.Normalize(mean=MEAN, std=STD),
])
api_transform = transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ToTensor(),
])
# ── Model ──────────────────────────────────────────────────────────────────────
device = "cpu"
model = create_binary_model().to(device)
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
model.eval()
# ── Helpers ────────────────────────────────────────────────────────────────────
def _logits_to_binary_mask(logits: torch.Tensor, threshold: float = 0.5) -> torch.Tensor:
probs = torch.sigmoid(logits)
return (probs > threshold).float()
def _overlay_mask(image: Image.Image, mask: torch.Tensor) -> np.ndarray:
"""Green-channel overlay of binary mask on the original image."""
img_np = api_transform(image).permute(1, 2, 0).numpy() # (H, W, 3) in [0,1]
mask_np = mask.squeeze().cpu().numpy() # (H, W) in {0,1}
overlay = img_np.copy()
overlay[:, :, 1] = np.maximum(overlay[:, :, 1], mask_np)
return (overlay * 255).clip(0, 255).astype(np.uint8)
# ── Inference ──────────────────────────────────────────────────────────────────
def inference(image: Image.Image, gt_mask: Image.Image = None):
img_tensor = inference_transform(image).unsqueeze(0) # (1, 3, H, W)
with torch.inference_mode():
logits = model(img_tensor) # (1, 1, H, W)
pred_mask = _logits_to_binary_mask(logits, THRESHOLD) # (1, 1, H, W)
pred_mask_np = pred_mask.squeeze().cpu().numpy() # (H, W) for Gradio display
overlay = _overlay_mask(image, pred_mask)
if gt_mask is None:
return pred_mask_np, overlay, "No ground truth provided → metrics unavailable"
gt_tensor = api_transform(gt_mask.convert("L")).unsqueeze(0) # (1, 1, H, W)
gt_binary = (gt_tensor > 0.5).float()
inter = (pred_mask * gt_binary).sum()
union = ((pred_mask + gt_binary) > 0).float().sum()
iou = (inter / (union + 1e-7)).item()
dice = (2 * inter / (pred_mask.sum() + gt_binary.sum() + 1e-7)).item()
metrics_text = f"IoU (Jaccard): {iou:.4f} Dice (F1): {dice:.4f}"
return pred_mask_np, overlay, metrics_text
# ── Gradio Interface ───────────────────────────────────────────────────────────
demo = gr.Interface(
fn=inference,
inputs=[
gr.Image(type="pil", label="Input Image"),
gr.Image(type="pil", label="Ground Truth Mask (optional)"),
],
outputs=[
gr.Image(label="Predicted Binary Mask"),
gr.Image(label="Overlay"),
gr.Textbox(label="Metrics"),
],
title="LinkNet Binary Text Segmentation — DocLayNet ChestNut 🌰",
description=(
"Upload a document image to obtain its binary text segmentation mask. "
"Optionally upload a ground-truth binary mask to compute IoU and Dice metrics."
),
examples=[
["examples/binary_1_img.png", "examples/binary_1_mask.png"],
["examples/binary_2_img.png", "examples/binary_2_mask.png"],
["examples/binary_3_img.png", "examples/binary_3_mask.png"],
],
)
if __name__ == "__main__":
demo.launch()
|