app.py

import gradio as gr
from transformers import pipeline

# 1) 이미지 캡셔닝 파이프라인 초기화
captioner = pipeline("image-captioning", model="Salesforce/blip-image-captioning-base")

# 2) (옵션) 스타일 변환 함수
def style_convert(raw_caption, style):
    if style == "TOEIC Speaking Part 1":
        # Part 1: “What do you see in the picture?”
        # 답변: “I see ~.” 한 문장
        return f"Q: What do you see in the picture?\nA: {raw_caption.capitalize()}."
    elif style == "IELTS Describe a Photo":
        return f"Describe the photo in two sentences:\n1. {raw_caption.capitalize()}.\n2. It also shows the context of daily life."
    else:
        return raw_caption

# 3) Gradio 인터페이스 함수
def generate_caption(image, style):
    # 3.1 이미지 캡셔닝
    result = captioner(image, max_length=30, num_beams=3)[0]["caption"]
    # 3.2 스타일 변환
    return style_convert(result, style)

# 4) Gradio Blocks 정의
with gr.Blocks() as demo:
    gr.Markdown("## 📸 이미지 캡셔닝 → English Test 스타일 문장 생성")
    with gr.Row():
        img_in = gr.Image(type="pil", label="Upload Image")
        style_sel = gr.Dropdown(
            choices=["Raw Caption", "TOEIC Speaking Part 1", "IELTS Describe a Photo"],
            value="TOEIC Speaking Part 1",
            label="시험 형식 선택"
        )
    output = gr.Textbox(label="Generated Caption", lines=4)
    btn = gr.Button("Generate")
    btn.click(fn=generate_caption, inputs=[img_in, style_sel], outputs=output)

# 5) 앱 실행
if __name__ == "__main__":
    demo.launch()