Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -17,10 +17,10 @@ import json
|
|
17 |
os.makedirs("models", exist_ok=True)
|
18 |
|
19 |
# Download model if it doesn't exist
|
20 |
-
model_path = "models/
|
21 |
if not os.path.exists(model_path):
|
22 |
-
url = "https://
|
23 |
-
print(f"Downloading model from {url}...")
|
24 |
r = requests.get(url)
|
25 |
with open(model_path, 'wb') as f:
|
26 |
f.write(r.content)
|
@@ -38,7 +38,6 @@ CLASSES = ["Caption", "Footnote", "Formula", "List-item", "Page-footer", "Page-h
|
|
38 |
VISUAL_ELEMENTS = ["Picture", "Caption", "Table", "Formula"]
|
39 |
|
40 |
# Define colors for visualization - Fix for ColorPalette issue
|
41 |
-
# Use the sv.ColorPalette directly or create a custom color palette based on supervision version
|
42 |
try:
|
43 |
# Try newer versions approach
|
44 |
COLORS = sv.ColorPalette.default()
|
@@ -141,7 +140,7 @@ def download_json(json_data):
|
|
141 |
|
142 |
# Set up the Gradio interface
|
143 |
with gr.Blocks() as demo:
|
144 |
-
gr.Markdown("#
|
145 |
gr.Markdown("Upload a document image to extract visual elements including diagrams, tables, formulas, and captions.")
|
146 |
|
147 |
with gr.Row():
|
@@ -189,11 +188,10 @@ with gr.Blocks() as demo:
|
|
189 |
|
190 |
gr.Markdown("## About")
|
191 |
gr.Markdown("""
|
192 |
-
This demo uses
|
193 |
-
|
194 |
""")
|
195 |
-
|
196 |
-
|
197 |
|
198 |
if __name__ == "__main__":
|
199 |
-
|
|
|
|
17 |
os.makedirs("models", exist_ok=True)
|
18 |
|
19 |
# Download model if it doesn't exist
|
20 |
+
model_path = "models/yolov8n-doclaynet.pt"
|
21 |
if not os.path.exists(model_path):
|
22 |
+
url = "https://huggingface.co/datasets/awsaf49/yolov8-doclaynet/resolve/main/yolov8n-doclaynet.pt"
|
23 |
+
print(f"Downloading smaller model from {url}...")
|
24 |
r = requests.get(url)
|
25 |
with open(model_path, 'wb') as f:
|
26 |
f.write(r.content)
|
|
|
38 |
VISUAL_ELEMENTS = ["Picture", "Caption", "Table", "Formula"]
|
39 |
|
40 |
# Define colors for visualization - Fix for ColorPalette issue
|
|
|
41 |
try:
|
42 |
# Try newer versions approach
|
43 |
COLORS = sv.ColorPalette.default()
|
|
|
140 |
|
141 |
# Set up the Gradio interface
|
142 |
with gr.Blocks() as demo:
|
143 |
+
gr.Markdown("# Document Layout Analysis for Visual Elements (YOLOv8n)")
|
144 |
gr.Markdown("Upload a document image to extract visual elements including diagrams, tables, formulas, and captions.")
|
145 |
|
146 |
with gr.Row():
|
|
|
188 |
|
189 |
gr.Markdown("## About")
|
190 |
gr.Markdown("""
|
191 |
+
This demo uses YOLOv8n for document layout analysis, with a focus on extracting visual elements.
|
192 |
+
The model is a smaller, more efficient version trained on the DocLayNet dataset.
|
193 |
""")
|
|
|
|
|
194 |
|
195 |
if __name__ == "__main__":
|
196 |
+
# Specify a lower queue_size and a maximum number of connections to limit memory use
|
197 |
+
demo.launch(share=True, max_threads=1, queue_size=5)
|