Rework the code, add json output

app.py

@@ -1,182 +1,71 @@
+import json
 
 import gradio as gr
 import cv2
-import requests
-import os
 import numpy as np
-from sahi.utils.yolov5 import (
-    download_yolov5s6_model,
-)
-
+from sahi.utils.yolov5 import download_yolov5s6_model
 
 # import required functions, classes
 from sahi import AutoDetectionModel
-from sahi.utils.cv import read_image
-from sahi.utils.file import download_from_url
-from sahi.predict import get_prediction, get_sliced_prediction, predict, visualize_object_predictions
-from IPython.display import Image
-
-from ultralytics import YOLO
-import gradio as gr
-import cv2
-import requests
-import os
-
-from ultralytics import YOLO
+from sahi.predict import get_sliced_prediction, visualize_object_predictions
 
 
-yolov5_model_path = 'best.pt'
+yolov5_model_path = "best.pt"
 download_yolov5s6_model(destination_path=yolov5_model_path)
 detection_model = AutoDetectionModel.from_pretrained(
-    model_type='yolov5',
+    model_type="yolov5",
     model_path=yolov5_model_path,
     confidence_threshold=0.01,
-    device="cpu", # or 'cuda:0'
+    device="cpu",  # or 'cuda:0'
 )
 
-#model = YOLO('/home/ubuntu/Receptacle_Detection_Demo/best.pt')
-
-
-
-demo = gr.Blocks()
-
 EXAMPLES = [
-    [ "test1.jpg"],
+    ["test1.jpg"],
     ["test2.jpg"],
     ["test3.jpg"],
     ["test4.jpg"],
 ]
 
 
-def with_labels(image_path):
+def do_detection(image_path, hide_labels=False):
     result = get_sliced_prediction(
-    image_path,
-    detection_model,
-    slice_height = 512,
-    slice_width = 512,
-    overlap_height_ratio = 0.12,
-    overlap_width_ratio = 0.12)
-    #result.export_visuals(export_dir="/home/ubuntu/Receptacle_Detection_Demo/")
-    #image = cv2.imread("/home/ubuntu/Receptacle_Detection_Demo/prediction_visual.png")
-    #img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
+        image_path,
+        detection_model,
+        slice_height=512,
+        slice_width=512,
+        overlap_height_ratio=0.12,
+        overlap_width_ratio=0.12,
+    )
 
     count = -1
-    new_list=[]
+    new_list = []
     for i in result.object_prediction_list:
         count += 1
-        print(i)
         score = i.score
         value = score.value
         category = i.category
         category_name = category.name
         if value > confidence_scores[category_name]:
-            print(value)
-            print(confidence_scores[category_name])
             new_list.append(result.object_prediction_list[count])
-            
-    img_converted = cv2.cvtColor(image_path, cv2.COLOR_BGR2RGB)
-    numpydata = np.asarray(img_converted)
-    visualize_object_predictions(
-    numpydata, 
-    object_prediction_list = new_list,
-    text_size=1,
-    text_th=1,
-    hide_labels = 0, 
-    rect_th=3,
-    output_dir='/home/ubuntu/Receptacle_Detection_Demo/',
-    file_name = 'result',
-    export_format = 'png')
-    image2 = cv2.imread("/home/ubuntu/Receptacle_Detection_Demo/result.png")
-    img_rgb = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB)
-
- 
-    class_counts = {}
-
-    
-    predictions = new_list
-    for i in predictions:
-        category = i.category
-        category_name = category.name
-        if category_name not in class_counts:
-            class_counts[category_name] = 1
-        else:
-            class_counts[category_name] += 1
-
- 
-    legend_text = 'Symbols Counted:'
-    for class_name, count in class_counts.items():
-        legend_text += f' {class_name}: {count} |'
-    
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    font_scale = 1
-    font_color = (255, 255, 255)  
-    font_thickness = 2
-    legend_bg_color = (131, 79, 0) 
-    legend_padding = 10
 
