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()