Spaces:
Sleeping
Sleeping
| 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.") |