ugolefoo commited on
Commit
4721608
Β·
verified Β·
1 Parent(s): b97942e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +34 -18
app.py CHANGED
@@ -4,6 +4,7 @@ import pytesseract
4
  import requests
5
  import pandas as pd
6
  import gradio as gr
 
7
  from io import BytesIO
8
 
9
  # ──────────────────────────────────────────────────────────────
@@ -54,10 +55,12 @@ def ocr_on_region(image: np.ndarray, box: tuple):
54
  Return the raw OCR text.
55
  """
56
  x, y, w, h = box
57
- cropped = image[y:y + h, x:x + w]
58
  gray_crop = cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY)
59
- _, thresh_crop = cv2.threshold(gray_crop, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
60
- custom_config = r'--oem 3 --psm 6'
 
 
61
  text = pytesseract.image_to_string(thresh_crop, config=custom_config)
62
  return text.strip()
63
 
@@ -84,7 +87,7 @@ def query_openlibrary(title_text: str, author_text: str = None):
84
  "title": doc.get("title", ""),
85
  "author_name": ", ".join(doc.get("author_name", [])),
86
  "publisher": ", ".join(doc.get("publisher", [])),
87
- "first_publish_year": doc.get("first_publish_year", "")
88
  }
89
  except Exception as e:
90
  print(f"OpenLibrary query failed: {e}")
@@ -97,7 +100,7 @@ def query_openlibrary(title_text: str, author_text: str = None):
97
  def process_image(image_file):
98
  """
99
  Gradio passes a PIL image or numpy array. Convert to OpenCV BGR, detect covers β†’ OCR β†’ OpenLibrary.
100
- Return a DataFrame and CSV bytes.
101
  """
102
  img = np.array(image_file)[:, :, ::-1].copy() # PIL to OpenCV BGR
103
  boxes = detect_book_regions(img)
@@ -116,20 +119,28 @@ def process_image(image_file):
116
  if meta:
117
  records.append(meta)
118
  else:
119
- records.append({
120
- "title": title_guess,
121
- "author_name": author_guess or "",
122
- "publisher": "",
123
- "first_publish_year": "",
124
- })
 
 
125
 
126
  if not records:
127
  df_empty = pd.DataFrame(columns=["title", "author_name", "publisher", "first_publish_year"])
128
- return df_empty, df_empty.to_csv(index=False).encode()
 
 
 
 
129
 
130
  df = pd.DataFrame(records)
131
  csv_bytes = df.to_csv(index=False).encode()
132
- return df, csv_bytes
 
 
133
 
134
  # ──────────────────────────────────────────────────────────────
135
  # 5. Build the Gradio Interface
@@ -151,15 +162,20 @@ def build_interface():
151
 
152
  output_table = gr.Dataframe(
153
  headers=["title", "author_name", "publisher", "first_publish_year"],
154
- label="Detected Books with Metadata"
 
155
  )
156
- download_btn = gr.Download(label="Download CSV")
157
 
158
  def on_run(image):
159
- df, csv_bytes = process_image(image)
160
- return df, csv_bytes
161
 
162
- run_button.click(fn=on_run, inputs=[img_in], outputs=[output_table, download_btn])
 
 
 
 
163
 
164
  return demo
165
 
 
4
  import requests
5
  import pandas as pd
6
  import gradio as gr
7
+ import io
8
  from io import BytesIO
9
 
10
  # ──────────────────────────────────────────────────────────────
 
55
  Return the raw OCR text.
56
  """
57
  x, y, w, h = box
58
+ cropped = image[y : y + h, x : x + w]
59
  gray_crop = cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY)
60
+ _, thresh_crop = cv2.threshold(
61
+ gray_crop, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
62
+ )
63
+ custom_config = r"--oem 3 --psm 6"
64
  text = pytesseract.image_to_string(thresh_crop, config=custom_config)
65
  return text.strip()
66
 
 
87
  "title": doc.get("title", ""),
88
  "author_name": ", ".join(doc.get("author_name", [])),
89
  "publisher": ", ".join(doc.get("publisher", [])),
90
+ "first_publish_year": doc.get("first_publish_year", ""),
91
  }
92
  except Exception as e:
93
  print(f"OpenLibrary query failed: {e}")
 
100
  def process_image(image_file):
101
  """
102
  Gradio passes a PIL image or numpy array. Convert to OpenCV BGR, detect covers β†’ OCR β†’ OpenLibrary.
103
+ Return a DataFrame and a (filename, BytesIO) tuple for CSV.
104
  """
105
  img = np.array(image_file)[:, :, ::-1].copy() # PIL to OpenCV BGR
106
  boxes = detect_book_regions(img)
 
119
  if meta:
120
  records.append(meta)
121
  else:
122
+ records.append(
123
+ {
124
+ "title": title_guess,
125
+ "author_name": author_guess or "",
126
+ "publisher": "",
127
+ "first_publish_year": "",
128
+ }
129
+ )
130
 
131
  if not records:
132
  df_empty = pd.DataFrame(columns=["title", "author_name", "publisher", "first_publish_year"])
133
+ # Build an empty CSV bytes buffer
134
+ empty_csv = df_empty.to_csv(index=False).encode()
135
+ buffer = io.BytesIO(empty_csv)
136
+ buffer.name = "books.csv"
137
+ return df_empty, buffer
138
 
139
  df = pd.DataFrame(records)
140
  csv_bytes = df.to_csv(index=False).encode()
141
+ buffer = io.BytesIO(csv_bytes)
142
+ buffer.name = "books.csv"
143
+ return df, buffer
144
 
145
  # ──────────────────────────────────────────────────────────────
146
  # 5. Build the Gradio Interface
 
162
 
163
  output_table = gr.Dataframe(
164
  headers=["title", "author_name", "publisher", "first_publish_year"],
165
+ label="Detected Books with Metadata",
166
+ datatype="pandas",
167
  )
168
+ download_file = gr.File(label="Download CSV")
169
 
170
  def on_run(image):
171
+ df, file_buffer = process_image(image)
172
+ return df, file_buffer
173
 
174
+ run_button.click(
175
+ fn=on_run,
176
+ inputs=[img_in],
177
+ outputs=[output_table, download_file],
178
+ )
179
 
180
  return demo
181