Update app.py
Browse files
app.py
CHANGED
@@ -11,7 +11,14 @@ from frontPrompt import main as main_f
|
|
11 |
model_path = "best.pt"
|
12 |
model = YOLO(model_path)
|
13 |
|
14 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
15 |
|
16 |
def preprocessing(image):
|
17 |
"""Apply three enhancement filters, including brightness reduction, and resize."""
|
@@ -84,9 +91,9 @@ def crop_image(image, bounding_boxes):
|
|
84 |
def vision_ai_api(image, doc_type):
|
85 |
|
86 |
if doc_type == "front":
|
87 |
-
results = main_f(image)
|
88 |
if doc_type == "back":
|
89 |
-
results = main_b(image)
|
90 |
|
91 |
return results
|
92 |
|
|
|
11 |
model_path = "best.pt"
|
12 |
model = YOLO(model_path)
|
13 |
|
14 |
+
model = AutoModel.from_pretrained(
|
15 |
+
path,
|
16 |
+
torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
|
17 |
+
# load_in_8bit=True,
|
18 |
+
low_cpu_mem_usage=True,
|
19 |
+
use_flash_attn=True,
|
20 |
+
trust_remote_code=True).eval()
|
21 |
+
tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
|
22 |
|
23 |
def preprocessing(image):
|
24 |
"""Apply three enhancement filters, including brightness reduction, and resize."""
|
|
|
91 |
def vision_ai_api(image, doc_type):
|
92 |
|
93 |
if doc_type == "front":
|
94 |
+
results = main_f(image,model,tokenizer)
|
95 |
if doc_type == "back":
|
96 |
+
results = main_b(image,model,tokenizer)
|
97 |
|
98 |
return results
|
99 |
|