lukiod commited on
Commit
f50b49c
1 Parent(s): 9a0cb2a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -1
app.py CHANGED
@@ -2,6 +2,7 @@ import streamlit as st
2
  import torch
3
  from PIL import Image
4
  import gc
 
5
  import os
6
  from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
7
  from qwen_vl_utils import process_vision_info
@@ -42,9 +43,14 @@ if image:
42
  st.write("Extracting text from image...")
43
  byaldi_model = load_byaldi_model()
44
 
 
 
 
 
 
45
  # Create a temporary index for the uploaded image
46
  with st.spinner("Processing image..."):
47
- byaldi_model.index(img, index_name="temp_index", overwrite=True)
48
 
49
  # Perform a dummy search to get the OCR results
50
  ocr_results = byaldi_model.search("Extract all text from the image", k=1)
@@ -62,6 +68,9 @@ if image:
62
  del byaldi_model
63
  clear_memory()
64
 
 
 
 
65
  # Text input field for question
66
  question = st.text_input("Ask a question about the image and extracted text")
67
 
 
2
  import torch
3
  from PIL import Image
4
  import gc
5
+ import tempfile
6
  import os
7
  from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
8
  from qwen_vl_utils import process_vision_info
 
43
  st.write("Extracting text from image...")
44
  byaldi_model = load_byaldi_model()
45
 
46
+ # Save the image to a temporary file
47
+ with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
48
+ img.save(temp_file, format="JPEG")
49
+ temp_file_path = temp_file.name
50
+
51
  # Create a temporary index for the uploaded image
52
  with st.spinner("Processing image..."):
53
+ byaldi_model.index(temp_file_path, index_name="temp_index", overwrite=True)
54
 
55
  # Perform a dummy search to get the OCR results
56
  ocr_results = byaldi_model.search("Extract all text from the image", k=1)
 
68
  del byaldi_model
69
  clear_memory()
70
 
71
+ # Remove the temporary file
72
+ os.unlink(temp_file_path)
73
+
74
  # Text input field for question
75
  question = st.text_input("Ask a question about the image and extracted text")
76