Spaces:
Sleeping
Sleeping
| import io, json, os | |
| import numpy as np | |
| from PIL import Image | |
| import gradio as gr | |
| import cv2 | |
| import pandas as pd | |
| from biz.segmentation import segment_cards | |
| from biz.gemini import extract_from_crop | |
| from biz.utils import crop_to_png_bytes, overlay_boxes, to_excel_file | |
| # --- helpers --- | |
| def np_to_pil(img_np) -> Image.Image: | |
| if isinstance(img_np, Image.Image): return img_np.convert("RGB") | |
| return Image.fromarray(img_np).convert("RGB") | |
| def pil_to_bgr(pil: Image.Image): | |
| return cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2BGR) | |
| # --- Gradio functions --- | |
| def do_segment(image_np): | |
| if image_np is None: | |
| return None, "[]", gr.update(visible=False), None | |
| pil = np_to_pil(image_np) | |
| bgr = pil_to_bgr(pil) | |
| boxes, W, H = segment_cards(bgr) | |
| overlay = overlay_boxes(pil, boxes) | |
| return overlay, json.dumps(boxes, ensure_ascii=False), gr.update(visible=True), None | |
| def do_extract(image_np, boxes_json): | |
| if image_np is None or not boxes_json: | |
| return pd.DataFrame(), None | |
| pil = np_to_pil(image_np) | |
| try: | |
| boxes = json.loads(boxes_json) | |
| except Exception: | |
| boxes = [] | |
| cards = [] | |
| for b in boxes: | |
| crop = crop_to_png_bytes(pil, b["x"], b["y"], b["w"], b["h"]) | |
| fields = extract_from_crop(crop, source_name="upload") | |
| fields["box_id"] = b["id"] | |
| cards.append(fields) | |
| df = pd.DataFrame(cards) | |
| xlsx_path = to_excel_file(cards) | |
| return df, xlsx_path | |
| def clear_all(): | |
| return None, None, "[]", gr.update(visible=False), pd.DataFrame(), None | |
| # --- UI --- | |
| with gr.Blocks(title="BizCards Extractor (Gradio)") as demo: | |
| gr.Markdown("## 💼 BizCards Extractor\nUpload → **Segment** → **Extract** → **Download Excel**") | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| in_img = gr.Image(type="numpy", label="Upload single or multi-card photo") | |
| with gr.Row(): | |
| btn_seg = gr.Button("Segment", variant="primary") | |
| btn_ext = gr.Button("Extract", variant="secondary") | |
| btn_clear = gr.Button("Clear") | |
| with gr.Column(scale=2): | |
| out_img = gr.Image(label="Segmented preview (boxes)", interactive=False) | |
| out_table = gr.Dataframe( | |
| headers=["box_id","company","person_romaji","person_kanji","person_kana", | |
| "title","department","email","phone","website","address_jp","notes","source_name"], | |
| wrap=True, height=350 | |
| ) | |
| dl = gr.File(label="Download Excel", visible=False) | |
| # hidden state for boxes in JSON | |
| boxes_state = gr.Textbox(label="boxes_json (debug)", visible=False, value="[]") | |
| # wiring | |
| btn_seg.click(fn=do_segment, inputs=[in_img], | |
| outputs=[out_img, boxes_state, dl, dl]) | |
| btn_ext.click(fn=do_extract, inputs=[in_img, boxes_state], | |
| outputs=[out_table, dl]) | |
| btn_clear.click(fn=clear_all, inputs=[], | |
| outputs=[in_img, out_img, boxes_state, dl, out_table, dl]) | |
| # show a warning if key missing | |
| if not os.getenv("GOOGLE_API_KEY") and not os.getenv("GOOGLE_GENAI_USE_VERTEXAI"): | |
| gr.Warning("GOOGLE_API_KEY is not set. Add it in Space → Settings → Variables & secrets.") | |
| if __name__ == "__main__": | |
| demo.queue(max_size=16).launch() | |