Update app.py
Browse files
app.py
CHANGED
@@ -1,4 +1,3 @@
|
|
1 |
-
# app.py
|
2 |
import streamlit as st
|
3 |
from transformers import ViTForImageClassification, ViTImageProcessor
|
4 |
from PIL import Image
|
@@ -6,19 +5,26 @@ import torch
|
|
6 |
import time
|
7 |
import gc
|
8 |
import logging
|
9 |
-
from knowledge_base import KNOWLEDGE_BASE, DAMAGE_TYPES
|
10 |
from rag_utils import RAGSystem
|
11 |
import structlog
|
12 |
from typing import Optional, Dict, Any
|
13 |
from functools import lru_cache
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
14 |
|
15 |
# Configure logging
|
16 |
logging.basicConfig(level=logging.INFO)
|
17 |
logger = structlog.get_logger()
|
18 |
|
19 |
# Constants
|
20 |
-
MAX_FILE_SIZE = 5 * 1024 * 1024 # 5MB
|
21 |
-
MAX_IMAGE_SIZE = 1024 # Maximum dimension
|
22 |
MODEL = None
|
23 |
PROCESSOR = None
|
24 |
RAG_SYSTEM = None
|
@@ -47,10 +53,12 @@ def init_session_state():
|
|
47 |
st.session_state.history = []
|
48 |
if 'dark_mode' not in st.session_state:
|
49 |
st.session_state.dark_mode = False
|
|
|
|
|
50 |
|
51 |
@st.cache_resource(show_spinner="Loading AI model...", ttl=3600*24)
|
52 |
def load_model():
|
53 |
-
"""Load and cache the model
|
54 |
try:
|
55 |
model_name = "google/vit-base-patch16-224"
|
56 |
processor = ViTImageProcessor.from_pretrained(model_name)
|
@@ -78,7 +86,7 @@ def validate_upload(file) -> bool:
|
|
78 |
return False
|
79 |
|
80 |
if file.size > MAX_FILE_SIZE:
|
81 |
-
st.error("File too large. Maximum size is
|
82 |
return False
|
83 |
|
84 |
if file.type not in ['image/jpeg', 'image/png']:
|
@@ -134,12 +142,49 @@ def analyze_damage(image: Image.Image, model: ViTForImageClassification,
|
|
134 |
st.error(f"Error analyzing image: {str(e)}")
|
135 |
return None
|
136 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
137 |
def display_analysis_results(predictions: torch.Tensor, analysis_time: float):
|
138 |
"""Display analysis results with damage details"""
|
139 |
st.markdown("### 📊 Analysis Results")
|
140 |
st.markdown(f"*Analysis completed in {analysis_time:.2f} seconds*")
|
141 |
|
|
|
142 |
detected = False
|
|
|
143 |
for idx, prob in enumerate(predictions):
|
144 |
confidence = float(prob) * 100
|
145 |
if confidence > 15:
|
@@ -165,9 +210,24 @@ def display_analysis_results(predictions: torch.Tensor, analysis_time: float):
|
|
165 |
with tabs[2]:
|
166 |
for safety in analysis['safety_considerations']:
|
167 |
st.warning(safety)
|
|
|
|
|
|
|
|
|
|
|
168 |
|
169 |
if not detected:
|
170 |
st.info("No significant structural damage detected. Regular maintenance recommended.")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
171 |
|
172 |
def main():
|
173 |
"""Main application function"""
|
@@ -208,6 +268,11 @@ def main():
|
|
208 |
RAG_SYSTEM = RAGSystem()
|
209 |
RAG_SYSTEM.initialize_knowledge_base(KNOWLEDGE_BASE)
|
210 |
|
|
|
|
|
|
|
|
|
|
|
211 |
# File upload
|
212 |
uploaded_file = st.file_uploader(
|
213 |
"Upload an image for analysis",
|
@@ -234,6 +299,7 @@ def main():
|
|
234 |
analysis_time = time.time() - start_time
|
235 |
display_analysis_results(predictions, analysis_time)
|
236 |
st.session_state.history.append(f"Analyzed {uploaded_file.name}")
|
|
|
237 |
|
238 |
except Exception as e:
|
239 |
logger.error("Error in main processing loop", error=str(e))
|
|
|
|
|
1 |
import streamlit as st
|
2 |
from transformers import ViTForImageClassification, ViTImageProcessor
|
3 |
from PIL import Image
|
|
|
5 |
import time
|
6 |
import gc
|
7 |
import logging
|
8 |
+
from knowledge_base import KNOWLEDGE_BASE, DAMAGE_TYPES, validate_knowledge_base
|
9 |
from rag_utils import RAGSystem
|
10 |
import structlog
|
11 |
from typing import Optional, Dict, Any
|
12 |
from functools import lru_cache
|
13 |
+
from dynaconf import Dynaconf
|
14 |
+
|
15 |
+
# Configure settings
|
16 |
+
settings = Dynaconf(
|
17 |
+
settings_files=['settings.yaml', '.secrets.yaml'],
|
18 |
+
environments=True
|
19 |
+
)
|
20 |
|
21 |
# Configure logging
|
22 |
logging.basicConfig(level=logging.INFO)
|
23 |
logger = structlog.get_logger()
|
24 |
|
25 |
# Constants
|
26 |
+
MAX_FILE_SIZE = settings.get('max_file_size', 5 * 1024 * 1024) # 5MB default
|
27 |
+
MAX_IMAGE_SIZE = settings.get('max_image_size', 1024) # Maximum dimension default
|
28 |
MODEL = None
|
29 |
PROCESSOR = None
|
30 |
RAG_SYSTEM = None
|
|
|
53 |
st.session_state.history = []
|
54 |
if 'dark_mode' not in st.session_state:
|
55 |
st.session_state.dark_mode = False
|
56 |
+
if 'analysis_count' not in st.session_state:
|
57 |
+
st.session_state.analysis_count = 0
|
58 |
|
59 |
@st.cache_resource(show_spinner="Loading AI model...", ttl=3600*24)
|
60 |
def load_model():
|
61 |
+
"""Load and cache the model with daily refresh"""
|
62 |
try:
|
63 |
model_name = "google/vit-base-patch16-224"
|
64 |
processor = ViTImageProcessor.from_pretrained(model_name)
|
|
|
86 |
return False
|
87 |
|
88 |
if file.size > MAX_FILE_SIZE:
|
89 |
+
st.error(f"File too large. Maximum size is {MAX_FILE_SIZE/1024/1024:.1f}MB.")
|
90 |
return False
|
91 |
|
92 |
if file.type not in ['image/jpeg', 'image/png']:
|
|
|
142 |
st.error(f"Error analyzing image: {str(e)}")
|
143 |
return None
|
144 |
|
145 |
+
def generate_downloadable_report(analysis_results: Dict):
|
146 |
+
"""Generate a downloadable PDF report"""
|
147 |
+
try:
|
148 |
+
import io
|
149 |
+
from reportlab.lib import colors
|
150 |
+
from reportlab.lib.pagesizes import letter
|
151 |
+
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
|
152 |
+
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
|
153 |
+
|
154 |
+
buffer = io.BytesIO()
|
155 |
+
doc = SimpleDocTemplate(buffer, pagesize=letter)
|
156 |
+
styles = getSampleStyleSheet()
|
157 |
+
story = []
|
158 |
+
|
159 |
+
# Add title
|
160 |
+
story.append(Paragraph("Structural Damage Analysis Report", styles['Title']))
|
161 |
+
story.append(Spacer(1, 12))
|
162 |
+
|
163 |
+
# Add analysis results
|
164 |
+
for damage_type, details in analysis_results.items():
|
165 |
+
story.append(Paragraph(f"Damage Type: {damage_type}", styles['Heading1']))
|
166 |
+
story.append(Paragraph(f"Confidence: {details['confidence']}%", styles['Normal']))
|
167 |
+
story.append(Paragraph("Recommendations:", styles['Heading2']))
|
168 |
+
for rec in details['recommendations']:
|
169 |
+
story.append(Paragraph(f"• {rec}", styles['Normal']))
|
170 |
+
story.append(Spacer(1, 12))
|
171 |
+
|
172 |
+
doc.build(story)
|
173 |
+
pdf = buffer.getvalue()
|
174 |
+
buffer.close()
|
175 |
+
return pdf
|
176 |
+
except Exception as e:
|
177 |
+
logger.error(f"Error generating report: {str(e)}")
|
178 |
+
return None
|
179 |
+
|
180 |
def display_analysis_results(predictions: torch.Tensor, analysis_time: float):
|
181 |
"""Display analysis results with damage details"""
|
182 |
st.markdown("### 📊 Analysis Results")
|
183 |
st.markdown(f"*Analysis completed in {analysis_time:.2f} seconds*")
|
184 |
|
185 |
+
analysis_results = {}
|
186 |
detected = False
|
187 |
+
|
188 |
for idx, prob in enumerate(predictions):
|
189 |
confidence = float(prob) * 100
|
190 |
if confidence > 15:
|
|
|
210 |
with tabs[2]:
|
211 |
for safety in analysis['safety_considerations']:
|
212 |
st.warning(safety)
|
213 |
+
|
214 |
+
analysis_results[damage_type] = {
|
215 |
+
'confidence': confidence,
|
216 |
+
'recommendations': analysis['expert_recommendations']
|
217 |
+
}
|
218 |
|
219 |
if not detected:
|
220 |
st.info("No significant structural damage detected. Regular maintenance recommended.")
|
221 |
+
else:
|
222 |
+
# Generate download button for report
|
223 |
+
pdf_report = generate_downloadable_report(analysis_results)
|
224 |
+
if pdf_report:
|
225 |
+
st.download_button(
|
226 |
+
label="Download Analysis Report",
|
227 |
+
data=pdf_report,
|
228 |
+
file_name="damage_analysis_report.pdf",
|
229 |
+
mime="application/pdf"
|
230 |
+
)
|
231 |
|
232 |
def main():
|
233 |
"""Main application function"""
|
|
|
268 |
RAG_SYSTEM = RAGSystem()
|
269 |
RAG_SYSTEM.initialize_knowledge_base(KNOWLEDGE_BASE)
|
270 |
|
271 |
+
# Validate knowledge base
|
272 |
+
if not validate_knowledge_base():
|
273 |
+
st.error("Knowledge base validation failed. Please check the logs.")
|
274 |
+
return
|
275 |
+
|
276 |
# File upload
|
277 |
uploaded_file = st.file_uploader(
|
278 |
"Upload an image for analysis",
|
|
|
299 |
analysis_time = time.time() - start_time
|
300 |
display_analysis_results(predictions, analysis_time)
|
301 |
st.session_state.history.append(f"Analyzed {uploaded_file.name}")
|
302 |
+
st.session_state.analysis_count += 1
|
303 |
|
304 |
except Exception as e:
|
305 |
logger.error("Error in main processing loop", error=str(e))
|