import os
import zipfile
import json
from PIL import Image
import gradio as gr

FIG_DIR = "FIG"
ZIP_FILE = "FIG.zip"

# 解压 FIG.zip（只做一次）
if not os.path.exists(FIG_DIR) and os.path.exists(ZIP_FILE):
    with zipfile.ZipFile(ZIP_FILE, 'r') as zip_ref:
        zip_ref.extractall(FIG_DIR)

# 图像提取主函数
def extract_figures():
    figures = []
    base_url = "https://hf.space/embed/SFEREWQW/poster/+/file/FIG"
    
    if not os.path.exists(FIG_DIR):
        return [{"error": "FIG folder not found."}]
    
    for filename in os.listdir(FIG_DIR):
        if filename.lower().endswith((".png", ".jpg", ".jpeg", ".gif")):
            figures.append({
                "filename": filename,
                "url": f"{base_url}/{filename}",
                "caption": f"Figure: {filename}"
            })
    return figures

# 创建 Gradio 接口
demo = gr.Interface(
    fn=extract_figures,
    inputs=[],
    outputs="json",
    title="Figure Extractor API",
    description="Click 'Submit' to get JSON from FIG folder"
)

demo.launch()