mirror of https://github.com/immich-app/immich.git
feat: ocr (#18836)
* feat: add OCR functionality and related configurations * chore: update labeler configuration for machine learning files * feat(i18n): enhance OCR model descriptions and add orientation classification and unwarping features * chore: update Dockerfile to include ccache for improved build performance * feat(ocr): enhance OCR model configuration with orientation classification and unwarping options, update PaddleOCR integration, and improve response structure * refactor(ocr): remove OCR_CLEANUP job from enum and type definitions * refactor(ocr): remove obsolete OCR entity and migration files, and update asset job status and schema to accommodate new OCR table structure * refactor(ocr): update OCR schema and response structure to use individual coordinates instead of bounding box, and adjust related service and repository files * feat: enhance OCR configuration and functionality - Updated OCR settings to include minimum detection box score, minimum detection score, and minimum recognition score. - Refactored PaddleOCRecognizer to utilize new scoring parameters. - Introduced new database tables for asset OCR data and search functionality. - Modified related services and repositories to support the new OCR features. - Updated translations for improved clarity in settings UI. * sql changes * use rapidocr * change dto * update web * update lock * update api * store positions as normalized floats * match column order in db * update admin ui settings descriptions fix max resolution key set min threshold to 0.1 fix bind * apply config correctly, adjust defaults * unnecessary model type * unnecessary sources * fix(ocr): switch RapidOCR lang type from LangDet to LangRec * fix(ocr): expose lang_type (LangRec.CH) and font_path on OcrOptions for RapidOCR * fix(ocr): make OCR text search case- and accent-insensitive using ILIKE + unaccent * fix(ocr): add OCR search fields * fix: Add OCR database migration and update ML prediction logic. * trigrams are already case insensitive * add tests * format * update migrations * wrong uuid function * linting * maybe fix medium tests * formatting * fix weblate check * openapi * sql * minor fixes * maybe fix medium tests part 2 * passing medium tests * format web * readd sql * format dart * disabled in e2e * chore: translation ordering --------- Co-authored-by: mertalev <101130780+mertalev@users.noreply.github.com> Co-authored-by: Alex Tran <alex.tran1502@gmail.com>pull/23287/head
Kang
GitHub
parent
c666dc6c67
commit
02b29046b3
@ -0,0 +1,86 @@
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from rapidocr.ch_ppocr_det import TextDetector as RapidTextDetector
|
||||
from rapidocr.inference_engine.base import FileInfo, InferSession
|
||||
from rapidocr.utils import DownloadFile, DownloadFileInput
|
||||
from rapidocr.utils.typings import EngineType, LangDet, OCRVersion, TaskType
|
||||
from rapidocr.utils.typings import ModelType as RapidModelType
|
||||
|
||||
from immich_ml.config import log
|
||||
from immich_ml.models.base import InferenceModel
|
||||
from immich_ml.models.transforms import decode_cv2
|
||||
from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
|
||||
from immich_ml.sessions.ort import OrtSession
|
||||
|
||||
from .schemas import OcrOptions, TextDetectionOutput
|
||||
|
||||
|
||||
class TextDetector(InferenceModel):
|
||||
depends = []
|
||||
identity = (ModelType.DETECTION, ModelTask.OCR)
|
||||
|
||||
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
|
||||
super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
|
||||
self.max_resolution = 736
|
||||
self.min_score = 0.5
|
||||
self.score_mode = "fast"
|
||||
self._empty: TextDetectionOutput = {
|
||||
"image": np.empty(0, dtype=np.float32),
|
||||
"boxes": np.empty(0, dtype=np.float32),
|
||||
"scores": np.empty(0, dtype=np.float32),
|
||||
}
|
||||
|
||||
def _download(self) -> None:
|
||||
model_info = InferSession.get_model_url(
|
||||
FileInfo(
|
||||
engine_type=EngineType.ONNXRUNTIME,
|
||||
ocr_version=OCRVersion.PPOCRV5,
|
||||
task_type=TaskType.DET,
|
||||
lang_type=LangDet.CH,
|
||||
model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
|
||||
)
|
||||
)
|
||||
download_params = DownloadFileInput(
|
||||
file_url=model_info["model_dir"],
|
||||
sha256=model_info["SHA256"],
|
||||
save_path=self.model_path,
|
||||
logger=log,
|
||||
)
|
||||
DownloadFile.run(download_params)
|
||||
|
||||
def _load(self) -> ModelSession:
|
||||
# TODO: support other runtime sessions
|
||||
session = OrtSession(self.model_path)
|
||||
self.model = RapidTextDetector(
|
||||
OcrOptions(
|
||||
session=session.session,
|
||||
limit_side_len=self.max_resolution,
|
||||
limit_type="min",
|
||||
box_thresh=self.min_score,
|
||||
score_mode=self.score_mode,
|
||||
)
|
||||
)
|
||||
return session
|
||||
|
||||
def _predict(self, inputs: bytes | Image.Image) -> TextDetectionOutput:
|
||||
results = self.model(decode_cv2(inputs))
|
||||
if results.boxes is None or results.scores is None or results.img is None:
|
||||
return self._empty
|
||||
return {
|
||||
"image": results.img,
|
||||
"boxes": np.array(results.boxes, dtype=np.float32),
|
||||
"scores": np.array(results.scores, dtype=np.float32),
|
||||
}
|
||||
|
||||
def configure(self, **kwargs: Any) -> None:
|
||||
if (max_resolution := kwargs.get("maxResolution")) is not None:
|
||||
self.max_resolution = max_resolution
|
||||
self.model.limit_side_len = max_resolution
|
||||
if (min_score := kwargs.get("minScore")) is not None:
|
||||
self.min_score = min_score
|
||||
self.model.postprocess_op.box_thresh = min_score
|
||||
if (score_mode := kwargs.get("scoreMode")) is not None:
|
||||
self.score_mode = score_mode
|
||||
self.model.postprocess_op.score_mode = score_mode
|
||||
@ -0,0 +1,117 @@
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from numpy.typing import NDArray
|
||||
from PIL.Image import Image
|
||||
from rapidocr.ch_ppocr_rec import TextRecInput
|
||||
from rapidocr.ch_ppocr_rec import TextRecognizer as RapidTextRecognizer
|
||||
from rapidocr.inference_engine.base import FileInfo, InferSession
|
||||
from rapidocr.utils import DownloadFile, DownloadFileInput
|
||||
from rapidocr.utils.typings import EngineType, LangRec, OCRVersion, TaskType
|
||||
from rapidocr.utils.typings import ModelType as RapidModelType
|
||||
|
||||
from immich_ml.config import log, settings
|
||||
from immich_ml.models.base import InferenceModel
|
||||
from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
|
||||
from immich_ml.sessions.ort import OrtSession
|
||||
|
||||
from .schemas import OcrOptions, TextDetectionOutput, TextRecognitionOutput
|
||||
|
||||
|
||||
class TextRecognizer(InferenceModel):
|
||||
depends = [(ModelType.DETECTION, ModelTask.OCR)]
|
||||
identity = (ModelType.RECOGNITION, ModelTask.OCR)
|
||||
|
||||
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
|
||||
self.min_score = model_kwargs.get("minScore", 0.9)
|
||||
self._empty: TextRecognitionOutput = {
|
||||
"box": np.empty(0, dtype=np.float32),
|
||||
"boxScore": np.empty(0, dtype=np.float32),
|
||||
"text": [],
|
||||
"textScore": np.empty(0, dtype=np.float32),
|
||||
}
|
||||
super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
|
||||
|
||||
def _download(self) -> None:
|
||||
model_info = InferSession.get_model_url(
|
||||
FileInfo(
|
||||
engine_type=EngineType.ONNXRUNTIME,
|
||||
ocr_version=OCRVersion.PPOCRV5,
|
||||
task_type=TaskType.REC,
|
||||
lang_type=LangRec.CH,
|
||||
model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
|
||||
)
|
||||
)
|
||||
download_params = DownloadFileInput(
|
||||
file_url=model_info["model_dir"],
|
||||
sha256=model_info["SHA256"],
|
||||
save_path=self.model_path,
|
||||
logger=log,
|
||||
)
|
||||
DownloadFile.run(download_params)
|
||||
|
||||
def _load(self) -> ModelSession:
|
||||
# TODO: support other runtimes
|
||||
session = OrtSession(self.model_path)
|
||||
self.model = RapidTextRecognizer(
|
||||
OcrOptions(
|
||||
session=session.session,
|
||||
rec_batch_num=settings.max_batch_size.text_recognition if settings.max_batch_size is not None else 6,
|
||||
rec_img_shape=(3, 48, 320),
|
||||
)
|
||||
)
|
||||
return session
|
||||
|
||||
def _predict(self, _: Image, texts: TextDetectionOutput) -> TextRecognitionOutput:
|
||||
boxes, img, box_scores = texts["boxes"], texts["image"], texts["scores"]
|
||||
if boxes.shape[0] == 0:
|
||||
return self._empty
|
||||
rec = self.model(TextRecInput(img=self.get_crop_img_list(img, boxes)))
|
||||
if rec.txts is None:
|
||||
return self._empty
|
||||
|
||||
height, width = img.shape[0:2]
|
||||
boxes[:, :, 0] /= width
|
||||
boxes[:, :, 1] /= height
|
||||
|
||||
text_scores = np.array(rec.scores)
|
||||
valid_text_score_idx = text_scores > self.min_score
|
||||
valid_score_idx_list = valid_text_score_idx.tolist()
|
||||
return {
|
||||
"box": boxes.reshape(-1, 8)[valid_text_score_idx].reshape(-1),
|
||||
"text": [rec.txts[i] for i in range(len(rec.txts)) if valid_score_idx_list[i]],
|
||||
"boxScore": box_scores[valid_text_score_idx],
|
||||
"textScore": text_scores[valid_text_score_idx],
|
||||
}
|
||||
|
||||
def get_crop_img_list(self, img: NDArray[np.float32], boxes: NDArray[np.float32]) -> list[NDArray[np.float32]]:
|
||||
img_crop_width = np.maximum(
|
||||
np.linalg.norm(boxes[:, 1] - boxes[:, 0], axis=1), np.linalg.norm(boxes[:, 2] - boxes[:, 3], axis=1)
|
||||
).astype(np.int32)
|
||||
img_crop_height = np.maximum(
|
||||
np.linalg.norm(boxes[:, 0] - boxes[:, 3], axis=1), np.linalg.norm(boxes[:, 1] - boxes[:, 2], axis=1)
|
||||
).astype(np.int32)
|
||||
pts_std = np.zeros((img_crop_width.shape[0], 4, 2), dtype=np.float32)
|
||||
pts_std[:, 1:3, 0] = img_crop_width[:, None]
|
||||
pts_std[:, 2:4, 1] = img_crop_height[:, None]
|
||||
|
||||
img_crop_sizes = np.stack([img_crop_width, img_crop_height], axis=1).tolist()
|
||||
imgs: list[NDArray[np.float32]] = []
|
||||
for box, pts_std, dst_size in zip(list(boxes), list(pts_std), img_crop_sizes):
|
||||
M = cv2.getPerspectiveTransform(box, pts_std)
|
||||
dst_img: NDArray[np.float32] = cv2.warpPerspective(
|
||||
img,
|
||||
M,
|
||||
dst_size,
|
||||
borderMode=cv2.BORDER_REPLICATE,
|
||||
flags=cv2.INTER_CUBIC,
|
||||
) # type: ignore
|
||||
dst_height, dst_width = dst_img.shape[0:2]
|
||||
if dst_height * 1.0 / dst_width >= 1.5:
|
||||
dst_img = np.rot90(dst_img)
|
||||
imgs.append(dst_img)
|
||||
return imgs
|
||||
|
||||
def configure(self, **kwargs: Any) -> None:
|
||||
self.min_score = kwargs.get("minScore", self.min_score)
|
||||
@ -0,0 +1,28 @@
|
||||
from typing import Any, Iterable
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
from rapidocr.utils.typings import EngineType, LangRec
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
|
||||
class TextDetectionOutput(TypedDict):
|
||||
image: npt.NDArray[np.float32]
|
||||
boxes: npt.NDArray[np.float32]
|
||||
scores: npt.NDArray[np.float32]
|
||||
|
||||
|
||||
class TextRecognitionOutput(TypedDict):
|
||||
box: npt.NDArray[np.float32]
|
||||
boxScore: npt.NDArray[np.float32]
|
||||
text: Iterable[str]
|
||||
textScore: npt.NDArray[np.float32]
|
||||
|
||||
|
||||
# RapidOCR expects `engine_type`, `lang_type`, and `font_path` to be attributes
|
||||
class OcrOptions(dict[str, Any]):
|
||||
def __init__(self, **options: Any) -> None:
|
||||
super().__init__(**options)
|
||||
self.engine_type = EngineType.ONNXRUNTIME
|
||||
self.lang_type = LangRec.CH
|
||||
self.font_path = None
|
||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,136 @@
|
||||
//
|
||||
// AUTO-GENERATED FILE, DO NOT MODIFY!
|
||||
//
|
||||
// @dart=2.18
|
||||
|
||||
// ignore_for_file: unused_element, unused_import
|
||||
// ignore_for_file: always_put_required_named_parameters_first
|
||||
// ignore_for_file: constant_identifier_names
|
||||
// ignore_for_file: lines_longer_than_80_chars
|
||||
|
||||
part of openapi.api;
|
||||
|
||||
class OcrConfig {
|
||||
/// Returns a new [OcrConfig] instance.
|
||||
OcrConfig({
|
||||
required this.enabled,
|
||||
required this.maxResolution,
|
||||
required this.minDetectionScore,
|
||||
required this.minRecognitionScore,
|
||||
required this.modelName,
|
||||
});
|
||||
|
||||
bool enabled;
|
||||
|
||||
/// Minimum value: 1
|
||||
int maxResolution;
|
||||
|
||||
/// Minimum value: 0.1
|
||||
/// Maximum value: 1
|
||||
double minDetectionScore;
|
||||
|
||||
/// Minimum value: 0.1
|
||||
/// Maximum value: 1
|
||||
double minRecognitionScore;
|
||||
|
||||
String modelName;
|
||||
|
||||
@override
|
||||
bool operator ==(Object other) => identical(this, other) || other is OcrConfig &&
|
||||
other.enabled == enabled &&
|
||||
other.maxResolution == maxResolution &&
|
||||
other.minDetectionScore == minDetectionScore &&
|
||||
other.minRecognitionScore == minRecognitionScore &&
|
||||
other.modelName == modelName;
|
||||
|
||||
@override
|
||||
int get hashCode =>
|
||||
// ignore: unnecessary_parenthesis
|
||||
(enabled.hashCode) +
|
||||
(maxResolution.hashCode) +
|
||||
(minDetectionScore.hashCode) +
|
||||
(minRecognitionScore.hashCode) +
|
||||
(modelName.hashCode);
|
||||
|
||||
@override
|
||||
String toString() => 'OcrConfig[enabled=$enabled, maxResolution=$maxResolution, minDetectionScore=$minDetectionScore, minRecognitionScore=$minRecognitionScore, modelName=$modelName]';
|
||||
|
||||
Map<String, dynamic> toJson() {
|
||||
final json = <String, dynamic>{};
|
||||
json[r'enabled'] = this.enabled;
|
||||
json[r'maxResolution'] = this.maxResolution;
|
||||
json[r'minDetectionScore'] = this.minDetectionScore;
|
||||
json[r'minRecognitionScore'] = this.minRecognitionScore;
|
||||
json[r'modelName'] = this.modelName;
|
||||
return json;
|
||||
}
|
||||
|
||||
/// Returns a new [OcrConfig] instance and imports its values from
|
||||
/// [value] if it's a [Map], null otherwise.
|
||||
// ignore: prefer_constructors_over_static_methods
|
||||
static OcrConfig? fromJson(dynamic value) {
|
||||
upgradeDto(value, "OcrConfig");
|
||||
if (value is Map) {
|
||||
final json = value.cast<String, dynamic>();
|
||||
|
||||
return OcrConfig(
|
||||
enabled: mapValueOfType<bool>(json, r'enabled')!,
|
||||
maxResolution: mapValueOfType<int>(json, r'maxResolution')!,
|
||||
minDetectionScore: (mapValueOfType<num>(json, r'minDetectionScore')!).toDouble(),
|
||||
minRecognitionScore: (mapValueOfType<num>(json, r'minRecognitionScore')!).toDouble(),
|
||||
modelName: mapValueOfType<String>(json, r'modelName')!,
|
||||
);
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
static List<OcrConfig> listFromJson(dynamic json, {bool growable = false,}) {
|
||||
final result = <OcrConfig>[];
|
||||
if (json is List && json.isNotEmpty) {
|
||||
for (final row in json) {
|
||||
final value = OcrConfig.fromJson(row);
|
||||
if (value != null) {
|
||||
result.add(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
return result.toList(growable: growable);
|
||||
}
|
||||
|
||||
static Map<String, OcrConfig> mapFromJson(dynamic json) {
|
||||
final map = <String, OcrConfig>{};
|
||||
if (json is Map && json.isNotEmpty) {
|
||||
json = json.cast<String, dynamic>(); // ignore: parameter_assignments
|
||||
for (final entry in json.entries) {
|
||||
final value = OcrConfig.fromJson(entry.value);
|
||||
if (value != null) {
|
||||
map[entry.key] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
return map;
|
||||
}
|
||||
|
||||
// maps a json object with a list of OcrConfig-objects as value to a dart map
|
||||
static Map<String, List<OcrConfig>> mapListFromJson(dynamic json, {bool growable = false,}) {
|
||||
final map = <String, List<OcrConfig>>{};
|
||||
if (json is Map && json.isNotEmpty) {
|
||||
// ignore: parameter_assignments
|
||||
json = json.cast<String, dynamic>();
|
||||
for (final entry in json.entries) {
|
||||
map[entry.key] = OcrConfig.listFromJson(entry.value, growable: growable,);
|
||||
}
|
||||
}
|
||||
return map;
|
||||
}
|
||||
|
||||
/// The list of required keys that must be present in a JSON.
|
||||
static const requiredKeys = <String>{
|
||||
'enabled',
|
||||
'maxResolution',
|
||||
'minDetectionScore',
|
||||
'minRecognitionScore',
|
||||
'modelName',
|
||||
};
|
||||
}
|
||||
|
||||
@ -0,0 +1,68 @@
|
||||
-- NOTE: This file is auto generated by ./sql-generator
|
||||
|
||||
-- OcrRepository.getById
|
||||
select
|
||||
"asset_ocr".*
|
||||
from
|
||||
"asset_ocr"
|
||||
where
|
||||
"asset_ocr"."id" = $1
|
||||
|
||||
-- OcrRepository.getByAssetId
|
||||
select
|
||||
"asset_ocr".*
|
||||
from
|
||||
"asset_ocr"
|
||||
where
|
||||
"asset_ocr"."assetId" = $1
|
||||
|
||||
-- OcrRepository.upsert
|
||||
with
|
||||
"deleted_ocr" as (
|
||||
delete from "asset_ocr"
|
||||
where
|
||||
"assetId" = $1
|
||||
),
|
||||
"inserted_ocr" as (
|
||||
insert into
|
||||
"asset_ocr" (
|
||||
"assetId",
|
||||
"x1",
|
||||
"y1",
|
||||
"x2",
|
||||
"y2",
|
||||
"x3",
|
||||
"y3",
|
||||
"x4",
|
||||
"y4",
|
||||
"text",
|
||||
"boxScore",
|
||||
"textScore"
|
||||
)
|
||||
values
|
||||
(
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
$11,
|
||||
$12,
|
||||
$13
|
||||
)
|
||||
),
|
||||
"inserted_search" as (
|
||||
insert into
|
||||
"ocr_search" ("assetId", "text")
|
||||
values
|
||||
($14, $15)
|
||||
on conflict ("assetId") do update
|
||||
set
|
||||
"text" = "excluded"."text"
|
||||
)
|
||||
select
|
||||
1 as "dummy"
|
||||
@ -0,0 +1,68 @@
|
||||
import { Injectable } from '@nestjs/common';
|
||||
import { Insertable, Kysely, sql } from 'kysely';
|
||||
import { InjectKysely } from 'nestjs-kysely';
|
||||
import { DummyValue, GenerateSql } from 'src/decorators';
|
||||
import { DB } from 'src/schema';
|
||||
import { AssetOcrTable } from 'src/schema/tables/asset-ocr.table';
|
||||
|
||||
@Injectable()
|
||||
export class OcrRepository {
|
||||
constructor(@InjectKysely() private db: Kysely<DB>) {}
|
||||
|
||||
@GenerateSql({ params: [DummyValue.UUID] })
|
||||
getById(id: string) {
|
||||
return this.db.selectFrom('asset_ocr').selectAll('asset_ocr').where('asset_ocr.id', '=', id).executeTakeFirst();
|
||||
}
|
||||
|
||||
@GenerateSql({ params: [DummyValue.UUID] })
|
||||
getByAssetId(id: string) {
|
||||
return this.db.selectFrom('asset_ocr').selectAll('asset_ocr').where('asset_ocr.assetId', '=', id).execute();
|
||||
}
|
||||
|
||||
deleteAll() {
|
||||
return this.db.transaction().execute(async (trx: Kysely<DB>) => {
|
||||
await sql`truncate ${sql.table('asset_ocr')}`.execute(trx);
|
||||
await sql`truncate ${sql.table('ocr_search')}`.execute(trx);
|
||||
});
|
||||
}
|
||||
|
||||
@GenerateSql({
|
||||
params: [
|
||||
DummyValue.UUID,
|
||||
[
|
||||
{
|
||||
assetId: DummyValue.UUID,
|
||||
x1: DummyValue.NUMBER,
|
||||
y1: DummyValue.NUMBER,
|
||||
x2: DummyValue.NUMBER,
|
||||
y2: DummyValue.NUMBER,
|
||||
x3: DummyValue.NUMBER,
|
||||
y3: DummyValue.NUMBER,
|
||||
x4: DummyValue.NUMBER,
|
||||
y4: DummyValue.NUMBER,
|
||||
text: DummyValue.STRING,
|
||||
boxScore: DummyValue.NUMBER,
|
||||
textScore: DummyValue.NUMBER,
|
||||
},
|
||||
],
|
||||
],
|
||||
})
|
||||
upsert(assetId: string, ocrDataList: Insertable<AssetOcrTable>[]) {
|
||||
let query = this.db.with('deleted_ocr', (db) => db.deleteFrom('asset_ocr').where('assetId', '=', assetId));
|
||||
if (ocrDataList.length > 0) {
|
||||
const searchText = ocrDataList.map((item) => item.text.trim()).join(' ');
|
||||
(query as any) = query
|
||||
.with('inserted_ocr', (db) => db.insertInto('asset_ocr').values(ocrDataList))
|
||||
.with('inserted_search', (db) =>
|
||||
db
|
||||
.insertInto('ocr_search')
|
||||
.values({ assetId, text: searchText })
|
||||
.onConflict((oc) => oc.column('assetId').doUpdateSet((eb) => ({ text: eb.ref('excluded.text') }))),
|
||||
);
|
||||
} else {
|
||||
(query as any) = query.with('deleted_search', (db) => db.deleteFrom('ocr_search').where('assetId', '=', assetId));
|
||||
}
|
||||
|
||||
return query.selectNoFrom(sql`1`.as('dummy')).execute();
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,16 @@
|
||||
import { Kysely, sql } from 'kysely';
|
||||
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
await sql`CREATE TABLE "asset_ocr" ("id" uuid NOT NULL DEFAULT uuid_generate_v4(), "assetId" uuid NOT NULL, "x1" real NOT NULL, "y1" real NOT NULL, "x2" real NOT NULL, "y2" real NOT NULL, "x3" real NOT NULL, "y3" real NOT NULL, "x4" real NOT NULL, "y4" real NOT NULL, "boxScore" real NOT NULL, "textScore" real NOT NULL, "text" text NOT NULL);`.execute(
|
||||
db,
|
||||
);
|
||||
await sql`ALTER TABLE "asset_ocr" ADD CONSTRAINT "asset_ocr_pkey" PRIMARY KEY ("id");`.execute(db);
|
||||
await sql`ALTER TABLE "asset_ocr" ADD CONSTRAINT "asset_ocr_assetId_fkey" FOREIGN KEY ("assetId") REFERENCES "asset" ("id") ON UPDATE CASCADE ON DELETE CASCADE;`.execute(
|
||||
db,
|
||||
);
|
||||
await sql`CREATE INDEX "asset_ocr_assetId_idx" ON "asset_ocr" ("assetId")`.execute(db);
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
await sql`DROP TABLE "asset_ocr";`.execute(db);
|
||||
}
|
||||
@ -0,0 +1,20 @@
|
||||
import { Kysely, sql } from 'kysely';
|
||||
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
await sql`CREATE TABLE "ocr_search" ("assetId" uuid NOT NULL, "text" text NOT NULL);`.execute(db);
|
||||
await sql`ALTER TABLE "ocr_search" ADD CONSTRAINT "ocr_search_pkey" PRIMARY KEY ("assetId");`.execute(db);
|
||||
await sql`ALTER TABLE "ocr_search" ADD CONSTRAINT "ocr_search_assetId_fkey" FOREIGN KEY ("assetId") REFERENCES "asset" ("id") ON UPDATE CASCADE ON DELETE CASCADE;`.execute(
|
||||
db,
|
||||
);
|
||||
await sql`CREATE INDEX "idx_ocr_search_text" ON "ocr_search" USING gin (f_unaccent("text") gin_trgm_ops);`.execute(
|
||||
db,
|
||||
);
|
||||
await sql`INSERT INTO "migration_overrides" ("name", "value") VALUES ('index_idx_ocr_search_text', '{"type":"index","name":"idx_ocr_search_text","sql":"CREATE INDEX \\"idx_ocr_search_text\\" ON \\"ocr_search\\" USING gin (f_unaccent(\\"text\\") gin_trgm_ops);"}'::jsonb);`.execute(
|
||||
db,
|
||||
);
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
await sql`DROP TABLE "ocr_search";`.execute(db);
|
||||
await sql`DELETE FROM "migration_overrides" WHERE "name" = 'index_idx_ocr_search_text';`.execute(db);
|
||||
}
|
||||
@ -0,0 +1,9 @@
|
||||
import { Kysely, sql } from 'kysely';
|
||||
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
await sql`ALTER TABLE "asset_job_status" ADD "ocrAt" timestamp with time zone;`.execute(db);
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
await sql`ALTER TABLE "asset_job_status" DROP COLUMN "ocrAt";`.execute(db);
|
||||
}
|
||||
@ -0,0 +1,45 @@
|
||||
import { AssetTable } from 'src/schema/tables/asset.table';
|
||||
import { Column, ForeignKeyColumn, Generated, PrimaryGeneratedColumn, Table } from 'src/sql-tools';
|
||||
|
||||
@Table('asset_ocr')
|
||||
export class AssetOcrTable {
|
||||
@PrimaryGeneratedColumn()
|
||||
id!: Generated<string>;
|
||||
|
||||
@ForeignKeyColumn(() => AssetTable, { onDelete: 'CASCADE', onUpdate: 'CASCADE' })
|
||||
assetId!: string;
|
||||
|
||||
// box positions are normalized, with values between 0 and 1
|
||||
@Column({ type: 'real' })
|
||||
x1!: number;
|
||||
|
||||
@Column({ type: 'real' })
|
||||
y1!: number;
|
||||
|
||||
@Column({ type: 'real' })
|
||||
x2!: number;
|
||||
|
||||
@Column({ type: 'real' })
|
||||
y2!: number;
|
||||
|
||||
@Column({ type: 'real' })
|
||||
x3!: number;
|
||||
|
||||
@Column({ type: 'real' })
|
||||
y3!: number;
|
||||
|
||||
@Column({ type: 'real' })
|
||||
x4!: number;
|
||||
|
||||
@Column({ type: 'real' })
|
||||
y4!: number;
|
||||
|
||||
@Column({ type: 'real' })
|
||||
boxScore!: number;
|
||||
|
||||
@Column({ type: 'real' })
|
||||
textScore!: number;
|
||||
|
||||
@Column({ type: 'text' })
|
||||
text!: string;
|
||||
}
|
||||
@ -0,0 +1,20 @@
|
||||
import { AssetTable } from 'src/schema/tables/asset.table';
|
||||
import { Column, ForeignKeyColumn, Index, Table } from 'src/sql-tools';
|
||||
|
||||
@Table('ocr_search')
|
||||
@Index({
|
||||
name: 'idx_ocr_search_text',
|
||||
using: 'gin',
|
||||
expression: 'f_unaccent("text") gin_trgm_ops',
|
||||
})
|
||||
export class OcrSearchTable {
|
||||
@ForeignKeyColumn(() => AssetTable, {
|
||||
onDelete: 'CASCADE',
|
||||
onUpdate: 'CASCADE',
|
||||
primary: true,
|
||||
})
|
||||
assetId!: string;
|
||||
|
||||
@Column({ type: 'text' })
|
||||
text!: string;
|
||||
}
|
||||
@ -0,0 +1,177 @@
|
||||
import { AssetVisibility, ImmichWorker, JobName, JobStatus } from 'src/enum';
|
||||
import { OcrService } from 'src/services/ocr.service';
|
||||
import { assetStub } from 'test/fixtures/asset.stub';
|
||||
import { systemConfigStub } from 'test/fixtures/system-config.stub';
|
||||
import { makeStream, newTestService, ServiceMocks } from 'test/utils';
|
||||
|
||||
describe(OcrService.name, () => {
|
||||
let sut: OcrService;
|
||||
let mocks: ServiceMocks;
|
||||
|
||||
beforeEach(() => {
|
||||
({ sut, mocks } = newTestService(OcrService));
|
||||
|
||||
mocks.config.getWorker.mockReturnValue(ImmichWorker.Microservices);
|
||||
});
|
||||
|
||||
it('should work', () => {
|
||||
expect(sut).toBeDefined();
|
||||
});
|
||||
|
||||
describe('handleQueueOcr', () => {
|
||||
it('should do nothing if machine learning is disabled', async () => {
|
||||
mocks.systemMetadata.get.mockResolvedValue(systemConfigStub.machineLearningDisabled);
|
||||
|
||||
await sut.handleQueueOcr({ force: false });
|
||||
|
||||
expect(mocks.database.setDimensionSize).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should queue the assets without ocr', async () => {
|
||||
mocks.assetJob.streamForOcrJob.mockReturnValue(makeStream([assetStub.image]));
|
||||
|
||||
await sut.handleQueueOcr({ force: false });
|
||||
|
||||
expect(mocks.job.queueAll).toHaveBeenCalledWith([{ name: JobName.Ocr, data: { id: assetStub.image.id } }]);
|
||||
expect(mocks.assetJob.streamForOcrJob).toHaveBeenCalledWith(false);
|
||||
});
|
||||
|
||||
it('should queue all the assets', async () => {
|
||||
mocks.assetJob.streamForOcrJob.mockReturnValue(makeStream([assetStub.image]));
|
||||
|
||||
await sut.handleQueueOcr({ force: true });
|
||||
|
||||
expect(mocks.job.queueAll).toHaveBeenCalledWith([{ name: JobName.Ocr, data: { id: assetStub.image.id } }]);
|
||||
expect(mocks.assetJob.streamForOcrJob).toHaveBeenCalledWith(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleOcr', () => {
|
||||
it('should do nothing if machine learning is disabled', async () => {
|
||||
mocks.systemMetadata.get.mockResolvedValue(systemConfigStub.machineLearningDisabled);
|
||||
|
||||
expect(await sut.handleOcr({ id: '123' })).toEqual(JobStatus.Skipped);
|
||||
|
||||
expect(mocks.asset.getByIds).not.toHaveBeenCalled();
|
||||
expect(mocks.machineLearning.encodeImage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should skip assets without a resize path', async () => {
|
||||
mocks.assetJob.getForOcr.mockResolvedValue({ visibility: AssetVisibility.Timeline, previewFile: null });
|
||||
|
||||
expect(await sut.handleOcr({ id: assetStub.noResizePath.id })).toEqual(JobStatus.Failed);
|
||||
|
||||
expect(mocks.ocr.upsert).not.toHaveBeenCalled();
|
||||
expect(mocks.machineLearning.ocr).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should save the returned objects', async () => {
|
||||
mocks.machineLearning.ocr.mockResolvedValue({
|
||||
box: [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150, 160],
|
||||
boxScore: [0.9, 0.8],
|
||||
text: ['One Two Three', 'Four Five'],
|
||||
textScore: [0.95, 0.85],
|
||||
});
|
||||
mocks.assetJob.getForOcr.mockResolvedValue({
|
||||
visibility: AssetVisibility.Timeline,
|
||||
previewFile: assetStub.image.files[1].path,
|
||||
});
|
||||
|
||||
expect(await sut.handleOcr({ id: assetStub.image.id })).toEqual(JobStatus.Success);
|
||||
|
||||
expect(mocks.machineLearning.ocr).toHaveBeenCalledWith(
|
||||
'/uploads/user-id/thumbs/path.jpg',
|
||||
expect.objectContaining({
|
||||
modelName: 'PP-OCRv5_mobile',
|
||||
minDetectionScore: 0.5,
|
||||
minRecognitionScore: 0.8,
|
||||
maxResolution: 736,
|
||||
}),
|
||||
);
|
||||
expect(mocks.ocr.upsert).toHaveBeenCalledWith(assetStub.image.id, [
|
||||
{
|
||||
assetId: assetStub.image.id,
|
||||
boxScore: 0.9,
|
||||
text: 'One Two Three',
|
||||
textScore: 0.95,
|
||||
x1: 10,
|
||||
y1: 20,
|
||||
x2: 30,
|
||||
y2: 40,
|
||||
x3: 50,
|
||||
y3: 60,
|
||||
x4: 70,
|
||||
y4: 80,
|
||||
},
|
||||
{
|
||||
assetId: assetStub.image.id,
|
||||
boxScore: 0.8,
|
||||
text: 'Four Five',
|
||||
textScore: 0.85,
|
||||
x1: 90,
|
||||
y1: 100,
|
||||
x2: 110,
|
||||
y2: 120,
|
||||
x3: 130,
|
||||
y3: 140,
|
||||
x4: 150,
|
||||
y4: 160,
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('should apply config settings', async () => {
|
||||
mocks.systemMetadata.get.mockResolvedValue({
|
||||
machineLearning: {
|
||||
enabled: true,
|
||||
ocr: {
|
||||
modelName: 'PP-OCRv5_server',
|
||||
enabled: true,
|
||||
minDetectionScore: 0.8,
|
||||
minRecognitionScore: 0.9,
|
||||
maxResolution: 1500,
|
||||
},
|
||||
},
|
||||
});
|
||||
mocks.machineLearning.ocr.mockResolvedValue({ box: [], boxScore: [], text: [], textScore: [] });
|
||||
mocks.assetJob.getForOcr.mockResolvedValue({
|
||||
visibility: AssetVisibility.Timeline,
|
||||
previewFile: assetStub.image.files[1].path,
|
||||
});
|
||||
|
||||
expect(await sut.handleOcr({ id: assetStub.image.id })).toEqual(JobStatus.Success);
|
||||
|
||||
expect(mocks.machineLearning.ocr).toHaveBeenCalledWith(
|
||||
'/uploads/user-id/thumbs/path.jpg',
|
||||
expect.objectContaining({
|
||||
modelName: 'PP-OCRv5_server',
|
||||
minDetectionScore: 0.8,
|
||||
minRecognitionScore: 0.9,
|
||||
maxResolution: 1500,
|
||||
}),
|
||||
);
|
||||
expect(mocks.ocr.upsert).toHaveBeenCalledWith(assetStub.image.id, []);
|
||||
});
|
||||
|
||||
it('should skip invisible assets', async () => {
|
||||
mocks.assetJob.getForOcr.mockResolvedValue({
|
||||
visibility: AssetVisibility.Hidden,
|
||||
previewFile: assetStub.image.files[1].path,
|
||||
});
|
||||
|
||||
expect(await sut.handleOcr({ id: assetStub.livePhotoMotionAsset.id })).toEqual(JobStatus.Skipped);
|
||||
|
||||
expect(mocks.machineLearning.ocr).not.toHaveBeenCalled();
|
||||
expect(mocks.ocr.upsert).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should fail if asset could not be found', async () => {
|
||||
mocks.assetJob.getForOcr.mockResolvedValue(void 0);
|
||||
|
||||
expect(await sut.handleOcr({ id: assetStub.image.id })).toEqual(JobStatus.Failed);
|
||||
|
||||
expect(mocks.machineLearning.ocr).not.toHaveBeenCalled();
|
||||
expect(mocks.ocr.upsert).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
@ -0,0 +1,86 @@
|
||||
import { Injectable } from '@nestjs/common';
|
||||
import { JOBS_ASSET_PAGINATION_SIZE } from 'src/constants';
|
||||
import { OnJob } from 'src/decorators';
|
||||
import { AssetVisibility, JobName, JobStatus, QueueName } from 'src/enum';
|
||||
import { OCR } from 'src/repositories/machine-learning.repository';
|
||||
import { BaseService } from 'src/services/base.service';
|
||||
import { JobItem, JobOf } from 'src/types';
|
||||
import { isOcrEnabled } from 'src/utils/misc';
|
||||
|
||||
@Injectable()
|
||||
export class OcrService extends BaseService {
|
||||
@OnJob({ name: JobName.OcrQueueAll, queue: QueueName.Ocr })
|
||||
async handleQueueOcr({ force }: JobOf<JobName.OcrQueueAll>): Promise<JobStatus> {
|
||||
const { machineLearning } = await this.getConfig({ withCache: false });
|
||||
if (!isOcrEnabled(machineLearning)) {
|
||||
return JobStatus.Skipped;
|
||||
}
|
||||
|
||||
if (force) {
|
||||
await this.ocrRepository.deleteAll();
|
||||
}
|
||||
|
||||
let jobs: JobItem[] = [];
|
||||
const assets = this.assetJobRepository.streamForOcrJob(force);
|
||||
|
||||
for await (const asset of assets) {
|
||||
jobs.push({ name: JobName.Ocr, data: { id: asset.id } });
|
||||
|
||||
if (jobs.length >= JOBS_ASSET_PAGINATION_SIZE) {
|
||||
await this.jobRepository.queueAll(jobs);
|
||||
jobs = [];
|
||||
}
|
||||
}
|
||||
|
||||
await this.jobRepository.queueAll(jobs);
|
||||
return JobStatus.Success;
|
||||
}
|
||||
|
||||
@OnJob({ name: JobName.Ocr, queue: QueueName.Ocr })
|
||||
async handleOcr({ id }: JobOf<JobName.Ocr>): Promise<JobStatus> {
|
||||
const { machineLearning } = await this.getConfig({ withCache: true });
|
||||
if (!isOcrEnabled(machineLearning)) {
|
||||
return JobStatus.Skipped;
|
||||
}
|
||||
|
||||
const asset = await this.assetJobRepository.getForOcr(id);
|
||||
if (!asset || !asset.previewFile) {
|
||||
return JobStatus.Failed;
|
||||
}
|
||||
|
||||
if (asset.visibility === AssetVisibility.Hidden) {
|
||||
return JobStatus.Skipped;
|
||||
}
|
||||
|
||||
const ocrResults = await this.machineLearningRepository.ocr(asset.previewFile, machineLearning.ocr);
|
||||
|
||||
await this.ocrRepository.upsert(id, this.parseOcrResults(id, ocrResults));
|
||||
|
||||
await this.assetRepository.upsertJobStatus({ assetId: id, ocrAt: new Date() });
|
||||
|
||||
this.logger.debug(`Processed ${ocrResults.text.length} OCR result(s) for ${id}`);
|
||||
return JobStatus.Success;
|
||||
}
|
||||
|
||||
private parseOcrResults(id: string, { box, boxScore, text, textScore }: OCR) {
|
||||
const ocrDataList = [];
|
||||
for (let i = 0; i < text.length; i++) {
|
||||
const boxOffset = i * 8;
|
||||
ocrDataList.push({
|
||||
assetId: id,
|
||||
x1: box[boxOffset],
|
||||
y1: box[boxOffset + 1],
|
||||
x2: box[boxOffset + 2],
|
||||
y2: box[boxOffset + 3],
|
||||
x3: box[boxOffset + 4],
|
||||
y3: box[boxOffset + 5],
|
||||
x4: box[boxOffset + 6],
|
||||
y4: box[boxOffset + 7],
|
||||
boxScore: boxScore[i],
|
||||
textScore: textScore[i],
|
||||
text: text[i],
|
||||
});
|
||||
}
|
||||
return ocrDataList;
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,243 @@
|
||||
import { Kysely } from 'kysely';
|
||||
import { AssetFileType, JobStatus } from 'src/enum';
|
||||
import { AssetJobRepository } from 'src/repositories/asset-job.repository';
|
||||
import { AssetRepository } from 'src/repositories/asset.repository';
|
||||
import { ConfigRepository } from 'src/repositories/config.repository';
|
||||
import { JobRepository } from 'src/repositories/job.repository';
|
||||
import { LoggingRepository } from 'src/repositories/logging.repository';
|
||||
import { MachineLearningRepository } from 'src/repositories/machine-learning.repository';
|
||||
import { OcrRepository } from 'src/repositories/ocr.repository';
|
||||
import { SystemMetadataRepository } from 'src/repositories/system-metadata.repository';
|
||||
import { DB } from 'src/schema';
|
||||
import { OcrService } from 'src/services/ocr.service';
|
||||
import { newMediumService } from 'test/medium.factory';
|
||||
import { getKyselyDB } from 'test/utils';
|
||||
|
||||
let defaultDatabase: Kysely<DB>;
|
||||
|
||||
const setup = (db?: Kysely<DB>) => {
|
||||
return newMediumService(OcrService, {
|
||||
database: db || defaultDatabase,
|
||||
real: [AssetRepository, AssetJobRepository, ConfigRepository, OcrRepository, SystemMetadataRepository],
|
||||
mock: [JobRepository, LoggingRepository, MachineLearningRepository],
|
||||
});
|
||||
};
|
||||
|
||||
beforeAll(async () => {
|
||||
defaultDatabase = await getKyselyDB();
|
||||
});
|
||||
|
||||
describe(OcrService.name, () => {
|
||||
it('should work', () => {
|
||||
const { sut } = setup();
|
||||
expect(sut).toBeDefined();
|
||||
});
|
||||
|
||||
it('should parse asset', async () => {
|
||||
const { sut, ctx } = setup();
|
||||
const { user } = await ctx.newUser();
|
||||
const { asset } = await ctx.newAsset({ ownerId: user.id });
|
||||
await ctx.newAssetFile({ assetId: asset.id, type: AssetFileType.Preview, path: 'preview.jpg' });
|
||||
|
||||
const machineLearningMock = ctx.getMock(MachineLearningRepository);
|
||||
machineLearningMock.ocr.mockResolvedValue({
|
||||
box: [10, 10, 50, 10, 50, 50, 10, 50],
|
||||
boxScore: [0.99],
|
||||
text: ['Test OCR'],
|
||||
textScore: [0.95],
|
||||
});
|
||||
|
||||
await expect(sut.handleOcr({ id: asset.id })).resolves.toBe(JobStatus.Success);
|
||||
|
||||
const ocrRepository = ctx.get(OcrRepository);
|
||||
await expect(ocrRepository.getByAssetId(asset.id)).resolves.toEqual([
|
||||
{
|
||||
assetId: asset.id,
|
||||
boxScore: 0.99,
|
||||
id: expect.any(String),
|
||||
text: 'Test OCR',
|
||||
textScore: 0.95,
|
||||
x1: 10,
|
||||
y1: 10,
|
||||
x2: 50,
|
||||
y2: 10,
|
||||
x3: 50,
|
||||
y3: 50,
|
||||
x4: 10,
|
||||
y4: 50,
|
||||
},
|
||||
]);
|
||||
await expect(
|
||||
ctx.database.selectFrom('ocr_search').selectAll().where('assetId', '=', asset.id).executeTakeFirst(),
|
||||
).resolves.toEqual({
|
||||
assetId: asset.id,
|
||||
text: 'Test OCR',
|
||||
});
|
||||
await expect(
|
||||
ctx.database
|
||||
.selectFrom('asset_job_status')
|
||||
.select('asset_job_status.ocrAt')
|
||||
.where('assetId', '=', asset.id)
|
||||
.executeTakeFirst(),
|
||||
).resolves.toEqual({ ocrAt: expect.any(Date) });
|
||||
});
|
||||
|
||||
it('should handle multiple boxes', async () => {
|
||||
const { sut, ctx } = setup();
|
||||
const { user } = await ctx.newUser();
|
||||
const { asset } = await ctx.newAsset({ ownerId: user.id });
|
||||
await ctx.newAssetFile({ assetId: asset.id, type: AssetFileType.Preview, path: 'preview.jpg' });
|
||||
|
||||
const machineLearningMock = ctx.getMock(MachineLearningRepository);
|
||||
machineLearningMock.ocr.mockResolvedValue({
|
||||
box: Array.from({ length: 8 * 5 }, (_, i) => i),
|
||||
boxScore: [0.7, 0.67, 0.65, 0.62, 0.6],
|
||||
text: ['One', 'Two', 'Three', 'Four', 'Five'],
|
||||
textScore: [0.9, 0.89, 0.88, 0.87, 0.86],
|
||||
});
|
||||
|
||||
await expect(sut.handleOcr({ id: asset.id })).resolves.toBe(JobStatus.Success);
|
||||
|
||||
const ocrRepository = ctx.get(OcrRepository);
|
||||
await expect(ocrRepository.getByAssetId(asset.id)).resolves.toEqual([
|
||||
{
|
||||
assetId: asset.id,
|
||||
boxScore: 0.7,
|
||||
id: expect.any(String),
|
||||
text: 'One',
|
||||
textScore: 0.9,
|
||||
x1: 0,
|
||||
y1: 1,
|
||||
x2: 2,
|
||||
y2: 3,
|
||||
x3: 4,
|
||||
y3: 5,
|
||||
x4: 6,
|
||||
y4: 7,
|
||||
},
|
||||
{
|
||||
assetId: asset.id,
|
||||
boxScore: 0.67,
|
||||
id: expect.any(String),
|
||||
text: 'Two',
|
||||
textScore: 0.89,
|
||||
x1: 8,
|
||||
y1: 9,
|
||||
x2: 10,
|
||||
y2: 11,
|
||||
x3: 12,
|
||||
y3: 13,
|
||||
x4: 14,
|
||||
y4: 15,
|
||||
},
|
||||
{
|
||||
assetId: asset.id,
|
||||
boxScore: 0.65,
|
||||
id: expect.any(String),
|
||||
text: 'Three',
|
||||
textScore: 0.88,
|
||||
x1: 16,
|
||||
y1: 17,
|
||||
x2: 18,
|
||||
y2: 19,
|
||||
x3: 20,
|
||||
y3: 21,
|
||||
x4: 22,
|
||||
y4: 23,
|
||||
},
|
||||
{
|
||||
assetId: asset.id,
|
||||
boxScore: 0.62,
|
||||
id: expect.any(String),
|
||||
text: 'Four',
|
||||
textScore: 0.87,
|
||||
x1: 24,
|
||||
y1: 25,
|
||||
x2: 26,
|
||||
y2: 27,
|
||||
x3: 28,
|
||||
y3: 29,
|
||||
x4: 30,
|
||||
y4: 31,
|
||||
},
|
||||
{
|
||||
assetId: asset.id,
|
||||
boxScore: 0.6,
|
||||
id: expect.any(String),
|
||||
text: 'Five',
|
||||
textScore: 0.86,
|
||||
x1: 32,
|
||||
y1: 33,
|
||||
x2: 34,
|
||||
y2: 35,
|
||||
x3: 36,
|
||||
y3: 37,
|
||||
x4: 38,
|
||||
y4: 39,
|
||||
},
|
||||
]);
|
||||
await expect(
|
||||
ctx.database.selectFrom('ocr_search').selectAll().where('assetId', '=', asset.id).executeTakeFirst(),
|
||||
).resolves.toEqual({
|
||||
assetId: asset.id,
|
||||
text: 'One Two Three Four Five',
|
||||
});
|
||||
await expect(
|
||||
ctx.database
|
||||
.selectFrom('asset_job_status')
|
||||
.select('asset_job_status.ocrAt')
|
||||
.where('assetId', '=', asset.id)
|
||||
.executeTakeFirst(),
|
||||
).resolves.toEqual({ ocrAt: expect.any(Date) });
|
||||
});
|
||||
|
||||
it('should handle no boxes', async () => {
|
||||
const { sut, ctx } = setup();
|
||||
const { user } = await ctx.newUser();
|
||||
const { asset } = await ctx.newAsset({ ownerId: user.id });
|
||||
await ctx.newAssetFile({ assetId: asset.id, type: AssetFileType.Preview, path: 'preview.jpg' });
|
||||
|
||||
const machineLearningMock = ctx.getMock(MachineLearningRepository);
|
||||
machineLearningMock.ocr.mockResolvedValue({ box: [], boxScore: [], text: [], textScore: [] });
|
||||
|
||||
await expect(sut.handleOcr({ id: asset.id })).resolves.toBe(JobStatus.Success);
|
||||
|
||||
const ocrRepository = ctx.get(OcrRepository);
|
||||
await expect(ocrRepository.getByAssetId(asset.id)).resolves.toEqual([]);
|
||||
await expect(
|
||||
ctx.database.selectFrom('ocr_search').selectAll().where('assetId', '=', asset.id).executeTakeFirst(),
|
||||
).resolves.toBeUndefined();
|
||||
await expect(
|
||||
ctx.database
|
||||
.selectFrom('asset_job_status')
|
||||
.select('asset_job_status.ocrAt')
|
||||
.where('assetId', '=', asset.id)
|
||||
.executeTakeFirst(),
|
||||
).resolves.toEqual({ ocrAt: expect.any(Date) });
|
||||
});
|
||||
|
||||
it('should update existing results', async () => {
|
||||
const { sut, ctx } = setup();
|
||||
const { user } = await ctx.newUser();
|
||||
const { asset } = await ctx.newAsset({ ownerId: user.id });
|
||||
await ctx.newAssetFile({ assetId: asset.id, type: AssetFileType.Preview, path: 'preview.jpg' });
|
||||
|
||||
const machineLearningMock = ctx.getMock(MachineLearningRepository);
|
||||
machineLearningMock.ocr.mockResolvedValue({
|
||||
box: [10, 10, 50, 10, 50, 50, 10, 50],
|
||||
boxScore: [0.99],
|
||||
text: ['Test OCR'],
|
||||
textScore: [0.95],
|
||||
});
|
||||
await expect(sut.handleOcr({ id: asset.id })).resolves.toBe(JobStatus.Success);
|
||||
|
||||
machineLearningMock.ocr.mockResolvedValue({ box: [], boxScore: [], text: [], textScore: [] });
|
||||
await expect(sut.handleOcr({ id: asset.id })).resolves.toBe(JobStatus.Success);
|
||||
|
||||
const ocrRepository = ctx.get(OcrRepository);
|
||||
await expect(ocrRepository.getByAssetId(asset.id)).resolves.toEqual([]);
|
||||
await expect(
|
||||
ctx.database.selectFrom('ocr_search').selectAll().where('assetId', '=', asset.id).executeTakeFirst(),
|
||||
).resolves.toBeUndefined();
|
||||
});
|
||||
});
|
||||
Loading…
Reference in New Issue