Spaces:
Sleeping
Sleeping
File size: 9,195 Bytes
8ab8e97 abcce98 8ab8e97 abcce98 8ab8e97 abcce98 8ab8e97 abcce98 8ab8e97 abcce98 8ab8e97 abcce98 8ab8e97 abcce98 8ab8e97 abcce98 8ab8e97 abcce98 8ab8e97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 |
import streamlit as st
import fitz # PyMuPDF
import io
from PIL import Image
import google.generativeai as genai
from dotenv import load_dotenv
import os
# Load environment variables
load_dotenv()
# Initialize Gemini model
genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
gemini_model = genai.GenerativeModel('gemini-2.0-flash-exp')
class SequentialDrawingAnalyzer:
def __init__(self):
self.drawings_list = [] # Store all extracted images info
self.analyzed_drawings = [] # Store analyzed results
def extract_page_as_image(self, page):
"""Extract entire page as an image"""
try:
# Get the page's pixmap
zoom = 2 # Increase resolution
mat = fitz.Matrix(zoom, zoom)
pix = page.get_pixmap(matrix=mat)
# Convert pixmap to PIL Image
img_data = pix.tobytes("png")
return img_data
except Exception as e:
st.warning(f"Could not extract page as image: {str(e)}")
return None
def extract_drawings_list(self, pdf_bytes):
"""First pass: Extract all drawings from PDF and create a list"""
try:
doc = fitz.open(stream=pdf_bytes, filetype="pdf")
for page_num in range(len(doc)):
page = doc[page_num]
image_list = page.get_images()
if len(image_list) > 0:
# Process individual images if found
for img_idx, img_info in enumerate(image_list):
try:
xref = img_info[0]
base_image = doc.extract_image(xref)
image_bytes = base_image["image"]
self.drawings_list.append({
'page': page_num + 1,
'drawing_number': img_idx + 1,
'xref': xref,
'image_bytes': image_bytes,
'type': 'embedded'
})
except Exception as img_error:
st.warning(f"Could not extract drawing {img_idx + 1} on page {page_num + 1}: {str(img_error)}")
else:
# If no images found, extract entire page as image
page_image = self.extract_page_as_image(page)
if page_image:
self.drawings_list.append({
'page': page_num + 1,
'drawing_number': 1,
'image_bytes': page_image,
'type': 'full_page'
})
st.info(f"Extracted page {page_num + 1} as full-page drawing")
doc.close()
return len(self.drawings_list)
except Exception as e:
st.error(f"Error extracting drawings: {str(e)}")
return 0
def analyze_drawing(self, drawing_info):
"""Analyze a single drawing"""
try:
image = Image.open(io.BytesIO(drawing_info['image_bytes']))
drawing_type = "full page" if drawing_info.get('type') == 'full_page' else "embedded"
engineering_prompt = f"""
Analyze this engineering drawing in detail ({drawing_type} drawing). Please provide:
1. Drawing Type and Purpose
- Identify the type of drawing (assembly, detail, section view, etc.)
- Main purpose and function of the depicted component/system
2. Dimensional Analysis
- Key dimensions and measurements
- Scale and proportions
- Tolerances if specified
3. Component Details
- List all visible components and parts
- Materials specifications if indicated
- Surface finish markings
4. Technical Specifications
- Any technical notes or special instructions
- Welding symbols or special instructions
- Reference standards mentioned
5. Critical Features
- Important geometric features
- Key interfaces or connections
- Safety-critical aspects
"""
response = gemini_model.generate_content([
engineering_prompt,
image
])
return {
'page': drawing_info['page'],
'drawing_number': drawing_info['drawing_number'],
'image': image,
'analysis': response.text,
'type': drawing_info.get('type', 'embedded')
}
except Exception as e:
st.error(f"Error analyzing drawing {drawing_info['drawing_number']}: {str(e)}")
return None
# Streamlit UI
st.title("Sequential Engineering Drawing Analyzer")
# Initialize session state
if "processed" not in st.session_state:
st.session_state.processed = False
if "analyzer" not in st.session_state:
st.session_state.analyzer = SequentialDrawingAnalyzer()
if "current_analysis_index" not in st.session_state:
st.session_state.current_analysis_index = 0
if "analyzed_drawings" not in st.session_state:
st.session_state.analyzed_drawings = []
# File upload
pdf_file = st.file_uploader("Upload PDF containing engineering drawings", type="pdf")
if pdf_file is not None:
# First pass: Extract all drawings if not already processed
if not st.session_state.processed:
try:
with st.spinner("Extracting drawings from PDF..."):
pdf_bytes = pdf_file.getvalue()
total_drawings = st.session_state.analyzer.extract_drawings_list(pdf_bytes)
st.session_state.processed = True
st.success(f"Found {total_drawings} drawings in the PDF!")
# Display list of all drawings
st.subheader("List of Extracted Drawings:")
for drawing in st.session_state.analyzer.drawings_list:
drawing_type = "Full Page" if drawing.get('type') == 'full_page' else "Embedded"
st.write(f"{drawing_type} Drawing {drawing['drawing_number']} on Page {drawing['page']}")
st.markdown("---")
except Exception as e:
st.error(f"Failed to process PDF: {str(e)}")
st.session_state.processed = False
# Process drawings sequentially
if st.session_state.processed:
remaining_drawings = min(5, len(st.session_state.analyzer.drawings_list)) - st.session_state.current_analysis_index
if remaining_drawings > 0:
st.subheader(f"Analyzing Drawing {st.session_state.current_analysis_index + 1} of {min(5, len(st.session_state.analyzer.drawings_list))}")
# Analyze current drawing
current_drawing = st.session_state.analyzer.drawings_list[st.session_state.current_analysis_index]
with st.spinner(f"Analyzing drawing {current_drawing['drawing_number']} from page {current_drawing['page']}..."):
analysis_result = st.session_state.analyzer.analyze_drawing(current_drawing)
if analysis_result:
# Store analysis result
st.session_state.analyzed_drawings.append(analysis_result)
# Increment counter
st.session_state.current_analysis_index += 1
# Auto-rerun to process next drawing
if remaining_drawings > 1:
st.rerun()
else:
st.success("Completed analysis of first 5 drawings!")
elif len(st.session_state.analyzer.drawings_list) > 5:
st.info("First 5 drawings have been analyzed. Reload the page to analyze a different set of drawings.")
# Display all analyzed drawings
if st.session_state.analyzed_drawings:
st.subheader("Analyzed Drawings:")
for analysis in st.session_state.analyzed_drawings:
col1, col2 = st.columns([1, 1])
with col1:
drawing_type = "Full Page" if analysis['type'] == 'full_page' else "Embedded"
st.image(analysis['image'],
use_column_width=True,
caption=f"{drawing_type} Drawing {analysis['drawing_number']} (Page {analysis['page']})")
with col2:
st.markdown("### Analysis Results")
st.markdown(analysis['analysis'])
st.markdown("---")
else:
st.info("Please upload a PDF file containing engineering drawings to begin analysis.") |