diff --git a/src/api/extract.ts b/src/api/extract.ts index 4d5de57..dfc92c8 100644 --- a/src/api/extract.ts +++ b/src/api/extract.ts @@ -57,25 +57,32 @@ const route = app.post('/', zValidator('form', extractSchema), async (c) => { }) } - // 3. Detect Bounding Boxes + // 2.5 Emit texture sheet to client + const dimensions = extractionService.getGridDimensions(analysis.assets.length) await stream.writeSSE({ - event: 'status', - data: JSON.stringify({ status: 'splitting', message: 'Detecting asset boundaries...' }), + event: 'texture_generated', + data: JSON.stringify({ + image: sheetBase64, + grid: dimensions, + }), }) - let boxes + // 3. Crop Assets (Splitting) + await stream.writeSSE({ + event: 'status', + data: JSON.stringify({ status: 'splitting', message: 'Splitting texture sheet...' }), + }) + + let crops try { - boxes = await extractionService.detectBoundingBoxes(sheetBase64, analysis.assets) + crops = await extractionService.cropAssets(sheetBase64, analysis.assets) } catch (err: unknown) { - throw new Error(`Detection failed: ${err instanceof Error ? err.message : String(err)}`, { + throw new Error(`Splitting failed: ${err instanceof Error ? err.message : String(err)}`, { cause: err, }) } - // 4. Crop Assets - const crops = await extractionService.cropAssets(sheetBase64, boxes) - - // 5. Refine & Save + // 4. Refine & Save await stream.writeSSE({ event: 'status', data: JSON.stringify({ status: 'refining', message: 'Refining and saving assets...' }), @@ -85,12 +92,12 @@ const route = app.post('/', zValidator('form', extractSchema), async (c) => { for (const crop of crops) { // Find matching asset from analysis to get category/description + // Since crop.name comes from analysis.assets, this should always find it. const originalMeta = analysis.assets.find((a) => a.item_name === crop.name) || { item_name: crop.name, description: 'Extracted asset', category: 'Others', sub_category: 'Others', - background_color: 'unknown', } // Refine image diff --git a/src/service/equipment.test.ts b/src/service/equipment.test.ts index dd2c715..5767b68 100644 --- a/src/service/equipment.test.ts +++ b/src/service/equipment.test.ts @@ -45,7 +45,10 @@ describe('EquipmentService', () => { description: 'New Description', }) - await equipmentService.updateEquipment('1', { name: 'New Name', description: 'New Description' }) + await equipmentService.updateEquipment('1', { + name: 'New Name', + description: 'New Description', + }) expect(mockPrisma.equipment.update).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/src/service/extraction.test.ts b/src/service/extraction.test.ts index f4cf3ef..5533c97 100644 --- a/src/service/extraction.test.ts +++ b/src/service/extraction.test.ts @@ -1,5 +1,4 @@ -import { describe, expect, test, mock } from 'bun:test' -import { extractionService } from './extraction' +import { describe, expect, test, mock, beforeAll } from 'bun:test' // Mock dependencies mock.module('./ai', () => ({ @@ -9,7 +8,6 @@ mock.module('./ai', () => ({ { item_name: 'Helmet', description: 'A rusty helmet', - background_color: 'Red', category: 'Headwear', sub_category: 'Helmet', }, @@ -20,17 +18,28 @@ mock.module('./ai', () => ({ })) mock.module('sharp', () => { - return () => ({ - metadata: () => Promise.resolve({ width: 100, height: 100 }), + const sharpInstance = { + metadata: () => Promise.resolve({ width: 300, height: 100 }), clone: () => ({ extract: () => ({ toBuffer: () => Promise.resolve(Buffer.from('mockCrop')), }), }), - }) + } + return { + default: () => sharpInstance, + } }) describe('ExtractionService', () => { + let extractionService: any + + beforeAll(async () => { + // Dynamic import to ensure mocks are applied + const mod = await import('./extraction') + extractionService = mod.extractionService + }) + test('analyzeImage calls AI service', async () => { const file = new File(['dummy'], 'test.png', { type: 'image/png' }) const result = await extractionService.analyzeImage(file) @@ -40,7 +49,29 @@ describe('ExtractionService', () => { test('generateTextureSheet calls AI service', async () => { const file = new File(['dummy'], 'test.png', { type: 'image/png' }) - const result = await extractionService.generateTextureSheet(file, []) + const assets = [ + { + item_name: 'Helmet', + description: 'desc', + category: 'Headwear', + }, + ] + const result = await extractionService.generateTextureSheet(file, assets) expect(result).toBe('data:image/png;base64,mockImage') }) + + test('cropAssets splits image based on grid', async () => { + const assets = [ + { + item_name: 'Helmet', + description: 'desc', + category: 'Headwear', + }, + ] + // 1 asset -> 1x3 grid. width 300 / 3 = 100. height 100 / 1 = 100. + const result = await extractionService.cropAssets('data:image/png;base64,mockSheet', assets) + expect(result.length).toBe(1) + expect(result[0].name).toBe('Helmet') + expect(result[0].base64).toContain('data:image/png;base64,') + }) }) diff --git a/src/service/extraction.ts b/src/service/extraction.ts index 5caf082..c2dc377 100644 --- a/src/service/extraction.ts +++ b/src/service/extraction.ts @@ -7,7 +7,6 @@ import { EQUIPMENT_CATEGORIES } from '../lib/categories' export interface ExtractedAsset { item_name: string description: string - background_color: string category: string sub_category?: string } @@ -16,18 +15,6 @@ export interface AnalysisResponse { assets: ExtractedAsset[] } -export interface BoundingBoxResponse { - grid_dimensions: { - rows: number - cols: number - } - assets: { - item_name: string - coordinates: [number, number, number, number] // ymin, xmin, ymax, xmax - bg_color_detected: string - }[] -} - export class ExtractionService { /** * Analyzes an image to identify extractable assets. @@ -40,9 +27,8 @@ Task: Character Asset Identification for Extraction Instructions: 1. Analyze the character in the image and identify all distinct, extractable equipment and clothing items. 2. Filter for items that are clearly visible. -3. For each identified item, assign one of the following high-contrast background colors for future segmentation: Red, Yellow, Green, White, Blue, Black, Magenta, Cyan, Orange, Gray. Avoid repeating colors if possible. -4. Limit to a maximum of 9 most prominent assets. -5. Assign a category and sub-category to each item based on the provided list. +3. Limit to a maximum of 9 most prominent assets. +4. Assign a category and sub-category to each item based on the provided list. - Choose the best fitting "main_category" and "sub_categories". - If no suitable category/sub-category is found, use "Others". @@ -57,7 +43,6 @@ Return ONLY a JSON object with a list of assets. assets: z.array( z.object({ item_name: z.string().describe('Name of the item'), - background_color: z.string().describe('Background color assigned'), description: z.string().describe('Description of the item'), category: z.string().describe('Main category of the item'), sub_category: z.string().optional().describe('Sub category of the item'), @@ -68,45 +53,42 @@ Return ONLY a JSON object with a list of assets. return aiService.generateText(prompt, [file], schema, 'step_analyze_model') } + public getGridDimensions(count: number): { rows: number; cols: number } { + if (count <= 3) return { rows: 1, cols: 3 } + if (count <= 4) return { rows: 2, cols: 2 } + if (count <= 6) return { rows: 2, cols: 3 } + return { rows: 3, cols: 3 } + } + /** * Generates a texture sheet with isolated assets on solid backgrounds. */ async generateTextureSheet(file: File, assets: ExtractedAsset[]): Promise { - const HEX_COLORS: Record = { - Red: '#FF0000', - Yellow: '#FFFF00', - Green: '#00FF00', - White: '#FFFFFF', - Blue: '#0000FF', - Black: '#000000', - Magenta: '#FF00FF', - Cyan: '#00FFFF', - Orange: '#FFA500', - Gray: '#808080', - } + const { rows, cols } = this.getGridDimensions(assets.length) - const mappingString = assets - .map((asset) => { - const colorName = asset.background_color - const hex = HEX_COLORS[colorName] || '#000000' - return `- ${asset.item_name}: ${colorName} Background (${hex}). ${asset.description}` - }) + const assetList = assets + .map((asset, index) => `${index + 1}. ${asset.item_name}: ${asset.description}`) .join('\n') const prompt = ` Task: Character Asset Extraction & Grid Generation Instructions: -Analyze the character in the provided image and extract the following ${assets.length} specific assets into a new, single image organized in a grid layout (e.g., 2x3, 3x3 depending on count). +Analyze the character in the provided image and extract the following ${assets.length} specific assets into a new, single image organized in a ${rows}x${cols} grid layout. Core Constraints: -1. Isolated Assets: Each item must be placed within its own distinct square tile. No overlapping. -2. Ghost Mannequin Style: Extract ONLY the items. Remove the character's body entirely. Clothing and armor must appear as empty, 3D shells as if worn by an invisible person. -3. Solid Chroma Key Backgrounds: Each tile must use the specific solid background color assigned below. Ensure zero shadows, highlights, or gradients on the background to facilitate clean digital segmentation. -4. Output Format: Generate a square texture sheet. Ignore the original image aspect ratio for the final output canvas. +1. Grid Layout: The output image must be a strict ${rows} rows x ${cols} columns grid. +2. Equal Cells: Every cell in the grid MUST have identical width and height. +3. Sequential Placement: Place the assets in the grid cells in the order listed below, starting from top-left, moving right, then down to the next row. +4. Empty Cells: If there are fewer assets than grid cells, leave the remaining cells at the end (bottom-right) completely empty (pure white). +5. Isolated Assets: Each item must be placed within its own distinct square tile. No overlapping. +6. Ghost Mannequin Style: Extract ONLY the items. Remove the character's body entirely. Clothing and armor must appear as empty, 3D shells as if worn by an invisible person. +7. Pure White Background: The entire background of the image and each cell must be pure white (#FFFFFF). No shadows, no gradients. +8. Output Format: Generate a square texture sheet. Ignore the original image aspect ratio for the final output canvas. +9. EXACT REPLICA: You MUST reproduce the items exactly as they appear in the original image. Maintain all textures, patterns, logos, weathering, materials, and small details. Do not simplify or stylize. -Item & Background Mapping: -${mappingString} +Assets to Extract (in order): +${assetList} Output Requirement: High-resolution texture sheet, flat lay presentation, sharp edges, and uniform lighting across all assets. ` @@ -115,76 +97,11 @@ Output Requirement: High-resolution texture sheet, flat lay presentation, sharp } /** - * Detects bounding boxes for assets in the texture sheet. - */ - async detectBoundingBoxes( - sheetBase64: string, - analyzedAssets: ExtractedAsset[], - ): Promise { - const buffer = Buffer.from(sheetBase64.replace(/^data:image\/\w+;base64,/, ''), 'base64') - const file = new File([buffer], 'sheet.png', { type: 'image/png' }) - - const assetContext = analyzedAssets - .map( - (a) => - `- Name: ${a.item_name}, Description: ${a.description}, Background Color: ${a.background_color}`, - ) - .join('\n') - - const prompt = ` -Task: Spatial Asset Detection & Normalized Grid Mapping - -Instructions: -1. Analyze the provided image consisting of character assets arranged in a grid. -2. Identify the bounding box for each distinct asset container (the colored square background tiles). -3. Return the coordinates for each box using a normalized scale of 0.0 to 1.0. - - [ymin, xmin, ymax, xmax] - - (0, 0) is the top-left corner; (1.0, 1.0) is the bottom-right corner. -4. Ensure the box coordinates strictly encompass the outer boundary of each colored tile. -5. Use the provided asset context to correctly identify which item is in which box. - -Asset Context: -${assetContext} - -Output Format (Strict JSON): -Return ONLY a JSON object. Provide the grid dimensions and a list of detected assets with their coordinates. -` - - // Create an enum from the analyzed asset names plus "unknown" - const assetNames = analyzedAssets.map((a) => a.item_name) - // Zod enum requires at least one value, and they must be unique - const uniqueNames = Array.from(new Set([...assetNames, 'unknown'])) as [string, ...string[]] - - const schema = z.object({ - grid_dimensions: z.object({ - rows: z.number(), - cols: z.number(), - }), - assets: z.array( - z.object({ - item_name: z.enum(uniqueNames), - coordinates: z - .tuple([z.number(), z.number(), z.number(), z.number()]) - .describe('[ymin, xmin, ymax, xmax]'), - bg_color_detected: z.string(), - }), - ), - }) - - return aiService.generateText( - prompt, - [file], - schema, - 'step_bounding_box_model', - ) - } - - /** - * Crops assets from the texture sheet based on bounding boxes. + * Crops assets from the texture sheet based on calculated grid. */ async cropAssets( sheetBase64: string, - boxes: BoundingBoxResponse, + assets: ExtractedAsset[], ): Promise<{ name: string; base64: string }[]> { const buffer = Buffer.from(sheetBase64.replace(/^data:image\/\w+;base64,/, ''), 'base64') const image = sharp(buffer) @@ -192,19 +109,23 @@ Return ONLY a JSON object. Provide the grid dimensions and a list of detected as if (!metadata.width || !metadata.height) throw new Error('Could not determine image dimensions') + const { rows, cols } = this.getGridDimensions(assets.length) + const cellWidth = Math.floor(metadata.width / cols) + const cellHeight = Math.floor(metadata.height / rows) + const croppedAssets = [] - for (const box of boxes.assets) { - if (box.item_name === 'unknown') continue + for (let i = 0; i < assets.length; i++) { + const row = Math.floor(i / cols) + const col = i % cols - const [ymin, xmin, ymax, xmax] = box.coordinates + // Calculate coordinates + const left = col * cellWidth + const top = row * cellHeight + const width = cellWidth + const height = cellHeight - const left = Math.floor(xmin * metadata.width) - const top = Math.floor(ymin * metadata.height) - const width = Math.floor((xmax - xmin) * metadata.width) - const height = Math.floor((ymax - ymin) * metadata.height) - - // Ensure we don't go out of bounds + // Ensure we don't go out of bounds (though math implies we fit) const safeLeft = Math.max(0, left) const safeTop = Math.max(0, top) const safeWidth = Math.min(width, metadata.width - safeLeft) @@ -218,7 +139,7 @@ Return ONLY a JSON object. Provide the grid dimensions and a list of detected as .toBuffer() croppedAssets.push({ - name: box.item_name, + name: assets[i].item_name, base64: `data:image/png;base64,${cropBuffer.toString('base64')}`, }) } @@ -238,10 +159,11 @@ Task: Asset Refinement Instructions: 1. Process the provided asset image. -2. Remove the existing colored background and replace it with a pure white background (#FFFFFF) (or transparent if possible, but white is safer for now). +2. Remove any remaining background artifacts and ensure a pure white or transparent background. 3. Upscale the image and enhance details for high-quality game asset presentation. 4. Ensure the object is centered and fully visible. -5. Return only the image. +5. EXTREME DETAIL: Preserve every existing detail, texture, and pattern from the input image. Do not change the design. Only sharpen and clarify. +6. Return only the image. ` // Using generateImage with image input return aiService.generateImage(prompt, [file], 'step_refine_model') diff --git a/src/service/file.test.ts b/src/service/file.test.ts index 94d693f..7779d19 100644 --- a/src/service/file.test.ts +++ b/src/service/file.test.ts @@ -1,18 +1,15 @@ -import { describe, expect, test, mock } from 'bun:test' -import { fileService } from './file' - -// Valid 1x1 PNG Base64 -const VALID_PNG_BASE64 = - 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==' -const VALID_PNG_BUFFER = Buffer.from(VALID_PNG_BASE64, 'base64') +import { describe, expect, test, mock, beforeAll } from 'bun:test' // Mock sharp mock.module('sharp', () => { - return () => ({ + const sharpInstance = { webp: () => ({ - toBuffer: () => Promise.resolve(new Uint8Array([1, 2, 3])), + toBuffer: () => Promise.resolve(Buffer.from([1, 2, 3])), }), - }) + } + return { + default: () => sharpInstance, + } }) // Mock Bun.write @@ -23,6 +20,19 @@ const mockWrite = mock((path: string, content: string | Uint8Array) => Bun.write = mockWrite describe('FileService', () => { + let fileService: any + + beforeAll(async () => { + // Dynamic import to ensure mocks are applied + const mod = await import('./file?v=' + Date.now()) + fileService = mod.fileService + }) + + // Valid 1x1 PNG Base64 + const VALID_PNG_BASE64 = + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==' + const VALID_PNG_BUFFER = Buffer.from(VALID_PNG_BASE64, 'base64') + test('saveFile saves a file and converts to webp if image', async () => { // Pass valid PNG buffer const file = new File([VALID_PNG_BUFFER], 'test.png', { type: 'image/png' }) @@ -30,8 +40,12 @@ describe('FileService', () => { expect(result.filename.endsWith('.webp')).toBe(true) expect(result.mimeType).toBe('image/webp') + // Check if write was called expect(mockWrite).toHaveBeenCalled() - expect(mockWrite.mock.calls[0][0]).toContain('data/files') + // The first call's first argument (path) should contain data/files + const writeCalls = mockWrite.mock.calls + const lastCall = writeCalls[writeCalls.length - 1] + expect(lastCall[0]).toContain('data/files') }) test('saveBase64Image saves a base64 string as webp', async () => { diff --git a/src/ui/components/EquipmentDetailsDialog.tsx b/src/ui/components/EquipmentDetailsDialog.tsx index 863e9d5..aefa7c2 100644 --- a/src/ui/components/EquipmentDetailsDialog.tsx +++ b/src/ui/components/EquipmentDetailsDialog.tsx @@ -1,13 +1,7 @@ import React, { useEffect, useState } from 'react' import { InferResponseType } from 'hono/client' import { client } from '../client' -import { - Dialog, - DialogContent, - DialogHeader, - DialogTitle, - DialogDescription, -} from './ui/dialog' +import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogDescription } from './ui/dialog' import { Button } from './ui/button' import { Input } from './ui/input' import { Textarea } from './ui/textarea' diff --git a/src/ui/components/ExtractionFlowDesigner.tsx b/src/ui/components/ExtractionFlowDesigner.tsx index 0444f74..912a549 100644 --- a/src/ui/components/ExtractionFlowDesigner.tsx +++ b/src/ui/components/ExtractionFlowDesigner.tsx @@ -1,4 +1,4 @@ -import { ArrowDown, BoxSelect, FileSearch, Image as ImageIcon, User, Wand2 } from 'lucide-react' +import { ArrowDown, FileSearch, Image as ImageIcon, User, Wand2 } from 'lucide-react' import { Select, SelectContent, @@ -43,14 +43,7 @@ export function ExtractionFlowDesigner({ icon: , type: 'image', }, - { - id: 'bbox', - key: 'step_bounding_box_model', - title: 'Detect Bounding Boxes', - description: 'Identifies regions for individual equipment pieces.', - icon: , - type: 'text', - }, + // Bounding Box Detection step removed as part of grid-based extraction refactor { id: 'refine', key: 'step_refine_model', diff --git a/src/ui/components/ExtractorDialog.tsx b/src/ui/components/ExtractorDialog.tsx index 4d43596..5a6576c 100644 --- a/src/ui/components/ExtractorDialog.tsx +++ b/src/ui/components/ExtractorDialog.tsx @@ -1,5 +1,5 @@ import { AlertCircle, Loader2 } from 'lucide-react' -import React, { useState } from 'react' +import React, { useEffect, useState } from 'react' import { client } from '../client' import { ExtractedAsset, ExtractionStatus } from '../types' import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from './ui/dialog' @@ -10,7 +10,10 @@ interface ExtractorDialogProps { onSuccess: (assets: ExtractedAsset[]) => void } -import { useEffect } from 'react' +interface GridDimensions { + rows: number + cols: number +} export const ExtractorDialog: React.FC = ({ open, @@ -23,6 +26,8 @@ export const ExtractorDialog: React.FC = ({ const [statusMessage, setStatusMessage] = useState('') const [error, setError] = useState(null) const [results, setResults] = useState([]) + const [textureSheet, setTextureSheet] = useState(null) + const [gridDimensions, setGridDimensions] = useState(null) const processFile = (f: File) => { setFile(f) @@ -30,6 +35,8 @@ export const ExtractorDialog: React.FC = ({ setResults([]) setError(null) setStatus('idle') + setTextureSheet(null) + setGridDimensions(null) } const handleFileChange = (e: React.ChangeEvent) => { @@ -60,6 +67,8 @@ export const ExtractorDialog: React.FC = ({ setError(null) setStatus('analyzing') setStatusMessage('Starting extraction...') + setTextureSheet(null) + setGridDimensions(null) try { const response = await client.extract.$post({ @@ -95,6 +104,9 @@ export const ExtractorDialog: React.FC = ({ if (event === 'status') { setStatus(data.status) setStatusMessage(data.message) + } else if (event === 'texture_generated') { + setTextureSheet(data.image) + setGridDimensions(data.grid) } else if (event === 'asset_refined') { setResults((prev) => [...prev, data.asset]) } else if (event === 'complete') { @@ -124,11 +136,13 @@ export const ExtractorDialog: React.FC = ({ setResults([]) setStatus('idle') setError(null) + setTextureSheet(null) + setGridDimensions(null) } return ( - + Equipment Extractor @@ -195,7 +209,7 @@ export const ExtractorDialog: React.FC = ({
{/* Source Image */} -
+
@@ -206,6 +220,42 @@ export const ExtractorDialog: React.FC = ({
+ {/* Generated Texture Sheet (with Grid Overlay) */} + {textureSheet && gridDimensions && ( +
+
+ + + {/* Grid Overlay */} +
+ {Array.from({ length: gridDimensions.rows * gridDimensions.cols }).map( + (_, i) => ( +
+ {i + 1} +
+ ), + )} +
+ +
+ + Generated Sheet + +
+
+
+ )} + {/* Controls & Status */}
{status === 'idle' || status === 'error' ? ( diff --git a/src/ui/components/LookDetailsDialog.tsx b/src/ui/components/LookDetailsDialog.tsx index 7b7a2ff..5645550 100644 --- a/src/ui/components/LookDetailsDialog.tsx +++ b/src/ui/components/LookDetailsDialog.tsx @@ -1,13 +1,7 @@ import React from 'react' import { InferResponseType } from 'hono/client' import { client } from '../client' -import { - Dialog, - DialogContent, - DialogHeader, - DialogTitle, - DialogDescription, -} from './ui/dialog' +import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogDescription } from './ui/dialog' import { ScrollArea } from './ui/scroll-area' type GenerationResponse = InferResponseType