Skip to content

Commit

Permalink
⚡️ perf: Optimize the image upload size for gpt-4-vision (lobehub#669)
Browse files Browse the repository at this point in the history
* ♻️ refactor: `FileItem` changed to `ImageFileItem`

* ⚡️ perf: Optimize the image upload size for `gpt-4-vision`

* 🔧 test: add Test Configuration

* 🧪 test: add case

* ✏️ chore: typos

---------

Co-authored-by: wuxh <wxh1220@gmail.com>
  • Loading branch information
mushan0x0 and Wxh16144 authored Dec 18, 2023
1 parent e709d02 commit d038d24
Show file tree
Hide file tree
Showing 9 changed files with 135 additions and 8 deletions.
3 changes: 2 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,8 @@
"typescript": "^5",
"unified": "^11",
"unist-util-visit": "^5",
"vitest": "0.34.6"
"vitest": "0.34.6",
"vitest-canvas-mock": "^0.3.3"
},
"publishConfig": {
"access": "public",
Expand Down
4 changes: 2 additions & 2 deletions src/components/FileList/EditableFileList.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import { useResponsive } from 'antd-style';
import { memo } from 'react';
import { Flexbox } from 'react-layout-kit';

import FileItem from '@/components/FileList/FileItem';
import ImageFileItem from '@/components/FileList/ImageFileItem';

interface EditableFileListProps {
alwaysShowClose?: boolean;
Expand All @@ -24,7 +24,7 @@ const EditableFileList = memo<EditableFileListProps>(
>
<ImageGallery>
{items.map((i) => (
<FileItem alwaysShowClose={alwaysShowClose} editable={editable} id={i} key={i} />
<ImageFileItem alwaysShowClose={alwaysShowClose} editable={editable} id={i} key={i} />
))}
</ImageGallery>
</Flexbox>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ interface FileItemProps {
onClick?: () => void;
style?: CSSProperties;
}
const FileItem = memo<FileItemProps>(({ editable, id, alwaysShowClose }) => {
const ImageFileItem = memo<FileItemProps>(({ editable, id, alwaysShowClose }) => {
const [useFetchFile, removeFile] = useFileStore((s) => [s.useFetchFile, s.removeFile]);
const IMAGE_SIZE = editable ? MIN_IMAGE_SIZE : '100%';
const { data, isLoading } = useFetchFile(id);
Expand Down Expand Up @@ -74,4 +74,4 @@ const FileItem = memo<FileItemProps>(({ editable, id, alwaysShowClose }) => {
);
});

export default FileItem;
export default ImageFileItem;
6 changes: 3 additions & 3 deletions src/components/FileList/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import { Flexbox } from 'react-layout-kit';
import { MAX_SIZE_DESKTOP, MAX_SIZE_MOBILE } from '@/components/FileList/style';

import FileGrid from './FileGrid';
import FileItem from './FileItem';
import ImageFileItem from './ImageFileItem';

interface FileListProps {
items: string[];
Expand Down Expand Up @@ -44,13 +44,13 @@ const FileList = memo<FileListProps>(({ items }) => {
<Flexbox gap={gap}>
<FileGrid col={firstRow.length} gap={gap} max={max}>
{firstRow.map((i) => (
<FileItem id={i} key={i} />
<ImageFileItem id={i} key={i} />
))}
</FileGrid>
{lastRow.length > 0 && (
<FileGrid col={lastRow.length > 2 ? 3 : lastRow.length} gap={gap} max={max}>
{lastRow.map((i) => (
<FileItem id={i} key={i} />
<ImageFileItem id={i} key={i} />
))}
</FileGrid>
)}
Expand Down
33 changes: 33 additions & 0 deletions src/services/file.ts
Original file line number Diff line number Diff line change
@@ -1,14 +1,47 @@
import { FileModel } from '@/database/models/file';
import { DB_File } from '@/database/schemas/files';
import { FilePreview } from '@/types/files';
import compressImage from '@/utils/compressImage';

class FileService {
private isImage(fileType: string) {
const imageRegex = /^image\//;
return imageRegex.test(fileType);
}
async uploadFile(file: DB_File) {
// 跳过图片上传测试
const isTestData = file.size === 1;
if (this.isImage(file.fileType) && !isTestData) {
return this.uploadImageFile(file);
}
// save to local storage
// we may want to save to a remote server later
return FileModel.create(file);
}

async uploadImageFile(file: DB_File) {
// 加载图片
const url = file.url || URL.createObjectURL(new Blob([file.data]));
const img = new Image();
img.src = url;
await (() =>
new Promise((resolve) => {
img.addEventListener('load', resolve);
}))();

// 压缩图片
const fileType = 'image/webp';
const base64String = compressImage({
img,
type: fileType,
});
const binaryString = atob(base64String.split('base64,')[1]);
const uint8Array = Uint8Array.from(binaryString, (char) => char.charCodeAt(0));
file.data = uint8Array.buffer;

return FileModel.create(file);
}

async removeFile(id: string) {
return FileModel.delete(id);
}
Expand Down
59 changes: 59 additions & 0 deletions src/utils/compressImage.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
import compressImage from './compressImage';

const getContextSpy = vi.spyOn(global.HTMLCanvasElement.prototype, 'getContext');
const drawImageSpy = vi.spyOn(CanvasRenderingContext2D.prototype, 'drawImage');

beforeEach(() => {
getContextSpy.mockClear();
drawImageSpy.mockClear();
});

describe('compressImage', () => {
it('should compress image with maxWidth', () => {
const img = document.createElement('img');
img.width = 3000;
img.height = 2000;

const r = compressImage({ img });

expect(r).toMatch(/^data:image\/webp;base64,/);

expect(getContextSpy).toBeCalledTimes(1);
expect(getContextSpy).toBeCalledWith('2d');

expect(drawImageSpy).toBeCalledTimes(1);
expect(drawImageSpy).toBeCalledWith(img, 0, 0, 3000, 2000, 0, 0, 2160, 1440);
});

it('should compress image with maxHeight', () => {
const img = document.createElement('img');
img.width = 2000;
img.height = 3000;

const r = compressImage({ img });

expect(r).toMatch(/^data:image\/webp;base64,/);

expect(getContextSpy).toBeCalledTimes(1);
expect(getContextSpy).toBeCalledWith('2d');

expect(drawImageSpy).toBeCalledTimes(1);
expect(drawImageSpy).toBeCalledWith(img, 0, 0, 2000, 3000, 0, 0, 1440, 2160);
});

it('should not compress image', () => {
const img = document.createElement('img');
img.width = 2000;
img.height = 2000;

const r = compressImage({ img });

expect(r).toMatch(/^data:image\/webp;base64,/);

expect(getContextSpy).toBeCalledTimes(1);
expect(getContextSpy).toBeCalledWith('2d');

expect(drawImageSpy).toBeCalledTimes(1);
expect(drawImageSpy).toBeCalledWith(img, 0, 0, 2000, 2000, 0, 0, 2000, 2000);
});
});
29 changes: 29 additions & 0 deletions src/utils/compressImage.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
const compressImage = ({ img, type = 'image/webp' }: { img: HTMLImageElement; type?: string }) => {
// 设置最大宽高
const maxWidth = 2160;
const maxHeight = 2160;
let width = img.width;
let height = img.height;

if (width > height && width > maxWidth) {
// 如果图片宽度大于高度且大于最大宽度限制
width = maxWidth;
height = Math.round((maxWidth / img.width) * img.height);
} else if (height > width && height > maxHeight) {
// 如果图片高度大于宽度且大于最大高度限制
height = maxHeight;
width = Math.round((maxHeight / img.height) * img.width);
}

const canvas = document.createElement('canvas');
const ctx = canvas.getContext('2d') as CanvasRenderingContext2D;

canvas.width = width;
canvas.height = height;

ctx.drawImage(img, 0, 0, img.width, img.height, 0, 0, width, height);

return canvas.toDataURL(type);
};

export default compressImage;
1 change: 1 addition & 0 deletions tests/setup.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { theme } from 'antd';
// refs: https://github.com/dumbmatter/fakeIndexedDB#dexie-and-other-indexeddb-api-wrappers
import 'fake-indexeddb/auto';
import React from 'react';
import 'vitest-canvas-mock';

// remove antd hash on test
theme.defaultConfig.hashed = false;
Expand Down
4 changes: 4 additions & 0 deletions vitest.config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@ export default defineConfig({
provider: 'v8',
reporter: ['text', 'json', 'lcov', 'text-summary'],
},
deps: {
inline: ['vitest-canvas-mock'],
},
// threads: false,
environment: 'jsdom',
globals: true,
setupFiles: './tests/setup.ts',
Expand Down

0 comments on commit d038d24

Please sign in to comment.