deb113 commited on
Commit
22a9afa
·
verified ·
1 Parent(s): f9fda49

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -10
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/yolov11x_best.pt"
21
  if not os.path.exists(model_path):
22
- url = "https://github.com/moured/YOLOv11-Document-Layout-Analysis/releases/download/doclaynet_weights/yolov11x_best.pt"
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("# YOLOv11x Document Layout Analysis for Visual Elements")
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 YOLOv11x for document layout analysis, with a focus on extracting visual elements.
193
- Model from [moured/YOLOv11-Document-Layout-Analysis](https://github.com/moured/YOLOv11-Document-Layout-Analysis)
194
  """)
195
-
196
-
197
 
198
  if __name__ == "__main__":
199
- demo.launch()
 
 
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)