Update app.py

app.py

@@ -49,10 +49,23 @@ def process_image(image):
     # Detect objects (e.g., helmets) in the frame
     items = detect_objects(frame)
 
+    # Check if helmet is detected (you can adapt this based on your model's labels)
+    helmet_detected = False
+    for score, label, box in items:
+        if model.config.id2label[label] == "helmet":  # Replace "helmet" with the actual class name in your model
+            helmet_detected = True
+
+    # If no helmet detected, show a traffic violation notification
+    if not helmet_detected:
+        violation_message = "Serious Traffic Violation: Rider not wearing a helmet!"
+    else:
+        violation_message = "Helmet detected: No violation."
+
+    # Save the image with detected items
     if items:  # If objects are detected, save the data
         save_data(frame, items)
 
-    return {"items_detected": items}
+    return {"items_detected": items, "violation_message": violation_message}
 
 def save_data(frame, items):
     """Save image and extract plate number."""
@@ -76,7 +89,7 @@ def save_to_database(image_filename, plate_number, items):
 # Define the Gradio interface using updated syntax
 interface = gr.Interface(fn=process_image,
                          inputs=gr.Image(type="pil"),
-                         outputs=gr.JSON(),
+                         outputs=[gr.JSON(), gr.Textbox()],
                          live=True)
 
 # Launch the Gradio app