Done

2026-04-01 13:24:03 +10:00
parent 3813fc4cbf
commit b24adb0957
3 changed files with 654 additions and 28 deletions
@@ -3,43 +3,108 @@ 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 replicate = new Replicate({ auth: process.env.REPLICATE_API_TOKEN });
+const removeWatermark = async () => {

-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."
+    const replicate = new Replicate({auth: process.env.REPLICATE_API_TOKEN});

-fs.mkdirSync(outputDirectory, { recursive: true });
+    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."

-const pngFiles = fs
-    .readdirSync(inputDirectory)
-    .filter((file) => file.toLowerCase().endsWith(".png"))
+    fs.mkdirSync(outputDirectory, {recursive: true});

-for (const file of pngFiles) {
-    const filePath = path.join(inputDirectory, file);
-    const outputPath = path.join(outputDirectory, file);
+    const pngFiles = fs
+        .readdirSync(inputDirectory)
+        .filter((file) => file.toLowerCase().endsWith(".png"))

-    if (fs.existsSync(outputPath)) {
-        continue;
+    for (const file of pngFiles) {
+        const filePath = path.join(inputDirectory, file);
+        const outputPath = path.join(outputDirectory, file);
+
+        if (fs.existsSync(outputPath)) {
+            console.log(`Converting: ${file} (webp → png)`);
+            const webpData = fs.readFileSync(outputPath);
+            const pngData = await sharp(webpData).png().toBuffer();
+            fs.writeFileSync(outputPath, pngData);
+            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}`);
    }
+}

-    console.log(`Processing: ${filePath}`);
-    const watermarkData = fs.readFileSync(watermarkImage);
+const removeBackground = async () => {
+    fs.mkdirSync(backgroundRemovedDirectory, { recursive: true });

-    const imageData = fs.readFileSync(filePath);
-    const watermarkBase64 = `data:image/png;base64,${watermarkData.toString("base64")}`;
+    const pngFiles = fs
+        .readdirSync(outputDirectory)
+        .filter((file) => file.toLowerCase().endsWith(".png"));

-    const imageBase64 = `data:image/png;base64,${imageData.toString("base64")}`;
+    // Load mask once
+    const { data: maskData } = await sharp("images/Mask.png")
+        .greyscale()
+        .raw()
+        .toBuffer({ resolveWithObject: true });

-    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);
+    for (const file of pngFiles) {
+        const filePath = path.join(outputDirectory, file);
+        const outputPath = path.join(backgroundRemovedDirectory, file);

-    console.log(`Saved: ${outputPath}`);
-}
+        if (fs.existsSync(outputPath)) {
+            console.log(`Skipping: ${file} (already processed)`);
+            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]();