demo.py

import contextlib
import time
from functools import wraps
from io import StringIO
from zipfile import ZipFile

import streamlit as st
from PIL import Image

import evaluator
from yolo_dataset import YoloDataset
from yolo_fire import detect


def main():
    # Header & Page Config.
    st.set_page_config(
        page_title="Fire&Smoke Detection",
        layout="centered")
    st.title("Fire&Smoke Detection:")

    detect_tab, evaluate_tab = st.tabs(["Detect", "Evaluate"])

    with evaluate_tab:
        buffer = st.file_uploader("Upload your Yolo Dataset here", type=["zip"])

        if buffer:
            with st.spinner('Wait for it...'):
                # Slider for changing confidence
                confidence = st.slider('Confidence Threshold', 0, 100, 30)
                yolo_dataset = YoloDataset.from_zip_file(ZipFile(buffer))
                capture_output(evaluator.evaluate)(coco_gt=yolo_dataset.to_coco(), loader=yolo_dataset.load_image,
                                                   confidence_threshold=confidence / 100.0)

    with detect_tab:
        # This will let you upload PNG, JPG & JPEG File
        buffer = st.file_uploader("Upload your Image here", type=["jpg", "png", "jpeg"])

        if buffer:
            # Object Detecting
            with st.spinner('Wait for it...'):
                # Slider for changing confidence
                confidence = st.slider('Confidence Threshold', 0, 100, 30)

                # Calculating time for detection
                t1 = time.time()
                im = Image.open(buffer)
                # im.save("saved_images/image.jpg")
                res_img = detect(im, confidence / 100.0)
                t2 = time.time()

            # Displaying the image
            st.image(res_img, use_column_width=True)

            # Printing Time
            st.write("\n")
            st.write("Time taken: ", t2 - t1, "sec.")


def capture_output(func):
    """Capture output from running a function and write using streamlit."""

    @wraps(func)
    def wrapper(*args, **kwargs):
        # Redirect output to string buffers
        stdout, stderr = StringIO(), StringIO()
        try:
            with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
                return func(*args, **kwargs)
        except Exception as err:
            st.write(f"Failure while executing: {err}")
        finally:
            if _stdout := stdout.getvalue():
                st.write("Execution stdout:")
                st.code(_stdout)
            if _stderr := stderr.getvalue():
                st.write("Execution stderr:")
                st.code(_stderr)

    return wrapper


if __name__ == '__main__':
    main()