Shakir60 commited on
Commit
8b2b9dd
·
verified ·
1 Parent(s): 5ec4050

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +72 -6
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 for images
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 and processor"""
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 5MB.")
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))