diff --git a/frontend/src/components/dataset/EpisodesDialog.test.tsx b/frontend/src/components/dataset/EpisodesDialog.test.tsx new file mode 100644 index 000000000..f3aa24f62 --- /dev/null +++ b/frontend/src/components/dataset/EpisodesDialog.test.tsx @@ -0,0 +1,112 @@ +// @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 EpisodesDialog from "./EpisodesDialog"; + +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( + {}} + onSubmit={async (indices, name) => { + moved.push([indices, name!]); + }} + />, + ); + }); + + const episodes = document.getElementById("episode-selection") 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(); + } +}); + +test("delete mode has no name field and deletes the selected episodes", async () => { + Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + const host = document.createElement("div"); + document.body.append(host); + const root = createRoot(host); + const deleted: [number[], string | null][] = []; + + try { + await act(async () => { + root.render( + {}} + onSubmit={async (indices, name) => { + deleted.push([indices, name]); + }} + />, + ); + }); + + expect(document.getElementById("moved-dataset-name")).toBeNull(); + const deleteButton = [...document.body.querySelectorAll("button")].find((b) => + b.textContent?.startsWith("Delete "), + ) as HTMLButtonElement; + expect(deleteButton.textContent).toBe("Delete 1 episode"); + await act(async () => deleteButton.click()); + expect(deleted).toEqual([[[2], null]]); + } finally { + await act(async () => root.unmount()); + host.remove(); + } +}); diff --git a/frontend/src/components/dataset/EpisodesDialog.tsx b/frontend/src/components/dataset/EpisodesDialog.tsx new file mode 100644 index 000000000..6b5e83944 --- /dev/null +++ b/frontend/src/components/dataset/EpisodesDialog.tsx @@ -0,0 +1,173 @@ +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"; + +/** "move" sends the episodes to a new dataset; "delete" drops them. */ +export type EpisodesDialogMode = "move" | "delete"; + +interface EpisodesDialogProps { + mode: EpisodesDialogMode; + repoId: string; + totalEpisodes: number; + currentEpisode: number | null; + open: boolean; + onOpenChange: (open: boolean) => void; + /** `outputName` is the new dataset's name in "move" mode, and null in "delete" mode. */ + onSubmit: (episodeIndices: number[], outputName: string | null) => Promise; +} + +export default function EpisodesDialog({ + mode, + repoId, + totalEpisodes, + currentEpisode, + open, + onOpenChange, + onSubmit, +}: EpisodesDialogProps) { + const [episodesText, setEpisodesText] = useState(""); + const [outputName, setOutputName] = useState(""); + const [working, setWorking] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + if (open) { + setEpisodesText(currentEpisode === null ? "" : String(currentEpisode)); + setOutputName(""); + setWorking(false); + setError(null); + } + }, [open, currentEpisode]); + + const isMove = mode === "move"; + const verb = isMove ? "Move" : "Delete"; + const episodes = parseEpisodeList(episodesText, totalEpisodes); + const takesEverything = episodes !== null && episodes.length >= totalEpisodes; + const nameValid = !isMove || /^[A-Za-z0-9._-]+$/.test(outputName); + const canSubmit = episodes !== null && !takesEverything && nameValid && !working; + + const handleOpenChange = (nextOpen: boolean) => { + if (working) return; + onOpenChange(nextOpen); + }; + + const handleSubmit = async () => { + if (!canSubmit || episodes === null) return; + setWorking(true); + try { + await onSubmit(episodes, isMove ? outputName : null); + onOpenChange(false); + } catch (caughtError) { + setError( + caughtError instanceof Error + ? caughtError.message + : `Could not ${verb.toLowerCase()} the episodes.`, + ); + } finally { + setWorking(false); + } + }; + + return ( + + + + + {isMove ? "Move episodes to a new dataset" : "Delete episodes"} + + + {isMove ? ( + <> + The episodes leave {repoId} and are saved locally as a new dataset. + + ) : ( + <> + The episodes are permanently removed from {repoId}. To keep a copy, + move them to a new dataset instead. + + )}{" "} + The episodes left in {repoId} are renumbered. + + + +
+
+ + { + setError(null); + setEpisodesText(event.target.value); + }} + placeholder="3, 7, 10-12" + disabled={working} + 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. +

+ )} + {takesEverything && ( +

At least one episode must stay in the dataset.

+ )} +
+ + {isMove && ( +
+ + { + setError(null); + setOutputName(event.target.value.replace(/[^A-Za-z0-9._-]/g, "_")); + }} + placeholder="rejected_takes" + disabled={working} + 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..5d15dfe47 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -16,6 +16,20 @@ export interface MergeDatasetsResult { total_frames: number; } +export interface MoveEpisodesResult { + success: true; + repo_id: string; + moved_episodes: number; + 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, @@ -55,3 +69,30 @@ 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", + }); +} + +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 53ec628a1..9df683b73 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -7,8 +7,10 @@ import { Database, Disc, ExternalLink, + ListX, Loader2, Merge, + Scissors, Trash2, Upload, } from "lucide-react"; @@ -17,6 +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 EpisodesDialog, { EpisodesDialogMode } from "@/components/dataset/EpisodesDialog"; import DatasetPicker from "@/components/landing/DatasetPicker"; import RecordingModal from "@/components/landing/RecordingModal"; import { useApi } from "@/contexts/ApiContext"; @@ -27,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 } 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"; @@ -55,6 +58,13 @@ const EditDataset = () => { const { toast } = useToast(); const [mergeOpen, setMergeOpen] = useState(false); const [deleteOpen, setDeleteOpen] = 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"), @@ -67,7 +77,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 +120,23 @@ const EditDataset = () => { setDataset(result.repo_id); toast({ title: "Datasets merged", description: `${result.num_episodes} episodes are ready to browse.` }); }; + const handleEpisodesEdit = async (episodeIndices: number[], outputName: string | null) => { + if (!repoId) return; + 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: outputName === null ? "Episodes deleted" : "Episodes moved", description }); + }; const setEpisode = (next: number) => { if (!repoId) return; setSearchParams(withApi({ dataset: repoId, episode: String(next) })); @@ -176,6 +203,28 @@ const EditDataset = () => { Merge )} + {repoId && selectedIsLocal && data?.success && data.total_episodes > 1 && ( + + )} + {repoId && selectedIsLocal && data?.success && data.total_episodes > 1 && ( + + )} {repoId && selectedSource === "both" && (