ysn-rfd/text-dataset-tiny-code-script-py-format

1import cv2
2import numpy as np
3import time
4from sklearn.neighbors import KNeighborsClassifier
5from collections import defaultdict, deque
6
7# Create background subtractor for motion detection
8back_sub = cv2.createBackgroundSubtractorKNN(history=500, dist2Threshold=400, detectShadows=True)
9cap = cv2.VideoCapture(0)
10
11# Store object traces
12object_traces = defaultdict(lambda: deque(maxlen=30))  # Last 30 points of each object
13object_last_seen = {}
14object_id_counter = 0
15
16# For real-time learning
17knn = KNeighborsClassifier(n_neighbors=3)
18features_set = []
19labels_set = []
20frame_count = 0
21learning_interval = 30
22
23# Timer for data collection
24start_time = time.time()
25learning_time_limit = 60  # 1 minute for data collection
26
27# Variable to avoid predicting before training
28is_trained = False
29
30def apply_noise_reduction(mask):
31    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
32    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=2)
33    mask = cv2.dilate(mask, kernel, iterations=1)
34    return mask
35
36def get_centroid(x, y, w, h):
37    return (int(x + w / 2), int(y + h / 2))
38
39def calculate_direction(trace):
40    if len(trace) < 2:
41        return "-"
42    dx = trace[-1][0] - trace[0][0]
43    dy = trace[-1][1] - trace[0][1]
44    if abs(dx) > abs(dy):
45        return "Left" if dx < 0 else "Right"
46    else:
47        return "Up" if dy < 0 else "Down"
48
49def calculate_speed(trace, duration):
50    if len(trace) < 2 or duration == 0:
51        return 0
52    dist = np.linalg.norm(np.array(trace[-1]) - np.array(trace[0]))
53    return dist / duration
54
55def count_direction_changes(trace):
56    changes = 0
57    for i in range(2, len(trace)):
58        dx1 = trace[i-1][0] - trace[i-2][0]
59        dx2 = trace[i][0] - trace[i-1][0]
60        if dx1 * dx2 < 0:  # Horizontal direction change
61            changes += 1
62    return changes
63
64while True:
65    ret, frame = cap.read()
66    if not ret:
67        break
68
69    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
70    fg_mask = back_sub.apply(frame)
71    fg_mask = apply_noise_reduction(fg_mask)
72
73    contours, _ = cv2.findContours(fg_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
74
75    current_ids = []
76    predicted = 1  # Default prediction value (if no prediction is made)
77    for cnt in contours:
78        area = cv2.contourArea(cnt)
79        if area < 150:
80            continue
81
82        x, y, w, h = cv2.boundingRect(cnt)
83        centroid = get_centroid(x, y, w, h)
84
85        # Identify or create a new ID for the object
86        matched_id = None
87        for oid, trace in object_traces.items():
88            if np.linalg.norm(np.array(trace[-1]) - np.array(centroid)) < 50:
89                matched_id = oid
90                break
91
92        if matched_id is None:
93            matched_id = object_id_counter
94            object_id_counter += 1
95
96        object_traces[matched_id].append(centroid)
97        object_last_seen[matched_id] = time.time()
98        current_ids.append(matched_id)
99
100        trace = object_traces[matched_id]
101        duration = time.time() - object_last_seen[matched_id] + 0.001
102        speed = calculate_speed(trace, duration)
103        direction = calculate_direction(trace)
104        direction_changes = count_direction_changes(trace)
105        total_move = sum(np.linalg.norm(np.array(trace[i]) - np.array(trace[i-1])) for i in range(1, len(trace)))
106
107        # Feature for the model
108        feature = [w, h, centroid[0], centroid[1], area, speed, direction_changes]
109        label = 1  # Default label: Normal
110
111        # Simple automatic labeling:
112        if speed > 100 or direction_changes > 4:
113            label = 2  # Suspicious
114
115        features_set.append(feature)
116        labels_set.append(label)
117
118        # Train the model only after enough data is collected
119        if time.time() - start_time < learning_time_limit:
120            # Still in data collection phase
121            continue
122        elif not is_trained:  # If the model hasn't been trained yet
123            if len(features_set) > 10:
124                knn.fit(features_set, labels_set)  # Train the model
125                is_trained = True  # Model is trained
126                print("Model updated.")
127
128        # Prediction only after training
129        if is_trained:
130            predicted = knn.predict([feature])[0]
131
132        # Draw information on the frame
133        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0) if label == 1 else (0, 0, 255), 2)
134        cv2.circle(frame, centroid, 4, (255, 255, 255), -1)
135        cv2.putText(frame, f"ID: {matched_id} | Direction: {direction} | Speed: {int(speed)}", (x, y - 25),
136                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1)
137        cv2.putText(frame, f"Behavior: {'Normal' if predicted == 1 else 'Suspicious'}", (x, y - 5),
138                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
139
140    frame_count += 1
141
142    # Remove old object IDs
143    for oid in list(object_last_seen):
144        if time.time() - object_last_seen[oid] > 2:
145            object_traces.pop(oid, None)
146            object_last_seen.pop(oid, None)
147
148    cv2.imshow("Behavioral Intelligence", frame)
149    if cv2.waitKey(1) & 0xFF == 27:
150        break
151
152cap.release()
153cv2.destroyAllWindows()
154