video-product-finder/src/product-detector.ts

import { generateObject } from 'ai';
import { createOpenAI } from '@ai-sdk/openai';
import { z } from 'zod';
import type { ExtractedFrame } from './frame-extractor.ts';
import type { ProductFrame } from './types.ts';
import { imageToBase64 } from './frame-extractor.ts';
import type { VisionConfig } from './index.ts';

// Pass 1: quick filter — discard frames that clearly have no product
const FilterSchema = z.object({
  keep: z.boolean(),
  reason: z.enum(['product_visible', 'content_only', 'hands_only', 'blur', 'transition', 'background_only']),
});

// Pass 2: comparative ranking across all candidates
const RankingSchema = z.object({
  bestFrameIndex: z.number().int(),
  description: z.string(),
  reasoning: z.string(),
  // normalized 0-1 relative to image dimensions: [x1, y1, x2, y2]
  boundingBox: z.tuple([z.number(), z.number(), z.number(), z.number()]),
});

const FILTER_PROMPT = `You are filtering frames from a TikTok/Douyin ecommerce product video.

Keep a frame (keep=true) ONLY if the HERO PRODUCT (the main item being sold) is at least partially visible as a recognizable object.
Discard (keep=false) if: only hands/texture/contents visible, motion blur, black/transition frame, or no product at all.

reason options: product_visible | content_only | hands_only | blur | transition | background_only`;

const RANKING_PROMPT = (count: number) => `You are selecting the single best product image from ${count} video frames for ecommerce image search.

The frames are numbered 0 to ${count - 1} in the order shown.

Pick the ONE frame where the HERO PRODUCT is:
1. Cleanest — fewest distractions, no hands blocking it, no clutter in foreground
2. Most complete — full product silhouette visible, no edges cropped
3. Most isolated — product stands out from background clearly
4. Empty/minimal load preferred — a product without contents (e.g. an empty rack) beats one stuffed with items if both show the full structure equally

Return:
- bestFrameIndex: 0-based index of chosen frame
- description: concise search query under 12 words (product type + material + color + key feature)
- reasoning: one sentence explaining why this frame was chosen
- boundingBox: tight bounding box of the HERO PRODUCT ONLY in the chosen frame as [x1, y1, x2, y2] normalized 0.0–1.0 (top-left origin). Exclude hands, background, and unrelated objects. The product is assumed to be near the center.`;

function createVisionModel(config: VisionConfig) {
  const provider = createOpenAI({ apiKey: config.apiKey, baseURL: config.baseURL });
  return provider(config.model);
}

async function filterFrame(
  frame: ExtractedFrame,
  model: ReturnType<ReturnType<typeof createOpenAI>>,
): Promise<boolean> {
  const base64Image = imageToBase64(frame.imagePath);

  const { object } = await generateObject({
    model,
    schema: FilterSchema,
    messages: [{
      role: 'user',
      content: [
        { type: 'image', image: `data:image/jpeg;base64,${base64Image}` },
        { type: 'text', text: FILTER_PROMPT },
      ],
    }],
  });

  return object.keep;
}

async function rankCandidates(
  candidates: ExtractedFrame[],
  model: ReturnType<ReturnType<typeof createOpenAI>>,
): Promise<{ bestFrame: ExtractedFrame; description: string; reasoning: string; boundingBox: [number, number, number, number] }> {
  const imageContent = candidates.map((f) => ({
    type: 'image' as const,
    image: `data:image/jpeg;base64,${imageToBase64(f.imagePath)}`,
  }));

  const { object } = await generateObject({
    model,
    schema: RankingSchema,
    mode: 'json',
    messages: [{
      role: 'user',
      content: [
        ...imageContent,
        { type: 'text', text: RANKING_PROMPT(candidates.length) },
      ],
    }],
  });

  const idx = Math.max(0, Math.min(object.bestFrameIndex, candidates.length - 1));
  return {
    bestFrame: candidates[idx],
    description: object.description,
    reasoning: object.reasoning,
    boundingBox: object.boundingBox,
  };
}

