Add cropping

jobs/queue.js

@@ -1,4 +1,5 @@
 import PQueue from 'p-queue';
+import sharp from 'sharp';
 import { generateText } from 'ai';
 import { createAnthropic } from '@ai-sdk/anthropic';
 import { createOpenAI } from '@ai-sdk/openai';
@@ -78,8 +79,140 @@ async function setStatus(job, status, extra = {}) {
   broadcast({ type: 'job_update', job: job.toJSON() });
 }
 
+/** Padding factor applied to each side of the detected bounding box. */
+const BBOX_PADDING = parseFloat(process.env.BBOX_PADDING ?? '0.15');
+
+/**
+ * Extract the raw image buffer + mime type from a data-URL string.
+ */
+function parseDataUrl(dataUrl) {
+  const match = dataUrl.match(/^data:([^;]+);base64,(.+)$/s);
+  if (!match) throw new Error('Invalid data URL');
+  return { mimeType: match[1], buffer: Buffer.from(match[2], 'base64') };
+}
+
+/**
+ * Build a data-URL from a buffer + mime type.
+ */
+function toDataUrl(buffer, mimeType) {
+  return `data:${mimeType};base64,${buffer.toString('base64')}`;
+}
+
 // ---------------------------------------------------------------------------
-// Main runner
+// Pass 1 — Ask the LLM to locate the subject described in the prompt
+// ---------------------------------------------------------------------------
+const BBOX_SYSTEM_PROMPT = `You are a vision assistant that locates objects in images.
+The user will give you an image and a description of what they are interested in.
+Your job is to identify the main subject described and return the bounding box
+coordinates for that subject as a proportion of the image dimensions (values 0-1).
+
+You MUST respond with ONLY a JSON object in this exact format, no other text:
+{
+  "subject": "<short description of what you identified>",
+  "bbox": {
+    "x": <left edge 0-1>,
+    "y": <top edge 0-1>,
+    "width": <box width 0-1>,
+    "height": <box height 0-1>
+  }
+}
+
+If the prompt doesn't clearly reference a specific subject in the image, or if the
+subject fills most of the image already, return the full image:
+{ "subject": "full image", "bbox": { "x": 0, "y": 0, "width": 1, "height": 1 } }`;
+
+async function detectSubjectBbox(model, imageDataUrl, prompt) {
+  const { text } = await generateText({
+    model,
+    system: BBOX_SYSTEM_PROMPT,
+    messages: [
+      {
+        role: 'user',
+        content: [
+          {
+            type: 'text',
+            text: `Locate the main subject referenced in this prompt and return its bounding box.\n\nPrompt: "${prompt}"`,
+          },
+          { type: 'image', image: imageDataUrl },
+        ],
+      },
+    ],
+    maxTokens: 256,
+  });
+
+  // Parse JSON from the response — tolerate markdown fences
+  const cleaned = text.replace(/```json\s*/g, '').replace(/```/g, '').trim();
+  const parsed = JSON.parse(cleaned);
+
+  // Validate and clamp values
+  const bbox = parsed.bbox;
+  if (
+    typeof bbox?.x !== 'number' ||
+    typeof bbox?.y !== 'number' ||
+    typeof bbox?.width !== 'number' ||
+    typeof bbox?.height !== 'number'
+  ) {
+    throw new Error('LLM returned invalid bbox structure');
+  }
+
+  const clamp = (v) => Math.max(0, Math.min(1, v));
+  return {
+    subject: parsed.subject ?? 'unknown',
+    bbox: {
+      x:      clamp(bbox.x),
+      y:      clamp(bbox.y),
+      width:  clamp(bbox.width),
+      height: clamp(bbox.height),
+    },
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Crop helper — takes a data-URL, bounding box (0-1), and padding factor
+// ---------------------------------------------------------------------------
+async function cropImage(imageDataUrl, bbox, padding = BBOX_PADDING) {
+  const { mimeType, buffer } = parseDataUrl(imageDataUrl);
+  const metadata = await sharp(buffer).metadata();
+  const imgW = metadata.width;
+  const imgH = metadata.height;
+
+  // Convert proportional bbox to pixel values
+  let left   = bbox.x * imgW;
+  let top    = bbox.y * imgH;
+  let width  = bbox.width * imgW;
+  let height = bbox.height * imgH;
+
+  // Add padding (proportional to the box dimensions)
+  const padX = width * padding;
+  const padY = height * padding;
+  left   = Math.max(0, left - padX);
+  top    = Math.max(0, top - padY);
+  width  = Math.min(imgW - left, width + padX * 2);
+  height = Math.min(imgH - top, height + padY * 2);
+
+  // Round to integers for sharp
+  left   = Math.round(left);
+  top    = Math.round(top);
+  width  = Math.round(width);
+  height = Math.round(height);
+
+  // Guard against degenerate crops
+  if (width < 1 || height < 1) {
+    return { croppedDataUrl: imageDataUrl, cropPixels: null };
+  }
+
+  const croppedBuf = await sharp(buffer)
+    .extract({ left, top, width, height })
+    .toBuffer();
+
+  return {
+    croppedDataUrl: toDataUrl(croppedBuf, mimeType),
+    cropPixels: { left, top, width, height, imgW, imgH },
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Main runner — two-pass: detect bbox → crop → analyse cropped image
 // ---------------------------------------------------------------------------
 async function runJob(jobId) {
   const job = await Job.findByPk(jobId);
@@ -89,15 +222,69 @@ async function runJob(jobId) {
 
   try {
     resolveModelMeta(job.model);
+    const model = resolveModel(job.model);
 
+    // ---- Pass 1: detect subject bounding box ----
+    let bboxResult;
+    try {
+      bboxResult = await detectSubjectBbox(model, job.imageDataUrl, job.prompt);
+    } catch (bboxErr) {
+      // If bbox detection fails, fall back to the full image
+      console.warn(`Bbox detection failed, using full image: ${bboxErr.message}`);
+      bboxResult = {
+        subject: 'full image (fallback)',
+        bbox: { x: 0, y: 0, width: 1, height: 1 },
+      };
+    }
+
+    // Persist the detected bbox for debugging / UI display
+    await job.update({
+      detectedSubject: bboxResult.subject,
+      bboxX:      bboxResult.bbox.x,
+      bboxY:      bboxResult.bbox.y,
+      bboxWidth:  bboxResult.bbox.width,
+      bboxHeight: bboxResult.bbox.height,
+    });
+    broadcast({ type: 'job_update', job: job.toJSON() });
+
+    // ---- Crop the image around the detected subject ----
+    const isFullImage =
+      bboxResult.bbox.x === 0 &&
+      bboxResult.bbox.y === 0 &&
+      bboxResult.bbox.width === 1 &&
+      bboxResult.bbox.height === 1;
+
+    let imageForAnalysis = job.imageDataUrl;
+
+    if (!isFullImage) {
+      const { croppedDataUrl, cropPixels } = await cropImage(
+        job.imageDataUrl,
+        bboxResult.bbox,
+      );
+      imageForAnalysis = croppedDataUrl;
+
+      // Optionally store the cropped image for UI preview
+      await job.update({ croppedImageDataUrl: croppedDataUrl });
+      broadcast({ type: 'job_update', job: job.toJSON() });
+
+      if (cropPixels) {
+        console.log(
+          `Cropped to ${cropPixels.width}×${cropPixels.height} ` +
+          `from ${cropPixels.imgW}×${cropPixels.imgH} ` +
+          `(subject: "${bboxResult.subject}")`,
+        );
+      }
+    }
+
+    // ---- Pass 2: run the original prompt against the (cropped) image ----
     const { text, usage } = await generateText({
-      model: resolveModel(job.model),
+      model,
       messages: [
         {
           role: 'user',
           content: [
             { type: 'text',  text: job.prompt },
-            { type: 'image', image: job.imageDataUrl },
+            { type: 'image', image: imageForAnalysis },
           ],
         },
       ],