zliang commited on
Commit
b1e4794
1 Parent(s): cff5fa2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -4
app.py CHANGED
@@ -5,13 +5,16 @@ import spaces
5
  from ultralytics import YOLOv10
6
 
7
  # Load the trained model
 
8
  model = YOLOv10("best.pt")
9
 
 
10
  # Define the class indices for figures and tables
11
  figure_class_index = 3 # class index for figures
12
  table_class_index = 4 # class index for tables
13
 
14
- # Function to perform inference on a single image and return bounding boxes for figures and tables
 
15
  def infer_image_and_get_boxes(image, confidence_threshold=0.6):
16
  results = model.predict(image)
17
  boxes = [
@@ -22,6 +25,7 @@ def infer_image_and_get_boxes(image, confidence_threshold=0.6):
22
  return boxes
23
 
24
  # Function to crop images from the boxes
 
25
  def crop_images_from_boxes(image, boxes, scale_factor):
26
  cropped_images = [
27
  image[int(y1 * scale_factor):int(y2 * scale_factor), int(x1 * scale_factor):int(x2 * scale_factor)]
@@ -44,7 +48,7 @@ def process_pdf(pdf_file):
44
 
45
  # Pre-cache all page pixmaps at low DPI
46
  low_res_pixmaps = [page.get_pixmap(dpi=low_dpi) for page in doc]
47
-
48
  # Loop through each page
49
  for page_num, low_res_pix in enumerate(low_res_pixmaps):
50
  low_res_img = np.frombuffer(low_res_pix.samples, dtype=np.uint8).reshape(low_res_pix.height, low_res_pix.width, 3)
@@ -74,5 +78,3 @@ iface = gr.Interface(
74
 
75
  # Launch the app
76
  iface.launch()
77
-
78
-
 
5
  from ultralytics import YOLOv10
6
 
7
  # Load the trained model
8
+
9
  model = YOLOv10("best.pt")
10
 
11
+
12
  # Define the class indices for figures and tables
13
  figure_class_index = 3 # class index for figures
14
  table_class_index = 4 # class index for tables
15
 
16
+ # Function to perform inference on an image and return bounding boxes for figures and tables
17
+
18
  def infer_image_and_get_boxes(image, confidence_threshold=0.6):
19
  results = model.predict(image)
20
  boxes = [
 
25
  return boxes
26
 
27
  # Function to crop images from the boxes
28
+
29
  def crop_images_from_boxes(image, boxes, scale_factor):
30
  cropped_images = [
31
  image[int(y1 * scale_factor):int(y2 * scale_factor), int(x1 * scale_factor):int(x2 * scale_factor)]
 
48
 
49
  # Pre-cache all page pixmaps at low DPI
50
  low_res_pixmaps = [page.get_pixmap(dpi=low_dpi) for page in doc]
51
+
52
  # Loop through each page
53
  for page_num, low_res_pix in enumerate(low_res_pixmaps):
54
  low_res_img = np.frombuffer(low_res_pix.samples, dtype=np.uint8).reshape(low_res_pix.height, low_res_pix.width, 3)
 
78
 
79
  # Launch the app
80
  iface.launch()