UniquePratham commited on
Commit
b0416c1
·
verified ·
1 Parent(s): 6bb168e

Update ocr_cpu.py

Browse files

Error during text extraction: eval() arg 1 must be a string, bytes or code object , Fixing this error

Files changed (1) hide show
  1. 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
- temp_file_path = 'temp_image.jpg'
33
  with open(temp_file_path, 'wb') as temp_file:
34
- temp_file.write(uploaded_file.read()) # Save file
 
 
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 the model without autocast (not necessary for CPU)
46
  for ocr_type in ocr_types:
47
  with torch.no_grad():
48
- outputs = model.chat(
49
- tokenizer, temp_file_path, ocr_type=ocr_type
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
- outputs = model.chat(
59
- tokenizer, temp_file_path, ocr_type=ocr_type, ocr_box=box
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
- outputs = model.chat(
70
- tokenizer, temp_file_path, ocr_type=ocr_type, ocr_color=color
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
- outputs = model.chat_crop(
81
- tokenizer, temp_file_path, ocr_type=ocr_type
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
- # If no text was extracted
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