mirror of
https://github.com/tcsenpai/pensieve.git
synced 2025-06-10 13:07:15 +00:00
feat(ocr): support ocrmac
This commit is contained in:
parent
a1a5e576c0
commit
65190f1e4c
@ -12,6 +12,7 @@ from concurrent.futures import ThreadPoolExecutor
|
|||||||
from functools import partial
|
from functools import partial
|
||||||
import yaml
|
import yaml
|
||||||
import io
|
import io
|
||||||
|
import platform
|
||||||
|
|
||||||
MAX_THUMBNAIL_SIZE = (1920, 1920)
|
MAX_THUMBNAIL_SIZE = (1920, 1920)
|
||||||
|
|
||||||
@ -73,14 +74,38 @@ def convert_ocr_results(results):
|
|||||||
return converted
|
return converted
|
||||||
|
|
||||||
|
|
||||||
|
def convert_ocr_data(ocr_data):
|
||||||
|
converted_data = []
|
||||||
|
for text, score, bbox in ocr_data:
|
||||||
|
x_min, y_min, x_max, y_max = bbox
|
||||||
|
dt_boxes = [
|
||||||
|
[x_min, y_min],
|
||||||
|
[x_max, y_min],
|
||||||
|
[x_max, y_max],
|
||||||
|
[x_min, y_max]
|
||||||
|
]
|
||||||
|
entry = {
|
||||||
|
'dt_boxes': dt_boxes,
|
||||||
|
'rec_txt': text,
|
||||||
|
'score': float(score)
|
||||||
|
}
|
||||||
|
converted_data.append(entry)
|
||||||
|
return converted_data
|
||||||
|
|
||||||
|
|
||||||
def predict_local(img_path):
|
def predict_local(img_path):
|
||||||
try:
|
try:
|
||||||
with Image.open(img_path) as img:
|
if platform.system() == 'Darwin': # Check if the OS is macOS
|
||||||
img = img.convert("RGB")
|
from ocrmac import ocrmac
|
||||||
img.thumbnail(MAX_THUMBNAIL_SIZE)
|
result = ocrmac.OCR(img_path, language_preference=['zh-Hans']).recognize(px=True)
|
||||||
img_array = np.array(img)
|
return convert_ocr_data(result)
|
||||||
results, _ = ocr(img_array)
|
else:
|
||||||
return convert_ocr_results(results)
|
with Image.open(img_path) as img:
|
||||||
|
img = img.convert("RGB")
|
||||||
|
img.thumbnail(MAX_THUMBNAIL_SIZE)
|
||||||
|
img_array = np.array(img)
|
||||||
|
results, _ = ocr(img_array)
|
||||||
|
return convert_ocr_results(results)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error processing image {img_path}: {str(e)}")
|
logger.error(f"Error processing image {img_path}: {str(e)}")
|
||||||
return None
|
return None
|
||||||
|
@ -36,6 +36,7 @@ dependencies = [
|
|||||||
"pyobjc; sys_platform == 'darwin'",
|
"pyobjc; sys_platform == 'darwin'",
|
||||||
"pyobjc-core; sys_platform == 'darwin'",
|
"pyobjc-core; sys_platform == 'darwin'",
|
||||||
"pyobjc-framework-Quartz; sys_platform == 'darwin'",
|
"pyobjc-framework-Quartz; sys_platform == 'darwin'",
|
||||||
|
"ocrmac; sys_platform == 'darwin'",
|
||||||
"sentence-transformers",
|
"sentence-transformers",
|
||||||
"torch",
|
"torch",
|
||||||
"numpy",
|
"numpy",
|
||||||
|
Loading…
x
Reference in New Issue
Block a user