mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-16 00:59:46 +00:00
add web interface example
This commit is contained in:
croissant
@@ -0,0 +1,481 @@
|
||||
body {
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
|
||||
background: #f5f5f5;
|
||||
}
|
||||
|
||||
main {
|
||||
min-height: 100vh;
|
||||
padding: 2rem;
|
||||
}
|
||||
|
||||
header {
|
||||
text-align: center;
|
||||
margin-bottom: 2rem;
|
||||
}
|
||||
|
||||
h1 {
|
||||
font-size: 2rem;
|
||||
font-weight: 600;
|
||||
color: #333;
|
||||
margin: 0;
|
||||
}
|
||||
|
||||
h2 {
|
||||
font-size: 1.25rem;
|
||||
font-weight: 600;
|
||||
color: #333;
|
||||
margin: 0 0 1rem 0;
|
||||
}
|
||||
|
||||
h3 {
|
||||
font-size: 0.875rem;
|
||||
font-weight: 600;
|
||||
color: #666;
|
||||
margin: 0 0 0.5rem 0;
|
||||
text-transform: uppercase;
|
||||
letter-spacing: 0.5px;
|
||||
}
|
||||
|
||||
.container {
|
||||
max-width: 1600px;
|
||||
margin: 0 auto;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 1.5rem;
|
||||
}
|
||||
|
||||
.panel {
|
||||
background: white;
|
||||
border-radius: 8px;
|
||||
padding: 1.5rem;
|
||||
box-shadow: 0 1px 3px rgba(0,0,0,0.1);
|
||||
}
|
||||
|
||||
.config-panel {
|
||||
border: 2px solid #e5e7eb;
|
||||
}
|
||||
|
||||
.config-header {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
align-items: center;
|
||||
cursor: pointer;
|
||||
user-select: none;
|
||||
padding: 0.5rem 0;
|
||||
}
|
||||
|
||||
.config-header:hover {
|
||||
opacity: 0.7;
|
||||
}
|
||||
|
||||
.toggle-icon {
|
||||
font-size: 1rem;
|
||||
color: #6b7280;
|
||||
transition: transform 0.2s;
|
||||
}
|
||||
|
||||
.config-content {
|
||||
margin-top: 1rem;
|
||||
padding-top: 1rem;
|
||||
border-top: 1px solid #e5e7eb;
|
||||
}
|
||||
|
||||
.robot-setup {
|
||||
margin-bottom: 0.5rem;
|
||||
}
|
||||
|
||||
.robot-status {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
padding: 1rem;
|
||||
border-radius: 6px;
|
||||
font-weight: 500;
|
||||
gap: 1rem;
|
||||
}
|
||||
|
||||
.robot-status.ready {
|
||||
background: linear-gradient(135deg, #d1fae5 0%, #a7f3d0 100%);
|
||||
color: #065f46;
|
||||
border: 1px solid #10b981;
|
||||
}
|
||||
|
||||
.robot-status.not-ready {
|
||||
background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
|
||||
color: #92400e;
|
||||
border: 1px solid #f59e0b;
|
||||
}
|
||||
|
||||
.btn-setup {
|
||||
background: #10b981;
|
||||
color: white;
|
||||
border: none;
|
||||
padding: 0.5rem 1rem;
|
||||
border-radius: 4px;
|
||||
font-size: 0.875rem;
|
||||
font-weight: 500;
|
||||
cursor: pointer;
|
||||
transition: background 0.2s;
|
||||
}
|
||||
|
||||
.btn-setup:hover:not(:disabled) {
|
||||
background: #059669;
|
||||
}
|
||||
|
||||
.btn-setup:disabled {
|
||||
background: #d1d5db;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
.btn-disconnect {
|
||||
background: #ef4444;
|
||||
color: white;
|
||||
border: none;
|
||||
padding: 0.5rem 1rem;
|
||||
border-radius: 4px;
|
||||
font-size: 0.875rem;
|
||||
font-weight: 500;
|
||||
cursor: pointer;
|
||||
transition: background 0.2s;
|
||||
}
|
||||
|
||||
.btn-disconnect:hover {
|
||||
background: #dc2626;
|
||||
}
|
||||
|
||||
.btn-refresh {
|
||||
background: #3b82f6;
|
||||
color: white;
|
||||
border: none;
|
||||
padding: 0.4rem 0.8rem;
|
||||
border-radius: 4px;
|
||||
font-size: 0.75rem;
|
||||
font-weight: 500;
|
||||
cursor: pointer;
|
||||
transition: background 0.2s;
|
||||
}
|
||||
|
||||
.btn-refresh:hover:not(:disabled) {
|
||||
background: #2563eb;
|
||||
}
|
||||
|
||||
.btn-refresh:disabled {
|
||||
background: #d1d5db;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
.control-panel {
|
||||
border: 2px solid #10b981;
|
||||
}
|
||||
|
||||
.status-banner {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 1rem;
|
||||
padding: 1rem 1.5rem;
|
||||
border-radius: 6px;
|
||||
margin-bottom: 1.5rem;
|
||||
font-weight: 500;
|
||||
font-size: 0.95rem;
|
||||
}
|
||||
|
||||
.status-banner.initializing {
|
||||
background: linear-gradient(135deg, #dbeafe 0%, #bfdbfe 100%);
|
||||
color: #1e40af;
|
||||
border-left: 4px solid #3b82f6;
|
||||
}
|
||||
|
||||
.status-banner.encoding {
|
||||
background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
|
||||
color: #92400e;
|
||||
border-left: 4px solid #f59e0b;
|
||||
}
|
||||
|
||||
.status-banner.uploading {
|
||||
background: linear-gradient(135deg, #e0e7ff 0%, #c7d2fe 100%);
|
||||
color: #3730a3;
|
||||
border-left: 4px solid #6366f1;
|
||||
}
|
||||
|
||||
.status-banner.success {
|
||||
background: linear-gradient(135deg, #d1fae5 0%, #a7f3d0 100%);
|
||||
color: #065f46;
|
||||
border-left: 4px solid #10b981;
|
||||
}
|
||||
|
||||
.status-banner.warning {
|
||||
background: linear-gradient(135deg, #fee2e2 0%, #fecaca 100%);
|
||||
color: #991b1b;
|
||||
border-left: 4px solid #ef4444;
|
||||
}
|
||||
|
||||
.spinner {
|
||||
width: 20px;
|
||||
height: 20px;
|
||||
border: 3px solid rgba(0, 0, 0, 0.1);
|
||||
border-top-color: currentColor;
|
||||
border-radius: 50%;
|
||||
animation: spin 0.8s linear infinite;
|
||||
}
|
||||
|
||||
@keyframes spin {
|
||||
to { transform: rotate(360deg); }
|
||||
}
|
||||
|
||||
.control-horizontal {
|
||||
display: grid;
|
||||
grid-template-columns: 2fr 1fr;
|
||||
gap: 2rem;
|
||||
align-items: start;
|
||||
}
|
||||
|
||||
.control-left {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 1rem;
|
||||
}
|
||||
|
||||
.control-right {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
}
|
||||
|
||||
.input-group {
|
||||
display: flex;
|
||||
gap: 0.5rem;
|
||||
margin-bottom: 0;
|
||||
}
|
||||
|
||||
input[type="text"] {
|
||||
flex: 1;
|
||||
padding: 0.75rem;
|
||||
border: 1px solid #ddd;
|
||||
border-radius: 4px;
|
||||
font-size: 1rem;
|
||||
}
|
||||
|
||||
input[type="text"]:disabled {
|
||||
background: #f5f5f5;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
input[type="text"]:focus {
|
||||
outline: none;
|
||||
border-color: #10b981;
|
||||
}
|
||||
|
||||
button {
|
||||
padding: 0.75rem 1.5rem;
|
||||
border: none;
|
||||
border-radius: 4px;
|
||||
font-size: 1rem;
|
||||
font-weight: 500;
|
||||
cursor: pointer;
|
||||
transition: all 0.2s;
|
||||
}
|
||||
|
||||
.btn-start {
|
||||
background: #10b981;
|
||||
color: white;
|
||||
}
|
||||
|
||||
.btn-start:hover:not(:disabled) {
|
||||
background: #059669;
|
||||
}
|
||||
|
||||
.btn-start:disabled {
|
||||
background: #d1d5db;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
.btn-stop {
|
||||
background: #ef4444;
|
||||
color: white;
|
||||
}
|
||||
|
||||
.btn-stop:hover {
|
||||
background: #dc2626;
|
||||
}
|
||||
|
||||
.btn-reset {
|
||||
padding: 0.5rem 1rem;
|
||||
background: #6b7280;
|
||||
color: white;
|
||||
font-size: 0.875rem;
|
||||
}
|
||||
|
||||
.btn-reset:hover {
|
||||
background: #4b5563;
|
||||
}
|
||||
|
||||
.status {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.75rem;
|
||||
padding: 1rem;
|
||||
border-radius: 4px;
|
||||
margin-bottom: 1rem;
|
||||
}
|
||||
|
||||
.status.recording {
|
||||
background: #fee2e2;
|
||||
color: #991b1b;
|
||||
}
|
||||
|
||||
.status.idle {
|
||||
background: #f3f4f6;
|
||||
color: #374151;
|
||||
}
|
||||
|
||||
.indicator {
|
||||
width: 12px;
|
||||
height: 12px;
|
||||
border-radius: 50%;
|
||||
background: #ef4444;
|
||||
animation: pulse 1.5s ease-in-out infinite;
|
||||
}
|
||||
|
||||
@keyframes pulse {
|
||||
0%, 100% { opacity: 1; }
|
||||
50% { opacity: 0.5; }
|
||||
}
|
||||
|
||||
.counter {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
align-items: center;
|
||||
gap: 0.75rem;
|
||||
padding: 1.5rem;
|
||||
background: linear-gradient(135deg, #f9fafb 0%, #f3f4f6 100%);
|
||||
border-radius: 8px;
|
||||
border: 2px solid #e5e7eb;
|
||||
min-width: 200px;
|
||||
}
|
||||
|
||||
.counter-label {
|
||||
font-size: 0.75rem;
|
||||
color: #6b7280;
|
||||
text-transform: uppercase;
|
||||
letter-spacing: 0.5px;
|
||||
font-weight: 600;
|
||||
}
|
||||
|
||||
.counter-value {
|
||||
font-size: 3rem;
|
||||
font-weight: 700;
|
||||
color: #10b981;
|
||||
line-height: 1;
|
||||
}
|
||||
|
||||
.time-display {
|
||||
font-size: 1.5rem;
|
||||
font-weight: 600;
|
||||
font-family: 'Courier New', monospace;
|
||||
}
|
||||
|
||||
.error-box {
|
||||
padding: 1rem;
|
||||
background: #fee2e2;
|
||||
color: #991b1b;
|
||||
border-radius: 4px;
|
||||
border-left: 4px solid #ef4444;
|
||||
font-size: 0.875rem;
|
||||
}
|
||||
|
||||
.config-section {
|
||||
margin-bottom: 1.5rem;
|
||||
}
|
||||
|
||||
.config-section:last-child {
|
||||
margin-bottom: 0;
|
||||
}
|
||||
|
||||
.config-grid {
|
||||
display: grid;
|
||||
grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
|
||||
gap: 1rem;
|
||||
}
|
||||
|
||||
label {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 0.5rem;
|
||||
font-size: 0.875rem;
|
||||
color: #374151;
|
||||
font-weight: 500;
|
||||
}
|
||||
|
||||
select {
|
||||
padding: 0.5rem;
|
||||
border: 1px solid #ddd;
|
||||
border-radius: 4px;
|
||||
font-size: 0.875rem;
|
||||
background: white;
|
||||
}
|
||||
|
||||
select:disabled {
|
||||
background: #f5f5f5;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
.camera-grid {
|
||||
display: grid;
|
||||
grid-template-columns: repeat(auto-fit, minmax(300px, 1fr));
|
||||
gap: 1rem;
|
||||
}
|
||||
|
||||
.camera {
|
||||
border: 1px solid #e5e7eb;
|
||||
border-radius: 4px;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.camera h3 {
|
||||
padding: 0.75rem;
|
||||
background: #f9fafb;
|
||||
border-bottom: 1px solid #e5e7eb;
|
||||
margin: 0;
|
||||
}
|
||||
|
||||
.camera img {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
display: block;
|
||||
background: #000;
|
||||
min-height: 300px;
|
||||
object-fit: cover;
|
||||
}
|
||||
|
||||
.camera-placeholder {
|
||||
text-align: center;
|
||||
padding: 4rem 2rem;
|
||||
background: #f9fafb;
|
||||
border-radius: 4px;
|
||||
border: 2px dashed #d1d5db;
|
||||
}
|
||||
|
||||
.camera-placeholder p {
|
||||
margin: 0.5rem 0;
|
||||
font-size: 1rem;
|
||||
color: #6b7280;
|
||||
}
|
||||
|
||||
.camera-placeholder p:first-child {
|
||||
font-size: 1.25rem;
|
||||
font-weight: 500;
|
||||
color: #374151;
|
||||
}
|
||||
|
||||
.hint {
|
||||
margin-top: 0.5rem;
|
||||
font-size: 0.75rem;
|
||||
color: #6b7280;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.5rem;
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,562 @@
|
||||
import { useState, useEffect, useCallback, useRef } from 'react';
|
||||
import './App.css';
|
||||
|
||||
const API_BASE = 'http://localhost:8000/api';
|
||||
|
||||
function App() {
|
||||
// State
|
||||
const [task, setTask] = useState('');
|
||||
const [isRecording, setIsRecording] = useState(false);
|
||||
const [isInitializing, setIsInitializing] = useState(false);
|
||||
const [isEncoding, setIsEncoding] = useState(false);
|
||||
const [isUploading, setIsUploading] = useState(false);
|
||||
const [robotsReady, setRobotsReady] = useState(false);
|
||||
const [elapsedTime, setElapsedTime] = useState(0);
|
||||
const [currentFps, setCurrentFps] = useState(0);
|
||||
const [episodeCount, setEpisodeCount] = useState(0);
|
||||
const [error, setError] = useState(null);
|
||||
const [statusMessage, setStatusMessage] = useState('Ready');
|
||||
const [uploadStatus, setUploadStatus] = useState(null);
|
||||
const [configExpanded, setConfigExpanded] = useState(false);
|
||||
|
||||
// Configuration
|
||||
const [config, setConfig] = useState({
|
||||
leader_left: 'can0',
|
||||
leader_right: 'can1',
|
||||
follower_left: 'can2',
|
||||
follower_right: 'can3',
|
||||
left_wrist: '/dev/video0',
|
||||
right_wrist: '/dev/video1',
|
||||
base: '/dev/video4'
|
||||
});
|
||||
|
||||
// Available options
|
||||
const [availableCameras, setAvailableCameras] = useState([]);
|
||||
const canInterfaces = ['can0', 'can1', 'can2', 'can3'];
|
||||
|
||||
const statusIntervalRef = useRef(null);
|
||||
const hasInitializedRef = useRef(false);
|
||||
|
||||
const loadConfig = () => {
|
||||
try {
|
||||
const saved = localStorage.getItem('openarms_config');
|
||||
if (saved) {
|
||||
setConfig(prev => ({ ...prev, ...JSON.parse(saved) }));
|
||||
}
|
||||
} catch (e) {
|
||||
console.error('Load config error:', e);
|
||||
}
|
||||
};
|
||||
|
||||
const saveConfig = (newConfig) => {
|
||||
try {
|
||||
localStorage.setItem('openarms_config', JSON.stringify(newConfig || config));
|
||||
} catch (e) {
|
||||
console.error('Save config error:', e);
|
||||
}
|
||||
};
|
||||
|
||||
// Fetch status periodically
|
||||
const fetchStatus = async () => {
|
||||
try {
|
||||
const response = await fetch(`${API_BASE}/status`);
|
||||
const data = await response.json();
|
||||
|
||||
setIsRecording(data.is_recording);
|
||||
setIsInitializing(data.is_initializing);
|
||||
setIsEncoding(data.is_encoding);
|
||||
setIsUploading(data.is_uploading);
|
||||
setRobotsReady(data.robots_ready);
|
||||
setElapsedTime(data.elapsed_time);
|
||||
setCurrentFps(data.current_fps || 0);
|
||||
setEpisodeCount(data.episode_count);
|
||||
setError(data.error);
|
||||
setStatusMessage(data.status_message || 'Ready');
|
||||
setUploadStatus(data.upload_status);
|
||||
|
||||
if (data.config) {
|
||||
// Only merge server config if we don't have a saved config (first load)
|
||||
if (!localStorage.getItem('openarms_config')) {
|
||||
setConfig(prev => {
|
||||
const merged = { ...data.config, ...prev };
|
||||
localStorage.setItem('openarms_config', JSON.stringify(merged));
|
||||
return merged;
|
||||
});
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
console.error('Failed to fetch status:', e);
|
||||
}
|
||||
};
|
||||
|
||||
const setupRobots = async () => {
|
||||
setError(null);
|
||||
try {
|
||||
const response = await fetch(`${API_BASE}/robots/setup`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify(config)
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const data = await response.json();
|
||||
throw new Error(data.detail || 'Failed to setup robots');
|
||||
}
|
||||
|
||||
await response.json();
|
||||
saveConfig(config);
|
||||
} catch (e) {
|
||||
setError(`Robot setup failed: ${e.message}`);
|
||||
}
|
||||
};
|
||||
|
||||
// Disconnect robots
|
||||
const disconnectRobots = async () => {
|
||||
try {
|
||||
await fetch(`${API_BASE}/robots/disconnect`, { method: 'POST' });
|
||||
setRobotsReady(false);
|
||||
} catch (e) {
|
||||
console.error('Failed to disconnect robots:', e);
|
||||
}
|
||||
};
|
||||
|
||||
// Discover cameras
|
||||
const discoverCameras = async () => {
|
||||
try {
|
||||
const response = await fetch(`${API_BASE}/cameras/discover`);
|
||||
const data = await response.json();
|
||||
const cameras = data.cameras || [];
|
||||
setAvailableCameras(cameras);
|
||||
|
||||
// Get list of valid camera IDs
|
||||
const validCameraIds = cameras.map(cam => String(cam.id));
|
||||
|
||||
// Auto-fix config if current values are invalid or not set
|
||||
const updated = { ...config };
|
||||
let changed = false;
|
||||
|
||||
// Auto-fix invalid camera config
|
||||
if (!config.left_wrist || !validCameraIds.includes(config.left_wrist)) {
|
||||
if (cameras.length >= 1) {
|
||||
updated.left_wrist = String(cameras[0].id);
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (!config.right_wrist || !validCameraIds.includes(config.right_wrist)) {
|
||||
if (cameras.length >= 2) {
|
||||
updated.right_wrist = String(cameras[1].id);
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (!config.base || !validCameraIds.includes(config.base)) {
|
||||
if (cameras.length >= 3) {
|
||||
updated.base = String(cameras[2].id);
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (changed) {
|
||||
setConfig(updated);
|
||||
saveConfig(updated);
|
||||
}
|
||||
|
||||
if (cameras.length === 0) {
|
||||
setError('No cameras detected! Please connect cameras and refresh.');
|
||||
}
|
||||
} catch (e) {
|
||||
console.error('Failed to discover cameras:', e);
|
||||
setError(`Camera discovery failed: ${e.message}`);
|
||||
}
|
||||
};
|
||||
|
||||
// Start recording
|
||||
const startRecording = async () => {
|
||||
if (!task.trim()) {
|
||||
setError('Please enter a task description');
|
||||
return;
|
||||
}
|
||||
|
||||
setError(null);
|
||||
|
||||
try {
|
||||
const response = await fetch(`${API_BASE}/recording/start`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ task, ...config })
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const data = await response.json();
|
||||
throw new Error(data.detail || 'Failed to start recording');
|
||||
}
|
||||
|
||||
await response.json();
|
||||
saveConfig(config);
|
||||
} catch (e) {
|
||||
setError(e.message);
|
||||
}
|
||||
};
|
||||
|
||||
// Stop recording
|
||||
const stopRecording = async () => {
|
||||
try {
|
||||
const response = await fetch(`${API_BASE}/recording/stop`, {
|
||||
method: 'POST'
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const data = await response.json();
|
||||
throw new Error(data.detail || 'Failed to stop recording');
|
||||
}
|
||||
|
||||
await response.json();
|
||||
setError(null);
|
||||
} catch (e) {
|
||||
setError(e.message);
|
||||
}
|
||||
};
|
||||
|
||||
// Reset counter
|
||||
const resetCounter = async () => {
|
||||
try {
|
||||
await fetch(`${API_BASE}/counter/reset`, { method: 'POST' });
|
||||
setEpisodeCount(0);
|
||||
} catch (e) {
|
||||
console.error('Failed to reset counter:', e);
|
||||
}
|
||||
};
|
||||
|
||||
// Format time as MM:SS
|
||||
const formatTime = (seconds) => {
|
||||
const mins = Math.floor(seconds / 60);
|
||||
const secs = Math.floor(seconds % 60);
|
||||
return `${mins.toString().padStart(2, '0')}:${secs.toString().padStart(2, '0')}`;
|
||||
};
|
||||
|
||||
// Update config and save
|
||||
const updateConfig = (key, value) => {
|
||||
const updated = { ...config, [key]: value };
|
||||
setConfig(updated);
|
||||
saveConfig(updated);
|
||||
};
|
||||
|
||||
// Initialize on mount only
|
||||
useEffect(() => {
|
||||
// Prevent double-initialization in development
|
||||
if (hasInitializedRef.current) {
|
||||
return;
|
||||
}
|
||||
hasInitializedRef.current = true;
|
||||
|
||||
loadConfig();
|
||||
discoverCameras();
|
||||
fetchStatus();
|
||||
statusIntervalRef.current = setInterval(fetchStatus, 1000);
|
||||
|
||||
return () => {
|
||||
if (statusIntervalRef.current) {
|
||||
clearInterval(statusIntervalRef.current);
|
||||
}
|
||||
};
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, []); // Run only once on mount
|
||||
|
||||
return (
|
||||
<main>
|
||||
<header>
|
||||
<h1>OpenArms Recording</h1>
|
||||
</header>
|
||||
|
||||
<div className="container">
|
||||
{/* Configuration Panel */}
|
||||
<section className="panel config-panel">
|
||||
<div
|
||||
className="config-header"
|
||||
onClick={() => setConfigExpanded(!configExpanded)}
|
||||
role="button"
|
||||
tabIndex={0}
|
||||
onKeyDown={(e) => e.key === 'Enter' && setConfigExpanded(!configExpanded)}
|
||||
>
|
||||
<h2>⚙️ Configuration</h2>
|
||||
<span className="toggle-icon">{configExpanded ? '▼' : '▶'}</span>
|
||||
</div>
|
||||
|
||||
{configExpanded && (
|
||||
<div className="config-content">
|
||||
{/* Robot Setup */}
|
||||
<div className="config-section">
|
||||
<h3>🤖 Robot Setup</h3>
|
||||
<div className="robot-setup">
|
||||
{robotsReady ? (
|
||||
<div className="robot-status ready">
|
||||
<span>✅ Robots Ready - Recording will start instantly</span>
|
||||
<button onClick={disconnectRobots} className="btn-disconnect">
|
||||
Disconnect Robots
|
||||
</button>
|
||||
</div>
|
||||
) : (
|
||||
<div className="robot-status not-ready">
|
||||
<span>⚠️ Robots not initialized - Recording will take ~10 seconds</span>
|
||||
<button
|
||||
onClick={setupRobots}
|
||||
disabled={isRecording || isInitializing}
|
||||
className="btn-setup"
|
||||
>
|
||||
🚀 Setup Robots
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* CAN Interfaces */}
|
||||
<div className="config-section">
|
||||
<h3>CAN Interfaces</h3>
|
||||
<div className="config-grid">
|
||||
<label>
|
||||
Leader Left
|
||||
<select
|
||||
value={config.leader_left}
|
||||
onChange={(e) => updateConfig('leader_left', e.target.value)}
|
||||
disabled={isRecording || robotsReady}
|
||||
>
|
||||
{canInterfaces.map((iface) => (
|
||||
<option key={iface} value={iface}>{iface}</option>
|
||||
))}
|
||||
</select>
|
||||
</label>
|
||||
|
||||
<label>
|
||||
Leader Right
|
||||
<select
|
||||
value={config.leader_right}
|
||||
onChange={(e) => updateConfig('leader_right', e.target.value)}
|
||||
disabled={isRecording || robotsReady}
|
||||
>
|
||||
{canInterfaces.map((iface) => (
|
||||
<option key={iface} value={iface}>{iface}</option>
|
||||
))}
|
||||
</select>
|
||||
</label>
|
||||
|
||||
<label>
|
||||
Follower Left
|
||||
<select
|
||||
value={config.follower_left}
|
||||
onChange={(e) => updateConfig('follower_left', e.target.value)}
|
||||
disabled={isRecording || robotsReady}
|
||||
>
|
||||
{canInterfaces.map((iface) => (
|
||||
<option key={iface} value={iface}>{iface}</option>
|
||||
))}
|
||||
</select>
|
||||
</label>
|
||||
|
||||
<label>
|
||||
Follower Right
|
||||
<select
|
||||
value={config.follower_right}
|
||||
onChange={(e) => updateConfig('follower_right', e.target.value)}
|
||||
disabled={isRecording || robotsReady}
|
||||
>
|
||||
{canInterfaces.map((iface) => (
|
||||
<option key={iface} value={iface}>{iface}</option>
|
||||
))}
|
||||
</select>
|
||||
</label>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Camera Configuration */}
|
||||
<div className="config-section">
|
||||
<div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center', marginBottom: '0.5rem' }}>
|
||||
<h3>Cameras {availableCameras.length > 0 && `(${availableCameras.length} detected)`}</h3>
|
||||
<button
|
||||
onClick={discoverCameras}
|
||||
className="btn-refresh"
|
||||
disabled={isRecording || robotsReady}
|
||||
>
|
||||
🔄 Refresh
|
||||
</button>
|
||||
</div>
|
||||
<div className="config-grid">
|
||||
<label>
|
||||
Left Wrist
|
||||
<select
|
||||
value={config.left_wrist}
|
||||
onChange={(e) => updateConfig('left_wrist', e.target.value)}
|
||||
disabled={isRecording || robotsReady}
|
||||
>
|
||||
{availableCameras.map((cam) => (
|
||||
<option key={cam.id} value={String(cam.id)}>
|
||||
{cam.name || `Camera @ ${cam.id}`}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
</label>
|
||||
|
||||
<label>
|
||||
Right Wrist
|
||||
<select
|
||||
value={config.right_wrist}
|
||||
onChange={(e) => updateConfig('right_wrist', e.target.value)}
|
||||
disabled={isRecording || robotsReady}
|
||||
>
|
||||
{availableCameras.map((cam) => (
|
||||
<option key={cam.id} value={String(cam.id)}>
|
||||
{cam.name || `Camera @ ${cam.id}`}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
</label>
|
||||
|
||||
<label>
|
||||
Base Camera
|
||||
<select
|
||||
value={config.base}
|
||||
onChange={(e) => updateConfig('base', e.target.value)}
|
||||
disabled={isRecording || robotsReady}
|
||||
>
|
||||
{availableCameras.map((cam) => (
|
||||
<option key={cam.id} value={String(cam.id)}>
|
||||
{cam.name || `Camera @ ${cam.id}`}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
</label>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</section>
|
||||
|
||||
{/* Control Panel */}
|
||||
<section className="panel control-panel">
|
||||
<h2>🎬 Recording Control</h2>
|
||||
|
||||
{/* Status Banner - Always show important statuses */}
|
||||
{isInitializing && (
|
||||
<div className="status-banner initializing">
|
||||
<div className="spinner"></div>
|
||||
<span>{statusMessage}</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{isEncoding && (
|
||||
<div className="status-banner encoding">
|
||||
<div className="spinner"></div>
|
||||
<span>📹 {statusMessage}</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{isUploading && (
|
||||
<div className="status-banner uploading">
|
||||
<div className="spinner"></div>
|
||||
<span>☁️ {statusMessage}</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{uploadStatus && !isRecording && !isEncoding && !isUploading && (
|
||||
<div className={`status-banner ${uploadStatus.startsWith('✓') ? 'success' : 'warning'}`}>
|
||||
<span>{uploadStatus}</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="control-horizontal">
|
||||
{/* Task Input and Status */}
|
||||
<div className="control-left">
|
||||
<div className="input-group">
|
||||
<input
|
||||
type="text"
|
||||
value={task}
|
||||
onChange={(e) => setTask(e.target.value)}
|
||||
placeholder="Task description (e.g., 'pick and place')"
|
||||
disabled={isRecording || isInitializing || isEncoding || isUploading}
|
||||
/>
|
||||
<button
|
||||
onClick={startRecording}
|
||||
disabled={isRecording || isInitializing || isEncoding || isUploading || !robotsReady}
|
||||
className="btn-start"
|
||||
title={!robotsReady ? 'Please setup robots first' : ''}
|
||||
>
|
||||
{isInitializing
|
||||
? '⏳ Initializing...'
|
||||
: isRecording
|
||||
? '⏺ Recording...'
|
||||
: robotsReady
|
||||
? '⏺ Start Recording'
|
||||
: '⏺ Setup Robots First'}
|
||||
</button>
|
||||
</div>
|
||||
|
||||
{/* Recording Status */}
|
||||
{isRecording && (
|
||||
<div className="status recording">
|
||||
<div className="indicator"></div>
|
||||
<span className="time-display">
|
||||
{formatTime(elapsedTime)} @ {currentFps.toFixed(1)} FPS
|
||||
</span>
|
||||
<button onClick={stopRecording} className="btn-stop">
|
||||
⏹ Stop
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Episode Counter */}
|
||||
<div className="control-right">
|
||||
<div className="counter">
|
||||
<div className="counter-label">Episodes Recorded</div>
|
||||
<div className="counter-value">{episodeCount}</div>
|
||||
<button onClick={resetCounter} className="btn-reset">
|
||||
Reset
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Error Display */}
|
||||
{error && (
|
||||
<div className="error-box">
|
||||
⚠️ {error}
|
||||
</div>
|
||||
)}
|
||||
</section>
|
||||
|
||||
{/* Camera Feeds */}
|
||||
<section className="panel cameras">
|
||||
<h2>📹 Camera Views</h2>
|
||||
{robotsReady || isRecording || isInitializing ? (
|
||||
<div className="camera-grid">
|
||||
<div className="camera">
|
||||
<h3>Left Wrist</h3>
|
||||
<img src={`${API_BASE}/camera/stream/left_wrist`} alt="Left Wrist Camera" />
|
||||
</div>
|
||||
|
||||
<div className="camera">
|
||||
<h3>Base</h3>
|
||||
<img src={`${API_BASE}/camera/stream/base`} alt="Base Camera" />
|
||||
</div>
|
||||
|
||||
<div className="camera">
|
||||
<h3>Right Wrist</h3>
|
||||
<img src={`${API_BASE}/camera/stream/right_wrist`} alt="Right Wrist Camera" />
|
||||
</div>
|
||||
</div>
|
||||
) : (
|
||||
<div className="camera-placeholder">
|
||||
<p>📷 Camera feeds will appear when robots are set up</p>
|
||||
<p className="hint">Click "Setup Robots" above to preview camera feeds</p>
|
||||
</div>
|
||||
)}
|
||||
</section>
|
||||
|
||||
</div>
|
||||
</main>
|
||||
);
|
||||
}
|
||||
|
||||
export default App;
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
# OpenArms Web Recording Interface
|
||||
|
||||
A web interface for recording OpenArms datasets.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
cd examples/openarms_web_interface
|
||||
npm install
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
**Start everything with one command:**
|
||||
|
||||
```bash
|
||||
./launch.sh
|
||||
```
|
||||
|
||||
This will:
|
||||
- Start the FastAPI backend on port 8000
|
||||
- Start the React frontend on port 5173
|
||||
- Show live logs from both services
|
||||
|
||||
Then open your browser to: **http://localhost:5173**
|
||||
|
||||
**Stop with:** `Ctrl+C`
|
||||
|
||||
---
|
||||
|
||||
## Workflow
|
||||
|
||||
1. **Configure CAN interfaces** and **camera paths** in the dropdowns
|
||||
2. Click **"Setup Robots"** to initialize (once at start)
|
||||
3. Enter a **task description**
|
||||
4. Click **"Start Recording"** to begin an episode
|
||||
5. Click **"Stop Recording"** when done
|
||||
6. Dataset is automatically encoded and uploaded to HuggingFace Hub as **private**
|
||||
7. Repeat steps 3-6 for more episodes (no need to re-setup robots!)
|
||||
|
||||
---
|
||||
@@ -0,0 +1,12 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<title>OpenArms Recording Interface</title>
|
||||
</head>
|
||||
<body>
|
||||
<div id="root"></div>
|
||||
<script type="module" src="/main.jsx"></script>
|
||||
</body>
|
||||
</html>
|
||||
Executable
+125
@@ -0,0 +1,125 @@
|
||||
#!/bin/bash
|
||||
|
||||
# OpenArms Web Interface Launcher
|
||||
# Starts Rerun viewer, FastAPI backend, and React frontend
|
||||
|
||||
set -e
|
||||
|
||||
# Colors for output
|
||||
GREEN='\033[0;32m'
|
||||
BLUE='\033[0;34m'
|
||||
YELLOW='\033[1;33m'
|
||||
RED='\033[0;31m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
# Get script directory
|
||||
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
echo -e "${BLUE}╔════════════════════════════════════════╗${NC}"
|
||||
echo -e "${BLUE}║ OpenArms Web Recording Interface ║${NC}"
|
||||
echo -e "${BLUE}╚════════════════════════════════════════╝${NC}"
|
||||
echo ""
|
||||
|
||||
# Function to cleanup on exit
|
||||
cleanup() {
|
||||
echo ""
|
||||
echo -e "${YELLOW}Shutting down services...${NC}"
|
||||
|
||||
# Kill all child processes
|
||||
pkill -P $$ 2>/dev/null || true
|
||||
|
||||
# Kill specific services by port
|
||||
lsof -ti:8000 | xargs kill -9 2>/dev/null || true # Backend
|
||||
lsof -ti:5173 | xargs kill -9 2>/dev/null || true # Frontend
|
||||
lsof -ti:9876 | xargs kill -9 2>/dev/null || true # Rerun (if spawned)
|
||||
|
||||
echo -e "${GREEN}✓ Services stopped${NC}"
|
||||
exit 0
|
||||
}
|
||||
|
||||
# Register cleanup on script exit
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
# Check if required commands exist
|
||||
command -v rerun >/dev/null 2>&1 || {
|
||||
echo -e "${RED}✗ Error: 'rerun' not found. Please install: pip install rerun-sdk${NC}"
|
||||
exit 1
|
||||
}
|
||||
|
||||
command -v python >/dev/null 2>&1 || {
|
||||
echo -e "${RED}✗ Error: 'python' not found${NC}"
|
||||
exit 1
|
||||
}
|
||||
|
||||
command -v npm >/dev/null 2>&1 || {
|
||||
echo -e "${RED}✗ Error: 'npm' not found${NC}"
|
||||
exit 1
|
||||
}
|
||||
|
||||
# Check if node_modules exists
|
||||
if [ ! -d "node_modules" ]; then
|
||||
echo -e "${YELLOW}⚠ node_modules not found. Running npm install...${NC}"
|
||||
npm install
|
||||
echo -e "${GREEN}✓ Dependencies installed${NC}"
|
||||
echo ""
|
||||
fi
|
||||
|
||||
echo -e "${GREEN}Starting services...${NC}"
|
||||
echo ""
|
||||
|
||||
# 1. Start FastAPI backend (Rerun will start when recording begins)
|
||||
echo -e "${BLUE}[1/2]${NC} Starting FastAPI backend on port 8000..."
|
||||
cd "$SCRIPT_DIR"
|
||||
python web_record_server.py > /tmp/openarms_backend.log 2>&1 &
|
||||
BACKEND_PID=$!
|
||||
sleep 3
|
||||
|
||||
if ps -p $BACKEND_PID > /dev/null; then
|
||||
echo -e "${GREEN}✓ Backend started${NC} (PID: $BACKEND_PID)"
|
||||
echo -e " URL: ${BLUE}http://localhost:8000${NC}"
|
||||
else
|
||||
echo -e "${RED}✗ Failed to start backend${NC}"
|
||||
echo -e "${YELLOW}Check logs: tail -f /tmp/openarms_backend.log${NC}"
|
||||
exit 1
|
||||
fi
|
||||
echo ""
|
||||
|
||||
# 2. Start React frontend
|
||||
echo -e "${BLUE}[2/2]${NC} Starting React frontend on port 5173..."
|
||||
cd "$SCRIPT_DIR"
|
||||
npm run dev > /tmp/openarms_frontend.log 2>&1 &
|
||||
FRONTEND_PID=$!
|
||||
sleep 3
|
||||
|
||||
if ps -p $FRONTEND_PID > /dev/null; then
|
||||
echo -e "${GREEN}✓ Frontend started${NC} (PID: $FRONTEND_PID)"
|
||||
echo -e " URL: ${BLUE}http://localhost:5173${NC}"
|
||||
else
|
||||
echo -e "${RED}✗ Failed to start frontend${NC}"
|
||||
echo -e "${YELLOW}Check logs: tail -f /tmp/openarms_frontend.log${NC}"
|
||||
exit 1
|
||||
fi
|
||||
echo ""
|
||||
|
||||
# Display status
|
||||
echo -e "${GREEN}╔════════════════════════════════════════╗${NC}"
|
||||
echo -e "${GREEN}║ All services running! 🚀 ║${NC}"
|
||||
echo -e "${GREEN}╚════════════════════════════════════════╝${NC}"
|
||||
echo ""
|
||||
echo -e "🔧 ${BLUE}Backend:${NC} http://localhost:8000"
|
||||
echo -e "🌐 ${BLUE}Frontend:${NC} http://localhost:5173"
|
||||
echo -e "📊 ${BLUE}Rerun:${NC} Will spawn automatically when recording starts"
|
||||
echo ""
|
||||
echo -e "${YELLOW}Open your browser to:${NC} ${BLUE}http://localhost:5173${NC}"
|
||||
echo ""
|
||||
echo -e "${YELLOW}Logs:${NC}"
|
||||
echo -e " • Backend: tail -f /tmp/openarms_backend.log"
|
||||
echo -e " • Frontend: tail -f /tmp/openarms_frontend.log"
|
||||
echo ""
|
||||
echo -e "${RED}Press Ctrl+C to stop all services${NC}"
|
||||
echo ""
|
||||
|
||||
# Keep script running and wait for any service to exit
|
||||
wait
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
import { createRoot } from 'react-dom/client'
|
||||
import App from './App.jsx'
|
||||
|
||||
createRoot(document.getElementById('root')).render(
|
||||
<App />
|
||||
)
|
||||
|
||||
+1955
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,21 @@
|
||||
{
|
||||
"name": "openarms-web-interface",
|
||||
"private": true,
|
||||
"version": "0.0.0",
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"dev": "vite",
|
||||
"build": "vite build",
|
||||
"preview": "vite preview"
|
||||
},
|
||||
"dependencies": {
|
||||
"react": "^18.3.1",
|
||||
"react-dom": "^18.3.1"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/react": "^18.3.12",
|
||||
"@types/react-dom": "^18.3.1",
|
||||
"@vitejs/plugin-react": "^4.3.4",
|
||||
"vite": "^6.0.1"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
import { defineConfig } from 'vite'
|
||||
import react from '@vitejs/plugin-react'
|
||||
|
||||
// https://vite.dev/config/
|
||||
export default defineConfig({
|
||||
plugins: [react()],
|
||||
server: {
|
||||
port: 5173,
|
||||
strictPort: false,
|
||||
host: true,
|
||||
open: false
|
||||
},
|
||||
build: {
|
||||
outDir: 'dist',
|
||||
sourcemap: true
|
||||
}
|
||||
})
|
||||
@@ -0,0 +1,867 @@
|
||||
"""
|
||||
OpenArms Web Recording Server
|
||||
|
||||
FastAPI backend for recording OpenArms datasets with gravity compensation.
|
||||
Provides camera streaming, robot control, and automatic HuggingFace upload.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import platform
|
||||
import re
|
||||
import shutil
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional, Any
|
||||
import threading
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
# LeRobot imports
|
||||
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
|
||||
from lerobot.robots.openarms.config_openarms_follower import OpenArmsFollowerConfig
|
||||
from lerobot.robots.openarms.openarms_follower import OpenArmsFollower
|
||||
from lerobot.teleoperators.openarms.config_openarms_leader import OpenArmsLeaderConfig
|
||||
from lerobot.teleoperators.openarms.openarms_leader import OpenArmsLeader
|
||||
|
||||
app = FastAPI(title="OpenArms Recording Server")
|
||||
|
||||
# Enable CORS
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Global state
|
||||
recording_state = {
|
||||
"is_recording": False,
|
||||
"is_initializing": False,
|
||||
"is_encoding": False,
|
||||
"is_uploading": False,
|
||||
"robots_ready": False,
|
||||
"start_time": None,
|
||||
"task": "",
|
||||
"episode_count": 0,
|
||||
"error": None,
|
||||
"status_message": "Ready",
|
||||
"upload_status": None,
|
||||
"current_fps": 0.0,
|
||||
"config": {
|
||||
"leader_left": "can0",
|
||||
"leader_right": "can1",
|
||||
"follower_left": "can2",
|
||||
"follower_right": "can3",
|
||||
"left_wrist": "/dev/video0",
|
||||
"right_wrist": "/dev/video1",
|
||||
"base": "/dev/video4",
|
||||
}
|
||||
}
|
||||
|
||||
# Global robot instances
|
||||
robot_instances = {
|
||||
"follower": None,
|
||||
"leader": None,
|
||||
"dataset": None,
|
||||
"dataset_features": None,
|
||||
}
|
||||
|
||||
# Camera frames from recording loop (for streaming during recording)
|
||||
camera_frames = {
|
||||
"left_wrist": None,
|
||||
"right_wrist": None,
|
||||
"base": None,
|
||||
}
|
||||
camera_frames_lock = threading.Lock()
|
||||
|
||||
# Recording control
|
||||
recording_thread = None
|
||||
stop_recording_flag = threading.Event()
|
||||
FPS = 30
|
||||
FRICTION_SCALE = 1.0
|
||||
|
||||
|
||||
class RecordingConfig(BaseModel):
|
||||
task: str
|
||||
leader_left: str
|
||||
leader_right: str
|
||||
follower_left: str
|
||||
follower_right: str
|
||||
left_wrist: str
|
||||
right_wrist: str
|
||||
base: str
|
||||
|
||||
|
||||
class RobotSetupConfig(BaseModel):
|
||||
"""Configuration for robot setup (no task required)."""
|
||||
leader_left: str
|
||||
leader_right: str
|
||||
follower_left: str
|
||||
follower_right: str
|
||||
left_wrist: str
|
||||
right_wrist: str
|
||||
base: str
|
||||
|
||||
|
||||
class CounterUpdate(BaseModel):
|
||||
value: int
|
||||
|
||||
|
||||
def discover_cameras_sync():
|
||||
"""Discover available OpenCV cameras."""
|
||||
try:
|
||||
from lerobot.cameras.opencv.camera_opencv import OpenCVCamera
|
||||
cameras = OpenCVCamera.find_cameras()
|
||||
print(f"[Cameras] Found {len(cameras)} camera(s)")
|
||||
return cameras
|
||||
except Exception as e:
|
||||
print(f"[Cameras] Error: {e}")
|
||||
return []
|
||||
|
||||
|
||||
@app.get("/api/cameras/discover")
|
||||
async def discover_cameras():
|
||||
"""Discover available cameras."""
|
||||
cameras = discover_cameras_sync()
|
||||
return {"cameras": cameras}
|
||||
|
||||
|
||||
@app.get("/api/can/interfaces")
|
||||
async def get_can_interfaces():
|
||||
"""Get available CAN interfaces."""
|
||||
return {"interfaces": ["can0", "can1", "can2", "can3"]}
|
||||
|
||||
|
||||
def initialize_robots_only(config: RobotSetupConfig):
|
||||
"""Initialize robots only (no dataset) for pre-setup."""
|
||||
global robot_instances, recording_state
|
||||
|
||||
print(f"[Setup] Initializing robots...")
|
||||
|
||||
# Update status: Configuring cameras
|
||||
recording_state["status_message"] = "Configuring cameras..."
|
||||
camera_config = {
|
||||
"left_wrist": OpenCVCameraConfig(index_or_path=config.left_wrist, width=640, height=480, fps=FPS),
|
||||
"right_wrist": OpenCVCameraConfig(index_or_path=config.right_wrist, width=640, height=480, fps=FPS),
|
||||
"base": OpenCVCameraConfig(index_or_path=config.base, width=640, height=480, fps=FPS),
|
||||
}
|
||||
|
||||
# Configure follower robot with cameras
|
||||
recording_state["status_message"] = "Configuring follower robot..."
|
||||
follower_config = OpenArmsFollowerConfig(
|
||||
port_left=config.follower_left,
|
||||
port_right=config.follower_right,
|
||||
can_interface="socketcan",
|
||||
id="openarms_follower",
|
||||
disable_torque_on_disconnect=True,
|
||||
max_relative_target=10.0,
|
||||
cameras=camera_config,
|
||||
)
|
||||
|
||||
# Configure leader teleoperator
|
||||
recording_state["status_message"] = "Configuring leader teleoperator..."
|
||||
leader_config = OpenArmsLeaderConfig(
|
||||
port_left=config.leader_left,
|
||||
port_right=config.leader_right,
|
||||
can_interface="socketcan",
|
||||
id="openarms_leader",
|
||||
manual_control=False,
|
||||
)
|
||||
|
||||
# Initialize and connect
|
||||
recording_state["status_message"] = "Connecting to follower robot..."
|
||||
follower = OpenArmsFollower(follower_config)
|
||||
follower.connect(calibrate=False) # Skip calibration in web mode
|
||||
|
||||
recording_state["status_message"] = "Connecting to leader teleoperator..."
|
||||
leader = OpenArmsLeader(leader_config)
|
||||
leader.connect(calibrate=False)
|
||||
|
||||
# Verify URDF is loaded
|
||||
recording_state["status_message"] = "Loading URDF model for gravity compensation..."
|
||||
if leader.pin_robot is None:
|
||||
raise RuntimeError("URDF model not loaded on leader. Gravity compensation not available.")
|
||||
|
||||
# Enable gravity compensation
|
||||
recording_state["status_message"] = "Enabling gravity compensation..."
|
||||
leader.bus_right.enable_torque()
|
||||
leader.bus_left.enable_torque()
|
||||
time.sleep(0.1)
|
||||
|
||||
robot_instances["follower"] = follower
|
||||
robot_instances["leader"] = leader
|
||||
|
||||
|
||||
def initialize_robot_systems(config: RecordingConfig):
|
||||
"""Initialize robot, leader, and dataset."""
|
||||
global robot_instances, recording_state
|
||||
|
||||
# Check if robots are already initialized
|
||||
if robot_instances.get("follower") and robot_instances.get("leader"):
|
||||
recording_state["status_message"] = "Using existing robots..."
|
||||
follower = robot_instances["follower"]
|
||||
leader = robot_instances["leader"]
|
||||
print(f"[Initialize] Reusing existing robots")
|
||||
else:
|
||||
# Full initialization required
|
||||
# Update status: Configuring cameras
|
||||
recording_state["status_message"] = "Configuring cameras..."
|
||||
camera_config = {
|
||||
"left_wrist": OpenCVCameraConfig(index_or_path=config.left_wrist, width=640, height=480, fps=FPS),
|
||||
"right_wrist": OpenCVCameraConfig(index_or_path=config.right_wrist, width=640, height=480, fps=FPS),
|
||||
"base": OpenCVCameraConfig(index_or_path=config.base, width=640, height=480, fps=FPS),
|
||||
}
|
||||
|
||||
# Configure follower robot with cameras
|
||||
recording_state["status_message"] = "Configuring follower robot..."
|
||||
follower_config = OpenArmsFollowerConfig(
|
||||
port_left=config.follower_left,
|
||||
port_right=config.follower_right,
|
||||
can_interface="socketcan",
|
||||
id="openarms_follower",
|
||||
disable_torque_on_disconnect=True,
|
||||
max_relative_target=10.0,
|
||||
cameras=camera_config,
|
||||
)
|
||||
|
||||
# Configure leader teleoperator
|
||||
recording_state["status_message"] = "Configuring leader teleoperator..."
|
||||
leader_config = OpenArmsLeaderConfig(
|
||||
port_left=config.leader_left,
|
||||
port_right=config.leader_right,
|
||||
can_interface="socketcan",
|
||||
id="openarms_leader",
|
||||
manual_control=False,
|
||||
)
|
||||
|
||||
# Initialize and connect
|
||||
recording_state["status_message"] = "Connecting to follower robot..."
|
||||
follower = OpenArmsFollower(follower_config)
|
||||
follower.connect(calibrate=False) # Skip calibration in web mode
|
||||
|
||||
recording_state["status_message"] = "Connecting to leader teleoperator..."
|
||||
leader = OpenArmsLeader(leader_config)
|
||||
leader.connect(calibrate=False)
|
||||
|
||||
# Verify URDF is loaded
|
||||
recording_state["status_message"] = "Loading URDF model for gravity compensation..."
|
||||
if leader.pin_robot is None:
|
||||
raise RuntimeError("URDF model not loaded on leader. Gravity compensation not available.")
|
||||
|
||||
# Enable gravity compensation
|
||||
recording_state["status_message"] = "Enabling gravity compensation..."
|
||||
leader.bus_right.enable_torque()
|
||||
leader.bus_left.enable_torque()
|
||||
time.sleep(0.1)
|
||||
|
||||
# Configure dataset features
|
||||
recording_state["status_message"] = "Configuring dataset features..."
|
||||
action_features_hw = {}
|
||||
for key, value in follower.action_features.items():
|
||||
if key.endswith(".pos"):
|
||||
action_features_hw[key] = value
|
||||
|
||||
action_features = hw_to_dataset_features(action_features_hw, "action")
|
||||
obs_features = hw_to_dataset_features(follower.observation_features, "observation")
|
||||
dataset_features = {**action_features, **obs_features}
|
||||
|
||||
# Create dataset
|
||||
recording_state["status_message"] = "Creating dataset..."
|
||||
now = datetime.now()
|
||||
date_str = now.strftime("%Y-%m-%d")
|
||||
time_str = now.strftime("%H-%M")
|
||||
|
||||
# Validate and sanitize task name
|
||||
task_raw = config.task.strip()
|
||||
if not task_raw:
|
||||
raise ValueError("Task description cannot be empty")
|
||||
|
||||
# Replace spaces with hyphens and convert to lowercase
|
||||
task = task_raw.replace(" ", "-").lower()
|
||||
|
||||
# Remove any leading/trailing hyphens or dots
|
||||
task = task.strip("-.")
|
||||
|
||||
# Ensure task name is valid (alphanumeric, hyphens, underscores, dots only)
|
||||
task = re.sub(r'[^a-z0-9\-_.]', '', task)
|
||||
|
||||
if not task:
|
||||
raise ValueError("Task description must contain alphanumeric characters")
|
||||
|
||||
dataset_name = f"{task}-{date_str}-{time_str}"
|
||||
repo_id = f"lerobot-data-collection/{dataset_name}"
|
||||
|
||||
print(f"[Initialize] Sanitized task: '{config.task}' -> '{task}'")
|
||||
print(f"[Initialize] Dataset name: {dataset_name}")
|
||||
print(f"[Initialize] Repo ID: {repo_id}")
|
||||
|
||||
# Remove existing dataset if it exists
|
||||
dataset_path = Path.home() / ".cache" / "huggingface" / "lerobot" / repo_id
|
||||
if dataset_path.exists():
|
||||
recording_state["status_message"] = "Removing existing dataset..."
|
||||
shutil.rmtree(dataset_path)
|
||||
|
||||
try:
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id=repo_id,
|
||||
fps=FPS,
|
||||
features=dataset_features,
|
||||
robot_type=follower.name,
|
||||
use_videos=True,
|
||||
image_writer_processes=0, # Use threads only (faster for single process)
|
||||
image_writer_threads=8, # More threads for 3 cameras (2-3 per camera)
|
||||
)
|
||||
print(f"[Dataset] Created: {repo_id}")
|
||||
except Exception as e:
|
||||
print(f"[Dataset] Error: {e}")
|
||||
raise RuntimeError(f"Failed to create dataset: {e}")
|
||||
|
||||
robot_instances["follower"] = follower
|
||||
robot_instances["leader"] = leader
|
||||
robot_instances["dataset"] = dataset
|
||||
robot_instances["dataset_features"] = dataset_features
|
||||
robot_instances["repo_id"] = repo_id
|
||||
|
||||
if robot_instances.get("dataset") is None:
|
||||
raise RuntimeError("Dataset was not stored!")
|
||||
|
||||
print(f"[Initialize] All systems ready")
|
||||
|
||||
recording_state["status_message"] = "Recording..."
|
||||
# Keep robots_ready = True so we can reuse them for next recording
|
||||
# recording_state["robots_ready"] = False # Don't mark as consumed anymore
|
||||
|
||||
return follower, leader, dataset, dataset_features
|
||||
|
||||
|
||||
def record_loop_with_compensation():
|
||||
"""Main recording loop with compensation."""
|
||||
global recording_state, stop_recording_flag
|
||||
|
||||
follower = robot_instances.get("follower")
|
||||
leader = robot_instances.get("leader")
|
||||
dataset = robot_instances.get("dataset")
|
||||
dataset_features = robot_instances.get("dataset_features")
|
||||
task = recording_state.get("task", "")
|
||||
|
||||
if follower is None or leader is None or dataset is None:
|
||||
recording_state["error"] = "Robot or dataset not initialized"
|
||||
print(f"[Recording] Error: Missing components")
|
||||
return
|
||||
|
||||
print(f"[Recording] Starting recording loop...")
|
||||
|
||||
dt = 1 / FPS
|
||||
episode_start_time = time.perf_counter()
|
||||
frame_count = 0
|
||||
last_fps_update = episode_start_time
|
||||
fps_frame_count = 0
|
||||
|
||||
# All joints (both arms)
|
||||
all_joints = []
|
||||
for motor in leader.bus_right.motors:
|
||||
all_joints.append(f"right_{motor}")
|
||||
for motor in leader.bus_left.motors:
|
||||
all_joints.append(f"left_{motor}")
|
||||
|
||||
try:
|
||||
while not stop_recording_flag.is_set():
|
||||
loop_start = time.perf_counter()
|
||||
elapsed = loop_start - episode_start_time
|
||||
|
||||
# Calculate actual FPS every second
|
||||
fps_frame_count += 1
|
||||
if elapsed - (last_fps_update - episode_start_time) >= 1.0:
|
||||
actual_fps = fps_frame_count / (elapsed - (last_fps_update - episode_start_time))
|
||||
recording_state["current_fps"] = round(actual_fps, 1)
|
||||
fps_frame_count = 0
|
||||
last_fps_update = loop_start
|
||||
|
||||
# Get leader state
|
||||
leader_action = leader.get_action()
|
||||
|
||||
# Extract positions and velocities
|
||||
leader_positions_deg = {}
|
||||
leader_velocities_deg_per_sec = {}
|
||||
|
||||
for motor in leader.bus_right.motors:
|
||||
pos_key = f"right_{motor}.pos"
|
||||
vel_key = f"right_{motor}.vel"
|
||||
if pos_key in leader_action:
|
||||
leader_positions_deg[f"right_{motor}"] = leader_action[pos_key]
|
||||
if vel_key in leader_action:
|
||||
leader_velocities_deg_per_sec[f"right_{motor}"] = leader_action[vel_key]
|
||||
|
||||
for motor in leader.bus_left.motors:
|
||||
pos_key = f"left_{motor}.pos"
|
||||
vel_key = f"left_{motor}.vel"
|
||||
if pos_key in leader_action:
|
||||
leader_positions_deg[f"left_{motor}"] = leader_action[pos_key]
|
||||
if vel_key in leader_action:
|
||||
leader_velocities_deg_per_sec[f"left_{motor}"] = leader_action[vel_key]
|
||||
|
||||
# Calculate gravity and friction torques
|
||||
leader_positions_rad = {k: np.deg2rad(v) for k, v in leader_positions_deg.items()}
|
||||
leader_gravity_torques_nm = leader._gravity_from_q(leader_positions_rad)
|
||||
|
||||
leader_velocities_rad_per_sec = {k: np.deg2rad(v) for k, v in leader_velocities_deg_per_sec.items()}
|
||||
leader_friction_torques_nm = leader._friction_from_velocity(
|
||||
leader_velocities_rad_per_sec,
|
||||
friction_scale=FRICTION_SCALE
|
||||
)
|
||||
|
||||
# Combine torques
|
||||
leader_total_torques_nm = {}
|
||||
for motor_name in leader_gravity_torques_nm:
|
||||
gravity = leader_gravity_torques_nm.get(motor_name, 0.0)
|
||||
friction = leader_friction_torques_nm.get(motor_name, 0.0)
|
||||
leader_total_torques_nm[motor_name] = gravity + friction
|
||||
|
||||
# Apply compensation to RIGHT arm
|
||||
for motor in leader.bus_right.motors:
|
||||
full_name = f"right_{motor}"
|
||||
position = leader_positions_deg.get(full_name, 0.0)
|
||||
torque = leader_total_torques_nm.get(full_name, 0.0)
|
||||
kd = leader.get_damping_kd(motor)
|
||||
|
||||
leader.bus_right._mit_control(
|
||||
motor=motor, kp=0.0, kd=kd,
|
||||
position_degrees=position,
|
||||
velocity_deg_per_sec=0.0,
|
||||
torque=torque,
|
||||
)
|
||||
|
||||
# Apply compensation to LEFT arm
|
||||
for motor in leader.bus_left.motors:
|
||||
full_name = f"left_{motor}"
|
||||
position = leader_positions_deg.get(full_name, 0.0)
|
||||
torque = leader_total_torques_nm.get(full_name, 0.0)
|
||||
kd = leader.get_damping_kd(motor)
|
||||
|
||||
leader.bus_left._mit_control(
|
||||
motor=motor, kp=0.0, kd=kd,
|
||||
position_degrees=position,
|
||||
velocity_deg_per_sec=0.0,
|
||||
torque=torque,
|
||||
)
|
||||
|
||||
# Send positions to follower
|
||||
follower_action = {}
|
||||
for joint in all_joints:
|
||||
pos_key = f"{joint}.pos"
|
||||
if pos_key in leader_action:
|
||||
follower_action[pos_key] = leader_action[pos_key]
|
||||
|
||||
if follower_action:
|
||||
follower.send_action(follower_action)
|
||||
|
||||
# Get observation
|
||||
observation = follower.get_observation()
|
||||
|
||||
# Store camera frames for streaming (no extra camera reads needed!)
|
||||
with camera_frames_lock:
|
||||
for cam_name in ["left_wrist", "right_wrist", "base"]:
|
||||
if cam_name in observation:
|
||||
camera_frames[cam_name] = observation[cam_name].copy()
|
||||
|
||||
# Add to dataset
|
||||
try:
|
||||
obs_frame = build_dataset_frame(dataset_features, observation, prefix="observation")
|
||||
action_frame = build_dataset_frame(dataset_features, follower_action, prefix="action")
|
||||
frame = {**obs_frame, **action_frame}
|
||||
frame["task"] = task
|
||||
dataset.add_frame(frame)
|
||||
frame_count += 1
|
||||
|
||||
# Log progress every 5 seconds (150 frames @ 30 FPS)
|
||||
if frame_count % 150 == 0:
|
||||
print(f"[Recording] {frame_count} frames @ {recording_state['current_fps']} FPS")
|
||||
except Exception as frame_error:
|
||||
print(f"[Recording] Frame error: {frame_error}")
|
||||
# Continue recording even if one frame fails
|
||||
|
||||
# Maintain loop rate - don't wait if behind, just continue
|
||||
loop_duration = time.perf_counter() - loop_start
|
||||
sleep_time = dt - loop_duration
|
||||
|
||||
# Only sleep if we're ahead of schedule (more than 2ms remaining)
|
||||
if sleep_time > 0.002:
|
||||
time.sleep(min(sleep_time, dt * 0.5)) # Never sleep more than half a frame time
|
||||
# If behind, just continue without sleeping - maintain real-time rate
|
||||
|
||||
except Exception as e:
|
||||
recording_state["error"] = str(e)
|
||||
print(f"[Recording] Error: {e}")
|
||||
finally:
|
||||
print(f"[Recording] Stopped. Total frames: {frame_count}")
|
||||
# Get fresh reference to dataset from global state
|
||||
current_dataset = robot_instances.get("dataset")
|
||||
if current_dataset:
|
||||
print(f"[Recording] Dataset exists, checking buffer...")
|
||||
print(f"[Recording] Episode buffer is None: {current_dataset.episode_buffer is None}")
|
||||
print(f"[Recording] Episode buffer value: {current_dataset.episode_buffer}")
|
||||
if current_dataset.episode_buffer:
|
||||
print(f"[Recording] Buffer size: {current_dataset.episode_buffer.get('size', 'NO SIZE KEY')}")
|
||||
else:
|
||||
print(f"[Recording] WARNING: Buffer is None despite {frame_count} frames!")
|
||||
else:
|
||||
print(f"[Recording] WARNING: Dataset is None in robot_instances!")
|
||||
print(f"[Recording] Local dataset variable: {dataset}")
|
||||
|
||||
|
||||
@app.post("/api/recording/start")
|
||||
async def start_recording(config: RecordingConfig):
|
||||
"""Start recording an episode."""
|
||||
global recording_state, recording_thread, stop_recording_flag
|
||||
|
||||
if recording_state["is_recording"] or recording_state["is_initializing"]:
|
||||
raise HTTPException(status_code=400, detail="Already recording or initializing")
|
||||
|
||||
# Check if robots are available (either from setup or previous recording)
|
||||
if not robot_instances.get("follower") or not robot_instances.get("leader"):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Robots not initialized. Please click 'Setup Robots' first."
|
||||
)
|
||||
|
||||
try:
|
||||
# Update config
|
||||
recording_state["config"] = config.model_dump()
|
||||
recording_state["task"] = config.task
|
||||
recording_state["error"] = None
|
||||
recording_state["is_initializing"] = True
|
||||
recording_state["status_message"] = "Initializing recording..."
|
||||
|
||||
# Initialize robot systems (will reuse pre-initialized robots if available)
|
||||
initialize_robot_systems(config)
|
||||
|
||||
if robot_instances.get("dataset") is None:
|
||||
raise RuntimeError("Dataset not created!")
|
||||
|
||||
# Start recording
|
||||
recording_state["is_initializing"] = False
|
||||
recording_state["is_recording"] = True
|
||||
recording_state["start_time"] = time.time()
|
||||
stop_recording_flag.clear() # Clear the stop flag
|
||||
|
||||
# Start recording in background thread
|
||||
recording_thread = threading.Thread(target=record_loop_with_compensation, daemon=True)
|
||||
recording_thread.start()
|
||||
time.sleep(0.1)
|
||||
|
||||
if not recording_thread.is_alive():
|
||||
raise RuntimeError("Recording thread failed to start")
|
||||
|
||||
print(f"[Start] Recording started for: {config.task}")
|
||||
return {"status": "started", "task": config.task}
|
||||
|
||||
except Exception as e:
|
||||
recording_state["is_recording"] = False
|
||||
recording_state["is_initializing"] = False
|
||||
recording_state["error"] = str(e)
|
||||
recording_state["status_message"] = f"Error: {str(e)}"
|
||||
# Clean up dataset if initialization failed, but keep robots
|
||||
cleanup_robot_systems(keep_robots=True)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
|
||||
|
||||
@app.post("/api/recording/stop")
|
||||
async def stop_recording():
|
||||
"""Stop recording, encode, and upload episode."""
|
||||
global recording_state, recording_thread, stop_recording_flag
|
||||
|
||||
if not recording_state["is_recording"]:
|
||||
raise HTTPException(status_code=400, detail="Not recording")
|
||||
|
||||
recording_state["is_recording"] = False
|
||||
recording_state["status_message"] = "Stopping recording..."
|
||||
|
||||
# Stop the recording thread
|
||||
stop_recording_flag.set()
|
||||
if recording_thread:
|
||||
recording_thread.join(timeout=5)
|
||||
|
||||
dataset = robot_instances.get("dataset")
|
||||
dataset_name = robot_instances.get("repo_id", "").split("/")[-1] if robot_instances.get("repo_id") else ""
|
||||
|
||||
print(f"[Stop] Recording stopped")
|
||||
print(f"[Stop] Dataset is None: {dataset is None}")
|
||||
if dataset:
|
||||
print(f"[Stop] Dataset type: {type(dataset)}")
|
||||
print(f"[Stop] Episode buffer is None: {dataset.episode_buffer is None}")
|
||||
print(f"[Stop] Episode buffer value: {dataset.episode_buffer}")
|
||||
if dataset.episode_buffer:
|
||||
print(f"[Stop] Episode buffer keys: {list(dataset.episode_buffer.keys())}")
|
||||
print(f"[Stop] Episode buffer type: {type(dataset.episode_buffer)}")
|
||||
print(f"[Stop] Episode buffer size: {dataset.episode_buffer.get('size', 'NO SIZE KEY')}")
|
||||
|
||||
# Process episode immediately (blocking)
|
||||
try:
|
||||
# Check buffer the same way as reference implementation
|
||||
if dataset is not None and dataset.episode_buffer is not None and dataset.episode_buffer.get("size", 0) > 0:
|
||||
buffer_size = dataset.episode_buffer.get("size", 0)
|
||||
print(f"[Stop] Buffer size: {buffer_size}")
|
||||
|
||||
# Encode videos
|
||||
recording_state["is_encoding"] = True
|
||||
recording_state["is_uploading"] = False # Ensure upload flag is clear
|
||||
recording_state["status_message"] = f"Encoding videos ({buffer_size} frames)..."
|
||||
recording_state["upload_status"] = None # Clear upload status
|
||||
print(f"[Stop] Saving episode...")
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
recording_state["is_encoding"] = False
|
||||
recording_state["status_message"] = "Encoding complete, uploading..."
|
||||
print(f"[Stop] Episode saved")
|
||||
|
||||
# Upload to hub
|
||||
recording_state["is_uploading"] = True
|
||||
recording_state["status_message"] = "Uploading to HuggingFace..."
|
||||
recording_state["upload_status"] = "Uploading..."
|
||||
print(f"[Stop] Uploading to hub...")
|
||||
dataset.push_to_hub(private=True)
|
||||
|
||||
recording_state["is_uploading"] = False
|
||||
recording_state["upload_status"] = "✓ Upload successful!"
|
||||
recording_state["status_message"] = "Ready"
|
||||
recording_state["episode_count"] += 1
|
||||
print(f"[Stop] Upload complete. Episode count: {recording_state['episode_count']}")
|
||||
else:
|
||||
recording_state["status_message"] = "No data"
|
||||
recording_state["upload_status"] = "No data"
|
||||
print(f"[Stop] No dataset or buffer (dataset={dataset is not None}, buffer={dataset.episode_buffer is not None if dataset else 'N/A'}, size={dataset.episode_buffer.get('size', 0) if dataset and dataset.episode_buffer else 'N/A'})")
|
||||
|
||||
# Clean up dataset, keep robots
|
||||
cleanup_robot_systems(keep_robots=True)
|
||||
|
||||
except Exception as e:
|
||||
recording_state["is_encoding"] = False
|
||||
recording_state["is_uploading"] = False
|
||||
recording_state["error"] = f"Upload failed: {str(e)}"
|
||||
recording_state["status_message"] = f"Error: {str(e)}"
|
||||
recording_state["upload_status"] = f"✗ Upload failed"
|
||||
cleanup_robot_systems(keep_robots=True)
|
||||
print(f"[Stop] Error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return {
|
||||
"status": "stopped",
|
||||
"dataset_name": dataset_name,
|
||||
"episode_count": recording_state["episode_count"]
|
||||
}
|
||||
|
||||
|
||||
def cleanup_robot_systems(keep_robots=False):
|
||||
"""Clean up robot systems."""
|
||||
global robot_instances
|
||||
|
||||
try:
|
||||
if not keep_robots:
|
||||
if robot_instances.get("leader"):
|
||||
robot_instances["leader"].bus_right.disable_torque()
|
||||
robot_instances["leader"].bus_left.disable_torque()
|
||||
time.sleep(0.1)
|
||||
robot_instances["leader"].disconnect()
|
||||
|
||||
if robot_instances.get("follower"):
|
||||
robot_instances["follower"].disconnect()
|
||||
|
||||
robot_instances["follower"] = None
|
||||
robot_instances["leader"] = None
|
||||
recording_state["robots_ready"] = False
|
||||
print(f"[Cleanup] Robots disconnected")
|
||||
|
||||
# Always clean up dataset
|
||||
robot_instances["dataset"] = None
|
||||
robot_instances["dataset_features"] = None
|
||||
robot_instances["repo_id"] = None
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Cleanup] Error: {e}")
|
||||
|
||||
|
||||
@app.post("/api/robots/setup")
|
||||
async def setup_robots(config: RobotSetupConfig):
|
||||
"""Pre-initialize robots for faster recording start."""
|
||||
global recording_state, robot_instances
|
||||
|
||||
if recording_state["is_recording"] or recording_state["is_initializing"]:
|
||||
raise HTTPException(status_code=400, detail="Cannot setup robots while recording")
|
||||
|
||||
if recording_state["robots_ready"] and robot_instances.get("follower") and robot_instances.get("leader"):
|
||||
raise HTTPException(status_code=400, detail="Robots already initialized")
|
||||
|
||||
# Clean up any existing robots first
|
||||
if robot_instances.get("follower") or robot_instances.get("leader"):
|
||||
cleanup_robot_systems()
|
||||
|
||||
try:
|
||||
recording_state["error"] = None
|
||||
recording_state["status_message"] = "Setting up robots..."
|
||||
|
||||
# Update config (add empty task for compatibility)
|
||||
config_dict = config.model_dump()
|
||||
config_dict["task"] = "" # Task not needed for setup
|
||||
recording_state["config"] = config_dict
|
||||
|
||||
# Initialize robot systems (without dataset creation)
|
||||
initialize_robots_only(config)
|
||||
|
||||
recording_state["robots_ready"] = True
|
||||
recording_state["status_message"] = "Robots ready for recording"
|
||||
|
||||
return {"status": "ready", "message": "Robots initialized successfully"}
|
||||
|
||||
except Exception as e:
|
||||
recording_state["error"] = str(e)
|
||||
recording_state["status_message"] = f"Setup failed"
|
||||
recording_state["robots_ready"] = False
|
||||
print(f"[Setup] Error: {e}")
|
||||
cleanup_robot_systems()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.post("/api/robots/disconnect")
|
||||
async def disconnect_robots():
|
||||
"""Disconnect pre-initialized robots."""
|
||||
global recording_state
|
||||
|
||||
if recording_state["is_recording"] or recording_state["is_initializing"]:
|
||||
raise HTTPException(status_code=400, detail="Cannot disconnect while recording")
|
||||
|
||||
recording_state["robots_ready"] = False
|
||||
recording_state["status_message"] = "Robots disconnected"
|
||||
cleanup_robot_systems(keep_robots=False) # Actually disconnect robots
|
||||
|
||||
return {"status": "disconnected"}
|
||||
|
||||
|
||||
@app.get("/api/status")
|
||||
async def get_status():
|
||||
"""Get current recording status."""
|
||||
elapsed = 0
|
||||
if recording_state["is_recording"] and recording_state["start_time"]:
|
||||
elapsed = time.time() - recording_state["start_time"]
|
||||
|
||||
return {
|
||||
"is_recording": recording_state["is_recording"],
|
||||
"is_initializing": recording_state["is_initializing"],
|
||||
"is_encoding": recording_state["is_encoding"],
|
||||
"is_uploading": recording_state["is_uploading"],
|
||||
"robots_ready": recording_state["robots_ready"],
|
||||
"elapsed_time": elapsed,
|
||||
"current_fps": recording_state["current_fps"],
|
||||
"task": recording_state["task"],
|
||||
"episode_count": recording_state["episode_count"],
|
||||
"error": recording_state["error"],
|
||||
"status_message": recording_state["status_message"],
|
||||
"upload_status": recording_state["upload_status"],
|
||||
"config": recording_state["config"]
|
||||
}
|
||||
|
||||
|
||||
@app.post("/api/counter/reset")
|
||||
async def reset_counter():
|
||||
"""Reset the episode counter."""
|
||||
recording_state["episode_count"] = 0
|
||||
return {"episode_count": 0}
|
||||
|
||||
|
||||
@app.post("/api/counter/set")
|
||||
async def set_counter(update: CounterUpdate):
|
||||
"""Set the episode counter value."""
|
||||
recording_state["episode_count"] = update.value
|
||||
return {"episode_count": update.value}
|
||||
|
||||
|
||||
@app.get("/api/camera/stream/{camera_name}")
|
||||
async def stream_camera(camera_name: str):
|
||||
"""Stream camera feed from robot.
|
||||
|
||||
During recording: streams frames from the recording loop (no extra camera reads = no contention!)
|
||||
When not recording: streams directly from cameras.
|
||||
"""
|
||||
def generate():
|
||||
try:
|
||||
while True:
|
||||
frame = None
|
||||
|
||||
# If recording, use frames from recording loop (zero contention!)
|
||||
if recording_state["is_recording"]:
|
||||
with camera_frames_lock:
|
||||
frame = camera_frames.get(camera_name)
|
||||
|
||||
if frame is not None:
|
||||
# Frame already in RGB from recording loop
|
||||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
_, buffer = cv2.imencode('.jpg', frame_rgb, [cv2.IMWRITE_JPEG_QUALITY, 70])
|
||||
yield (b'--frame\r\n'
|
||||
b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')
|
||||
time.sleep(0.033) # ~30 FPS (matches recording)
|
||||
else:
|
||||
# Wait a bit if frame not available yet
|
||||
time.sleep(0.01)
|
||||
continue
|
||||
else:
|
||||
# Not recording: stream directly from camera
|
||||
follower = robot_instances.get("follower")
|
||||
if not follower or not follower.cameras:
|
||||
break
|
||||
|
||||
if camera_name not in follower.cameras:
|
||||
break
|
||||
|
||||
camera = follower.cameras[camera_name]
|
||||
frame = camera.async_read(timeout_ms=50)
|
||||
|
||||
if frame is None:
|
||||
time.sleep(0.01)
|
||||
continue
|
||||
|
||||
# Convert BGR to RGB
|
||||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
_, buffer = cv2.imencode('.jpg', frame_rgb, [cv2.IMWRITE_JPEG_QUALITY, 60])
|
||||
|
||||
yield (b'--frame\r\n'
|
||||
b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')
|
||||
|
||||
time.sleep(0.05) # ~20 FPS when not recording
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Camera] Stream error: {e}")
|
||||
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
media_type="multipart/x-mixed-replace; boundary=frame",
|
||||
headers={
|
||||
"Cross-Origin-Resource-Policy": "cross-origin",
|
||||
"Access-Control-Allow-Origin": "*"
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown_event():
|
||||
"""Clean up resources on shutdown."""
|
||||
global stop_recording_flag
|
||||
|
||||
stop_recording_flag.set()
|
||||
cleanup_robot_systems()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
Reference in New Issue
Block a user