Upload 2 files

app.py

@@ -1,42 +1,75 @@
-import streamlit as st
-import torch
-import tempfile
-import os
-import cv2
-from PIL import Image
-import numpy as np
-
-st.set_page_config(page_title="Fire Ignition Detector", layout="centered")
-
-# 상단 로고와 제목
-st.image("logoall.jpg", use_column_width=True)
-st.markdown("<h2 style='text-align: center;'>화재 발화지점 탐지기</h2>", unsafe_allow_html=True)
-
-# 모델 로드
-@st.cache_resource
-def load_model(model_path):
-    model = torch.load(model_path, map_location=torch.device('cpu'))
-    model.eval()
-    return model
-
-uploaded_model = st.file_uploader("🔥 YOLOv5 모델(.pt) 파일 업로드", type=["pt"])
-if uploaded_model is not None:
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as tmp_model:
-        tmp_model.write(uploaded_model.read())
-        model_path = tmp_model.name
-        model = load_model(model_path)
-        st.success("모델이 성공적으로 업로드되었습니다.")
-
-    # 이미지 업로드
-    uploaded_images = st.file_uploader("🖼️ 분석할 이미지 업로드 (여러 장 가능)", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
-
-    if uploaded_images:
-        for img_file in uploaded_images:
-            file_bytes = np.asarray(bytearray(img_file.read()), dtype=np.uint8)
-            image = cv2.imdecode(file_bytes, 1)
-            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-            results = model(image_rgb)
-
-            result_img = results.render()[0]
-            st.image(result_img, caption=f"🔍 분석 결과: {img_file.name}", use_column_width=True)
+
+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)

requirements.txt

@@ -1,5 +1,6 @@
-streamlit
-torch
-opencv-python
-numpy
+
+streamlit==1.29.0
+torch>=2.0.0
 Pillow
+numpy
+requests