utils_pic.py

from pydantic import BaseModel
from fastapi import UploadFile
from PIL import Image
import io
import cv2
import numpy as np
import pytesseract
from ultralytics import YOLO
import networkx as nx
import os


class PicToUmlRequest(BaseModel):
    is_bpmn: bool = True  # True → BPMN (текстовое описание), False → UML (PlantUML)


class PicToUmlResponse(BaseModel):
    result: str  # описание или PlantUML код
    debug_info: dict = {}  # опционально для отладки


# Глобальные модели
yolo_model = None


def load_yolo():
    global yolo_model
    if yolo_model is None:
        weights_path = "/home/user/app/weights/bestt.pt"
        if not os.path.exists(weights_path):
            raise FileNotFoundError(f"YOLO weights not found: {weights_path}")

        yolo_model = YOLO(weights_path)
        print(f"YOLOv8n загружена из {weights_path}")
    return yolo_model


# ... (остальной код без изменений)

def preprocess_image(image_bytes: bytes, max_size=1024):
    img = Image.open(io.BytesIO(image_bytes))
    img.thumbnail((max_size, max_size))  # уже 1024
    img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    return img_cv, img



def run_ocr(pil_img):
    # psm 6 — Assume a single uniform block of text
    config = r'--oem 1 --psm 6'
    text = pytesseract.image_to_string(pil_img, lang='eng+rus', config=config)
    return text.strip()


def build_graph_from_detections(detections, ocr_text):
    G = nx.DiGraph()
    nodes = []

    if not detections or len(detections) == 0 or not detections[0].boxes:
        print("No detections — fallback to OCR only")
        return G, nodes

    for result in detections:
        for box, cls, conf in zip(result.boxes.xyxy, result.boxes.cls, result.boxes.conf):
            box_np = box.cpu().numpy()
            center = ((box_np[0] + box_np[2]) / 2, (box_np[1] + box_np[3]) / 2)
            nodes.append({
                "center": center,
                "class": int(cls),
                "conf": float(conf),
                "box": box_np
            })

    if not nodes:
        return G, nodes

    # Сортировка по y (сверху вниз), затем по x (слева направо)
    nodes.sort(key=lambda n: (n["center"][1], n["center"][0]))

    # Простые последовательные связи + возможные ветвления
    for i in range(len(nodes) - 1):
        G.add_edge(i, i + 1)

    # Можно добавить логику для gateways (если class == gateway — проверить ближайшие 2–3)

    return G, nodes


def generate_description_from_graph(G, nodes, ocr_text, is_bpmn):
    # Здесь LLM превращает граф + OCR в описание
    nodes_str = "\n".join([f"Node {i}: class={n['class']}, center={n['center']}" for i, n in enumerate(nodes)])
    edges_str = "\n".join([f"{u} → {v}" for u, v in G.edges()])

    prompt = f"""Ты эксперт по извлечению семантики из диаграмм.
На основе YOLO-детекции элементов, OCR-текста и построенного графа опиши процесс.

OCR текст:
{ocr_text}

Элементы (nodes):
{nodes_str}

Связи (edges):
{edges_str}

Тип диаграммы: {'BPMN' if is_bpmn else 'UML'}

Для BPMN: выведи последовательность в формате "Актор: действие", с (start) и (finish)
Для UML: сразу сгенерируй полный PlantUML код activity/sequence/class в зависимости от элементов.

Выводи ТОЛЬКО результат, без пояснений."""

    return prompt


def get_pic_prompt(ocr_text: str, nodes: list, edges: list, is_bpmn: bool) -> list[dict]:
    system = """Ты эксперт по восстановлению диаграмм из OCR и детекции объектов.
На основе извлечённого текста и списка элементов/связей восстанови структуру диаграммы.

Правила:
- Если is_bpmn: выводи описание в формате "Актор: действие", группируй по пулам/ролям, используй (start) и (finish), указывай ветвления.
- Если UML: выводи ТОЛЬКО полный PlantUML код (@startuml ... @enduml), без лишнего текста.
- Используй swimlanes (|Актор|) для BPMN, если роли видны.
- Выводи ТОЛЬКО результат — никаких пояснений."""

    messages = [{"role": "system", "content": system}]

    examples = [
        {
            "user": """OCR: Клиент: Потребность оформить заказ\nСистема Додо: Проверка адреса\n...
Nodes: class=0 (start), class=2 (task), class=3 (gateway)
Edges: 0→1, 1→2, 2→3
is_bpmn: true""",
            "assistant": """Клиент: (start Потребность оформить заказ)
Клиент: Адрес и способ доставки
Система Додо: Проверка адреса и доступности доставки
... (и т.д. до finish)"""
        },
        {
            "user": """OCR: User -> System: login\nSystem --> User: success
Nodes: class=actor, class=message
Edges: 0→1
is_bpmn: false""",
            "assistant": """@startuml
actor User
participant System
User -> System: login
System --> User: success
@enduml"""
        }
    ]

    for ex in examples:
        messages.append({"role": "user", "content": ex["user"]})
        messages.append({"role": "assistant", "content": ex["assistant"]})

    current = """OCR текст:
    """ + ocr_text + """

    Элементы (nodes):
    """ + '\n'.join(
        [f"Node {i}: class={n['class']}, center={n['center']}, conf={n['conf']:.2f}" for i, n in enumerate(nodes)]) + """

    Связи (edges):
    """ + '\n'.join([f"{u} → {v}" for u, v in edges]) + """

    Тип: """ + ("BPMN (группируй по пулам/акторам)" if is_bpmn else "UML (сгенерируй PlantUML код)")

    messages.append({"role": "user", "content": current})
    return messages