Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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
|
|
|
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()
|
|
|
|