diff --git a/.gitignore b/.gitignore index 39d8187c..62105eac 100644 --- a/.gitignore +++ b/.gitignore @@ -161,6 +161,7 @@ cython_debug/ /env.sh /models +/datasets /custom/* !/custom/.gitkeep /.tmp diff --git a/ui/src/app/api/caption/[...imagePath]/route.ts b/ui/src/app/api/caption/[...imagePath]/route.ts new file mode 100644 index 00000000..3a35b88b --- /dev/null +++ b/ui/src/app/api/caption/[...imagePath]/route.ts @@ -0,0 +1,42 @@ +// src/app/api/img/[imagePath]/route.ts +import { NextRequest, NextResponse } from 'next/server'; +import fs from 'fs'; +import path from 'path'; +import { getDatasetsRoot } from '@/app/api/datasets/utils'; + +export async function GET(request: NextRequest, { params }: { params: { imagePath: string } }) { + const { imagePath } = await params; + try { + // Decode the path + const filepath = decodeURIComponent(imagePath); + + // caption name is the filepath without extension but with .txt + const captionPath = filepath.replace(/\.[^/.]+$/, '') + '.txt'; + + // Get allowed directories + const allowedDir = await getDatasetsRoot(); + + // Security check: Ensure path is in allowed directory + const isAllowed = filepath.startsWith(allowedDir) && !filepath.includes('..'); + + if (!isAllowed) { + console.warn(`Access denied: ${filepath} not in ${allowedDir}`); + return new NextResponse('Access denied', { status: 403 }); + } + + // Check if file exists + if (!fs.existsSync(captionPath)) { + // send back blank string if caption file does not exist + return new NextResponse(''); + } + + // Read caption file + const caption = fs.readFileSync(captionPath, 'utf-8'); + + // Return caption + return new NextResponse(caption); + } catch (error) { + console.error('Error getting caption:', error); + return new NextResponse('Error getting caption', { status: 500 }); + } +} diff --git a/ui/src/app/api/datasets/create/route.tsx b/ui/src/app/api/datasets/create/route.tsx new file mode 100644 index 00000000..f552b57f --- /dev/null +++ b/ui/src/app/api/datasets/create/route.tsx @@ -0,0 +1,22 @@ +import { NextResponse } from 'next/server'; +import fs from 'fs'; +import path from 'path'; +import { getDatasetsRoot } from '@/app/api/datasets/utils'; + +export async function POST(request: Request) { + try { + const body = await request.json(); + const { name } = body; + let datasetsPath = await getDatasetsRoot(); + let datasetPath = path.join(datasetsPath, name); + + // if folder doesnt exist, create it + if (!fs.existsSync(datasetPath)) { + fs.mkdirSync(datasetPath); + } + + return NextResponse.json({ success: true }); + } catch (error) { + return NextResponse.json({ error: 'Failed to create dataset' }, { status: 500 }); + } +} \ No newline at end of file diff --git a/ui/src/app/api/datasets/list/route.ts b/ui/src/app/api/datasets/list/route.ts new file mode 100644 index 00000000..1d841681 --- /dev/null +++ b/ui/src/app/api/datasets/list/route.ts @@ -0,0 +1,25 @@ +import { NextResponse } from 'next/server'; +import fs from 'fs'; +import { getDatasetsRoot } from '@/app/api/datasets/utils'; + +export async function GET() { + try { + let datasetsPath = await getDatasetsRoot(); + + // if folder doesnt exist, create it + if (!fs.existsSync(datasetsPath)) { + fs.mkdirSync(datasetsPath); + } + + // find all the folders in the datasets folder + let folders = fs + .readdirSync(datasetsPath, { withFileTypes: true }) + .filter(dirent => dirent.isDirectory()) + .filter(dirent => !dirent.name.startsWith('.')) + .map(dirent => dirent.name); + + return NextResponse.json(folders); + } catch (error) { + return NextResponse.json({ error: 'Failed to fetch datasets' }, { status: 500 }); + } +} diff --git a/ui/src/app/api/datasets/listImages/route.ts b/ui/src/app/api/datasets/listImages/route.ts new file mode 100644 index 00000000..e563d65f --- /dev/null +++ b/ui/src/app/api/datasets/listImages/route.ts @@ -0,0 +1,67 @@ +import { NextResponse } from 'next/server'; +import fs from 'fs'; +import path from 'path'; +import { getDatasetsRoot } from '@/app/api/datasets/utils'; + +export async function POST(request: Request) { + const datasetsPath = await getDatasetsRoot(); + const body = await request.json(); + const { datasetName } = body; + const datasetFolder = path.join(datasetsPath, datasetName); + + try { + // Check if folder exists + if (!fs.existsSync(datasetFolder)) { + return NextResponse.json( + { error: `Folder '${datasetName}' not found` }, + { status: 404 } + ); + } + + // Find all images recursively + const imageFiles = findImagesRecursively(datasetFolder); + + // Format response + const result = imageFiles.map(imgPath => ({ + img_path: imgPath + })); + + return NextResponse.json({ images: result }); + } catch (error) { + console.error('Error finding images:', error); + return NextResponse.json( + { error: 'Failed to process request' }, + { status: 500 } + ); + } +} + +/** + * Recursively finds all image files in a directory and its subdirectories + * @param dir Directory to search + * @returns Array of absolute paths to image files + */ +function findImagesRecursively(dir: string): string[] { + const imageExtensions = ['.png', '.jpg', '.jpeg', '.webp']; + let results: string[] = []; + + const items = fs.readdirSync(dir); + + for (const item of items) { + const itemPath = path.join(dir, item); + const stat = fs.statSync(itemPath); + + if (stat.isDirectory()) { + // If it's a directory, recursively search it + results = results.concat(findImagesRecursively(itemPath)); + } else { + // If it's a file, check if it's an image + const ext = path.extname(itemPath).toLowerCase(); + if (imageExtensions.includes(ext)) { + results.push(itemPath); + } + } + } + + return results; +} \ No newline at end of file diff --git a/ui/src/app/api/datasets/utils.ts b/ui/src/app/api/datasets/utils.ts new file mode 100644 index 00000000..88361636 --- /dev/null +++ b/ui/src/app/api/datasets/utils.ts @@ -0,0 +1,17 @@ +import { PrismaClient } from '@prisma/client'; +import { defaultDatasetsFolder } from '@/paths'; + +const prisma = new PrismaClient(); + +export const getDatasetsRoot = async () => { + let row = await prisma.settings.findFirst({ + where: { + key: 'DATASETS_FOLDER', + }, + }); + let datasetsPath = defaultDatasetsFolder; + if (row?.value && row.value !== '') { + datasetsPath = row.value; + } + return datasetsPath; +}; \ No newline at end of file diff --git a/ui/src/app/api/img/[...imagePath]/route.ts b/ui/src/app/api/img/[...imagePath]/route.ts new file mode 100644 index 00000000..875cdd1a --- /dev/null +++ b/ui/src/app/api/img/[...imagePath]/route.ts @@ -0,0 +1,66 @@ +// src/app/api/img/[imagePath]/route.ts +import { NextRequest, NextResponse } from 'next/server'; +import fs from 'fs'; +import path from 'path'; +import { getDatasetsRoot } from '@/app/api/datasets/utils'; + +export async function GET(request: NextRequest, { params }: { params: { imagePath: string } }) { + const { imagePath } = await params; + try { + // Decode the path + const filepath = decodeURIComponent(imagePath); + console.log('Serving image:', filepath); + + // Get allowed directories + const allowedDir = await getDatasetsRoot(); + + // Security check: Ensure path is in allowed directory + const isAllowed = filepath.startsWith(allowedDir) && !filepath.includes('..'); + + if (!isAllowed) { + console.warn(`Access denied: ${filepath} not in ${allowedDir}`); + return new NextResponse('Access denied', { status: 403 }); + } + + // Check if file exists + if (!fs.existsSync(filepath)) { + console.warn(`File not found: ${filepath}`); + return new NextResponse('File not found', { status: 404 }); + } + + // Get file info + const stat = fs.statSync(filepath); + if (!stat.isFile()) { + return new NextResponse('Not a file', { status: 400 }); + } + + // Determine content type + const ext = path.extname(filepath).toLowerCase(); + const contentTypeMap: { [key: string]: string } = { + '.jpg': 'image/jpeg', + '.jpeg': 'image/jpeg', + '.png': 'image/png', + '.gif': 'image/gif', + '.webp': 'image/webp', + '.svg': 'image/svg+xml', + '.bmp': 'image/bmp', + }; + + const contentType = contentTypeMap[ext] || 'application/octet-stream'; + + // Read file as buffer + const fileBuffer = fs.readFileSync(filepath); + + // Return file with appropriate headers + return new NextResponse(fileBuffer, { + headers: { + 'Content-Type': contentType, + 'Content-Length': String(stat.size), + 'Cache-Control': 'public, max-age=86400', + }, + }); + } catch (error) { + console.error('Error serving image:', error); + return new NextResponse('Internal Server Error', { status: 500 }); + } +} diff --git a/ui/src/app/api/settings/route.ts b/ui/src/app/api/settings/route.ts index 8062afde..458c53c5 100644 --- a/ui/src/app/api/settings/route.ts +++ b/ui/src/app/api/settings/route.ts @@ -1,6 +1,6 @@ import { NextResponse } from 'next/server'; import { PrismaClient } from '@prisma/client'; -import { defaultTrainFolder } from '@/paths'; +import { defaultTrainFolder, defaultDatasetsFolder } from '@/paths'; const prisma = new PrismaClient(); @@ -15,6 +15,10 @@ export async function GET() { if (!settingsObject.TRAINING_FOLDER || settingsObject.TRAINING_FOLDER === '') { settingsObject.TRAINING_FOLDER = defaultTrainFolder; } + // if DATASETS_FOLDER is not set, use default + if (!settingsObject.DATASETS_FOLDER || settingsObject.DATASETS_FOLDER === '') { + settingsObject.DATASETS_FOLDER = defaultDatasetsFolder; + } return NextResponse.json(settingsObject); } catch (error) { return NextResponse.json({ error: 'Failed to fetch settings' }, { status: 500 }); @@ -24,7 +28,7 @@ export async function GET() { export async function POST(request: Request) { try { const body = await request.json(); - const { HF_TOKEN, TRAINING_FOLDER } = body; + const { HF_TOKEN, TRAINING_FOLDER, DATASETS_FOLDER } = body; // Upsert both settings await Promise.all([ @@ -38,6 +42,11 @@ export async function POST(request: Request) { update: { value: TRAINING_FOLDER }, create: { key: 'TRAINING_FOLDER', value: TRAINING_FOLDER }, }), + prisma.settings.upsert({ + where: { key: 'DATASETS_FOLDER' }, + update: { value: DATASETS_FOLDER }, + create: { key: 'DATASETS_FOLDER', value: DATASETS_FOLDER }, + }), ]); return NextResponse.json({ success: true }); diff --git a/ui/src/app/dashboard/page.tsx b/ui/src/app/dashboard/page.tsx index fd05095f..94394919 100644 --- a/ui/src/app/dashboard/page.tsx +++ b/ui/src/app/dashboard/page.tsx @@ -5,7 +5,7 @@ import GpuMonitor from '@/components/GPUMonitor'; export default function Dashboard() { return (
-

Dashboard

+

Dashboard

); diff --git a/ui/src/app/datasets/[datasetName]/page.tsx b/ui/src/app/datasets/[datasetName]/page.tsx new file mode 100644 index 00000000..43725fb7 --- /dev/null +++ b/ui/src/app/datasets/[datasetName]/page.tsx @@ -0,0 +1,73 @@ +'use client'; + +import { useEffect, useState, use } from 'react'; +import Card from '@/components/Card'; +import { Modal } from '@/components/Modal'; +import Link from 'next/link'; +import { TextInput } from '@/components/formInputs'; +import { useRouter } from 'next/router'; +import DatasetImageCard from '@/components/DatasetImageCard'; + +export default function DatasetPage({ params }: { params: { datasetName: string } }) { + const [imgList, setImgList] = useState<{ img_path: string }[]>([]); + const usableParams = use(params as any) as { datasetName: string }; + const datasetName = usableParams.datasetName; + const [newDatasetName, setNewDatasetName] = useState(''); + const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false); + const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle'); + + const refreshImageList = (dbName: string) => { + setStatus('loading'); + fetch('/api/datasets/listImages', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ datasetName: dbName }), + }) + .then(res => res.json()) + .then(data => { + console.log('Images:', data.images); + // sort + data.images.sort((a: { img_path: string }, b: { img_path: string }) => a.img_path.localeCompare(b.img_path)); + setImgList(data.images); + setStatus('success'); + }) + .catch(error => { + console.error('Error fetching images:', error); + setStatus('error'); + }); + }; + useEffect(() => { + if (datasetName) { + refreshImageList(datasetName); + } + }, [datasetName]); + return ( + <> +
+
+
+

Dataset: {datasetName}

+
+
+ + {status === 'loading' &&

Loading...

} + {status === 'error' &&

Error fetching images

} + {status === 'success' && ( +
+ {imgList.length === 0 &&

No images found

} + {imgList.map(img => ( + + ))} +
+ )} + +
+ + ); +} diff --git a/ui/src/app/datasets/page.tsx b/ui/src/app/datasets/page.tsx new file mode 100644 index 00000000..3d1d7a2b --- /dev/null +++ b/ui/src/app/datasets/page.tsx @@ -0,0 +1,124 @@ +'use client'; + +import { useEffect, useState } from 'react'; +import Card from '@/components/Card'; +import { Modal } from '@/components/Modal'; +import Link from 'next/link'; +import { TextInput } from '@/components/formInputs'; + +export default function Datasets() { + const [datasets, setDatasets] = useState([]); + const [newDatasetName, setNewDatasetName] = useState(''); + const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false); + const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle'); + + const refreshDatasets = () => { + setStatus('loading'); + fetch('/api/datasets/list') + .then(res => res.json()) + .then(data => { + console.log('Datasets:', data); + // sort + data.sort((a: string, b: string) => a.localeCompare(b)); + setDatasets(data); + setStatus('success'); + }) + .catch(error => { + console.error('Error fetching datasets:', error); + setStatus('error'); + }); + }; + useEffect(() => { + refreshDatasets(); + }, []); + return ( + <> +
+
+
+

