generators/generator.py - main - jonique/scanner_AI

generators/generator.py
123 строки4 KB
Евгений Щипунов
term_release: final version for term paper done
11 июн 2026, 22:32
11 июн 2026, 22:32ba51fe8
import os
import queue
from concurrent.futures import ProcessPoolExecutor
from utils.utils import clean_old_images, class_name_to_index, compact_generated_files
from generators.generator_2d import generate_2d_image
from generators.generator_3d import generate_3d_image
from config import config
import random
import threading

AVAILABLE_BW_CLASSES = { 1, 4, 5, 6, 10 }

class Generator:
    def __init__(self, script_number, bw_prob, classes_distr, obj_number_distr):
        self.executor = None
        self.cancelled = False
        self.script_number = script_number
        self.bw_prob = bw_prob
        self.classes_distr=classes_distr
        self.objects_distr=obj_number_distr

        self.progress_queue = queue.Queue()
        self.result_thread = None

    def generate_image(self, count):
        clean_old_images()

        self.executor = None
        if self.script_number == 1:
            workers = min(os.cpu_count() or 1, 8)
            self.executor = ProcessPoolExecutor(max_workers=workers)

            futures = []

            for i in range(count):
                scene = self.generate_scene()
                futures.append(
                    self.executor.submit(generate_2d_image, scene, i)
                )

            def result_collector():
                for f in futures:
                    if self.cancelled:
                        break

                    try:
                        f.result()
                        self.progress_queue.put(1)
                    except Exception as e:
                        print(f"Image generation failed: {e}")

                real_count = compact_generated_files(
                    config["imagesSaveDirectory"],
                    config["labelsSaveDirectory"]
                )

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

            threading.Thread(target=result_collector, daemon=True).start()
            return self.progress_queue
        else:
            def worker():
                success_index = 0

                for _ in range(count):
                    if self.cancelled:
                        break

                    scene = self.generate_scene()

                    try:
                        generate_3d_image(scene, success_index)
                        success_index += 1
                        self.progress_queue.put(1)
                    except Exception as e:
                        print(f"3D image generation failed: {e}")

                real_count = compact_generated_files(
                    config["imagesSaveDirectory"],
                    config["labelsSaveDirectory"]
                )

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

            threading.Thread(target=worker, daemon=True).start()
            return self.progress_queue
        
    def generate_scene(self):
        is_bw = random.randint(0, 100) < self.bw_prob
        if self.script_number != 1:
            is_bw = 0

        objects_count = random.choices(
            [0, 1, 2, 3],
            weights=self.objects_distr
        )[0]

        classes = list(self.classes_distr.keys())
        weights = list(self.classes_distr.values())

        if is_bw:
            filtered = [
                (cls, w)
                for cls, w in zip(classes, weights)
                if class_name_to_index(cls) in AVAILABLE_BW_CLASSES
            ]
            classes, weights = zip(*filtered)

        objects = []
        for _ in range(objects_count):
            obj_class = random.choices(classes, weights=weights)[0]
            objects.append({"class": class_name_to_index(obj_class)})

        return {
            "is_bw": is_bw,
            "objects": objects
        }

    def cancel(self):
        self.cancelled = True

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