159 lines
4.1 KiB
JavaScript
Executable File
159 lines
4.1 KiB
JavaScript
Executable File
#!/usr/bin/env node
|
|
|
|
/**
|
|
* Node.js background removal script using RMBG SDK
|
|
* Usage: node remove-bg-node.js <input> [options]
|
|
*/
|
|
|
|
const { rmbg, createBriaaiModel, createModnetModel, createU2netpModel } = require('rmbg')
|
|
const { readFileSync, writeFileSync, existsSync } = require('fs')
|
|
const { basename, extname } = require('path')
|
|
|
|
// Parse command line arguments
|
|
const args = process.argv.slice(2)
|
|
|
|
if (args.length === 0 || args.includes('--help') || args.includes('-h')) {
|
|
console.log(`
|
|
Background Removal using RMBG SDK
|
|
|
|
Usage: node remove-bg-node.js <input> [options]
|
|
|
|
Arguments:
|
|
input Input image file path (required)
|
|
|
|
Options:
|
|
-o, --output <path> Output file path (default: auto-generated)
|
|
-m, --model <name> Model: briaai, modnet, u2netp (default: modnet)
|
|
-r, --resolution <n> Max resolution in pixels (default: 2048)
|
|
-p, --progress Show progress information
|
|
-h, --help Show this help message
|
|
|
|
Examples:
|
|
node remove-bg-node.js photo.jpg
|
|
node remove-bg-node.js photo.jpg -m briaai -o output.png
|
|
node remove-bg-node.js photo.jpg -r 4096 -p
|
|
`)
|
|
process.exit(0)
|
|
}
|
|
|
|
// Parse options
|
|
const input = args[0]
|
|
let output = null
|
|
let model = 'modnet'
|
|
let maxResolution = 2048
|
|
let showProgress = false
|
|
|
|
for (let i = 1; i < args.length; i++) {
|
|
const arg = args[i]
|
|
if (arg === '-o' || arg === '--output') {
|
|
output = args[++i]
|
|
} else if (arg === '-m' || arg === '--model') {
|
|
model = args[++i]
|
|
} else if (arg === '-r' || arg === '--resolution') {
|
|
maxResolution = parseInt(args[++i], 10)
|
|
} else if (arg === '-p' || arg === '--progress') {
|
|
showProgress = true
|
|
}
|
|
}
|
|
|
|
// Validate input
|
|
if (!input) {
|
|
console.error('Error: Input file is required')
|
|
process.exit(1)
|
|
}
|
|
|
|
if (!existsSync(input)) {
|
|
console.error(`Error: Input file '${input}' not found`)
|
|
process.exit(1)
|
|
}
|
|
|
|
// Generate output filename if not provided
|
|
if (!output) {
|
|
const name = basename(input, extname(input))
|
|
output = `${name}-no-bg.png`
|
|
}
|
|
|
|
// Select model
|
|
let modelInstance
|
|
switch (model.toLowerCase()) {
|
|
case 'briaai':
|
|
modelInstance = createBriaaiModel()
|
|
break
|
|
case 'u2netp':
|
|
modelInstance = createU2netpModel()
|
|
break
|
|
case 'modnet':
|
|
default:
|
|
modelInstance = createModnetModel()
|
|
break
|
|
}
|
|
|
|
// Display configuration
|
|
console.log('Background Removal Configuration:')
|
|
console.log(` Input: ${input}`)
|
|
console.log(` Model: ${model}`)
|
|
console.log(` Output: ${output}`)
|
|
console.log(` Resolution: ${maxResolution}`)
|
|
console.log('')
|
|
|
|
// Remove background
|
|
async function removeBackground() {
|
|
try {
|
|
console.log('Processing...')
|
|
const startTime = Date.now()
|
|
|
|
const options = {
|
|
model: modelInstance,
|
|
maxResolution,
|
|
output
|
|
}
|
|
|
|
if (showProgress) {
|
|
options.onProgress = (progress, download, process) => {
|
|
const percent = Math.round(progress * 100)
|
|
const downloadPercent = Math.round(download * 100)
|
|
const processPercent = Math.round(process * 100)
|
|
|
|
process.stdout.write(
|
|
`\rProgress: ${percent}% | Download: ${downloadPercent}% | Process: ${processPercent}%`
|
|
)
|
|
}
|
|
}
|
|
|
|
await rmbg(input, options)
|
|
|
|
const duration = ((Date.now() - startTime) / 1000).toFixed(2)
|
|
|
|
if (showProgress) {
|
|
console.log('') // New line after progress
|
|
}
|
|
|
|
console.log('✓ Background removed successfully')
|
|
console.log(` Output: ${output}`)
|
|
console.log(` Duration: ${duration}s`)
|
|
|
|
// Display file sizes
|
|
const inputStats = require('fs').statSync(input)
|
|
const outputStats = require('fs').statSync(output)
|
|
|
|
console.log('')
|
|
console.log('File sizes:')
|
|
console.log(` Input: ${formatBytes(inputStats.size)}`)
|
|
console.log(` Output: ${formatBytes(outputStats.size)}`)
|
|
} catch (error) {
|
|
console.error('✗ Background removal failed:', error.message)
|
|
process.exit(1)
|
|
}
|
|
}
|
|
|
|
function formatBytes(bytes) {
|
|
if (bytes === 0) return '0 Bytes'
|
|
const k = 1024
|
|
const sizes = ['Bytes', 'KB', 'MB', 'GB']
|
|
const i = Math.floor(Math.log(bytes) / Math.log(k))
|
|
return Math.round(bytes / Math.pow(k, i) * 100) / 100 + ' ' + sizes[i]
|
|
}
|
|
|
|
// Run
|
|
removeBackground()
|