Mohiit007/BrainCache-OrbitalScan
0
1<<<<<<< HEAD2import streamlit as st
3from ultralytics import YOLO
4import cv2
5import numpy as np
6from PIL import Image
7import tempfile
8import os
9import matplotlib.pyplot as plt
10import seaborn as sns
11from fpdf import FPDF
12import json
13import requests
14
15# ---------------- Lottie Animation Loader ---------------- #
16def load_lottie(url):
17 r = requests.get(url)
18 if r.status_code != 200:
19 return None
20 return r.json()
21
22lottie_rocket = load_lottie("https://assets2.lottiefiles.com/packages/lf20_zrqthn6o.json")
23
24# ---------------- UI Styling ---------------- #
25st.set_page_config(page_title="BrainCache - Space Station Safety AI", layout="wide")
26st.markdown(
27 """
28 <style>
29 body {background-color: #0e1117; color: white;}
30 .main {background-color: #0e1117;}
31 h1, h2, h3, h4 {color: #4CAF50;}
32 .stButton>button {background-color: #4CAF50; color: white; font-size:18px;}
33 </style>
34 """, unsafe_allow_html=True
35)
36
37# ---------------- Load Model ---------------- #
38@st.cache_resource
39def load_model():
40 return YOLO("best.pt")
41
42model = load_model()
43
44st.title("🚀 BrainCache – Space Station Safety AI")
45st.write("AI-powered detection of **Toolbox, Oxygen Tank, Fire Extinguisher** to ensure astronaut safety.")
46
47# ---------------- Tabs ---------------- #
48tab1, tab2, tab3 = st.tabs(["🛰 Detection", "📊 Analytics", "ℹ About Us"])
49
50# ---------------- Detection ---------------- #
51with tab1:
52 st.subheader("Upload Image, Video, or Use Camera")
53 option = st.radio("Select Input Type", ("Image", "Video", "Live Camera"))
54
55 if option == "Image":
56 uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"])
57 if uploaded_file:
58 img = Image.open(uploaded_file)
59
60 with st.spinner("🚀 Detecting objects..."):
61 results = model.predict(source=np.array(img))
62
63 annotated_img = results[0].plot()
64 st.image(annotated_img, caption="Detections", use_column_width=True)
65
66 # Download annotated image
67 cv2.imwrite("annotated_image.jpg", annotated_img)
68 with open("annotated_image.jpg", "rb") as f:
69 st.download_button("Download Annotated Image", f, file_name="annotated_image.jpg")
70
71 # Confidence Scores
72 st.write("### Confidence Scores")
73 for box in results[0].boxes:
74 st.write(f"{results[0].names[int(box.cls)]}: {float(box.conf):.2f}")
75
76 elif option == "Video":
77 uploaded_video = st.file_uploader("Upload a Video", type=["mp4", "mov", "avi"])
78 if uploaded_video:
79 temp_file = tempfile.NamedTemporaryFile(delete=False)
80 temp_file.write(uploaded_video.read())
81 st.video(uploaded_video)
82
83 if st.button("Run Detection on Video"):
84 with st.spinner("Analyzing video..."):
85 output_path = "annotated_video.mp4"
86 results = model.predict(source=temp_file.name, save=True)
87 # YOLO saves video automatically in runs/detect/predict
88 st.success("Video Processed!")
89 st.video("runs/detect/predict/video.mp4")
90 with open("runs/detect/predict/video.mp4", "rb") as f:
91 st.download_button("Download Annotated Video", f, file_name="annotated_video.mp4")
92
93 elif option == "Live Camera":
94 camera_image = st.camera_input("Capture an Image")
95 if camera_image:
96 img = Image.open(camera_image)
97 results = model.predict(source=np.array(img))
98 annotated_img = results[0].plot()
99 st.image(annotated_img, caption="Live Detection", use_column_width=True)
100
101 # Download captured detection
102 cv2.imwrite("live_detect.jpg", annotated_img)
103 with open("live_detect.jpg", "rb") as f:
104 st.download_button("Download Live Detection", f, file_name="live_detect.jpg")
105
106# ---------------- Analytics ---------------- #
107with tab2:
108 st.subheader("Model Performance")
109 st.metric("mAP@0.5", "0.916")
110 st.metric("mAP@0.5-0.95", "0.792")
111
112 # Confusion Matrix
113 if st.button("Generate Confusion Matrix"):
114 st.write("📊 Generating confusion matrix...")
115 labels = ["Toolbox", "Oxygen Tank", "Fire Extinguisher"]
116 confusion = np.array([[65, 2, 0],
117 [1, 58, 1],
118 [0, 3, 76]])
119
120 fig, ax = plt.subplots()
121 sns.heatmap(confusion, annot=True, fmt="d", cmap="Blues", xticklabels=labels, yticklabels=labels)
122 plt.xlabel("Predicted")
123 plt.ylabel("Actual")
124 st.pyplot(fig)
125
126 # Save confusion matrix for report
127 fig.savefig("confusion_matrix.png")
128
129 # Generate PDF Report
130 if st.button("Generate PDF Report"):
131 pdf = FPDF()
132 pdf.set_font("Arial", size=12)
133 pdf.add_page()
134 pdf.cell(200, 10, "Performance Report - BrainCache", ln=True, align="C")
135 pdf.cell(200, 10, f"mAP@0.5: 0.916", ln=True)
136 pdf.cell(200, 10, f"mAP@0.5-0.95: 0.792", ln=True)
137 if os.path.exists("confusion_matrix.png"):
138 pdf.image("confusion_matrix.png", x=50, w=100)
139 pdf.output("Performance_Report.pdf")
140 with open("Performance_Report.pdf", "rb") as f:
141 st.download_button("Download Performance Report", f, file_name="Performance_Report.pdf")
142
143# ---------------- About Us ---------------- #
144with tab3:
145 st.subheader("Our Mission")
146 st.write("""
147 - **Team Name:** BrainCache
148 - **Members:** Swastika, Mohit, Uday, Rohit
149 - **Goal:** AI-driven safety monitoring for astronauts.
150 - **Hackathon:** BuildWithIndia 2.0
151 """)
152=======153import streamlit as st154from ultralytics import YOLO155import cv2156import numpy as np157from PIL import Image158import tempfile159import os160import matplotlib.pyplot as plt161import seaborn as sns162from fpdf import FPDF163import json164import requests165 166# ---------------- Lottie Animation Loader ---------------- #167def load_lottie(url):168 r = requests.get(url)169 if r.status_code != 200:170 return None171 return r.json()172 173lottie_rocket = load_lottie("https://assets2.lottiefiles.com/packages/lf20_zrqthn6o.json")174 175# ---------------- UI Styling ---------------- #176st.set_page_config(page_title="BrainCache - Space Station Safety AI", layout="wide")177st.markdown(178 """179 <style>180 body {background-color: #0e1117; color: white;}181 .main {background-color: #0e1117;}182 h1, h2, h3, h4 {color: #4CAF50;}183 .stButton>button {background-color: #4CAF50; color: white; font-size:18px;}184 </style>185 """, unsafe_allow_html=True186)187 188# ---------------- Load Model ---------------- #189@st.cache_resource190def load_model():191 return YOLO("best.pt")192 193model = load_model()194 195st.title("🚀 BrainCache – Space Station Safety AI")196st.write("AI-powered detection of **Toolbox, Oxygen Tank, Fire Extinguisher** to ensure astronaut safety.")197 198# ---------------- Tabs ---------------- #199tab1, tab2, tab3 = st.tabs(["🛰 Detection", "📊 Analytics", "ℹ About Us"])200 201# ---------------- Detection ---------------- #202with tab1:203 st.subheader("Upload Image, Video, or Use Camera")204 option = st.radio("Select Input Type", ("Image", "Video", "Live Camera"))205 206 if option == "Image":207 uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"])208 if uploaded_file:209 img = Image.open(uploaded_file)210 211 with st.spinner("🚀 Detecting objects..."):212 results = model.predict(source=np.array(img))213 214 annotated_img = results[0].plot()215 st.image(annotated_img, caption="Detections", use_column_width=True)216 217 # Download annotated image218 cv2.imwrite("annotated_image.jpg", annotated_img)219 with open("annotated_image.jpg", "rb") as f:220 st.download_button("Download Annotated Image", f, file_name="annotated_image.jpg")221 222 # Confidence Scores223 st.write("### Confidence Scores")224 for box in results[0].boxes:225 st.write(f"{results[0].names[int(box.cls)]}: {float(box.conf):.2f}")226 227 elif option == "Video":228 uploaded_video = st.file_uploader("Upload a Video", type=["mp4", "mov", "avi"])229 if uploaded_video:230 temp_file = tempfile.NamedTemporaryFile(delete=False)231 temp_file.write(uploaded_video.read())232 st.video(uploaded_video)233 234 if st.button("Run Detection on Video"):235 with st.spinner("Analyzing video..."):236 output_path = "annotated_video.mp4"237 results = model.predict(source=temp_file.name, save=True)238 # YOLO saves video automatically in runs/detect/predict239 st.success("Video Processed!")240 st.video("runs/detect/predict/video.mp4")241 with open("runs/detect/predict/video.mp4", "rb") as f:242 st.download_button("Download Annotated Video", f, file_name="annotated_video.mp4")243 244 elif option == "Live Camera":245 camera_image = st.camera_input("Capture an Image")246 if camera_image:247 img = Image.open(camera_image)248 results = model.predict(source=np.array(img))249 annotated_img = results[0].plot()250 st.image(annotated_img, caption="Live Detection", use_column_width=True)251 252 # Download captured detection253 cv2.imwrite("live_detect.jpg", annotated_img)254 with open("live_detect.jpg", "rb") as f:255 st.download_button("Download Live Detection", f, file_name="live_detect.jpg")256 257# ---------------- Analytics ---------------- #258with tab2:259 st.subheader("Model Performance")260 st.metric("mAP@0.5", "0.916")261 st.metric("mAP@0.5-0.95", "0.792")262 263 # Confusion Matrix264 if st.button("Generate Confusion Matrix"):265 st.write("📊 Generating confusion matrix...")266 labels = ["Toolbox", "Oxygen Tank", "Fire Extinguisher"]267 confusion = np.array([[65, 2, 0],268 [1, 58, 1],269 [0, 3, 76]])270 271 fig, ax = plt.subplots()272 sns.heatmap(confusion, annot=True, fmt="d", cmap="Blues", xticklabels=labels, yticklabels=labels)273 plt.xlabel("Predicted")274 plt.ylabel("Actual")275 st.pyplot(fig)276 277 # Save confusion matrix for report278 fig.savefig("confusion_matrix.png")279 280 # Generate PDF Report281 if st.button("Generate PDF Report"):282 pdf = FPDF()283 pdf.set_font("Arial", size=12)284 pdf.add_page()285 pdf.cell(200, 10, "Performance Report - BrainCache", ln=True, align="C")286 pdf.cell(200, 10, f"mAP@0.5: 0.916", ln=True)287 pdf.cell(200, 10, f"mAP@0.5-0.95: 0.792", ln=True)288 if os.path.exists("confusion_matrix.png"):289 pdf.image("confusion_matrix.png", x=50, w=100)290 pdf.output("Performance_Report.pdf")291 with open("Performance_Report.pdf", "rb") as f:292 st.download_button("Download Performance Report", f, file_name="Performance_Report.pdf")293 294# ---------------- About Us ---------------- #295with tab3:296 st.subheader("Our Mission")297 st.write("""298 - **Team Name:** BrainCache 299 - **Members:** Swastika, Mohit, Uday, Rohit 300 - **Goal:** AI-driven safety monitoring for astronauts. 301 - **Hackathon:** BuildWithIndia 2.0 302 """)303 304>>>>>>> f3a1803 (Initial commit with YOLO model)305