diff --git a/ai-worker/best.pt b/ai-worker/best.pt new file mode 100644 index 00000000..a53d7a2b Binary files /dev/null and b/ai-worker/best.pt differ diff --git a/ai-worker/best3.pt b/ai-worker/best3.pt new file mode 100644 index 00000000..f44d99f2 Binary files /dev/null and b/ai-worker/best3.pt differ diff --git a/ai-worker/best6.pt b/ai-worker/best6.pt new file mode 100644 index 00000000..6e4bf6ae Binary files /dev/null and b/ai-worker/best6.pt differ diff --git a/ai-worker/requirements.txt b/ai-worker/requirements.txt new file mode 100644 index 00000000..ed07d3b7 --- /dev/null +++ b/ai-worker/requirements.txt @@ -0,0 +1,3 @@ +opencv-python +python-socketio[client] +ultralytics diff --git a/ai-worker/test_video_model.py b/ai-worker/test_video_model.py new file mode 100644 index 00000000..83a4aa61 --- /dev/null +++ b/ai-worker/test_video_model.py @@ -0,0 +1,128 @@ +import argparse +import time +from pathlib import Path + +import cv2 +from ultralytics import YOLO + + +CLASS_COLORS = { + "full": (0, 0, 255), + "half": (0, 255, 255), + "empty": (0, 255, 0), +} + + +def parse_args(): + parser = argparse.ArgumentParser(description="Test trash-bin YOLO model on a local video file.") + parser.add_argument("video", nargs="?", default="trash_test.mp4", help="Path to video file. Default: trash_test.mp4") + parser.add_argument("--model", default="best.pt", help="Path to YOLO model. Default: best.pt") + parser.add_argument("--conf", type=float, default=0.5, help="Confidence threshold. Default: 0.5") + parser.add_argument("--imgsz", type=int, default=640, help="YOLO image size. Default: 640") + parser.add_argument("--save", default="", help="Optional output video path, for example result.mp4") + return parser.parse_args() + + +def draw_detections(frame, result, names, conf_threshold): + best_status = "unknown" + best_confidence = 0.0 + + for box in result.boxes: + confidence = float(box.conf[0]) + if confidence < conf_threshold: + continue + + class_id = int(box.cls[0]) + class_name = str(names[class_id]).lower() + x1, y1, x2, y2 = map(int, box.xyxy[0]) + color = CLASS_COLORS.get(class_name, (255, 255, 255)) + + label = f"{class_name} {confidence:.2f}" + cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2) + cv2.putText(frame, label, (x1, max(y1 - 10, 24)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2) + + if confidence > best_confidence: + best_status = class_name + best_confidence = confidence + + return best_status, best_confidence + + +def main(): + args = parse_args() + video_path = Path(args.video) + model_path = Path(args.model) + + if not video_path.exists(): + raise FileNotFoundError(f"Video not found: {video_path}") + if not model_path.exists(): + raise FileNotFoundError(f"Model not found: {model_path}") + + print(f"Loading model: {model_path}") + model = YOLO(str(model_path)) + print("Model classes:", model.names) + + cap = cv2.VideoCapture(str(video_path)) + if not cap.isOpened(): + raise RuntimeError(f"Cannot open video: {video_path}") + + source_fps = cap.get(cv2.CAP_PROP_FPS) or 25 + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + writer = None + if args.save: + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + writer = cv2.VideoWriter(args.save, fourcc, source_fps, (width, height)) + print(f"Saving result to: {args.save}") + + paused = False + prev_time = time.time() + + print("Controls: q = quit, space = pause/resume") + + while True: + if not paused: + ok, frame = cap.read() + if not ok: + break + + results = model(frame, verbose=False, imgsz=args.imgsz) + status, confidence = draw_detections(frame, results[0], model.names, args.conf) + + now = time.time() + fps = 1 / max(now - prev_time, 0.001) + prev_time = now + + header_color = CLASS_COLORS.get(status, (255, 255, 255)) + cv2.rectangle(frame, (0, 0), (width, 52), (0, 0, 0), -1) + cv2.putText( + frame, + f"Status: {status} | Confidence: {confidence:.2f} | FPS: {fps:.1f}", + (16, 34), + cv2.FONT_HERSHEY_SIMPLEX, + 0.9, + header_color, + 2, + ) + + if writer: + writer.write(frame) + + cv2.imshow("Trash model video test", frame) + key = cv2.waitKey(1 if not paused else 50) & 0xFF + + if key == ord("q"): + break + if key == ord(" "): + paused = not paused + + cap.release() + if writer: + writer.release() + cv2.destroyAllWindows() + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/ai-worker/yolo_multistream.py b/ai-worker/yolo_multistream.py new file mode 100644 index 00000000..cc11eb11 --- /dev/null +++ b/ai-worker/yolo_multistream.py @@ -0,0 +1,166 @@ +import cv2 +import base64 +import socketio +import time +import threading +from ultralytics import YOLO + +print("Загрузка YOLOv8...") +model = YOLO("best.pt") +print("Модель загружена успешно!") + +NODE_URL = "https://localhost:3000" +CONFIDENCE_THRESHOLD = 0.5 +ANALYSIS_INTERVAL_SECONDS = 300 + +sio = socketio.Client(reconnection=True, ssl_verify=False) + +@sio.event +def connect(): + print("Подключено к Node.js серверу!") + +@sio.event +def disconnect(): + print("Отключено от сервера!") + +# --------------------------------------------------------- +# НАСТРОЙКИ ДЛЯ 2 КАМЕР +# Обрати внимание на конец ссылки: Channels/101 и Channels/102 +# --------------------------------------------------------- +CAMERAS = [ + # Первая камера (старый пароль, IP .237) + {"id": 1, "url": "rtsp://admin:112233**gsc@172.16.11.237:554/Streaming/Channels/102"}, + + # Вторая камера (НОВЫЙ пароль, IP .236) + {"id": 2, "url": "rtsp://admin:88888888a@172.16.11.236:554/Streaming/Channels/101"} +] + +# Статус ИИ для 2 камер +ai_enabled = {1: True, 2: True} + +@sio.on("set_ai") +def set_ai(data): + cam_id = int(data.get("cameraId")) + ai_enabled[cam_id] = bool(data.get("enabled")) + print(f"Статус ИИ для камеры {cam_id} изменен на: {ai_enabled[cam_id]}") + +def connect_to_server(): + while True: + try: + sio.connect(NODE_URL, wait_timeout=10) + return + except Exception as e: + print(f"Ошибка подключения к серверу: {e}. Повтор через 3 сек.") + time.sleep(3) + +connect_to_server() + +def pick_best_detection(results): + best = None + for r in results: + for box in r.boxes: + conf = float(box.conf[0]) + if conf < CONFIDENCE_THRESHOLD: + continue + + cls = int(box.cls[0]) + class_name = str(model.names[cls]).lower() + if best is None or conf > best["confidence"]: + best = { + "class_name": class_name, + "confidence": conf, + "box": tuple(map(int, box.xyxy[0])) + } + return best + +def process_camera(cam_id, rtsp_url): + cap = cv2.VideoCapture(rtsp_url, cv2.CAP_FFMPEG) + cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) + print(f"Запуск потока камеры {cam_id}...") + + prev_time = time.time() + last_analysis_time = 0 + + while True: + ret, frame = cap.read() + if not ret: + print(f"Камера {cam_id} потеряла связь. Переподключение...") + cap.release() + time.sleep(2) + cap = cv2.VideoCapture(rtsp_url, cv2.CAP_FFMPEG) + cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) + continue + + frame = cv2.resize(frame, (640, 360)) + + current_time = time.time() + fps = int(1 / max(current_time - prev_time, 0.001)) + prev_time = current_time + + success, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 50]) + + if success: + frame_base64 = base64.b64encode(buffer).decode("utf-8") + sio.emit("camera_frame", { + "type": "live", + "cameraId": cam_id, + "image": frame_base64, + "fps": fps, + "aiEnabled": ai_enabled.get(cam_id, True) + }) + + should_analyze = ( + ai_enabled.get(cam_id, True) + and current_time - last_analysis_time >= ANALYSIS_INTERVAL_SECONDS + ) + + if should_analyze: + last_analysis_time = current_time + analysis_frame = frame.copy() + trash_status = "unknown" + confidence = 0.0 + + results = model(analysis_frame, verbose=False, imgsz=640) + detection = pick_best_detection(results) + + if detection: + trash_status = detection["class_name"] + confidence = detection["confidence"] + x1, y1, x2, y2 = detection["box"] + color = (0, 0, 255) if trash_status == "full" else (0, 255, 0) + cv2.rectangle(analysis_frame, (x1, y1), (x2, y2), color, 2) + cv2.putText(analysis_frame, f"{trash_status} {confidence:.2f}", (x1, y1 - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) + + success, buffer = cv2.imencode(".jpg", analysis_frame, [cv2.IMWRITE_JPEG_QUALITY, 70]) + if success: + frame_base64 = base64.b64encode(buffer).decode("utf-8") + sio.emit("camera_frame", { + "type": "analysis", + "cameraId": cam_id, + "image": frame_base64, + "trashStatus": trash_status, + "confidence": confidence, + "fps": fps, + "aiEnabled": ai_enabled.get(cam_id, True), + "analysisIntervalSeconds": ANALYSIS_INTERVAL_SECONDS, + "analyzedAt": int(current_time * 1000) + }) + + time.sleep(0.01) + +threads = [] +for cam in CAMERAS: + t = threading.Thread(target=process_camera, args=(cam["id"], cam["url"])) + t.daemon = True + t.start() + threads.append(t) + +try: + while True: + time.sleep(1) +except KeyboardInterrupt: + print("\nОстановка всех камер и отключение...") + sio.disconnect() + +#88888888a diff --git a/ai-worker/yolov8.py b/ai-worker/yolov8.py new file mode 100644 index 00000000..0f6eee4f --- /dev/null +++ b/ai-worker/yolov8.py @@ -0,0 +1,225 @@ +import cv2 +import base64 +import socketio +import time +from ultralytics import YOLO +import threading +import os +import urllib.parse + +os.environ.setdefault("OPENCV_FFMPEG_CAPTURE_OPTIONS", "rtsp_transport;tcp|stimeout;5000000") + +SERVER_URLS = [ + os.environ.get("AI_SERVER_URL"), + "http://127.0.0.1:3000", + "http://localhost:3000", + "http://172.16.11.190:3000", +] +SERVER_URLS = list(dict.fromkeys(url for url in SERVER_URLS if url)) +FRAME_SIZE = (640, 360) +ANALYSIS_INTERVAL_SECONDS = 300 + +CONF_THRESHOLD = 0.10 +IOU_THRESHOLD = 0.60 + +sio = socketio.Client(ssl_verify=False) + +print("Загрузка YOLOv8...") +model = YOLO("best6.pt") + +ai_status = {} +active_cameras = {} + +def emit_camera_status(cam_id, online, status, error=""): + sio.emit("camera_status", { + "cameraId": cam_id, + "online": online, + "status": status, + "error": error, + "updatedAt": int(time.time() * 1000) + }) + +def build_rtsp_url(data): + direct_url = str(data.get("url") or data.get("rtspUrl") or "").strip() + ip_or_url = str(data.get("ip") or "").strip() + + if direct_url.startswith("rtsp://"): + return direct_url + if ip_or_url.startswith("rtsp://"): + return ip_or_url + + usr = str(data.get("username", "admin")).strip() + raw_pwd = str(data.get("password", "")).strip() + ip = ip_or_url + port = str(data.get("port", "554")).strip() or "554" + channel = str(data.get("channel", "101")).strip().strip("/") + + if not ip: + raise ValueError("IP камеры пустой") + + safe_usr = urllib.parse.quote(usr, safe="") + safe_pwd = urllib.parse.quote(raw_pwd, safe="*") + auth_part = safe_usr + if safe_pwd: + auth_part += f":{safe_pwd}" + + if channel.lower().startswith("streaming/channels/"): + path = channel + elif "/" in channel: + path = channel + else: + path = f"Streaming/Channels/{channel}" + + return f"rtsp://{auth_part}@{ip}:{port}/{path}" + +def open_capture(url): + cap = cv2.VideoCapture(url, cv2.CAP_FFMPEG) + cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) + return cap + +@sio.event +def connect(): + print("[+] Успешно подключено к Node.js серверу!") + +@sio.event +def disconnect(): + print("[-] Отключено от Node.js сервера!") + +@sio.on("set_ai") +def set_ai(data): + cam_id = int(data.get("cameraId")) + ai_status[cam_id] = bool(data.get("enabled")) + print(f"[*] Камера {cam_id} | Статус ИИ: {ai_status.get(cam_id)}") + +@sio.on("init_cameras") +def on_init_cameras(cameras_list): + print(f"[*] Получен список из {len(cameras_list)} камер от сервера.") + for cam in cameras_list: + on_start_camera_stream(cam) + +@sio.on("start_camera_stream") +def on_start_camera_stream(data): + cam_id = int(data["id"]) + if active_cameras.get(cam_id): + return + + try: + rtsp_url = build_rtsp_url(data) + except Exception as exc: + message = f"Ошибка настроек камеры: {exc}" + print(f"[!] Камера {cam_id}: {message}") + emit_camera_status(cam_id, False, "offline", message) + return + + cam_info = {"id": cam_id, "url": rtsp_url} + ai_status[cam_id] = data.get("aiEnabled", True) + active_cameras[cam_id] = True + + t = threading.Thread(target=process_camera, args=(cam_info,)) + t.daemon = True + t.start() + +@sio.on("delete_camera") +def on_delete_camera(data): + cam_id = int(data["id"]) + if cam_id in active_cameras: + active_cameras[cam_id] = False + +def connect_to_server(): + while True: + for node_url in SERVER_URLS: + try: + sio.connect(node_url, wait_timeout=5) + return + except Exception: + pass + time.sleep(3) + +def process_camera(cam_info): + cam_id = cam_info["id"] + url = cam_info["url"] + + print(f"[*] AI Worker камеры {cam_id} запущен (Интервал {ANALYSIS_INTERVAL_SECONDS} сек).") + + last_analysis_time = 0 + was_online = False + + while active_cameras.get(cam_id, True): + current_time = time.time() + needs_analysis = ai_status.get(cam_id, True) and (current_time - last_analysis_time >= ANALYSIS_INTERVAL_SECONDS) + + if not needs_analysis: + time.sleep(1) + continue + + print(f"[>] Камера {cam_id}: Захват кадра для ИИ-анализа...") + cap = open_capture(url) + ret, frame = cap.read() + cap.release() # Освобождаем девайс немедленно + + if not ret or frame is None: + emit_camera_status(cam_id, False, "offline", "Нет сигнала") + was_online = False + time.sleep(5) + continue + + if not was_online: + emit_camera_status(cam_id, True, "online", "") + was_online = True + + last_analysis_time = time.time() + analysis_frame = cv2.resize(frame, FRAME_SIZE) + + results = model(analysis_frame, conf=CONF_THRESHOLD, iou=IOU_THRESHOLD, agnostic_nms=False, verbose=False, imgsz=640) + + status_priority = {"trash_around": 2, "bin_full": 1, "bin_empty": 0} + global_trash_status = "unknown" + max_priority = -1 + highest_confidence = 0.0 + + for r in results: + for box in r.boxes: + conf = float(box.conf[0]) + if conf < CONF_THRESHOLD: continue + + cls = int(box.cls[0]) + class_name = str(model.names[cls]).lower() + + if "trash" in class_name or class_name == "2": + color, clean_name = (0, 0, 255), "trash_around" + elif "full" in class_name or class_name == "1": + color, clean_name = (0, 165, 255), "bin_full" + else: + color, clean_name = (0, 255, 0), "bin_empty" + + x1, y1, x2, y2 = tuple(map(int, box.xyxy[0])) + cv2.rectangle(analysis_frame, (x1, y1), (x2, y2), color, 2) + cv2.putText(analysis_frame, f"{clean_name} {conf:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) + + priority = status_priority.get(clean_name, -1) + if priority > max_priority: + max_priority = priority + global_trash_status = clean_name + highest_confidence = conf + elif priority == max_priority and conf > highest_confidence: + highest_confidence = conf + + success, buffer = cv2.imencode(".jpg", analysis_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70]) + if success: + jpg_as_text = base64.b64encode(buffer).decode("utf-8") + sio.emit("camera_frame", { + "type": "analysis", + "cameraId": cam_id, + "image": jpg_as_text, + "trashStatus": global_trash_status, + "confidence": highest_confidence, + "fps": 0, + "aiEnabled": ai_status.get(cam_id, True), + "analysisIntervalSeconds": ANALYSIS_INTERVAL_SECONDS, + "analyzedAt": int(time.time() * 1000) + }) + print(f"[+] Камера {cam_id}: Анализ завершен. Ожидание {ANALYSIS_INTERVAL_SECONDS} сек.") + +if __name__ == '__main__': + connect_to_server() + sio.wait() diff --git a/ai-worker/yolov8n.pt b/ai-worker/yolov8n.pt new file mode 100644 index 00000000..0db4ca4b Binary files /dev/null and b/ai-worker/yolov8n.pt differ