120 lines
4.3 KiB
TypeScript
120 lines
4.3 KiB
TypeScript
import fs from "fs";
|
|
import path from "path";
|
|
|
|
const inputDirectory = "images/unprocessed";
|
|
const outputDirectory = "images/processed";
|
|
const backgroundRemovedDirectory = "images/background_removed";
|
|
const watermarkImage = "images/Watermark.png";
|
|
import Replicate from "replicate";
|
|
import sharp from "sharp";
|
|
|
|
const removeWatermark = async () => {
|
|
const replicate = new Replicate({ auth: process.env.REPLICATE_API_TOKEN });
|
|
|
|
const prompt = "Carefully remove the watermark in image 1 from image 2, leaving everything else unchanged. Preserve icon details while removing watermark text that says 'Airhex.com'. Save with a white background."
|
|
|
|
fs.mkdirSync(outputDirectory, { recursive: true });
|
|
|
|
const pngFiles = fs
|
|
.readdirSync(inputDirectory)
|
|
.filter((file) => file.toLowerCase().endsWith(".png"));
|
|
|
|
for (const file of pngFiles) {
|
|
const filePath = path.join(inputDirectory, file);
|
|
const outputPath = path.join(outputDirectory, file);
|
|
|
|
// Ensure the output directory exists before every write
|
|
fs.mkdirSync(path.dirname(outputPath), { recursive: true });
|
|
|
|
if (fs.existsSync(outputPath)) {
|
|
const fileBuffer = fs.readFileSync(outputPath);
|
|
const metadata = await sharp(fileBuffer).metadata();
|
|
|
|
if (metadata.format !== 'png') {
|
|
console.log(`About to write to: ${outputPath}`);
|
|
console.log(`Dir exists: ${fs.existsSync(path.dirname(outputPath))}`);
|
|
console.log(`Dir writable:`, (() => { try { fs.accessSync(path.dirname(outputPath), fs.constants.W_OK); return true; } catch { return false; } })());
|
|
console.log(`File writable:`, (() => { try { fs.accessSync(outputPath, fs.constants.W_OK); return true; } catch { return false; } })());
|
|
const pngData = await sharp(fileBuffer).png().toBuffer();
|
|
console.log(`PNG buffer size: ${pngData.length}`);
|
|
fs.writeFileSync(outputPath, pngData);
|
|
console.log(`Write succeeded`);
|
|
}
|
|
|
|
continue;
|
|
}
|
|
|
|
console.log(`Processing: ${filePath}`);
|
|
const watermarkData = fs.readFileSync(watermarkImage);
|
|
const imageData = fs.readFileSync(filePath);
|
|
const watermarkBase64 = `data:image/png;base64,${watermarkData.toString("base64")}`;
|
|
const imageBase64 = `data:image/png;base64,${imageData.toString("base64")}`;
|
|
|
|
const output = await replicate.run("qwen/qwen-image-edit-plus", {
|
|
input: {
|
|
image: [watermarkBase64, imageBase64],
|
|
prompt,
|
|
},
|
|
}) as { url(): string }[];
|
|
|
|
const resultBuffer = Buffer.from(await (await fetch(output[0].url())).arrayBuffer());
|
|
fs.writeFileSync(outputPath, resultBuffer);
|
|
|
|
console.log(`Saved: ${outputPath}`);
|
|
}
|
|
};
|
|
|
|
const removeBackground = async () => {
|
|
fs.mkdirSync(backgroundRemovedDirectory, { recursive: true });
|
|
|
|
const pngFiles = fs
|
|
.readdirSync(outputDirectory)
|
|
.filter((file) => file.toLowerCase().endsWith(".png"));
|
|
|
|
// Load mask once
|
|
const { data: maskData } = await sharp("images/Mask.png")
|
|
.greyscale()
|
|
.raw()
|
|
.toBuffer({ resolveWithObject: true });
|
|
|
|
for (const file of pngFiles) {
|
|
const filePath = path.join(outputDirectory, file);
|
|
const outputPath = path.join(backgroundRemovedDirectory, file);
|
|
|
|
if (fs.existsSync(outputPath)) {
|
|
continue;
|
|
}
|
|
|
|
console.log(`Processing: ${filePath}`);
|
|
|
|
const { data, info } = await sharp(filePath)
|
|
.ensureAlpha()
|
|
.raw()
|
|
.toBuffer({ resolveWithObject: true });
|
|
|
|
for (let i = 0; i < info.width * info.height; i++) {
|
|
data[i * 4 + 3] = maskData[i];
|
|
}
|
|
|
|
await sharp(data, {
|
|
raw: { width: info.width, height: info.height, channels: 4 }
|
|
}).png().resize(256, 256).toFile(outputPath);
|
|
|
|
console.log(`Saved: ${outputPath}`);
|
|
}
|
|
};
|
|
|
|
const commands: Record<string, () => Promise<void>> = {
|
|
removeWatermark,
|
|
removeBackground,
|
|
};
|
|
|
|
const command = process.argv[2];
|
|
|
|
if (!command || !commands[command]) {
|
|
console.error(`Usage: npx tsx process-images.ts <command>`);
|
|
console.error(`Available commands: ${Object.keys(commands).join(", ")}`);
|
|
process.exit(1);
|
|
}
|
|
|
|
await commands[command](); |