Datasets

+
+
+ +
+
+ + {status === 'loading' &&

Loading...

} + {status === 'error' &&

Error fetching datasets

} + {status === 'success' && ( +
+ {datasets.length === 0 &&

No datasets found

} + {datasets.map((dataset: string) => ( + + {dataset} + + ))} +
+ )} + +
+ setIsNewDatasetModalOpen(false)} + title="New Dataset" + size="md" + > +
+
{ + e.preventDefault(); + console.log('Creating new dataset'); + // make post with name to create new dataset + fetch('/api/datasets/create', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ name: newDatasetName }), + }) + .then(res => res.json()) + .then(data => { + console.log('New dataset created:', data); + refreshDatasets(); + setNewDatasetName(''); + setIsNewDatasetModalOpen(false); + }) + .catch(error => { + console.error('Error creating new dataset:', error); + }); + }} + > +
+ This will create a new folder with the name below in your dataset folder. +
+
+ setNewDatasetName(value)} /> +
+ +
+ + +
+
+
+ + + ); +} diff --git a/ui/src/app/settings/page.tsx b/ui/src/app/settings/page.tsx index 6dc5f59b..9e7a7097 100644 --- a/ui/src/app/settings/page.tsx +++ b/ui/src/app/settings/page.tsx @@ -6,6 +6,7 @@ export default function Settings() { const [settings, setSettings] = useState({ HF_TOKEN: '', TRAINING_FOLDER: '', + DATASETS_FOLDER: '', }); const [status, setStatus] = useState<'idle' | 'saving' | 'success' | 'error'>('idle'); @@ -17,6 +18,7 @@ export default function Settings() { setSettings({ HF_TOKEN: data.HF_TOKEN || '', TRAINING_FOLDER: data.TRAINING_FOLDER || '', + DATASETS_FOLDER: data.DATASETS_FOLDER || '', }); }) .catch(error => console.error('Error fetching settings:', error)); @@ -65,7 +67,8 @@ export default function Settings() { {' '} Huggingface - if you need to access gated/private models. + {' '} + if you need to access gated/private models. Training Folder Path
- We will store your training information here. Must be an absolute path. If blank, it will default to the output folder in the project root. + We will store your training information here. Must be an absolute path. If blank, it will default to the + output folder in the project root.
+ +
+ + +