Update app.py

app.py

@@ -1,70 +1,86 @@
 import gradio as gr
 import torch
-from ultralyticsplus import YOLO, render_result
-
+from sahi.prediction import ObjectPrediction
+from sahi.utils.cv import visualize_object_predictions, read_image
+from ultralyticsplus import YOLO
 
 # Images
-torch.hub.download_url_to_file('https://external-content.duckduckgo.com/iu/?u=https%3A%2F%2Ftexashafts.com%2Fwp-content%2Fuploads%2F2016%2F04%2Fconstruction-worker.jpg', 'one.jpg')
-torch.hub.download_url_to_file(
-    'https://www.pearsonkoutcherlaw.com/wp-content/uploads/2020/06/Construction-Workers.jpg', 'two.jpg')
-torch.hub.download_url_to_file(
-    'https://nssgroup.com/wp-content/uploads/2019/02/Building-maintenance-blog.jpg', 'three.jpg')
-
+torch.hub.download_url_to_file('https://raw.githubusercontent.com/kadirnar/dethub/main/data/images/highway.jpg', 'highway.jpg')
+torch.hub.download_url_to_file('https://user-images.githubusercontent.com/34196005/142742872-1fefcc4d-d7e6-4c43-bbb7-6b5982f7e4ba.jpg', 'highway1.jpg')
+torch.hub.download_url_to_file('https://raw.githubusercontent.com/obss/sahi/main/tests/data/small-vehicles1.jpeg', 'small-vehicles1.jpeg')
 
-def yoloV8_func(image: gr.inputs.Image = None,
-                image_size: gr.inputs.Slider = 640,
-                conf_threshold: gr.inputs.Slider = 0.4,
-                iou_threshold: gr.inputs.Slider = 0.50):
-    """_summary_
+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 (gr.inputs.Image, optional): _description_. Defaults to None.
-        image_size (gr.inputs.Slider, optional): _description_. Defaults to 640.
-        conf_threshold (gr.inputs.Slider, optional): _description_. Defaults to 0.4.
-        iou_threshold (gr.inputs.Slider, optional): _description_. Defaults to 0.50.
+        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_path = "best.pt"
-    model = YOLO("foduucom/table-detection-and-extraction")
-
-    results = model.predict(image,
-                            conf=conf_threshold,
-                            iou=iou_threshold,
-                            imgsz=image_size)
-
-    # observe results
-    box = results[0].boxes
-    print("Object type:", box.cls)
-    print("Coordinates:", box.xyxy)
-    print("Probability:", box.conf)
-    render = render_result(model=model, image=image, result=results[0])
-    return render
+    model = YOLO(model_path)
+    model.conf = conf_threshold
+    model.iou = iou_threshold
+    results = model.predict(image, imgsz=image_size, return_outputs=True)
+    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.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"),
+    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 = "YOLOv8 101: Custome Object Detection on Construction Workers "
-
-
-examples = [['one.jpg', 640, 0.5, 0.7],
-            ['two.jpg', 800, 0.5, 0.6],
-            ['three.jpg',  900, 0.5, 0.8]]
+title = "YoloTableExtract: Efficient Table Detection"
 
-yolo_app = gr.Interface(
-    fn=yoloV8_func,
+description = "🔍 YoloTableExtract is a powerful space that utilizes YOLOv8 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 = [['highway.jpg', 'foduucom/table-detection-and-extraction', 640, 0.25, 0.45], ['highway1.jpg', 'foduucom/table-detection-and-extraction', 640, 0.25, 0.45], ['small-vehicles1.jpeg', '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',
+    theme='huggingface',
 )
-yolo_app.launch(debug=True, enable_queue=True)
+demo_app.launch(debug=True, enable_queue=True)