export async function cropProduct(
  imagePath: string,
  boundingBox: [number, number, number, number],
  outputPath: string,
  paddingFactor = 0.05,
): Promise<string> {
  const sharp = (await import('sharp')).default;
  const meta = await sharp(imagePath).metadata();
  const W = meta.width!;
  const H = meta.height!;

  let [x1, y1, x2, y2] = boundingBox;

  // add padding
  const pw = (x2 - x1) * paddingFactor;
  const ph = (y2 - y1) * paddingFactor;
  x1 = Math.max(0, x1 - pw);
  y1 = Math.max(0, y1 - ph);
  x2 = Math.min(1, x2 + pw);
  y2 = Math.min(1, y2 + ph);

  const left = Math.round(x1 * W);
  const top = Math.round(y1 * H);
  const width = Math.round((x2 - x1) * W);
  const height = Math.round((y2 - y1) * H);

  await sharp(imagePath)
    .extract({ left, top, width, height })
    .jpeg({ quality: 95 })
    .toFile(outputPath);

  return outputPath;
}

async function withConcurrency<T>(
  tasks: (() => Promise<T>)[],
  limit: number,
): Promise<T[]> {
  const results: T[] = new Array(tasks.length);
  let next = 0;
  async function worker() {
    while (next < tasks.length) {
      const i = next++;
      results[i] = await tasks[i]();
    }
  }
  await Promise.all(Array.from({ length: Math.min(limit, tasks.length) }, worker));
  return results;
}

// Skips Pass 1 filter entirely — ranks all frames and always returns the best one.
// Evenly samples down to maxCandidates when there are too many frames.
export async function detectBestFrame(
  frames: ExtractedFrame[],
  concurrency: number = 10,
  visionConfig: VisionConfig,
  maxCandidates: number = 20,
): Promise<ProductFrame | null> {
  if (frames.length === 0) return null;

  const model = createVisionModel(visionConfig);

  let candidates = frames;
  if (frames.length > maxCandidates) {
    const step = frames.length / maxCandidates;
    candidates = Array.from({ length: maxCandidates }, (_, i) => frames[Math.floor(i * step)]);
  }

  const { bestFrame, description, reasoning, boundingBox } = await rankCandidates(candidates, model);

  const croppedPath = bestFrame.imagePath.replace(/\.jpg$/, '_cropped.jpg');
  await cropProduct(bestFrame.imagePath, boundingBox, croppedPath);

  return {
    frameIndex: bestFrame.frameIndex,
    timestampSeconds: bestFrame.timestampSeconds,
    imagePath: bestFrame.imagePath,
    croppedImagePath: croppedPath,
    confidence: 0.95,
    description,
    boundingHint: reasoning,
  };
}

export async function detectProductFrames(
  frames: ExtractedFrame[],
  minConfidence: number,
  concurrency: number = 10,
  visionConfig: VisionConfig,
): Promise<ProductFrame[]> {
  const model = createVisionModel(visionConfig);

  // Pass 1: all frames in parallel, bounded by concurrency
  const keepFlags = await withConcurrency(
    frames.map((f) => () => filterFrame(f, model).catch(() => false)),
    concurrency,
  );

  const candidates = frames.filter((_, i) => keepFlags[i]);
  if (candidates.length === 0) return [];

  // Pass 2: single comparative call — model sees all candidates at once
  const { bestFrame, description, reasoning, boundingBox } = await rankCandidates(candidates, model);

  const croppedPath = bestFrame.imagePath.replace(/\.jpg$/, '_cropped.jpg');
  await cropProduct(bestFrame.imagePath, boundingBox, croppedPath);

  return [{
    frameIndex: bestFrame.frameIndex,
    timestampSeconds: bestFrame.timestampSeconds,
    imagePath: bestFrame.imagePath,
    croppedImagePath: croppedPath,
    confidence: 0.95,
    description,
    boundingHint: reasoning,
  }];
}