Spaces:
Runtime error
Runtime error
File size: 3,877 Bytes
2fdba78 06a2085 f977726 2fdba78 684e6f5 2fdba78 f977726 2fdba78 f977726 2fdba78 095f7cc 684e6f5 095f7cc 684e6f5 06a2085 f977726 095f7cc f977726 095f7cc 0821896 095f7cc 06a2085 095f7cc f977726 2fdba78 095f7cc 2fdba78 095f7cc 1d08b5b 095f7cc 1d08b5b c973a85 06a2085 c973a85 095f7cc 684e6f5 095f7cc 684e6f5 095f7cc f977726 2fdba78 f977726 2fdba78 f977726 2fdba78 f977726 ec500f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import contextlib
import os
import time
from functools import wraps
from io import StringIO
from zipfile import ZipFile
from tempfile import mktemp
import streamlit as st
from PIL import Image
import evaluator
from yolo_dataset import YoloDataset
from yolo_model import YoloModel
from models.yolo_crack import YoloModel as CrackModel
fire_and_smoke = YoloModel("SHOU-ISD/fire-and-smoke", "yolov8n.pt")
crack = CrackModel("SHOU-ISD/yolo-cracks", "last4.pt", "SHOU-ISD/yolo-cracks", "best.pt")
coco = YoloModel("ultralyticsplus/yolov8s", "yolov8s.pt")
def main():
# Header & Page Config.
st.set_page_config(
page_title=f"Detection",
layout="centered")
model = None
with st.sidebar:
model_choice = st.radio("Select Model", ["Fire&Smoke", "Crack"])
if model_choice == "Fire&Smoke":
model = fire_and_smoke
elif model_choice == "Crack":
model = crack
elif model_choice == "Coco":
model = coco
st.title(f"{model_choice} Detection:")
detect_tab, evaluate_tab = st.tabs(["Detect", "Evaluate"])
with evaluate_tab:
evaluate(model)
with detect_tab:
detect(model)
def evaluate(model: YoloModel):
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.coco_evaluate)(model=model,
# dataset=yolo_dataset,
# confidence_threshold=confidence / 100.0)
with evaluator.yolo_evaluator(model, yolo_dataset) as metrics:
st.subheader("Metrics:")
st.write("Speed: ")
st.json(metrics.speed)
st.write("Results: ")
st.json(metrics.results_dict)
for pic in os.listdir(metrics.save_dir):
st.write(pic)
st.image(os.path.join(metrics.save_dir, pic), use_column_width=True)
def detect(model: YoloModel):
# 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()
filename = mktemp(suffix=buffer.name)
Image.open(buffer).save(filename)
res_img = model.preview_detect(filename, 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()
|