yolov8 accident detector

app.py

@@ -1,44 +1,137 @@
 import gradio as gr
 # import torch
 from ultralyticsplus import YOLO, render_result
+import cv2
+from email.mime.text import MIMEText
+from email.mime.multipart import MIMEMultipart
+from email.mime.base import MIMEBase
+from email import encoders
+import smtplib
+
+sender_email = "[email protected]"
+receiver_email = "[email protected]"
+sender_password = "pahifrmjjjwahdnr"
+smtp_port = 587
+smtp_server = "smtp.gmail.com"
+subject = "Accident detected"
 
 classes: ['car', 'bike', 'person', 'car_car_accident', 'car_bike_accident', 'car_person_accident', 'bike_bike_accidnet', 'bike_person_accident', 'car_object_accident', 'bike_object_accident']
 
-def yolov8_func(image):
-                #image_size: gr.inputs.Slider = 640,
-                #conf_threshold: gr.inputs.Slider = 0.4,
-                #iou_threshold: gr.inputs.Slider = 0.50):
+def send_email(accident_type):
+    body = accident_type
+
+    msg = MIMEMultipart()
+    msg['From'] = sender_email
+    msg['To'] = receiver_email
+    msg['Subject'] = subject
+
+    msg.attach(MIMEText(body, 'plain'))
+
+    filename = "res.png"
+    folder = "./result/"
+    fullFileName = folder + filename
+
+    attachment = open(fullFileName, 'rb')
+
+    attachment_package = MIMEBase('application', 'octet-stream')
+    attachment_package.set_payload((attachment).read())
+    encoders.encode_base64(attachment_package)
+    attachment_package.add_header('Content-Disposition', "attachment; filename= " + filename)
+    msg.attach(attachment_package)
+
+    text = msg.as_string()
+
+
+    print("Connecting to server")
+    gmail_server = smtplib.SMTP(smtp_server, smtp_port)
+    gmail_server.starttls()
+    gmail_server.login(sender_email, sender_password)
+    print("Successfully connected to server")
+    print()
+
+    print("Sending email to ", receiver_email)
+    gmail_server.sendmail(sender_email, receiver_email, text)
+    print("Email sent to ", receiver_email)
+    print()
+
+    gmail_server.quit()
+
+
+def check_acc(box):
+    # return format(box.cls)
+    res_index_list = box.cls.tolist()
+    result = ""
+
+    for index in res_index_list:
+        if index ==3:
+            result = "car car accident detected"
+        elif index==4:
+            result = "car bike accident detected"
+        elif index==5:
+            result = "car person accident detected"
+        elif index==6:
+            result = "bike bike accident detected"
+        elif index==7:
+            result = "bike person accident detected"
+        elif index ==8:
+            result = "car object accident detected"
+        elif index == 9:
+            result = "bike object accident detected"
     
+    return result
+
+def image_predict(image):
     model_path = "best.pt"
     model = YOLO(model_path)
-
     results = model.predict(image,
                             conf = 0.4,
                             iou = 0.6,
                             imgsz = 640)
-    
     box = results[0].boxes
-
-    print("Object type: ", box.cls)
-    # print("Coordinates: ", box.xyxy)
-    # print("Probability: ", box.conf)
+    res = check_acc(box)
+    print("object Type: ", res)
 
     render = render_result(model=model, image=image, result=results[0])
+    # if res == "car car accident detected":
+    #     cv2.imwrite("./result/res.png", render)
+    # print(render)
+    if res == "car car accident detected":
+        render.save("./result/res.png")
+        send_email(res)
+
+    return (res, render)
+
+def extract_frames(video):
+    vidcap = cv2.VideoCapture(video)
+    while vidcap.isOpened():
+        success, image = vidcap.read()
+        if success ==False:
+            break
+        res, render = image_predict(image)
 
-    return (render, box.cls)
+        if res == "car car accident detected":
+            return (res, render)
+    return ("", None)
+
+def take_input(image, video):
+    if(video != None):
+        res = extract_frames(video)
+    else:
+        res = image_predict(image)
+    return res
 
 with gr.Blocks(title="YOLOS Object Detection - ClassCat", css=".gradio-container {background:lightyellow;}") as demo:
     gr.HTML('<h1>Yolo Object Detection</h1>')
     gr.HTML("<br>")
     with gr.Row():
-        input_image = gr.Image(label="Input image", type="pil")
-        output_image = gr.Image(label="Output image", type="pil")
+        input_image = gr.Image(label="Input image")
+        input_video = gr.Video(label="Input video")
         output_label = gr.Text(label="output label")
+        output_image = gr.Image(label="Output image")
     gr.HTML("<br>")
     send_btn = gr.Button("Detect")
     gr.HTML("<br>")
 
-    send_btn.click(fn=yolov8_func, inputs=[input_image], outputs=[output_image, output_label])
-
+    send_btn.click(fn=take_input, inputs=[input_image, input_video], outputs=[output_label, output_image])
 
 demo.launch(debug=True)

requirements.txt

@@ -1,3 +1,4 @@
 gradio==4.16.0
-torch==2.1.2
+opencv_python==4.7.0.72
+opencv_python_headless==4.8.0.74
 ultralyticsplus==0.0.29