import streamlit as st
import torch
from PIL import Image
import numpy as np
import tempfile
import os
import requests
from io import BytesIO
from pathlib import Path

st.set_page_config(page_title="Ignition Point Detector", layout="centered")

st.image("logoall.jpg", use_container_width=True)

st.title("Ignition Point Detector 🔥")

st.markdown("AI 기반 화재 이미지 분석 웹앱입니다. 아래에서 모델 파일(.pt)과 이미지를 업로드하여 발화지점을 예측해보세요.")

# 모델 URL 입력
model_url = st.text_input("① .pt 모델 파일 URL을 입력하세요 (예: Google Drive 공유 링크):")

# 이미지 업로드
uploaded_images = st.file_uploader("② 분석할 이미지를 업로드하세요 (여러 장 가능)", type=["jpg", "jpeg", "png"], accept_multiple_files=True)

# 예측 버튼
predict_btn = st.button("🔥 예측 시작")

def load_model_from_url(url):
    response = requests.get(url)
    if response.status_code == 200:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as tmp_file:
            tmp_file.write(response.content)
            tmp_path = tmp_file.name
        try:
            model = torch.load(tmp_path, map_location=torch.device('cpu'))
            model.eval()
            return model
        except Exception as e:
            st.error(f"모델 로드 실패: {e}")
            return None
    else:
        st.error("모델 다운로드 실패: URL을 확인해주세요.")
        return None

def run_prediction(model, image):
    try:
        img = Image.open(image).convert('RGB')
        img_resized = img.resize((640, 640))
        img_array = np.array(img_resized) / 255.0
        img_tensor = torch.tensor(img_array).permute(2, 0, 1).unsqueeze(0).float()
        results = model(img_tensor)
        if isinstance(results, dict) and 'pred' in results:
            pred_boxes = results['pred'][0]
            for box in pred_boxes:
                x1, y1, x2, y2, conf, cls = box.tolist()
                st.write(f"🔥 예측 박스: 좌상단 ({x1:.0f}, {y1:.0f}), 우하단 ({x2:.0f}, {y2:.0f}), 신뢰도: {conf:.2f}")
            st.image(img, caption="업로드한 이미지", use_container_width=True)
        else:
            st.warning("모델 예측 결과가 올바르지 않습니다.")
    except Exception as e:
        st.error(f"예측 실패: {e}")

if predict_btn:
    if not model_url:
        st.warning("모델 URL을 입력해주세요.")
    elif not uploaded_images:
        st.warning("이미지를 업로드해주세요.")
    else:
        with st.spinner("모델을 다운로드하고 로드 중입니다..."):
            model = load_model_from_url(model_url)
        if model:
            for img in uploaded_images:
                st.subheader(f"📷 {img.name}")
                run_prediction(model, img)