From 6f84ddf7a917dcae792f347486eacec74e53a747 Mon Sep 17 00:00:00 2001 From: "ravindu.somawansa" Date: Mon, 19 Oct 2026 23:09:54 +0200 Subject: [PATCH 1/2] feat(datasets): move episodes out of a local dataset into a new one --- .../dataset/MoveEpisodesDialog.test.tsx | 74 +++++++++ .../components/dataset/MoveEpisodesDialog.tsx | 144 ++++++++++++++++++ frontend/src/hooks/useEpisodes.ts | 9 +- frontend/src/lib/datasetApi.ts | 19 +++ frontend/src/lib/replayApi.ts | 21 +++ frontend/src/pages/EditDataset.tsx | 41 ++++- lelab/datasets.py | 67 +++++++- lelab/server.py | 32 +++- tests/test_datasets.py | 139 +++++++++++++++++ 9 files changed, 535 insertions(+), 11 deletions(-) create mode 100644 frontend/src/components/dataset/MoveEpisodesDialog.test.tsx create mode 100644 frontend/src/components/dataset/MoveEpisodesDialog.tsx diff --git a/frontend/src/components/dataset/MoveEpisodesDialog.test.tsx b/frontend/src/components/dataset/MoveEpisodesDialog.test.tsx new file mode 100644 index 000000000..1a12bb9be --- /dev/null +++ b/frontend/src/components/dataset/MoveEpisodesDialog.test.tsx @@ -0,0 +1,74 @@ +// @vitest-environment jsdom +import React, { act } from "react"; +import { createRoot } from "react-dom/client"; +import { expect, test } from "vitest"; + +import { parseEpisodeList } from "@/lib/datasetApi"; + +import MoveEpisodesDialog from "./MoveEpisodesDialog"; + +test("parses episode lists and ranges", () => { + expect(parseEpisodeList("3, 1, 5-7, 3", 10)).toEqual([1, 3, 5, 6, 7]); + expect(parseEpisodeList(" 2 - 4 ,", 10)).toEqual([2, 3, 4]); + expect(parseEpisodeList("", 10)).toBeNull(); + expect(parseEpisodeList("10", 10)).toBeNull(); + expect(parseEpisodeList("4-2", 10)).toBeNull(); + expect(parseEpisodeList("1, x", 10)).toBeNull(); +}); + +test("opens on the browsed episode and refuses to move every episode", async () => { + Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + const host = document.createElement("div"); + document.body.append(host); + const root = createRoot(host); + const moved: [number[], string][] = []; + + try { + await act(async () => { + root.render( + {}} + onMove={async (indices, name) => { + moved.push([indices, name]); + }} + />, + ); + }); + + const episodes = document.getElementById("move-episodes") as HTMLInputElement; + const name = document.getElementById("moved-dataset-name") as HTMLInputElement; + const moveButton = () => + [...document.body.querySelectorAll("button")].find((b) => + b.textContent?.startsWith("Move "), + ) as HTMLButtonElement; + const type = async (input: HTMLInputElement, value: string) => { + const setter = Object.getOwnPropertyDescriptor(HTMLInputElement.prototype, "value")!.set!; + await act(async () => { + setter.call(input, value); + input.dispatchEvent(new Event("input", { bubbles: true })); + }); + }; + + expect(episodes.value).toBe("1"); + expect(moveButton().disabled).toBe(true); + + await type(name, "rejected takes"); + expect(name.value).toBe("rejected_takes"); + expect(moveButton().disabled).toBe(false); + + await type(episodes, "0-1"); + expect(document.body.textContent).toContain("At least one episode must stay in the dataset."); + expect(moveButton().disabled).toBe(true); + + await type(episodes, "1"); + await act(async () => moveButton().click()); + expect(moved).toEqual([[[1], "rejected_takes"]]); + } finally { + await act(async () => root.unmount()); + host.remove(); + } +}); diff --git a/frontend/src/components/dataset/MoveEpisodesDialog.tsx b/frontend/src/components/dataset/MoveEpisodesDialog.tsx new file mode 100644 index 000000000..273ae5c04 --- /dev/null +++ b/frontend/src/components/dataset/MoveEpisodesDialog.tsx @@ -0,0 +1,144 @@ +import { useEffect, useState } from "react"; +import { Loader2 } from "lucide-react"; + +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { parseEpisodeList } from "@/lib/datasetApi"; + +interface MoveEpisodesDialogProps { + repoId: string; + totalEpisodes: number; + currentEpisode: number | null; + open: boolean; + onOpenChange: (open: boolean) => void; + onMove: (episodeIndices: number[], outputName: string) => Promise; +} + +export default function MoveEpisodesDialog({ + repoId, + totalEpisodes, + currentEpisode, + open, + onOpenChange, + onMove, +}: MoveEpisodesDialogProps) { + const [episodesText, setEpisodesText] = useState(""); + const [outputName, setOutputName] = useState(""); + const [moving, setMoving] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + if (open) { + setEpisodesText(currentEpisode === null ? "" : String(currentEpisode)); + setOutputName(""); + setMoving(false); + setError(null); + } + }, [open, currentEpisode]); + + const episodes = parseEpisodeList(episodesText, totalEpisodes); + const movesEverything = episodes !== null && episodes.length >= totalEpisodes; + const canMove = + episodes !== null && !movesEverything && /^[A-Za-z0-9._-]+$/.test(outputName) && !moving; + + const handleOpenChange = (nextOpen: boolean) => { + if (moving) return; + onOpenChange(nextOpen); + }; + + const handleMove = async () => { + if (!canMove || episodes === null) return; + setMoving(true); + try { + await onMove(episodes, outputName); + onOpenChange(false); + } catch (caughtError) { + setError(caughtError instanceof Error ? caughtError.message : "Could not move the episodes."); + } finally { + setMoving(false); + } + }; + + return ( + + + + Move episodes to a new dataset + + The episodes leave {repoId} and are saved locally as a new dataset. The + episodes left in {repoId} are renumbered. + + + +
+
+ + { + setError(null); + setEpisodesText(event.target.value); + }} + placeholder="3, 7, 10-12" + disabled={moving} + className="border-gray-700 bg-gray-800 text-white" + /> + {episodesText.trim() && episodes === null && ( +

+ Use episode numbers from 0 to {totalEpisodes - 1}, e.g. 3, 7, 10-12. +

+ )} + {movesEverything && ( +

At least one episode must stay in the dataset.

+ )} +
+ +
+ + { + setError(null); + setOutputName(event.target.value.replace(/[^A-Za-z0-9._-]/g, "_")); + }} + placeholder="rejected_takes" + disabled={moving} + className="border-gray-700 bg-gray-800 text-white" + /> +

