<gh_stars>1-10
import math
from copy import deepcopy
from os import path
from typing import Union

import cv2
import numpy as np
import pytesseract
from pydantic import BaseModel, HttpUrl

import circuit_recognizer.annotations as annotations
from circuit_recognizer.connect import clean_image, get_start_lines
from circuit_recognizer.utils import get_annotation_source_image


class AnnotationImage(BaseModel):
    source: Union[HttpUrl, str]
    content: np.ndarray

    class Config:
        allow_mutation = False
        arbitrary_types_allowed = True


def image_snippet(image, annotation):
    region = annotation.target["selector"]["value"]
    x, y, w, h = [math.floor(float(i)) for i in region.split(":")[1].split(",")]
    snippet = image[y : y + h, x : x + w]
    return snippet


def image_label(annotation):
    return annotation.body[0]["value"]


def save_labeled_snippets(labeled_snippets):
    for (snippet, label) in labeled_snippets:
        cv2.imwrite(f"{label}.png", snippet)


def find_template_matches(image, template, threshold):
    templates = [
        (template, 0),
        (cv2.rotate(template, cv2.ROTATE_90_COUNTERCLOCKWISE), 90),
        (cv2.rotate(template, cv2.ROTATE_180), 180),
        (cv2.rotate(template, cv2.ROTATE_90_CLOCKWISE), 270),
    ]
    for (template, rotation) in templates:
        res = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
        loc = np.where(res >= threshold)
        for top_left in zip(*loc[::-1]):
            confidence = res[top_left[1], top_left[0]]
            yield (rotation, top_left, confidence)


def find_matches(annotation_image, annotation_snippets, threshold=0.72):
    matches = []
    for (annotation, template) in annotation_snippets:
        for (rotation, top_left, conf) in find_template_matches(
            annotation_image.content, template, threshold
        ):
            new_annotation = annotations.Annotation.from_annotation(
                annotation, confidence=conf
            )
            new_annotation.target["source"] = annotation_image.source
            new_annotation.rotate(rotation)
            new_annotation.move_to(*top_left)
            matches.append(new_annotation)
    return matches


def find_unique_matches(annotation_image, annotation_snippets, threshold=0.72):
    annotations = find_matches(annotation_image, annotation_snippets, threshold)
    if len(annotations) == 0:
        return []

    group = [annotations.pop(0)]
    grouped_annotations = [group]
    is_in_group = lambda group, ann: next(
        (True for member in group if ann.isclose(member)), False
    )

    for annotation in annotations:
        group = next(
            (group for group in grouped_annotations if is_in_group(group, annotation)),
            None,
        )
        if group is None:
            group = []
            grouped_annotations.append(group)

        group.append(annotation)

    results = [group[int(len(group) / 2)] for group in grouped_annotations]
    return results


def recognize_components(src_image, annotations, target_images=None):
    if target_images is None:
        target_images = [src_image]

    annotated_snippets = [
        (annotation, image_snippet(src_image.content, annotation))
        for annotation in annotations
    ]
    return (find_matches(image, annotated_snippets) for image in target_images)


def recognize_ports(src_image, annotation):
    image = clean_image(src_image)
    min_length = 5
    start_lines = get_start_lines(image, annotation, min_length)
    port_rects = []

    for line in start_lines:
        r_line = line.reverse()

        box = annotation.bounding_box(inclusive=True)
        r_line.remove_overlap_into(box)

        line.end = r_line.end
        rect = line.expand(4)
        for i in range(len(port_rects) - 1, -1, -1):
            existing = port_rects[i]
            if rect.intersect(existing) is not None:
                rect = rect.union(existing)
                port_rects.pop(i)

        port_rects.append(rect)

    ports = [
        annotations.Annotation.from_bbox(
            bbox=rect.to_bbox(),
            parent=annotation,
            creator={
                "id": "circuit_recognizer.recognize_ports",
                "type": "Agent",
                "name": "circuit_recognizer.recognize_ports",
            },
        )
        for rect in port_rects
    ]

    return ports


