forked from tangger/lerobot
55 lines
1.3 KiB
Python
55 lines
1.3 KiB
Python
import io
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
import torch
|
|
import tqdm
|
|
|
|
|
|
def download_and_extract_zip(url: str, destination_folder: Path) -> bool:
|
|
print(f"downloading from {url}")
|
|
response = requests.get(url, stream=True)
|
|
if response.status_code == 200:
|
|
total_size = int(response.headers.get("content-length", 0))
|
|
progress_bar = tqdm.tqdm(total=total_size, unit="B", unit_scale=True)
|
|
|
|
zip_file = io.BytesIO()
|
|
for chunk in response.iter_content(chunk_size=1024):
|
|
if chunk:
|
|
zip_file.write(chunk)
|
|
progress_bar.update(len(chunk))
|
|
|
|
progress_bar.close()
|
|
|
|
zip_file.seek(0)
|
|
|
|
with zipfile.ZipFile(zip_file, "r") as zip_ref:
|
|
zip_ref.extractall(destination_folder)
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
|
|
def yuv_to_rgb(frames):
|
|
assert frames.dtype == torch.uint8
|
|
assert frames.ndim == 4
|
|
assert frames.shape[1] == 3
|
|
|
|
frames = frames.cpu().to(torch.float)
|
|
y = frames[..., 0, :, :]
|
|
u = frames[..., 1, :, :]
|
|
v = frames[..., 2, :, :]
|
|
|
|
y /= 255
|
|
u = u / 255 - 0.5
|
|
v = v / 255 - 0.5
|
|
|
|
r = y + 1.14 * v
|
|
g = y + -0.396 * u - 0.581 * v
|
|
b = y + 2.029 * u
|
|
|
|
rgb = torch.stack([r, g, b], 1)
|
|
rgb = (rgb * 255).clamp(0, 255).to(torch.uint8)
|
|
return rgb
|