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;
  personExclusionApplied: boolean;
  personExclusionMaskBuffer?: Buffer;
  repairMaskBuffer: Buffer;
  sourceSubjectMaskBuffer: 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;
  exclusionAlpha?: Buffer | null;
  imageBase64: string;
}) {
  const { box, exclusionAlpha, imageBase64 } = params;
  const rawSubjectBuffer = 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 rawAlpha = await sharp(rawSubjectBuffer, { failOn: "none" })
    .extractChannel(3)
    .raw()
    .toBuffer();
  const alpha = Buffer.from(rawAlpha);
  if (exclusionAlpha?.length === rawAlpha.length) {
    let productPixelCount = 0;
    let overlapPixelCount = 0;
    for (let index = 0; index < alpha.length; index += 1) {
      if (alpha[index] >= 18) productPixelCount += 1;
      if (alpha[index] >= 18 && exclusionAlpha[index] >= 8) {
        overlapPixelCount += 1;
      }
    }
    const overlapLimit = Math.max(96, Math.round(productPixelCount * 0.004));
    if (overlapPixelCount >= overlapLimit) {
      throw new Error(
        "目标产品被人物遮挡，暂无法在保持产品完整的前提下移动，请选择未被人物遮挡的完整产品"
      );
    }
    for (let index = 0; index < alpha.length; index += 1) {
      if (exclusionAlpha[index] >= 8) alpha[index] = 0;
    }
  }
  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("比例调整无法识别完整产品，请重新完整框选产品");
  }
  const subjectBuffer = exclusionAlpha
    ? await sharp(
        await sharp(rawSubjectBuffer, { failOn: "none" })
          .removeAlpha()
          .toColorspace("srgb")
          .raw()
          .toBuffer(),
        {
          raw: { channels: 3, height: box.height, width: box.width },
        }
      )
        .joinChannel(alpha, {
          raw: { channels: 1, height: box.height, width: box.width },
        })
        .png({ compressionLevel: 6 })
        .toBuffer()
    : rawSubjectBuffer;
  return { alpha, subjectBuffer };
}

async function buildExpandedMaskPiece(params: {
  alpha: Buffer;
  box: PixelBox;
  exclusionAlpha?: Buffer | null;
}) {
  const { alpha, box, exclusionAlpha } = params;
  const radius = Math.max(2, Math.round(Math.min(box.width, box.height) * 0.06));
  const blurredAlphaRaw = await sharp(alpha, {
    raw: { channels: 1, height: box.height, width: box.width },
  })
    .blur(radius)
    .extractChannel(0)
    .raw()
    .toBuffer();
  const expandedAlphaRaw = Buffer.from(blurredAlphaRaw);
  for (let index = 0; index < expandedAlphaRaw.length; index += 1) {
    expandedAlphaRaw[index] = expandedAlphaRaw[index] >= 8 ? 255 : 0;
  }
  if (exclusionAlpha?.length === expandedAlphaRaw.length) {
    for (let index = 0; index < expandedAlphaRaw.length; index += 1) {
      if (exclusionAlpha[index] >= 8) expandedAlphaRaw[index] = 0;
    }
  }
  const expandedAlpha = await sharp(expandedAlphaRaw, {
    raw: { channels: 1, height: box.height, width: box.width },
  })
    .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();
}

async function buildSubjectMaskPiece(params: { alpha: Buffer; box: PixelBox }) {
  const alphaPng = await sharp(params.alpha, {
    raw: { channels: 1, height: params.box.height, width: params.box.width },
  })
    .png({ compressionLevel: 6 })
    .toBuffer();
  const whiteRgb = await sharp({
    create: {
      background: { b: 255, g: 255, r: 255 },
      channels: 3,
      height: params.box.height,
      width: params.box.width,
    },
  })
    .png({ compressionLevel: 6 })
    .toBuffer();
  return sharp(whiteRgb)
    .joinChannel(alphaPng)
    .png({ compressionLevel: 6 })
    .toBuffer();
}

