Harsimran19/DefectDetection

1#!/usr/bin/env python2# -*- coding: utf-8 -*-3"""4-------------------------------------------------5   @File Name:     app.py6   @Author:        yash mohite7   @Date:          2023/5/158   @Description:9-------------------------------------------------10"""11from pathlib import Path12from PIL import Image13import streamlit as st14 15import config16from utils import load_model, infer_uploaded_image, infer_uploaded_video, infer_uploaded_webcam17 18# setting page layout19st.set_page_config(20    page_title="Defect Detection with YOLOv8",21    page_icon="🤖",22    layout="wide",23    initial_sidebar_state="expanded"24    )25 26# main page heading27st.title("Defect Detection with YOLOv8")28 29# sidebar30st.sidebar.header("DL Model Config")31 32# model options33task_type = st.sidebar.selectbox(34    "Select Task",35    ["Rust","Scratch","Fabric"]36)37 38model_type = None39model_type = st.sidebar.selectbox(40    "Select Model",41    config.DETECTION_MODEL_LIST42)43 44confidence = float(st.sidebar.slider(45    "Select Model Confidence", 30, 100, 50)) / 10046 47model_path = ""48if model_type:49    model_path = Path(config.DETECTION_MODEL_DIR, str(model_type))50else:51    st.error("Please Select Model in Sidebar")52 53# load pretrained DL model54try:55    model = load_model(model_path)56except Exception as e:57    st.error(f"Unable to load model. Please check the specified path: {model_path}")58 59# image/video options60st.sidebar.header("Image/Video Config")61source_selectbox = st.sidebar.selectbox(62    "Select Source",63    config.SOURCES_LIST64)65 66source_img = None67if source_selectbox == config.SOURCES_LIST[0]: # Image68    infer_uploaded_image(confidence, model)69elif source_selectbox == config.SOURCES_LIST[1]: # Video70    infer_uploaded_video(confidence, model)71elif source_selectbox == config.SOURCES_LIST[2]: # Webcam72    infer_uploaded_webcam(confidence, model)73else:74    st.error("Currently only 'Image' and 'Video' source are implemented")