Add cropping

This commit is contained in:
khalid@traclabs.com
2026-04-30 15:30:52 -05:00
parent bfc32e78a9
commit 62e5e309ad
4 changed files with 216 additions and 5 deletions

View File

@@ -20,3 +20,8 @@ OLLAMA_API_KEY=your-ollama-api-key
# ── Server ────────────────────────────────────────────────────────────────── # ── Server ──────────────────────────────────────────────────────────────────
PORT=3000 PORT=3000
JOB_CONCURRENCY=3 JOB_CONCURRENCY=3
# ── Subject Crop ────────────────────────────────────────────────────────────
# Padding factor (0-1) added to each side of the detected bounding box.
# 0.15 = 15% of the box width/height added as margin on each side.
BBOX_PADDING=0.15

View File

@@ -49,8 +49,25 @@ export const Job = sequelize.define('Job', {
allowNull: false, allowNull: false,
defaultValue: 'qwen3.5:397b-cloud', defaultValue: 'qwen3.5:397b-cloud',
}, },
inputTokens: { type: DataTypes.INTEGER, allowNull: true }, inputTokens: { type: DataTypes.INTEGER, allowNull: true },
outputTokens: { type: DataTypes.INTEGER, allowNull: true }, outputTokens: { type: DataTypes.INTEGER, allowNull: true },
// ---- Bounding-box detection fields (Pass 1) ----
detectedSubject: {
type: DataTypes.STRING(256),
allowNull: true,
comment: 'Short description of what the LLM identified as the subject',
},
bboxX: { type: DataTypes.FLOAT, allowNull: true },
bboxY: { type: DataTypes.FLOAT, allowNull: true },
bboxWidth: { type: DataTypes.FLOAT, allowNull: true },
bboxHeight: { type: DataTypes.FLOAT, allowNull: true },
// Cropped image data URL (stored for UI preview / debugging)
croppedImageDataUrl: {
type: DataTypes.TEXT,
allowNull: true,
},
}, { }, {
tableName: 'jobs', tableName: 'jobs',
timestamps: true, timestamps: true,
@@ -59,5 +76,6 @@ export const Job = sequelize.define('Job', {
export async function initDb() { export async function initDb() {
const { mkdirSync } = await import('fs'); const { mkdirSync } = await import('fs');
mkdirSync(join(__dirname, '../data'), { recursive: true }); mkdirSync(join(__dirname, '../data'), { recursive: true });
await sequelize.sync(); // alter: true adds new columns to an existing table without dropping data
await sequelize.sync({ alter: true });
} }

View File

@@ -1,4 +1,5 @@
import PQueue from 'p-queue'; import PQueue from 'p-queue';
import sharp from 'sharp';
import { generateText } from 'ai'; import { generateText } from 'ai';
import { createAnthropic } from '@ai-sdk/anthropic'; import { createAnthropic } from '@ai-sdk/anthropic';
import { createOpenAI } from '@ai-sdk/openai'; import { createOpenAI } from '@ai-sdk/openai';
@@ -78,8 +79,140 @@ async function setStatus(job, status, extra = {}) {
broadcast({ type: 'job_update', job: job.toJSON() }); 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) { async function runJob(jobId) {
const job = await Job.findByPk(jobId); const job = await Job.findByPk(jobId);
@@ -89,15 +222,69 @@ async function runJob(jobId) {
try { try {
resolveModelMeta(job.model); 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({ const { text, usage } = await generateText({
model: resolveModel(job.model), model,
messages: [ messages: [
{ {
role: 'user', role: 'user',
content: [ content: [
{ type: 'text', text: job.prompt }, { type: 'text', text: job.prompt },
{ type: 'image', image: job.imageDataUrl }, { type: 'image', image: imageForAnalysis },
], ],
}, },
], ],

View File

@@ -26,6 +26,7 @@
"react": "^18.3.1", "react": "^18.3.1",
"react-dom": "^18.3.1", "react-dom": "^18.3.1",
"sequelize": "^6.37.3", "sequelize": "^6.37.3",
"sharp": "^0.33.5",
"sqlite3": "^5.1.7" "sqlite3": "^5.1.7"
}, },
"devDependencies": { "devDependencies": {