def ocr_on_selection(selection, threshold=0.6):
    """Given a selection (like a rectangular fragment in an image), perform OCR with default settings.

    Parameters
    ----------
    selection: dict
        The selection dictionary. Generally, it should contain information about the underlying image
        and the fragment selector to perform OCR on
    threshold: float, default=0.6
        The threshold value for confidence, results with confidence below threshold value are discarded

    Returns
    -------
    annotations.Annotation or None
        If the OCR is successful, return an annotation with populated text from the selection
        If the OCR wasn't successful, return None
    """
    annotation = annotations.Annotation.from_selection(selection)
    image = get_annotation_source_image(annotation)
    selected_snippet = image_snippet(image, annotation)
    ocr_results = snippet_ocr(selected_snippet, threshold=threshold)
    if ocr_results is not None:
        predicted_text = ocr_results[0].strip().replace("\n", "")
        confidence = ocr_results[1]
        return predicted_text, confidence


def _text_conf_from_tess_dict(data):
    """From the results of pytesseract data dictionary, return joined text and avg confidence."""
    text_chunks = []
    confs = []
    for conf, text in zip(data["conf"], data["text"]):
        if float(conf) < 0.0 or not text:
            continue
        else:
            text_chunks.append(text)
            confs.append(float(conf))
    return " ".join(text_chunks), sum(confs) / (100 * len(confs)) if len(confs) else 0.0


def snippet_ocr(
    image, langs=None, config=None, threshold=0.6, pad=True, retry=True, scale=1.0
):
    """Given an image-snippet perform OCR using tesseract.

    Parameters
    ----------
    image: np.ndarray
        The read image by cv2.imread function.
    langs: optional, str, default=None
        The languages to use. If not provided, English ("eng") and Greek ("ell") are used.
    config: optional, str, default=None
        The tesseract config string, if not provided, r"--psm 12" is used.
    threshold: optional, float, default=0.6
        The threshold value for confidence, results with confidence below threshold value are discarded
    pad: optional, bool, default=True
        If true extend the image boundaries by 20 white pixels
    retry: optional, bool, default=True
        If true, retry the OCR with image scaled with double the current scale
    scale: float, default=1.0
        Scale the image by this value before passing it to the OCR engine

    Returns
    -------
    tuple or None
        The pair (text, confidence) as a result of OCR on selection if the confidence > threshold else None
    """
    if pad:
        image_copy = cv2.copyMakeBorder(
            image, 20, 20, 20, 20, cv2.BORDER_CONSTANT, value=(255, 255, 255)
        )
    else:
        image_copy = image[:, :, :]

    if scale != 1.0:
        width = int(image_copy.shape[1] * scale)
        height = int(image_copy.shape[0] * scale)
        image_copy = cv2.resize(image_copy, (width, height))

    data = ocr(image_copy, config=config, langs=langs)

    text, confidence = _text_conf_from_tess_dict(data)
    if confidence > threshold:
        return text, confidence
    elif retry:
        return snippet_ocr(
            image, langs, config, threshold, pad=False, retry=False, scale=scale * 2.0
        )


def ocr(image, config=None, langs=None):
    """Perform OCR using tesseract."""
    return pytesseract.image_to_data(
        image,
        lang=langs or "ell+eng",
        config=config or r"--psm 12",
        output_type=pytesseract.Output.DICT,
    )

import math
from copy import deepcopy
from os import path
from typing import Union

import cv2
import numpy as np
import pytesseract
from pydantic import BaseModel, HttpUrl

import circuit_recognizer.annotations as annotations
from circuit_recognizer.connect import clean_image, get_start_lines
from circuit_recognizer.utils import get_annotation_source_image


class AnnotationImage(BaseModel):
    source: Union[HttpUrl, str]
    content: np.ndarray

    class Config:
        allow_mutation = False
        arbitrary_types_allowed = True


def image_snippet(image, annotation):
    region = annotation.target["selector"]["value"]
    x, y, w, h = [math.floor(float(i)) for i in region.split(":")[1].split(",")]
    snippet = image[y : y + h, x : x + w]
    return snippet


