# app.py
import streamlit as st
from PIL import Image
import torch
from transformers import Blip2Processor, Blip2ForConditionalGeneration, pipeline
import mediapipe as mp
import cv2
import numpy as np
import tempfile

# -------------------------
# 1. BLIP-2 모델 로딩
# -------------------------
st.title("📷 AI 포즈 촬영 가이드")
st.caption("업로드한 이미지를 분석하고 따라하기 위한 포즈 안내 및 촬영 가이드를 생성합니다")

@st.cache_resource
def load_blip2():
    processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
    model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", device_map="auto", torch_dtype=torch.float16)
    return processor, model

processor, model = load_blip2()

# -------------------------
# 2. DistilGPT2 텍스트 생성기 로딩
# -------------------------
gpt_generator = pipeline("text-generation", model="distilgpt2")

# -------------------------
# 3. 이미지 업로드 & 분석
# -------------------------
uploaded_file = st.file_uploader("이미지를 업로드하세요", type=["jpg", "jpeg", "png"])

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="업로드한 이미지", use_column_width=True)

    # BLIP-2로 설명 생성
    inputs = processor(image, return_tensors="pt").to(model.device, torch.float16)
    with torch.no_grad():
        generated_ids = model.generate(**inputs)
        caption = processor.decode(generated_ids[0], skip_special_tokens=True)

    st.subheader("🧠 이미지 설명")
    st.info(caption)

    # 질문 기능
    user_q = st.text_input("이미지에 대해 질문해보세요 (예: 이 사람의 옷 스타일은?)")
    if user_q:
        prompt = f"질문: {user_q}\n이미지 설명: {caption}\n답변:"
        ans = gpt_generator(prompt, max_length=100, do_sample=True)[0]["generated_text"].split("답변:")[-1]
        st.success(ans)

    # -------------------------
    # 4. 포즈 스켈레톤 시각화
    # -------------------------
    st.subheader("🕺 포즈 아웃라인")
    mp_pose = mp.solutions.pose
    pose = mp_pose.Pose(static_image_mode=True)

    with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
        image.save(tmp_file.name)
        img_bgr = cv2.imread(tmp_file.name)
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        results = pose.process(img_rgb)

    if results.pose_landmarks:
        annotated = img_rgb.copy()
        mp_drawing = mp.solutions.drawing_utils
        mp_drawing.draw_landmarks(
            annotated,
            results.pose_landmarks,
            mp_pose.POSE_CONNECTIONS,
            landmark_drawing_spec=mp_drawing.DrawingSpec(color=(255, 0, 0), thickness=2, circle_radius=3),
            connection_drawing_spec=mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2),
        )
        st.image(annotated, caption="포즈 시각화", use_column_width=True)
    else:
        st.warning("포즈를 추출하지 못했습니다.")

    # -------------------------
    # 5. 연출 가이드 생성
    # -------------------------
    st.subheader("🎯 촬영 가이드")
    prompt = f"이미지 설명: {caption}\n이 이미지를 참고해 비슷한 사진을 찍으려면 어떻게 포즈를 잡아야 할까?"
    instruction = gpt_generator(prompt, max_length=100, do_sample=True)[0]['generated_text'].split("찍으려면")[-1]
    st.markdown(f"**포즈 가이드:** {instruction.strip()}")