SHOU-ISD/object-detection

1import contextlib2import os3import time4from functools import wraps5from io import StringIO6from zipfile import ZipFile7from tempfile import mktemp8 9import streamlit as st10from PIL import Image11 12import evaluator13from yolo_dataset import YoloDataset14from yolo_model import YoloModel15from models.yolo_crack import YoloModel as CrackModel16 17fire_and_smoke = YoloModel("SHOU-ISD/fire-and-smoke", "yolov8n.pt")18crack = CrackModel("SHOU-ISD/yolo-cracks", "last4.pt", "SHOU-ISD/yolo-cracks", "best.pt")19coco = YoloModel("ultralyticsplus/yolov8s", "yolov8s.pt")20 21 22def main():23    # Header & Page Config.24    st.set_page_config(25        page_title=f"Detection",26        layout="centered")27 28    model = None29    with st.sidebar:30        model_choice = st.radio("Select Model", ["Fire&Smoke", "Crack"])31        if model_choice == "Fire&Smoke":32            model = fire_and_smoke33        elif model_choice == "Crack":34            model = crack35        elif model_choice == "Coco":36            model = coco37 38    st.title(f"{model_choice} Detection:")39 40    detect_tab, evaluate_tab = st.tabs(["Detect", "Evaluate"])41 42    with evaluate_tab:43        evaluate(model)44    with detect_tab:45        detect(model)46 47 48def evaluate(model: YoloModel):49    buffer = st.file_uploader("Upload your Yolo Dataset here", type=["zip"])50 51    if buffer:52        with st.spinner('Wait for it...'):53            # Slider for changing confidence54            # confidence = st.slider('Confidence Threshold', 0, 100, 30)55            yolo_dataset = YoloDataset.from_zip_file(ZipFile(buffer))56            # capture_output(evaluator.coco_evaluate)(model=model,57            #                                         dataset=yolo_dataset,58            #                                         confidence_threshold=confidence / 100.0)59            with evaluator.yolo_evaluator(model, yolo_dataset) as metrics:60                st.subheader("Metrics:")61                st.write("Speed: ")62                st.json(metrics.speed)63                st.write("Results: ")64                st.json(metrics.results_dict)65                for pic in os.listdir(metrics.save_dir):66                    st.write(pic)67                    st.image(os.path.join(metrics.save_dir, pic), use_column_width=True)68 69 70def detect(model: YoloModel):71    # This will let you upload PNG, JPG & JPEG File72    buffer = st.file_uploader("Upload your Image here", type=["jpg", "png", "jpeg"])73 74    if buffer:75        # Object Detecting76        with (st.spinner('Wait for it...')):77            # Slider for changing confidence78            confidence = st.slider('Confidence Threshold', 0, 100, 30)79 80            # Calculating time for detection81            t1 = time.time()82            filename = mktemp(suffix=buffer.name)83            Image.open(buffer).save(filename)84            res_img = model.preview_detect(filename, confidence / 100.0)85            t2 = time.time()86 87        # Displaying the image88        st.image(res_img, use_column_width=True)89 90        # Printing Time91        st.write("\n")92        st.write("Time taken: ", t2 - t1, "sec.")93 94 95def capture_output(func):96    """Capture output from running a function and write using streamlit."""97 98    @wraps(func)99    def wrapper(*args, **kwargs):100        # Redirect output to string buffers101        stdout, stderr = StringIO(), StringIO()102        try:103            with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):104                return func(*args, **kwargs)105        except Exception as err:106            st.write(f"Failure while executing: {err}")107        finally:108            if _stdout := stdout.getvalue():109                st.write("Execution stdout:")110                st.code(_stdout)111            if _stderr := stderr.getvalue():112                st.write("Execution stderr:")113                st.code(_stderr)114 115    return wrapper116 117 118if __name__ == '__main__':119    main()120