utils/utils.py - main - jonique/scanner_AI

utils/utils.py
172 строки4 KB
Евгений Щипунов
term_release: final version for term paper done
11 июн 2026, 22:32
11 июн 2026, 22:32ba51fe8
from config import config
import shutil
import os
from pathlib import Path

def clean_old_images():
    images_folder = config["imagesSaveDirectory"]
    labels_folder = config["labelsSaveDirectory"]

    if os.path.exists(images_folder):
        shutil.rmtree(images_folder)
    os.makedirs(images_folder, exist_ok=True)
    if os.path.exists(labels_folder):
        shutil.rmtree(labels_folder)
    os.makedirs(labels_folder, exist_ok=True)

def clean_old_predictions():
    prediction_folder = config["predictionsDirectory"]
    if os.path.exists(prediction_folder):
        shutil.rmtree(prediction_folder)
    os.makedirs(prediction_folder, exist_ok=True)

def class_name_to_russian(engilsh_name):
    russian_dict = { 
        "Baton": "Дубинка",
        "Plier": "Клещи",
        "Hammer": "Молоток",
        "Powerbank": "Пауэрбанк",
        "Scissors": "Ножницы",
        "Wrench": "Гаечный ключ",
        "Gun": "Пистолет",
        "Bullet": "Пуля",
        "Sprayer": "Распылитель",
        "HandCuffs": "Наручники",
        "Knife": "Нож",
        "Lighter": "Зажигалка"
    }

    return russian_dict[engilsh_name]

def index_to_classname(index):
    names_dict = { 
        0: "Baton",
        1: "Plier",
        2: "Hammer",
        3: "Powerbank",
        4: "Scissors",
        5: "Wrench",
        6: "Gun",
        7: "Bullet",
        8: "Sprayer",
        9: "HandCuffs",
        10: "Knife",
        11: "Lighter",
        12: "Bag",
        100: "Clothes",
        101: "Water_container",
        102: "Book",
        103: "Shoe",
        104: "Small_item",
        105: "Small_tech",
    }
    return names_dict[index]

def class_name_to_index(name):
    names_dict = { 
        "Baton": 0,
        "Plier": 1,
        "Hammer": 2,
        "Powerbank": 3,
        "Scissors": 4,
        "Wrench": 5,
        "Gun": 6,
        "Bullet": 7,
        "Sprayer": 8,
        "HandCuffs": 9,
        "Knife": 10,
        "Lighter": 11
    }
    return names_dict[name]

def read_yolo_image(path, w, h):
    items = []
    if not os.path.exists(path):
        return items

    with open(path) as f:
        for line in f:
            cls, x, y, bw, bh = map(float, line.split())

            x1 = int((x - bw / 2) * w)
            y1 = int((y - bh / 2) * h)
            x2 = int((x + bw / 2) * w)
            y2 = int((y + bh / 2) * h)

            items.append((int(cls), (x1, y1, x2, y2)))

    return items


def calculate_mean_iou(gt, pred):
    if not gt or not pred:
        return 0.0

    def iou(a, b):
        xA = max(a[0], b[0])
        yA = max(a[1], b[1])
        xB = min(a[2], b[2])
        yB = min(a[3], b[3])
        inter = max(0, xB - xA) * max(0, yB - yA)
        areaA = (a[2] - a[0]) * (a[3] - a[1])
        areaB = (b[2] - b[0]) * (b[3] - b[1])
        return inter / (areaA + areaB - inter + 1e-6)

    total = 0
    for g in gt:
        total += max(iou(g, p) for p in pred)

    return total / len(gt)

def iou(a, b):
    ax, ay, aw, ah = a
    bx, by, bw, bh = b

    ax1, ay1 = ax - aw/2, ay - ah/2
    ax2, ay2 = ax + aw/2, ay + ah/2

    bx1, by1 = bx - bw/2, by - bh/2
    bx2, by2 = bx + bw/2, by + bh/2

    inter_x1 = max(ax1, bx1)
    inter_y1 = max(ay1, by1)
    inter_x2 = min(ax2, bx2)
    inter_y2 = min(ay2, by2)

    inter = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)

    area_a = aw * ah
    area_b = bw * bh

    return inter / (area_a + area_b - inter + 1e-6)

def read_yolo(path):
    boxes = []
    if not os.path.exists(path):
        return boxes

    with open(path) as f:
        for line in f:
            cls, x, y, w, h = map(float, line.split())
            boxes.append((int(cls), (x, y, w, h)))
    return boxes

def compact_generated_files(images_dir, labels_dir):
    image_files = sorted(
        Path(images_dir).glob("*.png"),
        key=lambda p: int(p.stem)
    )

    for new_i, img_path in enumerate(image_files):
        old_i = img_path.stem

        new_img = Path(images_dir) / f"{new_i}.png"
        old_label = Path(labels_dir) / f"{old_i}.txt"
        new_label = Path(labels_dir) / f"{new_i}.txt"

        img_path.rename(new_img)

        if old_label.exists():
            old_label.rename(new_label)

    return len(image_files)