Function process_item

annotation_model/transform_annotation_data.py:27–66

This function will be executed by each worker process.

(args)

def process_item(args):
    """This function will be executed by each worker process."""
    uid, metadata, source_folder, lmdb_path, lock = args

    img_folder = os.path.join(source_folder, uid)
    if not os.path.exists(img_folder):
        return  # No images, skip

    images = sorted(glob.glob(f"{img_folder}/*.png"))
    if len(images) != 40:
        return  # Not enough images, skip

    img_list = []
    for img_path in images:
        try:
            img = Image.open(img_path).convert("RGB")
        except UnidentifiedImageError:
            print(f"Error opening image")
            return
        except Exception as e:
            print(f"Error")
            return

        img_byte_arr = io.BytesIO()
        img.save(img_byte_arr, format="PNG")
        img_list.append(img_byte_arr.getvalue())

    metadata_values = [int(m) for m in metadata]
    metadata_tensor = torch.tensor(metadata_values)

    key = uid.encode("ascii")
    value = pickle.dumps((uid, img_list, metadata_tensor))

    # Locking LMDB access to avoid race conditions
    with lock:
        map_size = int(1e12)  # Size of the database
        db = lmdb.open(lmdb_path, map_size=map_size)
        with db.begin(write=True) as txn:
            txn.put(key, value)
        db.close()


def create_lmdb_dataset(source_folder, lmdb_path, all_metadata):