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

1import cv2
2import torch
3import torch.nn as nn
4import torch.optim as optim
5import numpy as np
6
7# ======== AI MODEL (PyTorch) ========
8device = torch.device("cpu")
9
10label_map = {"Idle": 0, "Normal": 1, "Erratic": 2}
11reverse_label = {v: k for k, v in label_map.items()}
12
13class BehaviorAI(nn.Module):
14    def __init__(self):
15        super().__init__()
16        self.model = nn.Sequential(
17            nn.Linear(4, 16),
18            nn.ReLU(),
19            nn.Linear(16, 8),
20            nn.ReLU(),
21            nn.Linear(8, 3)
22        )
23        self.loss_fn = nn.CrossEntropyLoss()
24        self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
25    
26    def forward(self, x):
27        return self.model(x)
28
29    def predict_behavior(self, features):
30        self.model.eval()
31        with torch.no_grad():
32            x = torch.tensor([features], dtype=torch.float32).to(device)
33            logits = self.model(x)
34            pred = torch.argmax(logits, dim=-1).item()
35            return reverse_label[pred]
36
37    def learn_from(self, features, label):
38        self.model.train()
39        x = torch.tensor([features], dtype=torch.float32).to(device)
40        y = torch.tensor([label_map[label]], dtype=torch.long).to(device)
41        logits = self.model(x)
42        loss = self.loss_fn(logits, y)
43        self.optimizer.zero_grad()
44        loss.backward()
45        self.optimizer.step()
46
47# ======== FEATURE EXTRACTION ========
48def extract_features(trace):
49    if len(trace) < 2:
50        return [0, 0, 0, 0]
51
52    dx = trace[-1][0] - trace[0][0]
53    dy = trace[-1][1] - trace[0][1]
54    speeds = []
55    directions = []
56
57    for i in range(1, len(trace)):
58        x1, y1 = trace[i-1]
59        x2, y2 = trace[i]
60        dist = np.linalg.norm([x2 - x1, y2 - y1])
61        speeds.append(dist)
62        directions.append(np.arctan2(y2 - y1, x2 - x1))
63
64    avg_speed = np.mean(speeds)
65    direction_changes = np.sum(np.abs(np.diff(directions)))
66    return [dx, dy, avg_speed, direction_changes]
67
68# ======== MAIN REAL-TIME TRACKING ========
69cap = cv2.VideoCapture(0)  # یا 'video.mp4' برای فایل
70
71bg_subtractor = cv2.createBackgroundSubtractorMOG2()
72traces = {}
73next_id = 0
74ai = BehaviorAI()
75
76while True:
77    ret, frame = cap.read()
78    if not ret:
79        break
80
81    fgmask = bg_subtractor.apply(frame)
82    contours, _ = cv2.findContours(fgmask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
83
84    current_positions = []
85
86    for cnt in contours:
87        if cv2.contourArea(cnt) < 500:
88            continue
89
90        x, y, w, h = cv2.boundingRect(cnt)
91        cx, cy = x + w // 2, y + h // 2
92        current_positions.append((cx, cy))
93        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
94
95    new_traces = {}
96    matched_ids = set()
97
98    for cx, cy in current_positions:
99        min_dist = float('inf')
100        matched_id = None
101        for id, trace in traces.items():
102            if len(trace) == 0:
103                continue
104            prev_x, prev_y = trace[-1]
105            dist = np.linalg.norm([cx - prev_x, cy - prev_y])
106            if dist < 50 and id not in matched_ids:
107                min_dist = dist
108                matched_id = id
109
110        if matched_id is None:
111            matched_id = next_id
112            next_id += 1
113            new_traces[matched_id] = []
114
115        else:
116            new_traces[matched_id] = traces[matched_id]
117
118        new_traces[matched_id].append((cx, cy))
119        matched_ids.add(matched_id)
120
121    traces = new_traces
122
123    for id, trace in traces.items():
124        if len(trace) >= 2:
125            for i in range(1, len(trace)):
126                cv2.line(frame, trace[i-1], trace[i], (255, 0, 0), 2)
127
128            features = extract_features(trace)
129            behavior = ai.predict_behavior(features)
130
131            if len(trace) >= 10:
132                if features[2] < 2:
133                    label = "Idle"
134                elif features[3] > 4:
135                    label = "Erratic"
136                else:
137                    label = "Normal"
138                ai.learn_from(features, label)
139
140            cv2.putText(frame, f"ID:{id} AI:{behavior}", trace[-1], cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
141
142    cv2.imshow("Real-Time Tracker with AI", frame)
143    if cv2.waitKey(1) == 27:  # ESC
144        break
145
146cap.release()
147cv2.destroyAllWindows()
148