+ Saved locally as local/{outputName || "name"}. +

+
+ {error &&

{error}

} +
+ + + + + + + + ); +} diff --git a/frontend/src/hooks/useEpisodes.ts b/frontend/src/hooks/useEpisodes.ts index 4158ab4d6..791c6dddf 100644 --- a/frontend/src/hooks/useEpisodes.ts +++ b/frontend/src/hooks/useEpisodes.ts @@ -1,4 +1,4 @@ -import { useEffect, useState } from "react"; +import { useCallback, useEffect, useState } from "react"; import { useApi } from "@/contexts/ApiContext"; import { EpisodeDetail, @@ -13,6 +13,7 @@ export const useEpisodes = (repoId: string | null) => { const [data, setData] = useState(null); const [loading, setLoading] = useState(false); const [error, setError] = useState(null); + const [version, setVersion] = useState(0); useEffect(() => { setData(null); @@ -44,11 +45,13 @@ export const useEpisodes = (repoId: string | null) => { if (!controller.signal.aborted) setLoading(false); }); return () => controller.abort(); - }, [baseUrl, fetchWithHeaders, repoId]); + }, [baseUrl, fetchWithHeaders, repoId, version]); + + const refresh = useCallback(() => setVersion((v) => v + 1), []); // Effects run after render: do not let the new dataset auto-select an // episode from the old list before this effect clears it. - return { data: data?.repo_id === repoId ? data : null, loading, error }; + return { data: data?.repo_id === repoId ? data : null, loading, error, refresh }; }; /** One episode's detail. Either arg null → idle. */ diff --git a/frontend/src/lib/datasetApi.ts b/frontend/src/lib/datasetApi.ts index 4ac27a8af..1acfb5fa5 100644 --- a/frontend/src/lib/datasetApi.ts +++ b/frontend/src/lib/datasetApi.ts @@ -189,3 +189,22 @@ export function formatDuration(seconds: number | null | undefined): string { const s = Math.round(seconds % 60); return `${m}m ${String(s).padStart(2, "0")}s`; } + +/** + * Parse "3, 7, 10-12" into sorted episode indices. Returns null when the text + * is not a list of indices and ranges below `total`. + */ +export function parseEpisodeList(text: string, total: number): number[] | null { + const indices = new Set(); + for (const part of text.split(",")) { + const item = part.trim(); + if (!item) continue; + const match = /^(\d+)(?:\s*-\s*(\d+))?$/.exec(item); + if (!match) return null; + const start = Number(match[1]); + const end = match[2] === undefined ? start : Number(match[2]); + if (end < start || end >= total) return null; + for (let i = start; i <= end; i++) indices.add(i); + } + return indices.size > 0 ? [...indices].sort((a, b) => a - b) : null; +} diff --git a/frontend/src/lib/replayApi.ts b/frontend/src/lib/replayApi.ts index 21c591ec2..4c2190939 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -16,6 +16,13 @@ export interface MergeDatasetsResult { total_frames: number; } +export interface MoveEpisodesResult { + success: true; + repo_id: string; + moved_episodes: number; + remaining_episodes: number; +} + export async function listDatasets( baseUrl: string, fetcher: Fetcher, @@ -55,3 +62,17 @@ export async function mergeLocalDatasets( action: "Merge datasets", }); } + +export async function moveEpisodes( + baseUrl: string, + fetcher: Fetcher, + repoId: string, + episodeIndices: number[], + outputName: string, +): Promise { + return apiRequest(baseUrl, fetcher, "/datasets/move-episodes", { + method: "POST", + body: { repo_id: repoId, episode_indices: episodeIndices, output_name: outputName }, + action: "Move episodes", + }); +} diff --git a/frontend/src/pages/EditDataset.tsx b/frontend/src/pages/EditDataset.tsx index 53ec628a1..ba09b2686 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -9,6 +9,7 @@ import { ExternalLink, Loader2, Merge, + Scissors, Trash2, Upload, } from "lucide-react"; @@ -17,6 +18,7 @@ import DeleteDatasetDialog from "@/components/dataset/DeleteDatasetDialog"; import EpisodeList from "@/components/dataset/EpisodeList"; import EpisodeViewer from "@/components/dataset/EpisodeViewer"; import MergeDatasetsDialog from "@/components/dataset/MergeDatasetsDialog"; +import MoveEpisodesDialog from "@/components/dataset/MoveEpisodesDialog"; import DatasetPicker from "@/components/landing/DatasetPicker"; import RecordingModal from "@/components/landing/RecordingModal"; import { useApi } from "@/contexts/ApiContext"; @@ -27,7 +29,7 @@ import { useRobots } from "@/hooks/useRobots"; import { useToast } from "@/hooks/use-toast"; import { formatDuration } from "@/lib/datasetApi"; import { openHubViewer } from "@/lib/hubViewer"; -import { DatasetItem, mergeLocalDatasets } from "@/lib/replayApi"; +import { DatasetItem, mergeLocalDatasets, moveEpisodes } from "@/lib/replayApi"; // The robot_type LeRobot stores for the SO-101 follower. const FOLLOWER_ROBOT_TYPE = "so_follower"; @@ -55,6 +57,7 @@ const EditDataset = () => { const { toast } = useToast(); const [mergeOpen, setMergeOpen] = useState(false); const [deleteOpen, setDeleteOpen] = useState(false); + const [moveOpen, setMoveOpen] = useState(false); // A Hub-only dataset has no videos on disk to decode. const localDatasets = useMemo( () => datasets.filter((d) => d.source === "local" || d.source === "both"), @@ -67,7 +70,7 @@ const EditDataset = () => { const selectedSource = datasets.find((d) => d.repo_id === repoId)?.source; const selectedIsLocal = selectedSource === "local" || selectedSource === "both"; - const { data, loading, error } = useEpisodes(repoId); + const { data, loading, error, refresh: refreshEpisodes } = useEpisodes(repoId); const { detail, loading: detailLoading, @@ -110,6 +113,18 @@ const EditDataset = () => { setDataset(result.repo_id); toast({ title: "Datasets merged", description: `${result.num_episodes} episodes are ready to browse.` }); }; + const handleMove = async (episodeIndices: number[], outputName: string) => { + if (!repoId) return; + const result = await moveEpisodes(baseUrl, fetchWithHeaders, repoId, episodeIndices, outputName); + refresh(); + // The episodes left behind are renumbered, so reopen the dataset from its first one. + setSearchParams(withApi({ dataset: repoId })); + refreshEpisodes(); + toast({ + title: "Episodes moved", + description: `${result.moved_episodes} moved to ${result.repo_id}, ${result.remaining_episodes} left in ${repoId}.`, + }); + }; const setEpisode = (next: number) => { if (!repoId) return; setSearchParams(withApi({ dataset: repoId, episode: String(next) })); @@ -176,6 +191,17 @@ const EditDataset = () => { Merge )} + {repoId && selectedIsLocal && data?.success && data.total_episodes > 1 && ( + + )} {repoId && selectedSource === "both" && ( + + + + + ); +} diff --git a/frontend/src/components/dataset/MoveEpisodesDialog.tsx b/frontend/src/components/dataset/MoveEpisodesDialog.tsx deleted file mode 100644 index 273ae5c04..000000000 --- a/frontend/src/components/dataset/MoveEpisodesDialog.tsx +++ /dev/null @@ -1,144 +0,0 @@ -import { useEffect, useState } from "react"; -import { Loader2 } from "lucide-react"; - -import { Button } from "@/components/ui/button"; -import { - Dialog, - DialogContent, - DialogDescription, - DialogFooter, - DialogHeader, - DialogTitle, -} from "@/components/ui/dialog"; -import { Input } from "@/components/ui/input"; -import { Label } from "@/components/ui/label"; -import { parseEpisodeList } from "@/lib/datasetApi"; - -interface MoveEpisodesDialogProps { - repoId: string; - totalEpisodes: number; - currentEpisode: number | null; - open: boolean; - onOpenChange: (open: boolean) => void; - onMove: (episodeIndices: number[], outputName: string) => Promise; -} - -export default function MoveEpisodesDialog({ - repoId, - totalEpisodes, - currentEpisode, - open, - onOpenChange, - onMove, -}: MoveEpisodesDialogProps) { - const [episodesText, setEpisodesText] = useState(""); - const [outputName, setOutputName] = useState(""); - const [moving, setMoving] = useState(false); - const [error, setError] = useState(null); - - useEffect(() => { - if (open) { - setEpisodesText(currentEpisode === null ? "" : String(currentEpisode)); - setOutputName(""); - setMoving(false); - setError(null); - } - }, [open, currentEpisode]); - - const episodes = parseEpisodeList(episodesText, totalEpisodes); - const movesEverything = episodes !== null && episodes.length >= totalEpisodes; - const canMove = - episodes !== null && !movesEverything && /^[A-Za-z0-9._-]+$/.test(outputName) && !moving; - - const handleOpenChange = (nextOpen: boolean) => { - if (moving) return; - onOpenChange(nextOpen); - }; - - const handleMove = async () => { - if (!canMove || episodes === null) return; - setMoving(true); - try { - await onMove(episodes, outputName); - onOpenChange(false); - } catch (caughtError) { - setError(caughtError instanceof Error ? caughtError.message : "Could not move the episodes."); - } finally { - setMoving(false); - } - }; - - return ( - - - - Move episodes to a new dataset - - The episodes leave {repoId} and are saved locally as a new dataset. The - episodes left in {repoId} are renumbered. - - - -
-
- - { - setError(null); - setEpisodesText(event.target.value); - }} - placeholder="3, 7, 10-12" - disabled={moving} - className="border-gray-700 bg-gray-800 text-white" - /> - {episodesText.trim() && episodes === null && ( -

- Use episode numbers from 0 to {totalEpisodes - 1}, e.g. 3, 7, 10-12. -

- )} - {movesEverything && ( -

At least one episode must stay in the dataset.

- )} -
- -
- - { - setError(null); - setOutputName(event.target.value.replace(/[^A-Za-z0-9._-]/g, "_")); - }} - placeholder="rejected_takes" - disabled={moving} - className="border-gray-700 bg-gray-800 text-white" - /> -

