backend/models/drawing_processor.py - master - nd4wn/AIProjectControl

backend/models/drawing_processor.py
222 строки9 KB
nd4wn
Create: __init__.py, __init__.py, llama_client.py, !!!!ocr_processor.py, cache_manager.py, document_analyzer.py, drawing_processor.py, import_paddlex_data.py, lighton_ocr.py, local_remarks_generator.py, pdf_annotator.py, pdf_processor(old).py, pdf_processor.py, pp87_checker.py, report_generator.py, run_full_analysis.py, structure_aware_splitter.py, structure_learning.py, agent.py, batch_processor.py, cache_control.py, load_norms.py, norms_references.json, requirements.txt, run_full_analysis.py, start_normocontrol.ps1, test_agent.py, !!!!ocr_processor.py, cache_manager.py, document_analyzer.py, drawing_processor.py, import_paddlex_data.py, local_remarks_generator.py, pdf_annotator.py, pdf_processor.py, report_generator.py, run_full_analysis.py, structure_aware_splitter.py, structure_learning.py
19 июл 2026, 18:22
Верифицирован
19 июл 2026, 18:22ce9f7f3
# app/utils/drawing_processor.py - ПОЛНАЯ ИСПРАВЛЕННАЯ ВЕРСИЯ

import os
from pathlib import Path
from typing import List, Dict, Optional, Union
from PIL import Image, ImageDraw, ImageFont
from reportlab.lib.pagesizes import A4
from reportlab.lib import colors
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, Image as ReportLabImage
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import mm
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
import pdfplumber


def register_russian_font():
    """Регистрирует кириллический шрифт для ReportLab"""
    font_paths = [
        "C:/Windows/Fonts/arial.ttf",
        "C:/Windows/Fonts/times.ttf",
    ]
    for path in font_paths:
        if os.path.exists(path):
            try:
                pdfmetrics.registerFont(TTFont('RussianFont', path))
                return 'RussianFont'
            except:
                continue
    return 'Helvetica'


FONT_NAME = register_russian_font()


def normalize_bbox(bbox) -> Optional[List[int]]:
    """Нормализует bbox в формат [x0, y0, x1, y1] с целыми числами"""
    if not bbox:
        return None
    
    if isinstance(bbox, str):
        import re
        numbers = re.findall(r'\d+', bbox)
        if len(numbers) >= 4:
            return [int(numbers[0]), int(numbers[1]), int(numbers[2]), int(numbers[3])]
        return None
    
    if isinstance(bbox, (list, tuple)) and len(bbox) >= 4:
        result = []
        for val in bbox[:4]:
            try:
                result.append(int(float(val)))
            except (ValueError, TypeError):
                return None
        return result
    
    return None