-    legend_size, _ = cv2.getTextSize(legend_text, font, font_scale, font_thickness)
-    legend_bg_height = legend_size[1] + 2 * legend_padding
-    legend_bg_width = legend_size[0] + 2 * legend_padding
-
-    legend_bg = np.zeros((legend_bg_height, legend_bg_width, 3), dtype=np.uint8)
-    legend_bg[:] = legend_bg_color
-    cv2.putText(legend_bg, legend_text, (legend_padding, legend_padding + legend_size[1]), font,
-                font_scale, font_color, font_thickness)
-    
-    img_height, img_width, _ = img_rgb.shape
-    legend_x = img_width - legend_bg_width
-    legend_y = img_height - legend_bg_height
-
-    img_rgb[legend_y:, legend_x:, :] = legend_bg
-
-    result_image_path = '/home/ubuntu/Receptacle_Detection_Demo/result_with_legend.png'
-    cv2.imwrite(result_image_path, img_rgb)
-
-    return cv2.cvtColor(cv2.imread(result_image_path), cv2.COLOR_BGR2RGB)
-
-def without_labels(image_path):
-    result = get_sliced_prediction(
-    image_path,
-    detection_model,
-    slice_height = 512,
-    slice_width = 512,
-    overlap_height_ratio = 0.12,
-    overlap_width_ratio = 0.12)
-    #result.export_visuals(export_dir="/home/ubuntu/Receptacle_Detection_Demo/")
-    #image = cv2.imread("/home/ubuntu/Receptacle_Detection_Demo/prediction_visual.png")
-    #img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
-
-    count = -1
-    new_list=[]
-    for i in result.object_prediction_list:
-        count += 1
-        print(i)
-        score = i.score
-        value = score.value
-        category = i.category
-        category_name = category.name
-        if value > confidence_scores[category_name]:
-            print(value)
-            print(confidence_scores[category_name])
-            new_list.append(result.object_prediction_list[count])
-            
     img_converted = cv2.cvtColor(image_path, cv2.COLOR_BGR2RGB)
     numpydata = np.asarray(img_converted)
     visualize_object_predictions(
-    numpydata, 
-    object_prediction_list = new_list,
-    hide_labels = 1, 
-    rect_th=3,
-    output_dir='/home/ubuntu/Receptacle_Detection_Demo/',
-    file_name = 'result',
-    export_format = 'png')
+        numpydata,
+        object_prediction_list=new_list,
+        text_size=1,
+        text_th=1,
+        hide_labels=hide_labels,
+        rect_th=3,
+        output_dir="/home/ubuntu/Receptacle_Detection_Demo/",
+        file_name="result",
+        export_format="png",
+    )
     image2 = cv2.imread("/home/ubuntu/Receptacle_Detection_Demo/result.png")
     img_rgb = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB)
 
- 
     class_counts = {}
 
-    
     predictions = new_list
     for i in predictions:
         category = i.category
@@ -186,16 +75,18 @@ def without_labels(image_path):
         else:
             class_counts[category_name] += 1
 
- 
-    legend_text = 'Symbols Counted:'
+    legend_text = "Symbols Counted:"
     for class_name, count in class_counts.items():
-        legend_text += f' {class_name}: {count} |'
-    
+        legend_text += f" {class_name}: {count} |"
+
     font = cv2.FONT_HERSHEY_SIMPLEX
-    font_scale = 1.5
-    font_color = (255, 255, 255)  
+    if hide_labels:
+        font_scale = 1.5
+    else:
+        font_scale = 1
+    font_color = (255, 255, 255)
     font_thickness = 2
-    legend_bg_color = (131, 79, 0) 
+    legend_bg_color = (131, 79, 0)
     legend_padding = 10
 
     legend_size, _ = cv2.getTextSize(legend_text, font, font_scale, font_thickness)
@@ -204,52 +95,62 @@ def without_labels(image_path):
 
     legend_bg = np.zeros((legend_bg_height, legend_bg_width, 3), dtype=np.uint8)
     legend_bg[:] = legend_bg_color
