Upload 2 files

app.py

@@ -1,38 +1,44 @@
-import gradio as gr
+import streamlit as st
 import torch
 import cv2
-import numpy as np
-from PIL import Image
 import tempfile
-import requests
 import os
-
-# 모델 다운로드 (Google Drive 또는 외부 URL에서 best.pt를 로드)
-MODEL_URL = "https://huggingface.co/Daebin/fire-ignition-detector/resolve/main/best.pt"
+import numpy as np
+from PIL import Image
+from pathlib import Path
 
 @st.cache_resource
-def load_model():
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as tmp:
-        tmp.write(requests.get(MODEL_URL).content)
-        tmp_path = tmp.name
-    model = torch.hub.load("ultralytics/yolov5", "custom", path=tmp_path, force_reload=True)
-    return model
-
-model = load_model()
-
-def detect_fire(image):
-    image = Image.fromarray(image)
-    results = model(image)
-    results.render()
-    return Image.fromarray(results.ims[0])
-
-iface = gr.Interface(
-    fn=detect_fire,
-    inputs=gr.Image(type="numpy"),
-    outputs=gr.Image(type="pil"),
-    title="Fire Ignition Detector",
-    description="Upload an image to detect ignition point using YOLOv5."
-)
+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__":
-    iface.launch()
+    main()

requirements.txt

@@ -1,7 +1,5 @@
+streamlit
 torch
-torchvision
-opencv-python
+opencv-python-headless
 numpy
 Pillow
-requests
-gradio