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.")