Upload 3 files

Dockerfile

@@ -0,0 +1,20 @@
+FROM tensorflow/tensorflow:2.15.0-gpu
+
+RUN apt-get update && apt-get install -y libgl1-mesa-glx libglib2.0-0 && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Copy code
+COPY . .
+
+# Install dependencies
+RUN pip install opencv-python fastapi uvicorn[standard] websockets
+
+ENV MODEL_PATH="model"
+ENV RESOLUTION="172"
+
+# Mở port 8000
+EXPOSE 8000
+
+# Chạy Server FastAPI
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

inference.py

@@ -0,0 +1,103 @@
+# inference.py
+import cv2
+import tensorflow as tf
+import time
+import base64
+import datetime
+import os
+
+# Load model 1 lần duy nhất khi import
+MODEL_PATH = os.getenv('MODEL_PATH', 'model')
+RESOLUTION = int(os.getenv('RESOLUTION', 172))
+CONFIDENCE_THRESHOLD = 0.65
+
+print(f"Loading MoViNet from {MODEL_PATH}...")
+model = tf.saved_model.load(MODEL_PATH)
+infer = model.signatures['serving_default']
+print("Model loaded!")
+
+def get_init_states():
+    dummy = tf.zeros([1, 1, RESOLUTION, RESOLUTION, 3], dtype=tf.float32)
+    return model.init_states(tf.shape(dummy))
+
+class VideoProcessor:
+    def __init__(self):
+        self.running = False
+
+    def start_processing(self, rtsp_url, result_queue):
+        self.running = True
+        cap = cv2.VideoCapture(rtsp_url)
+        if not cap.isOpened():
+            result_queue.put({"error": "Cannot open RTSP URL"})
+            return
+
+        states = get_init_states()
+        
+        # Logic quản lý sự kiện (Cooldown)
+        in_event = False
+        event_start_time = None
+        cooldown_counter = 0
+        COOLDOWN_LIMIT = 30 # Frames (~2-3s)
+
+        while self.running:
+            ret, frame = cap.read()
+            if not ret:
+                break
+
+            # 1. Inference
+            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            resized = tf.image.resize(rgb, [RESOLUTION, RESOLUTION])
+            input_tensor = tf.cast(resized, tf.float32) / 255.0
+            input_tensor = input_tensor[tf.newaxis, tf.newaxis, ...]
+
+            outputs = infer(image=input_tensor, **states)
+            logits = outputs['logits']
+            states = {k: v for k, v in outputs.items() if k != 'logits'}
+            probs = tf.nn.softmax(logits, axis=-1)[0]
+            
+            fight_conf = float(probs[0])
+            norm_conf = float(probs[1])
+            is_violence = (fight_conf > norm_conf) and (fight_conf > CONFIDENCE_THRESHOLD)
+
+            # 2. Logic xử lý kết quả để gửi về
+            current_time = datetime.datetime.now()
+            msg = None
+
+            if is_violence:
+                cooldown_counter = 0
+                if not in_event:
+                    in_event = True
+                    event_start_time = current_time
+                    
+                    # START: Gửi ảnh bằng chứng
+                    small_frame = cv2.resize(frame, (640, 360))
+                    _, buffer = cv2.imencode('.jpg', small_frame)
+                    img_base64 = base64.b64encode(buffer).decode('utf-8')
+                    
+                    msg = {
+                        "type": "START",
+                        "timestamp": current_time.isoformat(),
+                        "score": fight_conf,
+                        "image": img_base64
+                    }
+            else:
+                if in_event:
+                    cooldown_counter += 1
+                    if cooldown_counter >= COOLDOWN_LIMIT:
+                        # END: Gửi thời lượng
+                        duration = (current_time - event_start_time).total_seconds()
+                        msg = {
+                            "type": "END",
+                            "timestamp": current_time.isoformat(),
+                            "duration": duration
+                        }
+                        in_event = False
+            
+            # Nếu có tin quan trọng thì đẩy vào hàng đợi gửi về Laptop
+            if msg:
+                result_queue.put(msg)
+            
+            # Tùy chọn: Gửi Heartbeat mỗi 5s để biết model vẫn sống (nếu cần)
+
+        cap.release()
+        result_queue.put({"status": "Stream stopped"})

main.py

@@ -0,0 +1,55 @@
+# main.py
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from inference import VideoProcessor
+import threading
+import queue
+import asyncio
+import json
+
+app = FastAPI()
+
+@app.websocket("/ws/inference")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    processor = VideoProcessor()
+    result_queue = queue.Queue()
+    
+    try:
+        # 1. Nhận link RTSP từ Laptop
+        data = await websocket.receive_text()
+        config = json.loads(data)
+        rtsp_url = config.get("url")
+        
+        print(f"Received request for: {rtsp_url}")
+        
+        # 2. Chạy Model trong luồng riêng (Background Thread)
+        # Để không chặn WebSocket
+        process_thread = threading.Thread(
+            target=processor.start_processing, 
+            args=(rtsp_url, result_queue)
+        )
+        process_thread.start()
+
+        # 3. Vòng lặp gửi kết quả về Laptop
+        while True:
+            # Kiểm tra queue xem có kết quả mới từ Model không
+            try:
+                # Non-blocking get
+                result = result_queue.get_nowait()
+                await websocket.send_json(result)
+            except queue.Empty:
+                # Nếu không có kết quả, ngủ xíu để đỡ tốn CPU
+                await asyncio.sleep(0.1)
+                
+            # Kiểm tra nếu client ngắt kết nối thì dừng model
+            if not process_thread.is_alive():
+                await websocket.send_json({"status": "Processing finished"})
+                break
+
+    except WebSocketDisconnect:
+        print("Client disconnected")
+        processor.running = False # Ra lệnh dừng model
+        process_thread.join()
+    except Exception as e:
+        print(f"Error: {e}")
+        processor.running = False