model_runner/model_runner.py - main - jonique/scanner_AI

model_runner/model_runner.py
84 строки2 KB
Евгений Щипунов
term_release: final version for term paper done
11 июн 2026, 22:32
11 июн 2026, 22:32ba51fe8
import os
import queue
import threading
from concurrent.futures import ProcessPoolExecutor
from .run_worker import predict
import math
from utils.utils import clean_old_predictions

def chunkify(lst, n):
    if not lst:
        return []

    k = max(1, min(math.ceil(len(lst) / n), 10))
    return [lst[i:i + k] for i in range(0, len(lst), k)]


class ModelRunner:
    def __init__(self, model_path, images_dir):
        self.model_path = model_path
        self.images_dir = images_dir

        self.executor = None
        self.cancelled = False
        self.progress_queue = queue.Queue()
        self.result_thread = None

    def start(self, count):
        clean_old_predictions()

        if count <= 0:
            self.progress_queue.put(("done", 0))
            return self.progress_queue

        count_workers = os.cpu_count() or 1
        workers = min(count_workers, 8, count)

        self.executor = ProcessPoolExecutor(max_workers=workers)

        indexes = list(range(count))
        chunks = chunkify(indexes, workers)

        futures = [
            self.executor.submit(
                predict,
                self.model_path,
                self.images_dir,
                chunk
            )
            for chunk in chunks
        ]

        def result_collector():
            real_processed = 0

            for f in futures:
                if self.cancelled:
                    break

                try:
                    processed = f.result()
                    real_processed += processed

                    for _ in range(processed):
                        self.progress_queue.put(1)

                except Exception as e:
                    print(f"Model inference chunk failed: {e}")

            self.progress_queue.put(("done", real_processed))

        self.result_thread = threading.Thread(
            target=result_collector,
            daemon=True
        )
        self.result_thread.start()

        return self.progress_queue

    def cancel(self):
        self.cancelled = True

        if self.executor:
            self.executor.shutdown(wait=False, cancel_futures=True)
            self.executor = None