-    cv2.putText(legend_bg, legend_text, (legend_padding, legend_padding + legend_size[1]), font,
-                font_scale, font_color, font_thickness)
-    
+    cv2.putText(
+        legend_bg,
+        legend_text,
+        (legend_padding, legend_padding + legend_size[1]),
+        font,
+        font_scale,
+        font_color,
+        font_thickness,
+    )
+
     img_height, img_width, _ = img_rgb.shape
     legend_x = img_width - legend_bg_width
     legend_y = img_height - legend_bg_height
 
     img_rgb[legend_y:, legend_x:, :] = legend_bg
 
-    result_image_path = '/home/ubuntu/Receptacle_Detection_Demo/result_with_legend.png'
+    result_image_path = "/home/ubuntu/Receptacle_Detection_Demo/result_with_legend.png"
     cv2.imwrite(result_image_path, img_rgb)
 
-    return cv2.cvtColor(cv2.imread(result_image_path), cv2.COLOR_BGR2RGB)
+    return (
+        cv2.cvtColor(cv2.imread(result_image_path), cv2.COLOR_BGR2RGB),
+        result.to_coco_predictions(),
+    )
 
-def choose_function(choice, input_text):
-    if choice == "With Labels":
-        return with_labels(input_text)
-    else:
-        return without_labels(input_text)
-print("Starting the demo...")
 
+def update_duplex(val):
+    confidence_scores["Duplex - Standard"] = val
+    return "updated!"
 
 
-def update_duplex(val):
-    confidence_scores['Duplex - Standard'] = val
-    return 'updated!'
 def update_single(val):
-    confidence_scores['Singleplex - Standard'] = val
-    return 'updated!'
+    confidence_scores["Singleplex - Standard"] = val
+    return "updated!"
+
 
 def update_triplex(val):
-    confidence_scores['Triplex - Standard'] = val
-    return 'updated!'
+    confidence_scores["Triplex - Standard"] = val
+    return "updated!"
+
 
 def update_quadruplex(val):
-    confidence_scores['Quadruplex - Standard'] = val
-    return 'updated!'
+    confidence_scores["Quadruplex - Standard"] = val
+    return "updated!"
+
 
 def update_gfci(val):
-    confidence_scores['Duplex - GFCI'] = val
-    return 'updated!'
+    confidence_scores["Duplex - GFCI"] = val
+    return "updated!"
+
 
 def update_gfciwp(val):
-    confidence_scores['Duplex - Weatherproof-GFCI'] = val
-    return 'updated!'
+    confidence_scores["Duplex - Weatherproof-GFCI"] = val
+    return "updated!"
 
+
+demo = gr.Blocks()
 theme = gr.themes.Soft()
 
 with gr.Blocks(theme=theme) as demo:
