| import io, cv2, torch, gradio as gr |
| import matplotlib.pyplot as plt |
| from matplotlib.patches import Rectangle |
| from pathlib import Path |
| from PIL import Image |
| from transformers.models.fast.image_processing_fast import FastImageProcessor |
| from transformers.models.fast.modeling_fast import FastForSceneTextRecognition |
|
|
| THIS_DIR = Path(__file__).parent |
| model_dir = THIS_DIR / "converted_fast_base" |
|
|
| processor = FastImageProcessor.from_pretrained(model_dir) |
| model = FastForSceneTextRecognition.from_pretrained(model_dir).eval() |
|
|
| def draw_detections(img, dets): |
| plt.figure(figsize=(8, 8)); plt.imshow(img); ax = plt.gca() |
| for box in dets["boxes"]: |
| if len(box) == 5: |
| xc, yc, w, h, angle = box |
| pts = cv2.boxPoints(((xc, yc), (w, h), angle)).astype(int).tolist()+[()] |
| xs, ys = zip(*pts); ax.plot(xs, ys, "-r", lw=2) |
| elif len(box) == 4 and isinstance(box[0], (list, tuple)): |
| pts = box+[box[0]]; xs=[p[0] for p in pts]; ys=[p[1] for p in pts] |
| ax.plot(xs, ys, "-b", lw=2) |
| elif len(box) == 4: |
| xmin, ymin, xmax, ymax = box |
| ax.add_patch(Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, |
| fill=False, lw=2, ec="r")) |
| elif len(box) == 8 and all(isinstance(x, (int, float)) for x in box): |
| xs = list(box[0::2]) + [box[0]] |
| ys = list(box[1::2]) + [box[1]] |
| ax.plot(xs, ys, "-g", lw=2) |
| elif len(box) > 8 and all(isinstance(x, (int, float)) for x in box): |
| xs = list(box[0::2]) + [box[0]] |
| ys = list(box[1::2]) + [box[1]] |
| ax.plot(xs, ys, "-g", lw=2) |
| else: |
| raise ValueError(f"Unrecognized box format: {box!r}") |
| ax.axis("off"); plt.tight_layout() |
| buf=io.BytesIO(); plt.savefig(buf, format="png", bbox_inches="tight"); buf.seek(0) |
| out = Image.open(buf).convert("RGB"); plt.close(); return out |
|
|
| def run(image, mode="boxes"): |
| inputs = processor(images=image, return_tensors="pt") |
| with torch.no_grad(): |
| outputs = model(**inputs) |
| dets = processor.post_process_text_detection( |
| outputs, target_sizes=[image.size[::-1]], output_type=mode)[0] |
| return draw_detections(image, dets) |
|
|
| demo = gr.Interface( |
| fn=run, |
| inputs=[gr.Image(type="pil", label="Upload an image with text", sources=["upload", "clipboard"]), gr.Radio(["boxes", "polygons"], value="boxes")], |
| outputs=gr.Image(type="pil"), title="FAST Text Detection Demo", |
| description="Detect text in images using the FAST model from Hugging Face Transformers" |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
