基于TensorFlow与CV2的自定义目标检测触发动作实现问询

Hey there! Nice work building your custom object detection pipeline with TensorFlow and OpenCV. Adding that trigger logic—where detecting a specific object triggers an action like sending an Android app notification via Firebase—is totally doable, and I’ll walk you through exactly how to pull it off.

Step 1: Add Target Detection Trigger Logic #

First, we need to modify your existing detection loop to check when your target object is detected (with sufficient confidence).

In your code, after running the detection sess.run() call, you’ll have access to:

• classes: The class IDs of detected objects
• scores: Confidence scores for each detection
• category_index: Maps class IDs to human-readable labels

We’ll loop through these results, filter for your target object (e.g., "baseball bat") and a minimum confidence threshold (like 0.7), then trigger the action when we get a match.

Example: Print a Message on Detection #

Here’s how to add a simple print trigger first, to test the logic:

# After running sess.run() and before visualization
# Loop through detections to check for target object
target_class_name = "baseball bat"  # Match this to your label_map.pbtxt entry
min_confidence = 0.7

# Squeeze the arrays to remove extra dimensions
detected_classes = np.squeeze(classes).astype(np.int32)
detected_scores = np.squeeze(scores)

for idx, class_id in enumerate(detected_classes):
    if detected_scores[idx] >= min_confidence:
        class_name = category_index[class_id]['name']
        if class_name == target_class_name:
            print("let's try to make a home run")
            # We'll replace this print with Firebase logic next

Step 2: Integrate Firebase Cloud Messaging (FCM) #

To send messages to your Android app, we’ll use Firebase Admin SDK for Python. Here’s how to set it up:

1. Install Firebase Admin #

Run this command in your terminal:

pip install firebase-admin

2. Get Firebase Service Account Key #

• Go to your Firebase Console → Project Settings → Service Accounts → Generate New Private Key
• Download the JSON file, save it in your project directory (e.g., firebase-service-account.json)

3. Initialize Firebase and Write Notification Function #

Add this code at the top of your script (after imports):

import firebase_admin
from firebase_admin import credentials
from firebase_admin import messaging

# Initialize Firebase Admin SDK
cred = credentials.Certificate('firebase-service-account.json')
firebase_admin.initialize_app(cred)

def send_firebase_notification(target_token, title, body):
    # Create a notification message
    message = messaging.Message(
        notification=messaging.Notification(
            title=title,
            body=body
        ),
        token=target_token,
    )

    # Send the message
    try:
        response = messaging.send(message)
        print(f"Successfully sent message: {response}")
    except Exception as e:
        print(f"Error sending message: {e}")

• target_token: This is the FCM token of your Android device (you’ll get this from your Android app code)
• You can also send to topics instead of individual tokens if you prefer—adjust the message structure accordingly.

Step 3: Full Modified Code #

Here’s your original code with the trigger logic and Firebase integration added, plus a flag to avoid spamming notifications:

# coding: utf-8
# Object Detection Demo
import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
import cv2
# Firebase imports
import firebase_admin
from firebase_admin import credentials
from firebase_admin import messaging

# Initialize Firebase
cred = credentials.Certificate('firebase-service-account.json')
firebase_admin.initialize_app(cred)

cap = cv2.VideoCapture(0)

# Env setup
sys.path.append("..")
sys.path.append("D:\work\O_detection\models")
sys.path.append("D:\work\O_detection\models\slim")

# Object detection imports
from utils import label_map_util
from utils import visualization_utils as vis_util

# Model preparation
MODEL_NAME = 'gun_detection_graph'
PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'
PATH_TO_LABELS = os.path.join('training', 'object-detection.pbtxt')
NUM_CLASSES = 1

detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

# Loading label map
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)

# Helper code
def load_image_into_numpy_array(image):
    (im_width, im_height) = image.size
    return np.array(image.getdata()).reshape(
        (im_height, im_width, 3)).astype(np.uint8)

def send_firebase_notification(target_token, title, body):
    message = messaging.Message(
        notification=messaging.Notification(
            title=title,
            body=body
        ),
        token=target_token,
    )
    try:
        response = messaging.send(message)
        print(f"Notification sent successfully: {response}")
    except Exception as e:
        print(f"Failed to send notification: {str(e)}")

# Detection
with detection_graph.as_default():
    with tf.Session(graph=detection_graph) as sess:
        # Configure your target object and FCM token here
        TARGET_CLASS = "gun"  # Replace with your target object name from label_map
        MIN_CONFIDENCE = 0.7
        ANDROID_DEVICE_TOKEN = "YOUR_ANDROID_DEVICE_FCM_TOKEN"  # Replace with your device's token
        
        # To avoid spamming notifications, track if we've already sent one for the current detection
        notification_sent = False
        
        while True:
            ret, image_np = cap.read()
            if not ret:
                break
                
            image_np_expanded = np.expand_dims(image_np, axis=0)
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
            scores = detection_graph.get_tensor_by_name('detection_scores:0')
            classes = detection_graph.get_tensor_by_name('detection_classes:0')
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')

            # Actual detection
            (boxes, scores, classes, num_detections) = sess.run(
                [boxes, scores, classes, num_detections],
                feed_dict={image_tensor: image_np_expanded})
            
            # Check for target object
            detected_classes = np.squeeze(classes).astype(np.int32)
            detected_scores = np.squeeze(scores)
            
            target_detected = False
            for idx, class_id in enumerate(detected_classes):
                if detected_scores[idx] >= MIN_CONFIDENCE:
                    class_name = category_index[class_id]['name']
                    if class_name == TARGET_CLASS:
                        target_detected = True
                        break
            
            # Send notification if target is detected and we haven't sent one yet
            if target_detected and not notification_sent:
                send_firebase_notification(
                    ANDROID_DEVICE_TOKEN,
                    "Target Detected!",
                    f"{TARGET_CLASS} has been detected."
                )
                notification_sent = True
            elif not target_detected:
                # Reset flag if target is no longer detected
                notification_sent = False

            # Visualization
            vis_util.visualize_boxes_and_labels_on_image_array(
                image_np,
                np.squeeze(boxes),
                np.squeeze(classes).astype(np.int32),
                np.squeeze(scores),
                category_index,
                use_normalized_coordinates=True,
                line_thickness=8)

            cv2.imshow('object detection', cv2.resize(image_np,(800,600)))
            if cv2.waitKey(25) & 0xFF == ord('q'):
                cv2.destroyAllWindows()
                break

Key Notes: #

• Avoid Spamming: The notification_sent flag ensures we don’t send hundreds of notifications per second while the target is in frame. We only send one when the target first appears, and reset the flag when it’s gone.
• FCM Token: Make sure you replace YOUR_ANDROID_DEVICE_FCM_TOKEN with the actual token retrieved from your Android app. You can get this by calling FirebaseMessaging.getInstance().getToken() in your Android code.
• Confidence Threshold: Adjust MIN_CONFIDENCE based on your model’s accuracy—higher values mean fewer false positives.

内容的提问来源于stack exchange，提问作者Aadhar Bhatt