wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[pre-commit.ci] auto fixes from pre-commit.com hooks 

for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2025-04-03 03:55:07 +00:00
parent 4b35bb22e1
commit ca88603c4b
2 changed files with 34 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
lerobot/scripts/allowed_hosts.py
View File

@ -1,7 +1,7 @@
# List of allowed schemes and hosts for external requests # List of allowed schemes and hosts for external requests
ALLOWED_SCHEMES = {'http', 'https'} ALLOWED_SCHEMES = {"http", "https"}
ALLOWED_HOSTS = { ALLOWED_HOSTS = {
'localhost', "localhost",
'127.0.0.1', "127.0.0.1",
# Add other trusted hosts here as needed # Add other trusted hosts here as needed
} }

64
lerobot/scripts/visualize_dataset_html.py
View File

@ -15,29 +15,25 @@
import argparse import argparse
import base64 import base64
import json
import os import os
import sys import sys
import tempfile import tempfile
import urllib.parse import urllib.parse
from pathlib import Path from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union from typing import Dict, List, Tuple, Union
import cv2 import cv2
import flask
import numpy as np import numpy as np
import requests import requests
from flask import Flask, Response, jsonify, request from allowed_hosts import ALLOWED_HOSTS, ALLOWED_SCHEMES
from flask import Flask, jsonify, request
from flask_cors import CORS from flask_cors import CORS
from PIL import Image
from lerobot.data.dataset import Dataset from lerobot.data.dataset import Dataset
from lerobot.data.episode import Episode from lerobot.data.episode import Episode
from lerobot.data.frame import Frame from lerobot.data.frame import Frame
from lerobot.data.utils import get_dataset_path from lerobot.data.utils import get_dataset_path
from allowed_hosts import ALLOWED_SCHEMES, ALLOWED_HOSTS
app = Flask(__name__) app = Flask(__name__)
CORS(app) CORS(app)
@ -45,15 +41,15 @@ CORS(app)
def validate_url(url): def validate_url(url):
"""Validate URL against allowed schemes and hosts.""" """Validate URL against allowed schemes and hosts."""
parsed_url = urllib.parse.urlparse(url) parsed_url = urllib.parse.urlparse(url)
# Check if scheme is allowed # Check if scheme is allowed
if parsed_url.scheme not in ALLOWED_SCHEMES: if parsed_url.scheme not in ALLOWED_SCHEMES:
return False return False
# Check if host is allowed # Check if host is allowed
if parsed_url.netloc not in ALLOWED_HOSTS: if parsed_url.netloc not in ALLOWED_HOSTS:
return False return False
return True return True
@ -203,52 +199,52 @@ def index():
<body> <body>
<div class="container"> <div class="container">
<h1>Dataset Viewer</h1> <h1>Dataset Viewer</h1>
<div class="episode-selector"> <div class="episode-selector">
<label for="episode-id">Episode ID:</label> <label for="episode-id">Episode ID:</label>
<input type="text" id="episode-id" placeholder="Enter episode ID"> <input type="text" id="episode-id" placeholder="Enter episode ID">
<button onclick="loadEpisode()">Load Episode</button> <button onclick="loadEpisode()">Load Episode</button>
</div> </div>
<div class="frame-counter"> <div class="frame-counter">
Frame: <span id="current-frame">0</span> / <span id="total-frames">0</span> Frame: <span id="current-frame">0</span> / <span id="total-frames">0</span>
</div> </div>
<div class="frame-viewer"> <div class="frame-viewer">
<div class="frame-container"> <div class="frame-container">
<h3>RGB Image</h3> <h3>RGB Image</h3>
<img id="rgb-image" class="frame-image" src="" alt="RGB Image"> <img id="rgb-image" class="frame-image" src="" alt="RGB Image">
</div> </div>
<div class="frame-container"> <div class="frame-container">
<h3>Depth Image</h3> <h3>Depth Image</h3>
<img id="depth-image" class="frame-image" src="" alt="Depth Image"> <img id="depth-image" class="frame-image" src="" alt="Depth Image">
</div> </div>
</div> </div>
<div class="frame-info" id="frame-info"> <div class="frame-info" id="frame-info">
<h3>Frame Information</h3> <h3>Frame Information</h3>
<pre id="frame-data"></pre> <pre id="frame-data"></pre>
</div> </div>
<div class="navigation"> <div class="navigation">
<button id="prev-button" onclick="prevFrame()" disabled>Previous Frame</button> <button id="prev-button" onclick="prevFrame()" disabled>Previous Frame</button>
<button id="next-button" onclick="nextFrame()" disabled>Next Frame</button> <button id="next-button" onclick="nextFrame()" disabled>Next Frame</button>
</div> </div>
</div> </div>
<script> <script>
let currentEpisode = null; let currentEpisode = null;
let currentFrameIndex = 0; let currentFrameIndex = 0;
let frames = []; let frames = [];
function loadEpisode() { function loadEpisode() {
const episodeId = document.getElementById('episode-id').value; const episodeId = document.getElementById('episode-id').value;
if (!episodeId) { if (!episodeId) {
alert('Please enter an episode ID'); alert('Please enter an episode ID');
return; return;
} }
fetch(`/api/episode/${episodeId}`) fetch(`/api/episode/${episodeId}`)
.then(response => response.json()) .then(response => response.json())
.then(data => { .then(data => {
@ -264,10 +260,10 @@ def index():
alert('Error loading episode. Please check the episode ID and try again.'); alert('Error loading episode. Please check the episode ID and try again.');
}); });
} }
function loadFrame(frameIndex) { function loadFrame(frameIndex) {
if (!currentEpisode) return; if (!currentEpisode) return;
fetch(`/api/episode/${currentEpisode.episode_id}/frame/${frameIndex}`) fetch(`/api/episode/${currentEpisode.episode_id}/frame/${frameIndex}`)
.then(response => response.json()) .then(response => response.json())
.then(data => { .then(data => {
@ -277,14 +273,14 @@ def index():
} else { } else {
document.getElementById('rgb-image').src = ''; document.getElementById('rgb-image').src = '';
} }
// Update depth image // Update depth image
if (data.depth) { if (data.depth) {
document.getElementById('depth-image').src = `data:image/jpeg;base64,${data.depth}`; document.getElementById('depth-image').src = `data:image/jpeg;base64,${data.depth}`;
} else { } else {
document.getElementById('depth-image').src = ''; document.getElementById('depth-image').src = '';
} }
// Update frame info // Update frame info
const frameInfo = { const frameInfo = {
frame_id: data.frame_id, frame_id: data.frame_id,
@ -292,10 +288,10 @@ def index():
state: data.state state: data.state
}; };
document.getElementById('frame-data').textContent = JSON.stringify(frameInfo, null, 2); document.getElementById('frame-data').textContent = JSON.stringify(frameInfo, null, 2);
// Update current frame counter // Update current frame counter
document.getElementById('current-frame').textContent = frameIndex + 1; document.getElementById('current-frame').textContent = frameIndex + 1;
// Update navigation buttons // Update navigation buttons
document.getElementById('prev-button').disabled = frameIndex === 0; document.getElementById('prev-button').disabled = frameIndex === 0;
document.getElementById('next-button').disabled = frameIndex >= currentEpisode.num_frames - 1; document.getElementById('next-button').disabled = frameIndex >= currentEpisode.num_frames - 1;
@ -305,14 +301,14 @@ def index():
alert('Error loading frame data.'); alert('Error loading frame data.');
}); });
} }
function prevFrame() { function prevFrame() {
if (currentFrameIndex > 0) { if (currentFrameIndex > 0) {
currentFrameIndex--; currentFrameIndex--;
loadFrame(currentFrameIndex); loadFrame(currentFrameIndex);
} }
} }
function nextFrame() { function nextFrame() {
if (currentEpisode && currentFrameIndex < currentEpisode.num_frames - 1) { if (currentEpisode && currentFrameIndex < currentEpisode.num_frames - 1) {
currentFrameIndex++; currentFrameIndex++;
@ -371,14 +367,16 @@ def proxy():
# Make the request but don't forward headers from the original request # Make the request but don't forward headers from the original request
# to prevent header injection # to prevent header injection
response = requests.get(url, timeout=5) response = requests.get(url, timeout=5)
# Don't return the actual response to the user, just a success message # Don't return the actual response to the user, just a success message
# This prevents SSRF attacks where the response might contain sensitive information # This prevents SSRF attacks where the response might contain sensitive information
return jsonify({ return jsonify(
"status": "success", {
"message": "Request completed successfully", "status": "success",
"status_code": response.status_code "message": "Request completed successfully",
}) "status_code": response.status_code,
}
)
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 return jsonify({"error": str(e)}), 500