async function buildPersonExclusionAlpha(params: {
  sourceBuffer: Buffer;
  sourceHeight: number;
  sourceImageUrl: string;
  sourceWidth: number;
  userId: string;
}) {
  try {
    const person = await executeRmbgJob({
      foregroundType: "person",
      imageBuffer: params.sourceBuffer,
      prompt: "比例调整人物保护蒙版",
      skipGallerySave: true,
      sourceImageUrl: params.sourceImageUrl,
      userId: params.userId,
    });
    const alpha = await sharp(Buffer.from(person.image_base64, "base64"), {
      failOn: "none",
    })
      .rotate()
      .resize(params.sourceWidth, params.sourceHeight, { fit: "fill" })
      .ensureAlpha()
      .extractChannel(3)
      .raw()
      .toBuffer();
    const visiblePixels = alpha.reduce(
      (count, value) => count + (value >= 18 ? 1 : 0),
      0
    );
    if (visiblePixels < Math.max(32, params.sourceWidth * params.sourceHeight * 0.001)) {
      return null;
    }
    return sharp(alpha, {
      raw: {
        channels: 1,
        height: params.sourceHeight,
        width: params.sourceWidth,
      },
    })
      .blur(Math.max(1, Math.min(params.sourceWidth, params.sourceHeight) * 0.002))
      .threshold(8)
      .extractChannel(0)
      .raw()
      .toBuffer();
  } catch {
    return null;
  }
}

async function extractAlphaRegion(params: {
  alpha: Buffer | null;
  box: PixelBox;
  sourceHeight: number;
  sourceWidth: number;
}) {
  if (!params.alpha) return null;
  return sharp(params.alpha, {
    raw: {
      channels: 1,
      height: params.sourceHeight,
      width: params.sourceWidth,
    },
  })
    .extract(params.box)
    .extractChannel(0)
    .raw()
    .toBuffer();
}

export async function buildScaleAdjustProductLockAssets(params: {
  protectPeople?: boolean;
  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 sourceSubjectMaskInputs: sharp.OverlayOptions[] = [];
  const subjectReferenceUrls: string[] = [];
  const personExclusionAlpha = params.protectPeople
    ? await buildPersonExclusionAlpha({
        sourceBuffer,
        sourceHeight,
        sourceImageUrl,
        sourceWidth,
        userId,
      })
    : null;

  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 regionPersonExclusionAlpha = await extractAlphaRegion({
      alpha: personExclusionAlpha,
      box: sourceBox,
      sourceHeight,
      sourceWidth,
    });
    const { alpha, subjectBuffer } = await normalizeSubjectBuffer({
      box: sourceBox,
      exclusionAlpha: regionPersonExclusionAlpha,
      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,
        exclusionAlpha: regionPersonExclusionAlpha,
      }),
      left: sourceBox.left,
      top: sourceBox.top,
    });
    sourceSubjectMaskInputs.push({
      input: await buildSubjectMaskPiece({ 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, sourceSubjectMaskBuffer] =
    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(),
    sharp(transparentCanvas)
      .composite(sourceSubjectMaskInputs)
      .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 personExclusionMaskBuffer = personExclusionAlpha
    ? await sharp(personExclusionAlpha, {
        raw: { channels: 1, height: sourceHeight, width: sourceWidth },
      })
        .png({ compressionLevel: 6 })
        .toBuffer()
    : undefined;
  const backgroundPlateBuffer = hasTransparentBackground
    ? await sharp(transparentCanvas)
        .png({ compressionLevel: 6 })
        .toBuffer()
    : undefined;

  return {
    ...(backgroundPlateBuffer ? { backgroundPlateBuffer } : {}),
    modelMaskBuffer,
    overlayBuffer,
    personExclusionApplied: Boolean(personExclusionAlpha),
    ...(personExclusionMaskBuffer ? { personExclusionMaskBuffer } : {}),
    repairMaskBuffer,
    sourceSubjectMaskBuffer,
    subjectReferenceUrls,
  };
}