Upload 2 files

app.py

@@ -1,44 +1,42 @@
-import streamlit as st
-import torch
-import cv2
-import tempfile
-import os
-import numpy as np
-from PIL import Image
-from pathlib import Path
-
-@st.cache_resource
-def load_model(model_path):
-    return torch.hub.load('ultralytics/yolov5', 'custom', path=model_path, force_reload=False)
-
-def detect_image(img, model):
-    img_array = np.array(img.convert('RGB'))
-    img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
-    result = model(img_bgr)
-    result.render()
-    for img in result.ims:
-        yield Image.fromarray(img)
-
-def main():
-    st.title("🔥 Ignition Point Detector")
-    st.write("AI 기반 화재 발화지점 탐지 웹앱입니다.")
-    uploaded_model = st.file_uploader("📦 YOLOv5 모델 파일 (.pt)", type=['pt'])
-    uploaded_images = st.file_uploader("📷 이미지 업로드 (여러 장 가능)", type=['jpg', 'jpeg', 'png'], accept_multiple_files=True)
-
-    if uploaded_model and uploaded_images:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as tmp_model:
-            tmp_model.write(uploaded_model.read())
-            tmp_model_path = tmp_model.name
-
-        model = load_model(tmp_model_path)
-
-        for uploaded_image in uploaded_images:
-            image = Image.open(uploaded_image)
-            st.image(image, caption="🖼️ 원본 이미지", use_container_width=True)
-
-            st.markdown("🔍 분석 결과:")
-            for result_img in detect_image(image, model):
-                st.image(result_img, caption="🔥 탐지된 이미지", use_container_width=True)
-
-if __name__ == "__main__":
-    main()
+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)

requirements.txt

@@ -1,5 +1,5 @@
 streamlit
 torch
-opencv-python-headless
+opencv-python
 numpy
 Pillow