|
import gradio as gr |
|
import torch |
|
from sahi.prediction import ObjectPrediction |
|
from sahi.utils.cv import visualize_object_predictions, read_image |
|
from ultralyticsplus import YOLO |
|
|
|
|
|
torch.hub.download_url_to_file('https://huggingface.co/spaces/foduucom/table-extraction-yolov8/resolve/main/test/table1.jpg', 'document1.jpg') |
|
torch.hub.download_url_to_file('https://huggingface.co/spaces/foduucom/table-extraction-yolov8/resolve/main/test/table2.jpg', 'document2.jpg') |
|
torch.hub.download_url_to_file('https://huggingface.co/spaces/foduucom/table-extraction-yolov8/resolve/main/test/table3.jpg', 'document3.jpg') |
|
|
|
def yolov8_inference( |
|
image: gr.inputs.Image = None, |
|
model_path: gr.inputs.Dropdown = None, |
|
image_size: gr.inputs.Slider = 640, |
|
conf_threshold: gr.inputs.Slider = 0.25, |
|
iou_threshold: gr.inputs.Slider = 0.45, |
|
): |
|
""" |
|
YOLOv8 inference function |
|
Args: |
|
image: Input image |
|
model_path: Path to the model |
|
image_size: Image size |
|
conf_threshold: Confidence threshold |
|
iou_threshold: IOU threshold |
|
Returns: |
|
Rendered image |
|
""" |
|
model = YOLO(model_path) |
|
model.conf = conf_threshold |
|
model.iou = iou_threshold |
|
results = model.predict(image, imgsz=image_size) |
|
object_prediction_list = [] |
|
for _, image_results in enumerate(results): |
|
if len(image_results)!=0: |
|
image_predictions_in_xyxy_format = image_results['det'] |
|
for pred in image_predictions_in_xyxy_format: |
|
x1, y1, x2, y2 = ( |
|
int(pred[0]), |
|
int(pred[1]), |
|
int(pred[2]), |
|
int(pred[3]), |
|
) |
|
bbox = [x1, y1, x2, y2] |
|
score = pred[4] |
|
category_name = model.model.names[int(pred[5])] |
|
category_id = pred[5] |
|
object_prediction = ObjectPrediction( |
|
bbox=bbox, |
|
category_id=int(category_id), |
|
score=score, |
|
category_name=category_name, |
|
) |
|
object_prediction_list.append(object_prediction) |
|
|
|
image = read_image(image) |
|
output_image = visualize_object_predictions(image=image, object_prediction_list=object_prediction_list) |
|
return output_image['image'] |
|
|
|
|
|
inputs = [ |
|
gr.inputs.Image(type="filepath", label="Input Image"), |
|
gr.inputs.Dropdown(["foduucom/table-detection-and-extraction"], |
|
default="foduucom/table-detection-and-extraction", label="Model"), |
|
gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"), |
|
gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"), |
|
gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"), |
|
] |
|
|
|
outputs = gr.outputs.Image(type="filepath", label="Output Image") |
|
title = "YoloTableExtract: Efficient Table Detection" |
|
|
|
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 [email protected]. If you find this space helpful, please show your appreciation by liking it. ❤️👍🏼" |
|
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]] |
|
demo_app = gr.Interface( |
|
fn=yolov8_inference, |
|
inputs=inputs, |
|
outputs=outputs, |
|
title=title, |
|
description=description, |
|
examples=examples, |
|
cache_examples=True, |
|
theme='huggingface', |
|
) |
|
demo_app.launch(debug=True, enable_queue=True) |
