rishabhv471 commited on
Commit
36ab27c
β€’
1 Parent(s): 5e87164
Files changed (1) hide show
  1. app.py +66 -0
app.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from sahi.prediction import ObjectPrediction
4
+ from sahi.utils.cv import visualize_object_predictions, read_image
5
+ from ultralyticsplus import YOLO, render_result
6
+
7
+ # Images
8
+ torch.hub.download_url_to_file('https://huggingface.co/spaces/foduucom/table-extraction-yolov8/resolve/main/test/table1.jpg', 'document1.jpg')
9
+ torch.hub.download_url_to_file('https://huggingface.co/spaces/foduucom/table-extraction-yolov8/resolve/main/test/table2.jpg', 'document2.jpg')
10
+ torch.hub.download_url_to_file('https://huggingface.co/spaces/foduucom/table-extraction-yolov8/resolve/main/test/table3.jpg', 'document3.jpg')
11
+
12
+ def yolov8_inference(
13
+ image: gr.inputs.Image = None,
14
+ model_path: gr.inputs.Dropdown = None,
15
+ image_size: gr.inputs.Slider = 640,
16
+ conf_threshold: gr.inputs.Slider = 0.25,
17
+ iou_threshold: gr.inputs.Slider = 0.45,
18
+ ):
19
+ """
20
+ YOLOv8 inference function
21
+ Args:
22
+ image: Input image
23
+ model_path: Path to the model
24
+ image_size: Image size
25
+ conf_threshold: Confidence threshold
26
+ iou_threshold: IOU threshold
27
+ Returns:
28
+ Rendered image
29
+ """
30
+ model = YOLO(model_path)
31
+ model.overrides['conf'] = conf_threshold
32
+ model.overrides['iou']= iou_threshold
33
+ model.overrides['agnostic_nms'] = False # NMS class-agnostic
34
+ model.overrides['max_det'] = 1000
35
+ image = read_image(image)
36
+ results = model.predict(image)
37
+ render = render_result(model=model, image=image, result=results[0])
38
+
39
+ return render
40
+
41
+
42
+ inputs = [
43
+ gr.inputs.Image(type="filepath", label="Input Image"),
44
+ gr.inputs.Dropdown(["foduucom/table-detection-and-extraction"],
45
+ default="foduucom/table-detection-and-extraction", label="Model"),
46
+ gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
47
+ gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
48
+ gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
49
+ ]
50
+
51
+ outputs = gr.outputs.Image(type="filepath", label="Output Image")
52
+ title = "YoloTableExtract: Efficient Table Detection"
53
+
54
+ description = "πŸ” YoloTableExtract is a powerful space that utilizes YOLOv8s for accurate table detection and extraction. Whether tables are bordered or borderless, this space can effectively identify and extract them from images. For further assistance and support related to documentation or data-related issues, feel free to contact info@foduu.com. If you find this space helpful, please show your appreciation by liking it. β€οΈπŸ‘πŸΌ"
55
+ examples = [['document1.jpg', 'foduucom/table-detection-and-extraction', 640, 0.25, 0.45], ['document2.jpg', 'foduucom/table-detection-and-extraction', 640, 0.25, 0.45], ['document3.jpg', 'foduucom/table-detection-and-extraction', 1280, 0.25, 0.45]]
56
+ demo_app = gr.Interface(
57
+ fn=yolov8_inference,
58
+ inputs=inputs,
59
+ outputs=outputs,
60
+ title=title,
61
+ description=description,
62
+ examples=examples,
63
+ cache_examples=True,
64
+ theme='huggingface',
65
+ )
66
+ demo_app.launch(debug=True, enable_queue=True)