Plastic-Objects-Detection-System-DIP/app.py at main · zeeshan020dev/Plastic-Objects-Detection-System-DIP · GitHub

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
import streamlit as st
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image

# ----------------------------------
# PAGE CONFIG
# ----------------------------------
st.set_page_config(page_title="Plastic Detection", layout="wide")

# ----------------------------------
# CUSTOM CSS
# ----------------------------------
st.markdown("""
    <style>
        .main { background-color: #111827; color: white; }
        h1, h2, h3, h4 { color: white; }
        .stButton button {
            border-radius: 10px;
            padding: 0.6rem 1.2rem;
            font-size: 1rem;
        }
    </style>
""", unsafe_allow_html=True)

# ----------------------------------
# LOAD YOLO MODEL
# ----------------------------------
model = YOLO("best.pt")

# ----------------------------------
# APP TITLE
# ----------------------------------
st.title("♻️ Plastic Object Detection (Image + Real-Time Webcam)")


# ----------------------------------
# SIDEBAR SETTINGS
# ----------------------------------
st.sidebar.header("⚙️ Settings")
mode = st.sidebar.radio("Select Mode:", ["Image Upload", "Real-Time Webcam"])
conf = st.sidebar.slider("Confidence Threshold", 0.1, 1.0, 0.5)


# ----------------------------------
# IMAGE UPLOAD MODE
# ----------------------------------
if mode == "Image Upload":
    uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])

    if uploaded_file:
        img = Image.open(uploaded_file)
        st.image(img, caption="Original Image", use_column_width=True)

        with st.spinner("Detecting objects..."):
            results = model.predict(img, conf=conf)
            result_img = results[0].plot()

        st.image(result_img, caption="Detection Result", use_column_width=True)


# ----------------------------------
# REAL-TIME WEBCAM MODE
# ----------------------------------
elif mode == "Real-Time Webcam":

    st.subheader("📸 Live Camera Detection")
    start_cam = st.checkbox("Start Live Camera")

    FRAME_WINDOW = st.image([])

    if start_cam:
        cap = cv2.VideoCapture(0)

        while start_cam:
            ret, frame = cap.read()
            if not ret:
                st.warning("⚠️ Unable to access camera")
                break

            # YOLO detection
            results = model(frame, conf=conf)
            frame = results[0].plot()

            # Convert BGR → RGB
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            # Display frame
            FRAME_WINDOW.image(frame)

            # Check checkbox state (to stop loop)
            start_cam = st.session_state.get("Start Live Camera", True)

        cap.release()
        st.info("Camera stopped.")