Vinay15 commited on
Commit
e8ba698
·
verified ·
1 Parent(s): 9c1f1ad

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +17 -19
app.py CHANGED
@@ -5,36 +5,34 @@ import gradio as gr
5
  import tempfile
6
 
7
  # Load the OCR model and tokenizer
8
- tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2.0',
9
- trust_remote_code=True,
10
- token="<your_access_token>")
11
- model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2.0',
12
- trust_remote_code=True,
13
- low_cpu_mem_usage=True,
14
- pad_token_id=tokenizer.eos_token_id).eval()
15
 
16
  # Check if GPU is available and use it, else use CPU
17
- # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
18
- # model = model.to(device) # Uncomment if you want to use GPU
19
 
20
  # Function to perform OCR on the image
21
  def perform_ocr(image):
 
22
  with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
23
- image.save(temp_file.name)
24
- temp_image_path = temp_file.name
25
 
26
- # Perform OCR using the model
27
- result = model.chat(tokenizer, temp_image_path, ocr_type='ocr')
28
- return result
29
 
30
- # Create the Gradio interface
31
  interface = gr.Interface(
32
- fn=perform_ocr,
33
- inputs=gr.Image(type="pil"),
34
- outputs=gr.Textbox(),
35
  title="OCR Web App",
36
  description="Upload an image to extract text using the GOT-OCR2.0 model."
37
  )
38
 
39
  # Launch the app
40
- interface.launch()
 
5
  import tempfile
6
 
7
  # Load the OCR model and tokenizer
8
+ tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
9
+ model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True,
10
+ low_cpu_mem_usage=True,
11
+ pad_token_id=tokenizer.eos_token_id).eval()
 
 
 
12
 
13
  # Check if GPU is available and use it, else use CPU
14
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
15
+ model = model.to(device)
16
 
17
  # Function to perform OCR on the image
18
  def perform_ocr(image):
19
+ # Save the image to a temporary file
20
  with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
21
+ image.save(temp_file.name) # Save the image
22
+ temp_image_path = temp_file.name # Get the file path for the saved image
23
 
24
+ # Perform OCR using the model
25
+ result = model.chat(tokenizer, temp_image_path, ocr_type='ocr')
26
+ return result
27
 
28
+ # Create the Gradio interface using the new syntax
29
  interface = gr.Interface(
30
+ fn=perform_ocr,
31
+ inputs=gr.Image(type="pil"), # Updated to gr.Image
32
+ outputs=gr.Textbox(), # Updated to gr.Textbox
33
  title="OCR Web App",
34
  description="Upload an image to extract text using the GOT-OCR2.0 model."
35
  )
36
 
37
  # Launch the app
38
+ interface.launch()