Create app.py

app.py

@@ -0,0 +1,261 @@
+#import semua library yang dibutuhkan
+import os
+import io
+import cv2
+import numpy as np
+from PIL import Image, ImageDraw
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+import streamlit as st
+import torch
+import time
+import pandas as pd
+
+# Setting page layout
+st.set_page_config(
+    page_title="YoloS Helmet Detection",
+    page_icon="🤗",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+def input_image_setup(uploaded_file):
+    if uploaded_file is not None:
+        bytes_data = uploaded_file.getvalue()
+        image = Image.open(io.BytesIO(bytes_data))  # Convert bytes data to PIL image
+        return image
+    else:
+        raise FileNotFoundError("No file uploaded")
+
+# Function to convert OpenCV image to PIL image
+def cv2_to_pil(image):
+    return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+def draw_bounding_boxes(image, results, model, confidence):
+    draw = ImageDraw.Draw(image)
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        if score.item() >= confidence:
+            box = [int(i) for i in box.tolist()]
+            draw.rectangle(box, outline="purple", width=2)
+            label_text = f"{model.config.id2label[label.item()]} ({round(score.item(), 2)})"
+            draw.text((box[0], box[1]), label_text, fill="white")
+    return image
+
+def process_image(image, model, processor, confidence):
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, threshold=confidence, target_sizes=target_sizes)[0]
+    return results
+
+def detection_results_to_dict(results, model):
+    detection_dict = {
+        "objects": []
+    }
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = [round(i, 2) for i in box.tolist()]
+        detection_dict["objects"].append({
+            "label": model.config.id2label[label.item()],
+            "confidence": round(score.item(), 3),
+            "box": box
+        })
+    return detection_dict
+
+def convert_dict_to_csv(detection_dict_list):
+    combined_results = []
+    for detection_dict in detection_dict_list:
+        combined_results.extend(detection_dict["objects"])
+    df = pd.DataFrame(combined_results)
+    return df.to_csv(index=False).encode('utf-8')
+
+def clear_detection_results():
+    st.session_state.detection_dict_list = []
+
+# Initialize session state to store detection results
+if 'detection_dict_list' not in st.session_state:
+    st.session_state.detection_dict_list = []
+
+# Streamlit App Configuration
+st.header("Helmet Rider Detection")
+
+# Sidebar for Model Selection and Confidence Slider
+st.sidebar.header("ML Model Config")
+models = ["gghsgn/final200" ,"gghsgn/final100", "/gghsgn/final50"]
+
+model_name = st.sidebar.selectbox("Select model", models)
+confidence = st.sidebar.slider("Select Model Confidence", 25, 100, 40, step=5) / 100
+
+# Load Model and Processor
+processor = AutoImageProcessor.from_pretrained(model_name)
+model = AutoModelForObjectDetection.from_pretrained(model_name)
+
+# Option to select Real-Time or Upload Image
+mode = st.sidebar.selectbox("Select Input Mode", ["Upload Image", "Real-Time Webcam", "RTSP Video"])
+
+# Option if select Upload Image
+if mode == "Upload Image":
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image.", use_column_width=True)
+    else:
+        image = None
+    submit = st.button("Detect Objects")
+    if submit and image is not None:
+        try:
+            image_data = input_image_setup(uploaded_file)
+            st.subheader("The response is..")
+            
+            results = process_image(image, model, processor, confidence)
+            drawn_image = draw_bounding_boxes(image.copy(), results, model, confidence)
+            st.image(drawn_image, caption="Detected Objects", use_column_width=True)
+            
+            st.subheader("List of Objects:")
+            for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+                box = [round(i, 2) for i in box.tolist()]
+                st.write(
+                    f"Detected :orange[{model.config.id2label[label.item()]}] with confidence "
+                    f":green[{round(score.item(), 3)}] at location :violet[{box}]"
+                )
+
+            detected_objects = {model.config.id2label[label.item()]: 0 for label in results["labels"]}
+            for label in results["labels"]:
+                detected_objects[model.config.id2label[label.item()]] += 1
+            for obj, count in detected_objects.items():
+                st.write(f"Class :orange[{obj}] detected {count} time(s)")
+
+            detection_dict = detection_results_to_dict(results, model)
+            #st.write(detection_dict)
+
+            csv_data = convert_dict_to_csv([detection_dict])
+            st.download_button(
+                label="Download Results as CSV",
+                data=csv_data,
+                file_name="detection_results.csv",
+                mime="text/csv"
+            )
+        
+        except Exception as e:
+            st.error(f"Error: {e}")
+
+    elif submit and image is None:
+        st.error("Please upload an image before trying to detect objects.")
+
+# Option if select Realtime Webcam
+elif mode == "Real-Time Webcam":
+    run = st.checkbox("Run Webcam")
+    FRAME_WINDOW = st.image([])
+
+    if run:
+        cap = cv2.VideoCapture(0)
+        while run:
+            ret, frame = cap.read()
+            if not ret:
+                st.error("Failed to capture image from webcam")
+                break
+            
+            frame_pil = cv2_to_pil(frame)
+            
+            try:
+                results = process_image(frame_pil, model, processor, confidence)
+                drawn_image = draw_bounding_boxes(frame_pil.copy(), results, model, confidence)
+                FRAME_WINDOW.image(drawn_image, caption="Detected Objects", use_column_width=True)
+                
+                st.subheader("List of Objects:")
+                for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+                    box = [round(i, 2) for i in box.tolist()]
+                    st.write(
+                        f"Detected :orange[{model.config.id2label[label.item()]}] with confidence "
+                        f":green[{round(score.item(), 3)}] at location :violet[{box}]"
+                    )
+
+                detected_objects = {model.config.id2label[label.item()]: 0 for label in results["labels"]}
+                for label in results["labels"]:
+                    detected_objects[model.config.id2label[label.item()]] += 1
+                for obj, count in detected_objects.items():
+                    st.write(f"Class :orange[{obj}] detected {count} time(s)")
+
+                detection_dict = detection_results_to_dict(results, model)
+                st.session_state.detection_dict_list.append(detection_dict)
+                #st.write(detection_dict)
+                
+            
+            except Exception as e:
+                st.error(f"Error: {e}")
+
+            time.sleep(0.1)  # Delay for the next frame capture to create an illusion of real-time
+        cap.release()
+
+    if not run and st.session_state.detection_dict_list:
+        st.write("Detection stopped.")
+        csv_data = convert_dict_to_csv(st.session_state.detection_dict_list)
+        st.download_button(
+            label="Download All Results as CSV",
+            data=csv_data,
+            file_name="all_detection_results.csv",
+            mime="text/csv"
+        )
+        st.button("Clear Results", on_click=clear_detection_results)
+
+# Option if select Realtime RTSP Video
+elif mode == "RTSP Video":
+    rtsp_url = st.text_input("RTSP URL")
+    run = st.checkbox("Run RTSP Video")
+    FRAME_WINDOW = st.image([])
+
+    if rtsp_url and run:
+        cap = cv2.VideoCapture(rtsp_url)
+        st.subheader("List of Objects:")
+        while run:
+            ret, frame = cap.read()
+            if not ret:
+                st.error("Failed to capture image from RTSP stream")
+                break
+            
+            frame_pil = cv2_to_pil(frame)
+            
+            try:
+                results = process_image(frame_pil, model, processor, confidence)
+                drawn_image = draw_bounding_boxes(frame_pil.copy(), results, model, confidence)
+                FRAME_WINDOW.image(drawn_image, caption="Detected Objects", use_column_width=True)
+                st.subheader("List of Objects:")
+                for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+                    box = [round(i, 2) for i in box.tolist()]
+                    st.write(
+                        f"Detected :orange[{model.config.id2label[label.item()]}] with confidence "
+                        f":green[{round(score.item(), 3)}] at location :violet[{box}]"
+                    )
+
+                detected_objects = {model.config.id2label[label.item()]: 0 for label in results["labels"]}
+                for label in results["labels"]:
+                    detected_objects[model.config.id2label[label.item()]] += 1
+                for obj, count in detected_objects.items():
+                    st.write(f"Class :orange[{obj}] detected {count} time(s)")
+
+                detection_dict = detection_results_to_dict(results, model)
+                st.session_state.detection_dict_list.append(detection_dict)
+                #st.write(detection_dict)
+            
+            except Exception as e:
+                st.error(f"Error: {e}")
+
+            time.sleep(0.1)  # Delay for the next frame capture to create an illusion of real-time
+
+        cap.release()
+
+    if not run and st.session_state.detection_dict_list:
+        st.write("Detection stopped.")
+        csv_data = convert_dict_to_csv(st.session_state.detection_dict_list)
+        st.download_button(
+            label="Download All Results as CSV",
+            data=csv_data,
+            file_name="all_detection_results.csv",
+            mime="text/csv"
+        )
+        st.button("Clear Results", on_click=clear_detection_results)
+    elif run:
+        st.error("Please provide a valid RTSP URL before running the stream.")
+
+    # Ensure the video capture object is released if the checkbox is unchecked
+    if 'cap' in locals():
+        cap.release()