Spaces:
Running
Running
Update ocr_cpu.py
Browse filesError during text extraction: eval() arg 1 must be a string, bytes or code object , Fixing this error
- ocr_cpu.py +33 -22
ocr_cpu.py
CHANGED
@@ -21,19 +21,20 @@ model = AutoModel.from_pretrained(
|
|
21 |
# Ensure the model is in evaluation mode and loaded on CPU
|
22 |
device = torch.device("cpu")
|
23 |
dtype = torch.float32 # Use float32 on CPU
|
24 |
-
model = model.eval()
|
25 |
|
26 |
# OCR function
|
27 |
-
|
28 |
-
|
29 |
def extract_text_got(uploaded_file):
|
30 |
"""Use GOT-OCR2.0 model to extract text from the uploaded image."""
|
|
|
|
|
31 |
try:
|
32 |
-
|
33 |
with open(temp_file_path, 'wb') as temp_file:
|
34 |
-
temp_file.write(uploaded_file.read())
|
|
|
|
|
35 |
|
36 |
-
# OCR attempts
|
37 |
ocr_types = ['ocr', 'format']
|
38 |
fine_grained_options = ['ocr', 'format']
|
39 |
color_options = ['red', 'green', 'blue']
|
@@ -42,12 +43,15 @@ def extract_text_got(uploaded_file):
|
|
42 |
|
43 |
results = []
|
44 |
|
45 |
-
# Run
|
46 |
for ocr_type in ocr_types:
|
47 |
with torch.no_grad():
|
48 |
-
|
49 |
-
|
50 |
-
|
|
|
|
|
|
|
51 |
if isinstance(outputs, list) and outputs[0].strip():
|
52 |
return outputs[0].strip() # Return if successful
|
53 |
results.append(outputs[0].strip() if outputs else "No result")
|
@@ -55,9 +59,11 @@ def extract_text_got(uploaded_file):
|
|
55 |
# Try FINE-GRAINED OCR with box options
|
56 |
for ocr_type in fine_grained_options:
|
57 |
with torch.no_grad():
|
58 |
-
|
59 |
-
|
60 |
-
|
|
|
|
|
61 |
if isinstance(outputs, list) and outputs[0].strip():
|
62 |
return outputs[0].strip() # Return if successful
|
63 |
results.append(outputs[0].strip() if outputs else "No result")
|
@@ -66,25 +72,28 @@ def extract_text_got(uploaded_file):
|
|
66 |
for ocr_type in fine_grained_options:
|
67 |
for color in color_options:
|
68 |
with torch.no_grad():
|
69 |
-
|
70 |
-
|
71 |
-
|
|
|
|
|
72 |
if isinstance(outputs, list) and outputs[0].strip():
|
73 |
return outputs[0].strip() # Return if successful
|
74 |
-
results.append(outputs[0].strip()
|
75 |
-
if outputs else "No result")
|
76 |
|
77 |
# Try MULTI-CROP OCR
|
78 |
for ocr_type in multi_crop_types:
|
79 |
with torch.no_grad():
|
80 |
-
|
81 |
-
|
82 |
-
|
|
|
|
|
83 |
if isinstance(outputs, list) and outputs[0].strip():
|
84 |
return outputs[0].strip() # Return if successful
|
85 |
results.append(outputs[0].strip() if outputs else "No result")
|
86 |
|
87 |
-
#
|
88 |
if all(not text for text in results):
|
89 |
return "No text extracted."
|
90 |
else:
|
@@ -94,5 +103,7 @@ def extract_text_got(uploaded_file):
|
|
94 |
return f"Error during text extraction: {str(e)}"
|
95 |
|
96 |
finally:
|
|
|
97 |
if os.path.exists(temp_file_path):
|
98 |
os.remove(temp_file_path)
|
|
|
|
21 |
# Ensure the model is in evaluation mode and loaded on CPU
|
22 |
device = torch.device("cpu")
|
23 |
dtype = torch.float32 # Use float32 on CPU
|
24 |
+
model = model.eval().to(device)
|
25 |
|
26 |
# OCR function
|
|
|
|
|
27 |
def extract_text_got(uploaded_file):
|
28 |
"""Use GOT-OCR2.0 model to extract text from the uploaded image."""
|
29 |
+
temp_file_path = 'temp_image.jpg'
|
30 |
+
|
31 |
try:
|
32 |
+
# Save the uploaded file temporarily
|
33 |
with open(temp_file_path, 'wb') as temp_file:
|
34 |
+
temp_file.write(uploaded_file.read())
|
35 |
+
|
36 |
+
print(f"Processing image from path: {temp_file_path}") # Debug info
|
37 |
|
|
|
38 |
ocr_types = ['ocr', 'format']
|
39 |
fine_grained_options = ['ocr', 'format']
|
40 |
color_options = ['red', 'green', 'blue']
|
|
|
43 |
|
44 |
results = []
|
45 |
|
46 |
+
# Run basic OCR types
|
47 |
for ocr_type in ocr_types:
|
48 |
with torch.no_grad():
|
49 |
+
print(f"Running basic OCR with type: {ocr_type}") # Debug info
|
50 |
+
outputs = model.chat(tokenizer, temp_file_path, ocr_type=ocr_type)
|
51 |
+
|
52 |
+
# Debug outputs
|
53 |
+
print(f"Outputs for {ocr_type}: {outputs}")
|
54 |
+
|
55 |
if isinstance(outputs, list) and outputs[0].strip():
|
56 |
return outputs[0].strip() # Return if successful
|
57 |
results.append(outputs[0].strip() if outputs else "No result")
|
|
|
59 |
# Try FINE-GRAINED OCR with box options
|
60 |
for ocr_type in fine_grained_options:
|
61 |
with torch.no_grad():
|
62 |
+
print(f"Running fine-grained OCR with box, type: {ocr_type}") # Debug info
|
63 |
+
outputs = model.chat(tokenizer, temp_file_path, ocr_type=ocr_type, ocr_box=box)
|
64 |
+
|
65 |
+
print(f"Outputs for {ocr_type} with box: {outputs}")
|
66 |
+
|
67 |
if isinstance(outputs, list) and outputs[0].strip():
|
68 |
return outputs[0].strip() # Return if successful
|
69 |
results.append(outputs[0].strip() if outputs else "No result")
|
|
|
72 |
for ocr_type in fine_grained_options:
|
73 |
for color in color_options:
|
74 |
with torch.no_grad():
|
75 |
+
print(f"Running fine-grained OCR with color {color}, type: {ocr_type}") # Debug info
|
76 |
+
outputs = model.chat(tokenizer, temp_file_path, ocr_type=ocr_type, ocr_color=color)
|
77 |
+
|
78 |
+
print(f"Outputs for {ocr_type} with color {color}: {outputs}")
|
79 |
+
|
80 |
if isinstance(outputs, list) and outputs[0].strip():
|
81 |
return outputs[0].strip() # Return if successful
|
82 |
+
results.append(outputs[0].strip() if outputs else "No result")
|
|
|
83 |
|
84 |
# Try MULTI-CROP OCR
|
85 |
for ocr_type in multi_crop_types:
|
86 |
with torch.no_grad():
|
87 |
+
print(f"Running multi-crop OCR with type: {ocr_type}") # Debug info
|
88 |
+
outputs = model.chat_crop(tokenizer, temp_file_path, ocr_type=ocr_type)
|
89 |
+
|
90 |
+
print(f"Outputs for multi-crop {ocr_type}: {outputs}")
|
91 |
+
|
92 |
if isinstance(outputs, list) and outputs[0].strip():
|
93 |
return outputs[0].strip() # Return if successful
|
94 |
results.append(outputs[0].strip() if outputs else "No result")
|
95 |
|
96 |
+
# Return combined results or no text found message
|
97 |
if all(not text for text in results):
|
98 |
return "No text extracted."
|
99 |
else:
|
|
|
103 |
return f"Error during text extraction: {str(e)}"
|
104 |
|
105 |
finally:
|
106 |
+
# Clean up temporary file
|
107 |
if os.path.exists(temp_file_path):
|
108 |
os.remove(temp_file_path)
|
109 |
+
print(f"Temporary file {temp_file_path} removed.") # Debug info
|