- Saved locally as local/{outputName || "name"}. -

-
- {error &&

{error}

} -
- - - - - - - - ); -} diff --git a/frontend/src/lib/replayApi.ts b/frontend/src/lib/replayApi.ts index 4c2190939..5d15dfe47 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -23,6 +23,13 @@ export interface MoveEpisodesResult { remaining_episodes: number; } +export interface DeleteEpisodesResult { + success: true; + repo_id: string; + deleted_episodes: number; + remaining_episodes: number; +} + export async function listDatasets( baseUrl: string, fetcher: Fetcher, @@ -76,3 +83,16 @@ export async function moveEpisodes( action: "Move episodes", }); } + +export async function deleteEpisodes( + baseUrl: string, + fetcher: Fetcher, + repoId: string, + episodeIndices: number[], +): Promise { + return apiRequest(baseUrl, fetcher, "/datasets/delete-episodes", { + method: "POST", + body: { repo_id: repoId, episode_indices: episodeIndices }, + action: "Delete episodes", + }); +} diff --git a/frontend/src/pages/EditDataset.tsx b/frontend/src/pages/EditDataset.tsx index ba09b2686..9df683b73 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -7,6 +7,7 @@ import { Database, Disc, ExternalLink, + ListX, Loader2, Merge, Scissors, @@ -18,7 +19,7 @@ import DeleteDatasetDialog from "@/components/dataset/DeleteDatasetDialog"; import EpisodeList from "@/components/dataset/EpisodeList"; import EpisodeViewer from "@/components/dataset/EpisodeViewer"; import MergeDatasetsDialog from "@/components/dataset/MergeDatasetsDialog"; -import MoveEpisodesDialog from "@/components/dataset/MoveEpisodesDialog"; +import EpisodesDialog, { EpisodesDialogMode } from "@/components/dataset/EpisodesDialog"; import DatasetPicker from "@/components/landing/DatasetPicker"; import RecordingModal from "@/components/landing/RecordingModal"; import { useApi } from "@/contexts/ApiContext"; @@ -29,7 +30,7 @@ import { useRobots } from "@/hooks/useRobots"; import { useToast } from "@/hooks/use-toast"; import { formatDuration } from "@/lib/datasetApi"; import { openHubViewer } from "@/lib/hubViewer"; -import { DatasetItem, mergeLocalDatasets, moveEpisodes } from "@/lib/replayApi"; +import { DatasetItem, deleteEpisodes, mergeLocalDatasets, moveEpisodes } from "@/lib/replayApi"; // The robot_type LeRobot stores for the SO-101 follower. const FOLLOWER_ROBOT_TYPE = "so_follower"; @@ -57,7 +58,13 @@ const EditDataset = () => { const { toast } = useToast(); const [mergeOpen, setMergeOpen] = useState(false); const [deleteOpen, setDeleteOpen] = useState(false); - const [moveOpen, setMoveOpen] = useState(false); + const [episodesOpen, setEpisodesOpen] = useState(false); + // Kept after close, so the dialog does not change layout while it fades out. + const [episodesMode, setEpisodesMode] = useState("move"); + const openEpisodesDialog = (mode: EpisodesDialogMode) => { + setEpisodesMode(mode); + setEpisodesOpen(true); + }; // A Hub-only dataset has no videos on disk to decode. const localDatasets = useMemo( () => datasets.filter((d) => d.source === "local" || d.source === "both"), @@ -113,17 +120,22 @@ const EditDataset = () => { setDataset(result.repo_id); toast({ title: "Datasets merged", description: `${result.num_episodes} episodes are ready to browse.` }); }; - const handleMove = async (episodeIndices: number[], outputName: string) => { + const handleEpisodesEdit = async (episodeIndices: number[], outputName: string | null) => { if (!repoId) return; - const result = await moveEpisodes(baseUrl, fetchWithHeaders, repoId, episodeIndices, outputName); + const description = + outputName === null + ? await deleteEpisodes(baseUrl, fetchWithHeaders, repoId, episodeIndices).then( + (result) => `${result.deleted_episodes} deleted, ${result.remaining_episodes} left in ${repoId}.`, + ) + : await moveEpisodes(baseUrl, fetchWithHeaders, repoId, episodeIndices, outputName).then( + (result) => + `${result.moved_episodes} moved to ${result.repo_id}, ${result.remaining_episodes} left in ${repoId}.`, + ); refresh(); // The episodes left behind are renumbered, so reopen the dataset from its first one. setSearchParams(withApi({ dataset: repoId })); refreshEpisodes(); - toast({ - title: "Episodes moved", - description: `${result.moved_episodes} moved to ${result.repo_id}, ${result.remaining_episodes} left in ${repoId}.`, - }); + toast({ title: outputName === null ? "Episodes deleted" : "Episodes moved", description }); }; const setEpisode = (next: number) => { if (!repoId) return; @@ -195,13 +207,24 @@ const EditDataset = () => { )} + {repoId && selectedIsLocal && data?.success && data.total_episodes > 1 && ( + + )} {repoId && selectedSource === "both" && (