Update app.py

app.py

@@ -1,115 +1,66 @@
-
-import torch
-
-from transformers import AutoImageProcessor, AutoModelForObjectDetection
-#from transformers import pipeline
-
-from PIL import Image
-import matplotlib.pyplot as plt
-import matplotlib.patches as patches
-
-import io
-from random import choice
-
-
-image_processor_tiny = AutoImageProcessor.from_pretrained("foduucom/object_detection")
-model_tiny = AutoModelForObjectDetection.from_pretrained("foduucom/object_detection")
-
 import gradio as gr
-
-
-COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",
-            "#7f7fff", "#7fbfff", "#7fffff", "#7fffbf",
-            "#7fff7f", "#bfff7f", "#ffff7f", "#ffbf7f"]
-
-fdic = {
-    "family" : "DejaVu Serif",
-    "style" : "normal",
-    "size" : 18,
-    "color" : "yellow",
-    "weight" : "bold"
-}
-
-
-def get_figure(in_pil_img, in_results):
-    plt.figure(figsize=(16, 10))
-    plt.imshow(in_pil_img)
-    ax = plt.gca()
-
-    for score, label, box in zip(in_results["scores"], in_results["labels"], in_results["boxes"]):
-        selected_color = choice(COLORS)
-
-        box_int = [i.item() for i in torch.round(box).to(torch.int32)]
-        x, y, w, h = box_int[0], box_int[1], box_int[2]-box_int[0], box_int[3]-box_int[1]
-        #x, y, w, h = torch.round(box[0]).item(), torch.round(box[1]).item(), torch.round(box[2]-box[0]).item(), torch.round(box[3]-box[1]).item()
-
-        ax.add_patch(plt.Rectangle((x, y), w, h, fill=False, color=selected_color, linewidth=3, alpha=0.8))
-        ax.text(x, y, f"{model_tiny.config.id2label[label.item()]}: {round(score.item()*100, 2)}%", fontdict=fdic, alpha=0.8)
-
-    plt.axis("off")
-
-    return plt.gcf()
-
-
-def infer(in_pil_img, in_model="object_detection", in_threshold=0.9):
-    target_sizes = torch.tensor([in_pil_img.size[::-1]])
-    inputs = image_processor_tiny(images=in_pil_img, return_tensors="pt")
-    outputs = model_tiny(**inputs)
-
-        # convert outputs (bounding boxes and class logits) to COCO API
-    results = image_processor_tiny.post_process_object_detection(outputs, threshold=in_threshold, target_sizes=target_sizes)[0]
-
-    figure = get_figure(in_pil_img, results)
-
-    buf = io.BytesIO()
-    figure.savefig(buf, bbox_inches='tight')
-    buf.seek(0)
-    output_pil_img = Image.open(buf)
-
-    return output_pil_img
-
-
-with gr.Blocks(title="YOLOS Object Detection - ClassCat",
-                    css=".gradio-container {background:lightyellow;}"
-               ) as demo:
-    #sample_index = gr.State([])
-
-    gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:royalblue;">YOLOS Object Detection</div>""")
-
-    gr.HTML("""<h4 style="color:navy;">1. Select a model.</h4>""")
-
-    model = gr.Radio(["object_detection"], value="object_detection", label="Model name")
-
-    gr.HTML("""<br/>""")
-    gr.HTML("""<h4 style="color:navy;">2-a. Select an example by clicking a thumbnail below.</h4>""")
-    gr.HTML("""<h4 style="color:navy;">2-b. Or upload an image by clicking on the canvas.</h4>""")
-
-    with gr.Row():
-        input_image = gr.Image(label="Input image", type="pil")
-        output_image = gr.Image(label="Output image with predicted instances", type="pil")
-
-    gr.Examples(['samples/1.jpeg', 'samples/2.JPG'], inputs=input_image)
-
-    gr.HTML("""<br/>""")
-    gr.HTML("""<h4 style="color:navy;">3. Set a threshold value (default to 0.9)</h4>""")
-
-    threshold = gr.Slider(0, 1.0, value=0.9, label='threshold')
-
-    gr.HTML("""<br/>""")
-    gr.HTML("""<h4 style="color:navy;">4. Then, click "Infer" button to predict object instances.</h4>""")
-
-    send_btn = gr.Button("Infer")
-    send_btn.click(fn=infer, inputs=[input_image, model, threshold], outputs=[output_image])
-
-    gr.HTML("""<br/>""")
-    gr.HTML("""<h4 style="color:navy;">Reference</h4>""")
-    gr.HTML("""<ul>""")
-    gr.HTML("""<li><a href="https://huggingface.co/docs/transformers/model_doc/yolos" target="_blank">Hugging Face Transformers - YOLOv8n</a>""")
-    gr.HTML("""</ul>""")
-
-
-#demo.queue()
-demo.launch(debug=True)
-
-
-### EOF ###
+import torch
+from sahi.prediction import ObjectPrediction
+from sahi.utils.cv import visualize_object_predictions, read_image
+from ultralyticsplus import YOLO, render_result
+
+# Images
+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.overrides['conf'] = conf_threshold
+    model.overrides['iou']= iou_threshold
+    model.overrides['agnostic_nms'] = False  # NMS class-agnostic
+    model.overrides['max_det'] = 1000 
+    image = read_image(image)
+    results = model.predict(image)
+    render = render_result(model=model, image=image, result=results[0])
+    
+    return render
+        
+
+inputs = [
+    gr.inputs.Image(type="filepath", label="Input Image"),
+    gr.inputs.Dropdown(["foduucom/object_detection"], 
+                       default="foduucom/object_detection", 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 = "YOLOobjectdetection: Efficient ObjectDetection"
+
+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 = [['1.jpeg', "foduucom/object_detection", 640, 0.25, 0.45], ['2.JPG', "foduucom/object_detection", 640, 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)