database/metadata.py

import pickle
import os


class MetaData:
    def __init__(self):
        self.records = []

    def create(self, image_paths):
        self.records = []
        self.add(image_paths)

    def add(self, image_paths):
        for image_path in image_paths:
            self.records.append({"image_path": os.path.basename(image_path)})

    def get_record(self, index):
        return self.records[index]

    def get_records(self, indices):
        return [self.records[index] for index in indices if index != -1]

    def save(self, file_path):
        with open(file_path, "wb") as f:
            pickle.dump(self.records, f)

    def load(self, file_path):
        if os.path.exists(file_path):
            with open(file_path, "rb") as f:
                self.records = pickle.load(f)
        else:
            self.records = []

    def size(self):
        return len(self.records)

    def get_image_paths(self):
        return {record['image_path'] for record in self.records}