Vinay15 commited on
Commit
448d752
·
verified ·
1 Parent(s): dec293d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +17 -17
app.py CHANGED
@@ -5,34 +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', 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()
 
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',
10
+ trust_remote_code=True,
11
+ low_cpu_mem_usage=True,
12
+ pad_token_id=tokenizer.eos_token_id).eval()
13
 
14
  # Check if GPU is available and use it, else use CPU
15
+ # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
16
+ # model = model.to(device) # Uncomment if you want to use GPU
17
 
18
  # Function to perform OCR on the image
19
  def perform_ocr(image):
 
20
  with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
21
+ image.save(temp_file.name)
22
+ temp_image_path = temp_file.name
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
29
  interface = gr.Interface(
30
+ fn=perform_ocr,
31
+ inputs=gr.Image(type="pil"),
32
+ outputs=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()