Spaces:
Sleeping
Sleeping
File size: 3,318 Bytes
6548988 |
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 |
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()
|