def image_label(annotation):
    return annotation.body[0]["value"]


def save_labeled_snippets(labeled_snippets):
    for (snippet, label) in labeled_snippets:
        cv2.imwrite(f"{label}.png", snippet)


def find_template_matches(image, template, threshold):
    templates = [
        (template, 0),
        (cv2.rotate(template, cv2.ROTATE_90_COUNTERCLOCKWISE), 90),
        (cv2.rotate(template, cv2.ROTATE_180), 180),
        (cv2.rotate(template, cv2.ROTATE_90_CLOCKWISE), 270),
    ]
    for (template, rotation) in templates:
        res = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
        loc = np.where(res >= threshold)
        for top_left in zip(*loc[::-1]):
            confidence = res[top_left[1], top_left[0]]
            yield (rotation, top_left, confidence)


def find_matches(annotation_image, annotation_snippets, threshold=0.72):
    matches = []
    for (annotation, template) in annotation_snippets:
        for (rotation, top_left, conf) in find_template_matches(
            annotation_image.content, template, threshold
        ):
            new_annotation = annotations.Annotation.from_annotation(
                annotation, confidence=conf
            )
            new_annotation.target["source"] = annotation_image.source
            new_annotation.rotate(rotation)
            new_annotation.move_to(*top_left)
            matches.append(new_annotation)
    return matches


def find_unique_matches(annotation_image, annotation_snippets, threshold=0.72):
    annotations = find_matches(annotation_image, annotation_snippets, threshold)
    if len(annotations) == 0:
        return []

    group = [annotations.pop(0)]
    grouped_annotations = [group]
    is_in_group = lambda group, ann: next(
        (True for member in group if ann.isclose(member)), False
    )

    for annotation in annotations:
        group = next(
            (group for group in grouped_annotations if is_in_group(group, annotation)),
            None,
        )
        if group is None:
            group = []
            grouped_annotations.append(group)

        group.append(annotation)

    results = [group[int(len(group) / 2)] for group in grouped_annotations]
    return results


def recognize_components(src_image, annotations, target_images=None):
    if target_images is None:
        target_images = [src_image]

    annotated_snippets = [
        (annotation, image_snippet(src_image.content, annotation))
        for annotation in annotations
    ]
    return (find_matches(image, annotated_snippets) for image in target_images)


def recognize_ports(src_image, annotation):
    image = clean_image(src_image)
    min_length = 5
    start_lines = get_start_lines(image, annotation, min_length)
    port_rects = []

    for line in start_lines:
        r_line = line.reverse()

        box = annotation.bounding_box(inclusive=True)
        r_line.remove_overlap_into(box)

        line.end = r_line.end
        rect = line.expand(4)
        for i in range(len(port_rects) - 1, -1, -1):
            existing = port_rects[i]
            if rect.intersect(existing) is not None:
                rect = rect.union(existing)
                port_rects.pop(i)

        port_rects.append(rect)

    ports = [
        annotations.Annotation.from_bbox(
            bbox=rect.to_bbox(),
            parent=annotation,
            creator={
                "id": "circuit_recognizer.recognize_ports",
                "type": "Agent",
                "name": "circuit_recognizer.recognize_ports",
            },
        )
        for rect in port_rects
    ]

    return ports


def ocr_on_selection(selection, threshold=0.6):
    """Given a selection (like a rectangular fragment in an image), perform OCR with default settings.

    Parameters
    ----------
    selection: dict
        The selection dictionary. Generally, it should contain information about the underlying image
        and the fragment selector to perform OCR on
    threshold: float, default=0.6
        The threshold value for confidence, results with confidence below threshold value are discarded

    Returns
    -------
    annotations.Annotation or None
        If the OCR is successful, return an annotation with populated text from the selection
        If the OCR wasn't successful, return None
    """
    annotation = annotations.Annotation.from_selection(selection)
    image = get_annotation_source_image(annotation)
    selected_snippet = image_snippet(image, annotation)
    ocr_results = snippet_ocr(selected_snippet, threshold=threshold)
    if ocr_results is not None:
        predicted_text = ocr_results[0].strip().replace("\n", "")
        confidence = ocr_results[1]
        return predicted_text, confidence


