backend/prepare_data.py - develop - kuklindal/RealEstateAnalyzer

backend/prepare_data.py

87 строк2 KB

Daniil Kuklin

refactor: improve code style, async handling and project documentation

08 май 2026, 13:19

08 май 2026, 13:19f87acd4

import json
from ocr import extract_text


def clean_text(text):
    text = text.lower()
    text = text.replace("\n", " ")

    text = text.replace(",", "")
    text = text.replace(".", "")
    text = text.replace(":", "")
    text = text.replace(";", "")

    text = text.replace("ул.", "ул")
    text = text.replace("д.", "д")
    text = text.replace("кв.", "кв")

    text = text.replace("  ", " ")

    return text.strip()


def normalize_value(label, value):
    value = str(value).lower()

    value = value.replace(",", "")
    value = value.replace(".", "")
    value = value.replace(":", "")
    value = value.replace(";", "")

    value = value.replace("ул.", "ул")
    value = value.replace("д.", "д")
    value = value.replace("кв.", "кв")

    value = value.replace("  ", " ")

    return value.strip()


def find_span(text, value, label=None):
    text_clean = clean_text(text)
    value_clean = clean_text(str(value))

    start = text_clean.find(value_clean)

    if start != -1:
        return (start, start + len(value_clean))

    return None


def convert_dataset(json_path, images_path):
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    dataset = []

    for item in data:
        # ✅ ИСПРАВЛЕНО ЗДЕСЬ
        image_path = f"{images_path}/{item['image']}"

        text = extract_text(image_path)
        clean = clean_text(text)

        entities = []

        for label, value in item["entities"].items():
            span = find_span(text, value, label.upper())

            if span:
                entities.append((span[0], span[1], label.upper()))
            else:
                print(f"⚠️ НЕ НАЙДЕНО: {label} -> {value}")

        dataset.append((clean, {"entities": entities}))

    return dataset


if __name__ == "__main__":
    train_data = convert_dataset(
        "data/annotations/train.json",
        "data/images/train"
    )

    print("\n=== SAMPLE ===")
    print(train_data[0])