diff --git a/.env.example b/.env.example index 6f43079..d77ae60 100644 --- a/.env.example +++ b/.env.example @@ -20,3 +20,8 @@ OLLAMA_API_KEY=your-ollama-api-key # ── Server ────────────────────────────────────────────────────────────────── PORT=3000 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 diff --git a/db/models.js b/db/models.js index bea0e48..f966173 100644 --- a/db/models.js +++ b/db/models.js @@ -49,8 +49,25 @@ export const Job = sequelize.define('Job', { allowNull: false, defaultValue: 'qwen3.5:397b-cloud', }, - inputTokens: { type: DataTypes.INTEGER, allowNull: true }, + inputTokens: { 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', timestamps: true, @@ -59,5 +76,6 @@ export const Job = sequelize.define('Job', { export async function initDb() { const { mkdirSync } = await import('fs'); 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 }); } diff --git a/jobs/queue.js b/jobs/queue.js index 27f0037..f73d808 100644 --- a/jobs/queue.js +++ b/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": "", + "bbox": { + "x": , + "y": , + "width": , + "height": + } +} + +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 }, ], }, ], diff --git a/package.json b/package.json index 548b793..fe588c3 100644 --- a/package.json +++ b/package.json @@ -26,6 +26,7 @@ "react": "^18.3.1", "react-dom": "^18.3.1", "sequelize": "^6.37.3", + "sharp": "^0.33.5", "sqlite3": "^5.1.7" }, "devDependencies": {