def _text_conf_from_tess_dict(data):
    """From the results of pytesseract data dictionary, return joined text and avg confidence."""
    text_chunks = []
    confs = []
    for conf, text in zip(data["conf"], data["text"]):
        if float(conf) < 0.0 or not text:
            continue
        else:
            text_chunks.append(text)
            confs.append(float(conf))
    return " ".join(text_chunks), sum(confs) / (100 * len(confs)) if len(confs) else 0.0


def snippet_ocr(
    image, langs=None, config=None, threshold=0.6, pad=True, retry=True, scale=1.0
):
    """Given an image-snippet perform OCR using tesseract.

    Parameters
    ----------
    image: np.ndarray
        The read image by cv2.imread function.
    langs: optional, str, default=None
        The languages to use. If not provided, English ("eng") and Greek ("ell") are used.
    config: optional, str, default=None
        The tesseract config string, if not provided, r"--psm 12" is used.
    threshold: optional, float, default=0.6
        The threshold value for confidence, results with confidence below threshold value are discarded
    pad: optional, bool, default=True
        If true extend the image boundaries by 20 white pixels
    retry: optional, bool, default=True
        If true, retry the OCR with image scaled with double the current scale
    scale: float, default=1.0
        Scale the image by this value before passing it to the OCR engine

    Returns
    -------
    tuple or None
        The pair (text, confidence) as a result of OCR on selection if the confidence > threshold else None
    """
    if pad:
        image_copy = cv2.copyMakeBorder(
            image, 20, 20, 20, 20, cv2.BORDER_CONSTANT, value=(255, 255, 255)
        )
    else:
        image_copy = image[:, :, :]

    if scale != 1.0:
        width = int(image_copy.shape[1] * scale)
        height = int(image_copy.shape[0] * scale)
        image_copy = cv2.resize(image_copy, (width, height))

    data = ocr(image_copy, config=config, langs=langs)

    text, confidence = _text_conf_from_tess_dict(data)
    if confidence > threshold:
        return text, confidence
    elif retry:
        return snippet_ocr(
            image, langs, config, threshold, pad=False, retry=False, scale=scale * 2.0
        )


def ocr(image, config=None, langs=None):
    """Perform OCR using tesseract."""
    return pytesseract.image_to_data(
        image,
        lang=langs or "ell+eng",
        config=config or r"--psm 12",
        output_type=pytesseract.Output.DICT,
    )
Given a selection (like a rectangular fragment in an image), perform OCR with default settings.      Parameters     ----------     selection: dict         The selection dictionary. Generally, it should contain information about the underlying image         and the fragment selector to perform OCR on     threshold: float, default=0.6         The threshold value for confidence, results with confidence below threshold value are discarded      Returns     -------     annotations.Annotation or None         If the OCR is successful, return an annotation with populated text from the selection         If the OCR wasn't successful, return None From the results of pytesseract data dictionary, return joined text and avg confidence. Given an image-snippet perform OCR using tesseract.      Parameters     ----------     image: np.ndarray         The read image by cv2.imread function.     langs: optional, str, default=None         The languages to use. If not provided, English ("eng") and Greek ("ell") are used.     config: optional, str, default=None         The tesseract config string, if not provided, r"--psm 12" is used.     threshold: optional, float, default=0.6         The threshold value for confidence, results with confidence below threshold value are discarded     pad: optional, bool, default=True         If true extend the image boundaries by 20 white pixels     retry: optional, bool, default=True         If true, retry the OCR with image scaled with double the current scale     scale: float, default=1.0         Scale the image by this value before passing it to the OCR engine      Returns     -------     tuple or None         The pair (text, confidence) as a result of OCR on selection if the confidence > threshold else None Perform OCR using tesseract.