Create app.py

app.py

@@ -0,0 +1,202 @@
+import numpy as np
+import cv2
+import tensorflow as tf
+from tensorflow.keras.applications import ResNet50
+from tensorflow.keras.applications.resnet50 import preprocess_input
+from sklearn.metrics.pairwise import cosine_similarity
+from filterpy.kalman import KalmanFilter
+import gradio as gr
+
+# Load the frozen inference graph
+frozen_graph_path = "frozen_inference_graph.pb"
+
+# Load the frozen TensorFlow model
+with tf.io.gfile.GFile(frozen_graph_path, "rb") as f:
+    graph_def = tf.compat.v1.GraphDef()
+    graph_def.ParseFromString(f.read())
+
+# Convert the frozen graph to a function
+def wrap_frozen_graph(graph_def, inputs, outputs):
+    def _imports_graph_def():
+        tf.compat.v1.import_graph_def(graph_def, name="")
+    wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])
+    return wrapped_import.prune(
+        tf.nest.map_structure(wrapped_import.graph.as_graph_element, inputs),
+        tf.nest.map_structure(wrapped_import.graph.as_graph_element, outputs))
+
+# Define input and output tensors
+inputs = ["image_tensor:0"]
+outputs = ["detection_boxes:0", "detection_scores:0", "detection_classes:0", "num_detections:0"]
+
+# Get the detection function
+detection_fn = wrap_frozen_graph(graph_def, inputs, outputs)
+
+# TensorFlow function for detection
+@tf.function(input_signature=[tf.TensorSpec(shape=[None, None, None, 3], dtype=tf.uint8)])
+def detect_objects(image):
+    return detection_fn(image)
+
+# Load ResNet50 for feature extraction
+resnet_model = ResNet50(weights="imagenet", include_top=False, pooling="avg")
+
+# Initialize variables to store features and identities
+person_features = []
+person_identities = []
+person_colors = {}
+kalman_filters = {}
+next_person_id = 1  # Starting unique ID for persons
+
+# Function to generate unique colors based on person ID
+def get_color(person_id):
+    np.random.seed(person_id)  # Ensure color is unique for each person_id
+    color = tuple(np.random.randint(0, 256, size=3))  # Generates RGB tuple
+    return (int(color[0]), int(color[1]), int(color[2]))  # Ensure the color is a tuple of ints
+
+def extract_features(person_roi):
+    # Resize and preprocess the ROI for ResNet50 input
+    person_roi_resized = cv2.resize(person_roi, (224, 224))
+    person_roi_preprocessed = preprocess_input(person_roi_resized)
+
+    # Add batch dimension for ResNet50 input
+    input_tensor = np.expand_dims(person_roi_preprocessed, axis=0)
+
+    # Extract features using ResNet50
+    features = resnet_model.predict(input_tensor)
+    return features
+
+def initialize_kalman_filter(bbox):
+    kf = KalmanFilter(dim_x=7, dim_z=4)
+    kf.F = np.array([[1, 0, 0, 0, 1, 0, 0],
+                     [0, 1, 0, 0, 0, 1, 0],
+                     [0, 0, 1, 0, 0, 0, 1],
+                     [0, 0, 0, 1, 0, 0, 0],
+                     [0, 0, 0, 0, 1, 0, 0],
+                     [0, 0, 0, 0, 0, 1, 0],
+                     [0, 0, 0, 0, 0, 0, 1]])
+    kf.H = np.array([[1, 0, 0, 0, 0, 0, 0],
+                     [0, 1, 0, 0, 0, 0, 0],
+                     [0, 0, 0, 1, 0, 0, 0],
+                     [0, 0, 0, 0, 0, 1, 0]])
+    kf.R[2:, 2:] *= 10.
+    kf.P[4:, 4:] *= 1000.
+    kf.P *= 10.
+    kf.Q[-1, -1] *= 0.01
+    kf.Q[4:, 4:] *= 0.01
+    kf.x[:4] = bbox.reshape((4, 1))
+    return kf
+
+def predict_bbox(kf):
+    kf.predict()
+    return kf.x[:4].reshape((4,))
+
+def update_kalman_filter(kf, bbox):
+    kf.update(bbox.reshape((4, 1)))
+    return kf
+
+def match_and_identify(features, bbox):
+    global next_person_id
+
+    # Flag to check if a match is found
+    matched = False
+
+    # Iterate over existing identities to check for matches
+    for idx, (feat, identity) in enumerate(zip(person_features, person_identities)):
+        # Compute cosine similarity between features
+        similarity = cosine_similarity(
+            np.array(feat).reshape(1, -1),
+            np.array(features).reshape(1, -1)
+        )[0][0]
+
+        # If similarity is above threshold, consider them as the same person
+        similarity_threshold = 0.7  # Adjust as needed
+        if similarity > similarity_threshold:
+            # Assign color if not already assigned
+            if identity in person_colors:
+                color = person_colors[identity]
+            else:
+                color = get_color(identity)
+                person_colors[identity] = color
+
+            # Update Kalman filter
+            kalman_filters[identity] = update_kalman_filter(kalman_filters[identity], bbox)
+
+            # Set matched flag to True
+            matched = True
+            return identity, color
+
+    # If no match found, add new identity
+    if not matched:
+        person_features.append(features)
+        person_identities.append(next_person_id)
+        color = get_color(next_person_id)
+        person_colors[next_person_id] = color
+
+        # Initialize Kalman filter
+        kalman_filters[next_person_id] = initialize_kalman_filter(bbox)
+
+        identity = next_person_id
+        next_person_id += 1
+
+        return identity, color
+
+def process_image(image):
+    # Prepare the image tensor
+    image_np = np.array(image)
+    input_tensor = np.expand_dims(image_np, axis=0)
+
+    # Run inference
+    detections = detect_objects(input_tensor)
+
+    # Extract output tensors and convert to numpy arrays
+    boxes = detections[0].numpy()[0]
+    scores = detections[1].numpy()[0]
+    classes = detections[2].numpy()[0]
+    num_detections = int(detections[3].numpy()[0])
+
+    # Filter detections for 'person' class
+    threshold = 0.3  # Adjust this threshold as needed
+    for i in range(num_detections):
+        class_id = int(classes[i])
+        score = scores[i]
+        box = boxes[i]
+
+        if class_id == 1 and score > threshold:
+            h, w, _ = image.shape
+            ymin, xmin, ymax, xmax = box
+            left, right, top, bottom = int(xmin * w), int(xmax * w), int(ymin * h), int(ymax * h)
+            
+            # Extract person ROI
+            person_roi = image[top:bottom, left:right]
+
+            # Extract features
+            features = extract_features(person_roi)
+
+            # Predict bbox using Kalman filter
+            predicted_bbox = np.array([xmin, ymin, xmax, ymax])
+
+            # Match and identify
+            identity, color = match_and_identify(features, predicted_bbox)
+
+            # Draw bounding box
+            left, top, right, bottom = int(predicted_bbox[0] * w), int(predicted_bbox[1] * h), int(predicted_bbox[2] * w), int(predicted_bbox[3] * h)
+            cv2.rectangle(image, (left, top), (right, bottom), color, 2)
+            cv2.putText(image, f'Person {identity}', (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+    return image
+
+def gradio_interface(input_image):
+    # Process the input image
+    output_image = process_image(input_image)
+    return output_image
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Image(),
+    outputs=gr.Image(),
+    title="Person Detection and Tracking",
+    description="Upload an image to detect and track persons."
+)
+
+# Launch the interface
+iface.launch()