@@ -259,80 +160,87 @@ with gr.Blocks(theme=theme) as demo:
         """
     )
 
-    gr.Markdown("### Step 1: Upload an image")
-    
     with gr.Row():
         input_image = gr.Image(
-            label="Upload an image here.", source="upload", interactive=True,
+            label="Upload an image here.",
+            source="upload",
+            interactive=True,
         )
         examples = gr.Examples(
-            examples=EXAMPLES, inputs=[input_image], examples_per_page=4, label="Examples to use.",
+            examples=EXAMPLES,
+            inputs=[input_image],
+            examples_per_page=4,
+            label="Examples to use.",
+        )
+
+    hide_labels = gr.Checkbox(label="Hide labels")
+    with gr.Accordion("Visualization Confidence Thresholds", open=False):
+        filter_name1 = gr.Slider(
+            minimum=0.1,
+            maximum=1,
+            value=0.53,
+            interactive=True,
+            label="Singleplex",
+        )
+        filter_name2 = gr.Slider(
+            minimum=0.1,
+            maximum=1,
+            value=0.66,
+            interactive=True,
+            label="Duplex",
+        )
+        filter_name3 = gr.Slider(
+            minimum=0.1,
+            maximum=1,
+            value=0.65,
+            interactive=True,
+            label="Triplex",
+        )
+        filter_name4 = gr.Slider(
+            minimum=0.1,
+            maximum=1,
+            value=0.63,
+            interactive=True,
+            label="Quadruplex",
+        )
+        filter_name5 = gr.Slider(
+            minimum=0.1,
+            maximum=1,
+            value=0.31,
+            interactive=True,
+            label="GFCI",
+        )
+        filter_name6 = gr.Slider(
+            minimum=0.1,
+            maximum=1,
+            value=0.33,
+            interactive=True,
+            label="GFCI/WP",
         )
 
-    
-    gr.Markdown("### Step 2: Choose either \n With labels: See receptacles detected with type detected/confidence score included\n Without labels: See only bounding boxes")
-    filter_name = gr.Dropdown(
-        choices=["With Labels", "Without Labels"], label="With/Without Labels", interactive=True
-    )
-    gr.Markdown("### Step 3: Choose confidence score levels for each symbol detected (default are optimal scores)")
-    filter_name1 = gr.Slider(
-        minimum = .1,
-        maximum = 1,
-        value = .53,
-        interactive = True,
-        label = 'Singleplex',
-    )
-    filter_name2 = gr.Slider(
-        minimum = .1,
-        maximum = 1,
-        value = .66,
-        interactive = True,
-        label = 'Duplex',
-    )
-    filter_name3 = gr.Slider(
-        minimum = .1,
-        maximum = 1,
-        value = .65,
-        interactive = True,
-        label = 'Triplex',
-    )
-    filter_name4 = gr.Slider(
-        minimum = .1,
-        maximum = 1,
-        value = .63,
-        interactive = True,
-        label = 'Quadruplex',
-    )
-    filter_name5 = gr.Slider(
-        minimum = .1,
-        maximum = 1,
-        value = .31,
-        interactive = True,
-        label = 'GFCI',
-    )
-    filter_name6 = gr.Slider(
-        minimum = .1,
-        maximum = 1,
-        value = .33,
-        interactive = True,
-        label = 'GFCI/WP',
-    )
-    
-    filter_name2.change(fn=update_duplex, inputs=filter_name2)
     filter_name1.change(fn=update_single, inputs=filter_name1)
+    filter_name2.change(fn=update_duplex, inputs=filter_name2)
     filter_name3.change(fn=update_triplex, inputs=filter_name3)
     filter_name4.change(fn=update_quadruplex, inputs=filter_name4)
     filter_name5.change(fn=update_gfci, inputs=filter_name5)
     filter_name6.change(fn=update_gfciwp, inputs=filter_name6)
-    confidence_scores = {'Triplex - Standard': filter_name3.value,'Duplex - Standard': filter_name2.value,'Singleplex - Standard': filter_name1.value,'Duplex - GFCI': filter_name5.value,'Duplex - Weatherproof-GFCI':filter_name6.value,'Quadruplex - Standard': filter_name4.value}
- 
-    gr.Markdown("### Step 4: See results with number of symbols counted in the bottom right corner")
-    results_button = gr.Button("See Results")
+    confidence_scores = {
+        "Triplex - Standard": filter_name3.value,
+        "Duplex - Standard": filter_name2.value,
+        "Singleplex - Standard": filter_name1.value,
+        "Duplex - GFCI": filter_name5.value,
+        "Duplex - Weatherproof-GFCI": filter_name6.value,
+        "Quadruplex - Standard": filter_name4.value,
+    }
+
+    results_button = gr.Button("Submit")
     results_button.click(
-        choose_function,
-        inputs = [filter_name,input_image],
-        outputs = [gr.components.Image(type="numpy", label="Output Image")]
-        
+        do_detection,
+        inputs=[input_image, hide_labels],
+        outputs=[
+            gr.Image(type="numpy", label="Output Image"),
+            gr.Json(),
+        ],
     )
-   
+
 demo.launch()