flandre/flandre/utils/archive.py
2025-06-10 20:35:01 +08:00

115 lines
3.2 KiB
Python

import hashlib
import shutil
import subprocess
import zipfile
from pathlib import Path
import zstd
TEMP_FOLDER = Path("/mnt/16T/private_dataset/New Folder/temp")
def folder_to_zip(folder: Path):
shutil.rmtree(TEMP_FOLDER)
TEMP_FOLDER.mkdir(parents=True, exist_ok=True)
from flandre.utils.RfMeta import RfFrameMeta
for i, file in enumerate(folder.glob("*")):
file = Path(file)
b2b = hashlib.blake2b(file.read_bytes(), digest_size=4).hexdigest()
meta = RfFrameMeta.from_path(file)
meta.blake2b = b2b
print(i, b2b, file, meta.json_str)
src = file
dst = TEMP_FOLDER / f"{i}.zst"
dstj = TEMP_FOLDER / f"{i}.json"
dstj.write_text(meta.json_str)
subprocess.run(["zstd", "-f", src, "-o", dst])
subprocess.run(
[
"zip",
"-0",
"-j",
"-r",
TEMP_FOLDER.parent / f"{folder.name}.zip",
TEMP_FOLDER,
]
)
def to_zip(
li: list[tuple[Path, "RfFrameMeta", list[tuple[Path, str]]]],
temp_dst: Path,
zipdst: Path,
):
shutil.rmtree(temp_dst)
temp_dst.mkdir(parents=True, exist_ok=True)
for i, (file, meta, farr) in enumerate(li):
b2b = hashlib.blake2b(file.read_bytes(), digest_size=4).hexdigest()
meta.blake2b = b2b
src = file
dst = temp_dst / f"{i}.zst"
subprocess.run(["zstd", "-f", src, "-o", dst])
dstj = temp_dst / f"{i}.meta"
dstj.write_text(meta.name)
for srcf, suffix in farr:
dstf = temp_dst / f"{i}{suffix}"
shutil.copy(srcf, dstf)
subprocess.run(["zip", "-0", "-j", "-r", zipdst, temp_dst])
def zip_to_bytes(file: Path, name: int):
return zstd.loads(zipfile.ZipFile(file).open(f"{name}.zst").read())
def zip_to_bytes2(file: Path, name: str):
return zstd.loads(zipfile.ZipFile(file).open(name).read())
if __name__ == "__main__":
from flandre.utils.RfMeta import RfFrameMeta
# folder_to_zip()
arr = []
for i, f in enumerate(
Path("/mnt/16T/private_dataset/New Folder/T1,U=30,M=FMC,S=(256 256 4502)").glob(
"*"
)
):
m = RfFrameMeta.from_path(f)
m.encoder = 0
arr.append((Path(f), m, []))
arr.sort(key=lambda item: item[1].robot_x)
for i, e in enumerate(arr):
e[1].sequence_id = i
to_zip(
arr,
TEMP_FOLDER,
Path("/mnt/16T/private_dataset/T1,U=30,M=TFM,S=(256 256 4502)"),
)
# to_zip()
# pass
# cli()
# folder_to_zip(Path('/mnt/16T/private_dataset/us/R1,U=30,M=PWI,S=(256 1502)'))
# r = get(Path('/mnt/16T/private_dataset/New Folder/ST,U=60,M=FMC,S=(256 256 3002).zip'), 3)
# print(r.__len__())
# r = Path('/mnt/16T/private_dataset/New Folder/ST,U=60,M=FMC,S=(256 256 3002)/X=-204.bin').read_bytes()
# print(r.__len__())
# p = Path('/mnt/16T/private_dataset/New Folder/R1,U=90,M=PWI,S=(256 4502)')
# #
# archive = zipfile.ZipFile(p)
# for item in archive.infolist():
# print(int(Path(item.filename).stem))
# print()
# # imgfile = archive.open('img_01.png')
# print(archive.filelist)
# p2 = archive.open('X=20,Y=395.bin')
# print(p2.read().__len__())