Create main.py

main.py

@@ -0,0 +1,57 @@
+import streamlit as st
+import torch
+from PIL import Image
+from torchvision.transforms import functional as F
+from yolov5.utils.general import non_max_suppression, scale_coords
+from yolov5.models.experimental import attempt_load
+from yolov5.utils.plots import plot_one_box
+import cv2
+
+@st.cache(allow_output_mutation=True)
+def load_model():
+    # Load your pre-trained YOLOv5 model
+    model = attempt_load('best.pt', map_location=torch.device('cpu'))
+    return model
+
+def detect_objects(image, model, confidence=0.4, iou=0.5):
+    img = Image.fromarray(image.astype('uint8')).convert('RGB')
+    img_tensor = F.to_tensor(img)
+    img_tensor, _ = model.preprocess(img_tensor, None, None)
+    pred = model(img_tensor)[0]
+    pred = non_max_suppression(pred, confidence, iou)[0]
+
+    if pred is not None and len(pred):
+        pred[:, :4] = scale_coords(img_tensor.shape[2:], pred[:, :4], img.size).round()
+
+    return pred
+
+def main():
+    st.title("Real-time Object Detection with YOLOv5")
+
+    # Choose between image upload or video stream
+    option = st.radio("Choose Input Type:", ("Image Upload", "Video Stream"))
+
+    if option == "Image Upload":
+        uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+        if uploaded_image is not None:
+            image = Image.open(uploaded_image)
+            st.image(image, caption="Uploaded Image", use_column_width=True)
+            st.write("")
+
+            if st.button("Detect Objects"):
+                st.write("Detecting...")
+                model = load_model()
+                with st.spinner('Wait for it...'):
+                    pred = detect_objects(image, model)
+                    if pred is not None and len(pred):
+                        for *xyxy, conf, cls in pred:
+                            label = f'{model.names[int(cls)]} {conf:.2f}'
+                            plot_one_box(xyxy, image, label=label, color=(255, 0, 0), line_thickness=2)
+                        st.image(image, caption="Result", use_column_width=True)
+
+    elif option == "Video Stream":
+        st.write("Video stream functionality is not implemented yet.")
+
+if __name__ == '__main__':
+    main()