class DrawingProcessor:
    """Обработчик чертежей с аннотациями на основе bbox"""
    
    COLORS = {
        "critical": (220, 38, 38, 100),
        "critical_border": (220, 38, 38),
        "warning": (249, 115, 22, 100),
        "warning_border": (249, 115, 22)
    }
    
    def __init__(self, output_dir: str = "./output/drawings"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
    
    def extract_page_image(self, pdf_path: str, page_num: int, dpi: int = 150) -> Optional[Image.Image]:
        """Извлекает страницу как изображение"""
        try:
            with pdfplumber.open(pdf_path) as pdf:
                if page_num >= len(pdf.pages):
                    return None
                return pdf.pages[page_num].to_image(resolution=dpi).original
        except Exception as e:
            print(f"   Ошибка извлечения страницы {page_num + 1}: {e}")
            return None
    
    def draw_annotation(self, image: Image.Image, violations: List[Dict]) -> Image.Image:
        """Рисует прямоугольники и сноски на чертеже"""
        draw = ImageDraw.Draw(image, 'RGBA')
        
        try:
            font = ImageFont.truetype("C:/Windows/Fonts/arial.ttf", 12)
            font_bold = ImageFont.truetype("C:/Windows/Fonts/arialbd.ttf", 12)
        except:
            font = ImageFont.load_default()
            font_bold = font
        
        img_width, img_height = image.size
        
        for i, v in enumerate(violations):
            bbox = normalize_bbox(v.get('bbox'))
            
            if not bbox:
                # Fallback: создаём область в зависимости от индекса
                x0 = 50 + (i * 30) % (img_width - 200)
                y0 = 100 + (i * 50) % (img_height - 150)
                x1 = x0 + 150
                y1 = y0 + 30
                bbox = [x0, y0, x1, y1]
                print(f"      ⚠️ Fallback bbox для нарушения {i+1}: {bbox}")
            
            x0, y0, x1, y1 = bbox
            
            severity = v.get('severity', v.get('критичность', 'warning'))
            if severity in ('critical', 'критично'):
                fill = self.COLORS["critical"]
                border = self.COLORS["critical_border"]
            else:
                fill = self.COLORS["warning"]
                border = self.COLORS["warning_border"]
            
            # Рисуем полупрозрачный прямоугольник
            draw.rectangle([x0, y0, x1, y1], fill=fill, outline=border, width=2)
            
            # Рисуем метку
            rule_id = v.get('rule_id', v.get('id', f'ERR-{i+1:03d}'))
            label = f"[{rule_id}]"
            draw.text((x0, y0 - 15), label, fill=border, font=font_bold)
            
            # Линия-выноска
            draw.line([(x0, y0), (x0, y0 - 8)], fill=border, width=1)
        
        return image
    
    def generate_drawing_report(self, drawing: Image.Image, violations: List[Dict],
                                 page_num: int, section_info: str) -> str:
        """Генерирует PDF-отчёт для одного чертежа"""
        output_path = self.output_dir / f"drawing_page_{page_num:03d}.pdf"
        temp_path = self.output_dir / f"temp_drawing_{page_num}.png"
        
        # Сохраняем временное изображение
        drawing.save(temp_path, "PNG")
        
        # Создаём PDF
        doc = SimpleDocTemplate(str(output_path), pagesize=A4)
        story = []
        styles = getSampleStyleSheet()
        
        # Стили с кириллическим шрифтом
        title_style = ParagraphStyle('Title', parent=styles['Heading2'], fontName=FONT_NAME, fontSize=12)
        normal_style = ParagraphStyle('Normal', parent=styles['Normal'], fontName=FONT_NAME, fontSize=9)
        
        # Заголовок страницы
        story.append(Paragraph(f"Чертёж - Страница {page_num}", title_style))
        story.append(Paragraph(section_info, normal_style))
        story.append(Spacer(1, 10))
        
        # Вставляем изображение
        story.append(ReportLabImage(str(temp_path), width=170*mm, height=200*mm))
        story.append(Spacer(1, 10))
        
        # Таблица нарушений
        if violations:
            data = [["ID", "Нарушение", "Нормативный документ"]]
            for v in violations:
                data.append([
                    v.get('rule_id', '?'),
                    v.get('нарушение', '-')[:60],
                    v.get('пункт_норматива', '-')
                ])
            
            table = Table(data, colWidths=[30*mm, 100*mm, 45*mm])
            table.setStyle(TableStyle([
                ('BACKGROUND', (0, 0), (-1, 0), colors.grey),
                ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
                ('FONTNAME', (0, 0), (-1, -1), FONT_NAME),
                ('GRID', (0, 0), (-1, -1), 0.5, colors.black),
                ('FONTSIZE', (0, 0), (-1, 0), 9),
                ('FONTSIZE', (0, 1), (-1, -1), 8),
            ]))
            story.append(table)
        
        # Сборка PDF
        doc.build(story)
        
        # Удаляем временный файл
        if temp_path.exists():
            os.unlink(temp_path)
        
        print(f"      ✅ Отчёт создан: {output_path}")
        return str(output_path)
    
    def process_drawings(self, pdf_path: str, violations_by_page: Dict[int, List[Dict]], 
                          doc_structure) -> List[str]:
        """Обрабатывает все чертежи с аннотациями"""
        reports = []
        
        for page_num, violations in violations_by_page.items():
            print(f"   🎨 Страница {page_num}: {len(violations)} нарушений")
            
            # Проверяем наличие bbox
            has_bbox = any(v.get('bbox') for v in violations)
            if not has_bbox:
                print(f"      ⚠️ Нет bbox для нарушений на странице {page_num}")
            
            # Извлекаем изображение страницы
            img = self.extract_page_image(pdf_path, page_num - 1)
            if img is None:
                print(f"      ⚠️ Не удалось извлечь страницу {page_num}")
                continue
            
            # Рисуем аннотации
            img = self.draw_annotation(img, violations)
            
            # Генерируем отчёт
            report_path = self.generate_drawing_report(img, violations, page_num, f"Страница: {page_num}")
            reports.append(report_path)
        
        return reports


def get_drawing_processor(output_dir: str = "./output/drawings") -> DrawingProcessor:
    """Фабрика для получения экземпляра DrawingProcessor"""
    return DrawingProcessor(output_dir)