ai-chatbot/lib/canvas/scale-adjust-product-lock.ts

import sharp from "sharp";

import { executeRmbgJob } from "@/lib/ai/rmbg/service";
export type CanvasScaleAdjustRect = {
  height: number;
  width: number;
  x: number;
  y: number;
};

export type CanvasScaleAdjustIntentRegion = {
  index: number;
  rect: CanvasScaleAdjustRect;
  targetRect: CanvasScaleAdjustRect;
};

type PixelBox = {
  height: number;
  left: number;
  top: number;
  width: number;
};

export type CanvasScaleAdjustProductLockAssets = {
  backgroundPlateBuffer?: Buffer;
  modelMaskBuffer: Buffer;
  overlayBuffer: Buffer;
  repairMaskBuffer: Buffer;
  subjectReferenceUrls: string[];
};

function clamp(value: number, min: number, max: number) {
  return Math.min(max, Math.max(min, value));
}

function rectToPixelBox(
  rect: CanvasScaleAdjustRect,
  sourceWidth: number,
  sourceHeight: number
): PixelBox {
  const left = clamp(Math.floor(rect.x * sourceWidth), 0, sourceWidth - 1);
  const top = clamp(Math.floor(rect.y * sourceHeight), 0, sourceHeight - 1);
  const right = clamp(
    Math.ceil((rect.x + rect.width) * sourceWidth),
    left + 1,
    sourceWidth
  );
  const bottom = clamp(
    Math.ceil((rect.y + rect.height) * sourceHeight),
    top + 1,
    sourceHeight
  );
  return {
    height: bottom - top,
    left,
    top,
    width: right - left,
  };
}

async function normalizeSubjectBuffer(params: {
  box: PixelBox;
  imageBase64: string;
}) {
  const { box, imageBase64 } = params;
  const subjectBuffer = await sharp(Buffer.from(imageBase64, "base64"), {
    failOn: "none",
  })
    .rotate()
    .resize(box.width, box.height, {
      fit: "fill",
      kernel: sharp.kernel.lanczos3,
    })
    .ensureAlpha()
    .png({ compressionLevel: 6 })
    .toBuffer();
  const alpha = await sharp(subjectBuffer, { failOn: "none" })
    .extractChannel(3)
    .raw()
    .toBuffer();
  const visiblePixels = alpha.reduce(
    (count, value) => count + (value >= 18 ? 1 : 0),
    0
  );
  if (visiblePixels < Math.max(32, box.width * box.height * 0.008)) {
    throw new Error("比例调整无法识别完整产品，请重新完整框选产品");
  }
  return { alpha, subjectBuffer };
}

async function buildExpandedMaskPiece(params: {
  alpha: Buffer;
  box: PixelBox;
}) {
  const { alpha, box } = params;
  const radius = Math.max(2, Math.round(Math.min(box.width, box.height) * 0.06));
  const expandedAlpha = await sharp(alpha, {
    raw: { channels: 1, height: box.height, width: box.width },
  })
    .blur(radius)
    .threshold(8)
    .png({ compressionLevel: 6 })
    .toBuffer();
  const whiteRgb = await sharp({
    create: {
      background: { b: 255, g: 255, r: 255 },
      channels: 3,
      height: box.height,
      width: box.width,
    },
  })
    .png({ compressionLevel: 6 })
    .toBuffer();
  return sharp(whiteRgb).joinChannel(expandedAlpha).png({ compressionLevel: 6 }).toBuffer();
}

export async function buildScaleAdjustProductLockAssets(params: {
  regions: CanvasScaleAdjustIntentRegion[];
  sourceBuffer: Buffer;
  sourceHeight: number;
  sourceImageUrl: string;
  sourceWidth: number;
  userId: string;
}): Promise<CanvasScaleAdjustProductLockAssets> {
  const {
    regions,
    sourceBuffer,
    sourceHeight,
    sourceImageUrl,
    sourceWidth,
    userId,
  } = params;
  const overlayInputs: sharp.OverlayOptions[] = [];
  const maskInputs: sharp.OverlayOptions[] = [];
  const subjectReferenceUrls: string[] = [];

  for (const region of regions) {
    const sourceBox = rectToPixelBox(
      region.rect,
      sourceWidth,
      sourceHeight
    );
    const targetBox = rectToPixelBox(
      region.targetRect,
      sourceWidth,
      sourceHeight
    );
    const cropBuffer = await sharp(sourceBuffer, { failOn: "none" })
      .extract(sourceBox)
      .png({ compressionLevel: 6 })
      .toBuffer();
    const subject = await executeRmbgJob({
      imageBuffer: cropBuffer,
      prompt: "比例调整完整产品蒙版",
      skipGallerySave: true,
      sourceImageUrl,
      userId,
    });
    const { alpha, subjectBuffer } = await normalizeSubjectBuffer({
      box: sourceBox,
      imageBase64: subject.image_base64,
    });
    const scaledSubjectBuffer = await sharp(subjectBuffer, { failOn: "none" })
      .resize(targetBox.width, targetBox.height, {
        fit: "fill",
        kernel: sharp.kernel.lanczos3,
      })
      .png({ compressionLevel: 6 })
      .toBuffer();
    overlayInputs.push({
      input: scaledSubjectBuffer,
      left: targetBox.left,
      top: targetBox.top,
    });
    maskInputs.push({
      input: await buildExpandedMaskPiece({ alpha, box: sourceBox }),
      left: sourceBox.left,
      top: sourceBox.top,
    });
    subjectReferenceUrls.push(subject.original_url);
  }

  const transparentCanvas = {
    create: {
      background: { alpha: 0, b: 0, g: 0, r: 0 },
      channels: 4 as const,
      height: sourceHeight,
      width: sourceWidth,
    },
  };
  const [modelMaskBuffer, overlayBuffer, repairMaskBuffer] = await Promise.all([
    sharp({
      create: {
        background: { b: 0, g: 0, r: 0 },
        channels: 3,
        height: sourceHeight,
        width: sourceWidth,
      },
    })
      .composite(maskInputs)
      .flatten({ background: "#000000" })
      .png({ compressionLevel: 6 })
      .toBuffer(),
    sharp(transparentCanvas)
      .composite(overlayInputs)
      .png({ compressionLevel: 6 })
      .toBuffer(),
    sharp(transparentCanvas)
      .composite(maskInputs)
      .png({ compressionLevel: 6 })
      .toBuffer(),
  ]);
  const sourceAlpha = await sharp(sourceBuffer, { failOn: "none" })
    .rotate()
    .resize(sourceWidth, sourceHeight, { fit: "fill" })
    .ensureAlpha()
    .extractChannel(3)
    .raw()
    .toBuffer();
  const transparentPixelCount = sourceAlpha.reduce(
    (count, value) => count + (value < 250 ? 1 : 0),
    0
  );
  const hasTransparentBackground =
    transparentPixelCount >= sourceWidth * sourceHeight * 0.05;
  const backgroundPlateBuffer = hasTransparentBackground
    ? await sharp(transparentCanvas)
        .png({ compressionLevel: 6 })
        .toBuffer()
    : undefined;

  return {
    ...(backgroundPlateBuffer ? { backgroundPlateBuffer } : {}),
    modelMaskBuffer,
    overlayBuffer,
    repairMaskBuffer,
    subjectReferenceUrls,
  };
}