VI_SAFE / engine.py
anakvyas's picture
Deploy Vi-SAFE browser camera
41dd2e6
Raw
History Blame Contribute Delete
11 kB
# engine.py
import cv2
import torch
import numpy as np
import base64
import re
from collections import deque
from ultralytics import YOLO
from torchvision import transforms
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
import torch.nn as nn
import time
import os
import json
from datetime import datetime
try:
DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"
except Exception:
DEVICE = "cpu"
FRAME_BUFFER_SIZE = 16
VIOLENCE_THRESHOLD = 0.55
YOLO_CONFIDENCE = 0.4
MOTION_THRESHOLD = 0.35
MOTION_SUPPRESS = 0.85
ALERT_COOLDOWN = 10
ALERTS_LOG = "alerts.log"
ALERTS_JSONL = "alerts.jsonl"
# Model
class QuickViolenceNet(nn.Module):
def __init__(self):
super().__init__()
base = mobilenet_v2(weights=MobileNet_V2_Weights.IMAGENET1K_V1)
self.features = base.features
self.pool = nn.AdaptiveAvgPool2d(1)
self.lstm = nn.LSTM(1280, 128, num_layers=2, batch_first=True, dropout=0.3)
self.dropout = nn.Dropout(0.4)
self.fc = nn.Linear(128, 2)
def forward(self, x):
B, T, C, H, W = x.shape
x = x.view(B * T, C, H, W)
x = self.pool(self.features(x)).squeeze(-1).squeeze(-1)
x = x.view(B, T, -1)
out, _ = self.lstm(x)
return self.fc(self.dropout(out[:, -1]))
# Load models once
yolo = YOLO("yolov8n.pt")
classifier = QuickViolenceNet().to(DEVICE)
if os.path.exists("violence_classifier.pt"):
classifier.load_state_dict(torch.load("violence_classifier.pt", map_location=DEVICE))
else:
print("WARNING: violence_classifier.pt not found. Predictions may be inaccurate.")
classifier.eval()
transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((112, 112)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
def _decode_data_url(image_data):
if "," in image_data:
image_data = image_data.split(",", 1)[1]
image_data = re.sub(r"\s+", "", image_data)
raw = base64.b64decode(image_data)
arr = np.frombuffer(raw, dtype=np.uint8)
frame = cv2.imdecode(arr, cv2.IMREAD_COLOR)
if frame is None:
raise ValueError("Could not decode image frame")
return frame
def _analyze_frame(frame, state, location):
curr_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
if state["prev_gray"] is not None:
flow = cv2.calcOpticalFlowFarneback(
state["prev_gray"], curr_gray, None, 0.5, 3, 15, 3, 5, 1.2, 0
)
mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])
state["motion_mag"] = float(np.mean(mag))
state["prev_gray"] = curr_gray.copy()
results = yolo(frame, classes=[0], conf=YOLO_CONFIDENCE, verbose=False)
roi = frame
boxes = results[0].boxes
person_count = 0
if boxes is not None and len(boxes) > 0:
person_count = len(boxes)
xyxy = boxes.xyxy.cpu().numpy()
areas = [(b[2] - b[0]) * (b[3] - b[1]) for b in xyxy]
best = xyxy[np.argmax(areas)].astype(int)
x1, y1, x2, y2 = best
x1, y1 = max(0, x1), max(0, y1)
x2, y2 = min(frame.shape[1], x2), min(frame.shape[0], y2)
if x2 > x1 and y2 > y1:
roi = frame[y1:y2, x1:x2]
rgb = cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)
tensor = transform(rgb)
state["buffer"].append(tensor)
if len(state["buffer"]) == FRAME_BUFFER_SIZE:
with torch.no_grad():
clip = torch.stack(list(state["buffer"])).unsqueeze(0).to(DEVICE)
logits = classifier(clip)
probs = torch.softmax(logits, dim=1)
raw_score = probs[0][1].item()
if state["motion_mag"] < MOTION_THRESHOLD:
raw_score *= MOTION_SUPPRESS
state["score"] = raw_score
if state["score"] > VIOLENCE_THRESHOLD:
if state["violence_start_t"] is None:
state["violence_start_t"] = time.time()
else:
state["violence_start_t"] = None
duration = 0.0
if state["violence_start_t"] is not None:
duration = time.time() - state["violence_start_t"]
return {
"location": location,
"score": round(float(state["score"]), 4),
"is_violent": bool(state["score"] > VIOLENCE_THRESHOLD),
"motion": round(float(state["motion_mag"]), 4),
"duration_seconds": round(float(duration), 1),
"threshold": VIOLENCE_THRESHOLD,
"person_count": int(person_count),
"frames_collected": len(state["buffer"]),
"frames_required": FRAME_BUFFER_SIZE,
}
class BrowserFrameAnalyzer:
def __init__(self, location="Browser Camera"):
self.location = location
self.state = {
"buffer": deque(maxlen=FRAME_BUFFER_SIZE),
"score": 0.0,
"prev_gray": None,
"motion_mag": 1.0,
"violence_start_t": None,
}
def analyze_data_url(self, image_data):
frame = _decode_data_url(image_data)
return _analyze_frame(frame, self.state, self.location)
class CameraEngine:
def __init__(self, cam_id, location="Camera"):
self.cam_id = cam_id
self.cap = cv2.VideoCapture(cam_id)
self.buffer = deque(maxlen=FRAME_BUFFER_SIZE)
self.score = 0.0
self.location = location
self.prev_gray = None
self.motion_mag = 1.0
self.violence_start_t = None
self.last_alert_time = 0.0
self.last_alert_record = None
self.is_open = self.cap.isOpened()
def _ensure_camera_open(self):
if self.cap is not None and self.cap.isOpened():
self.is_open = True
return True
# Retry opening camera if it was disconnected or busy before.
self.cap = cv2.VideoCapture(self.cam_id)
self.is_open = self.cap.isOpened()
return self.is_open
def _trigger_alert(self, confidence, duration_seconds):
now = time.time()
if now - self.last_alert_time < ALERT_COOLDOWN:
return
self.last_alert_time = now
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
msg = (
f"[ALERT] {timestamp} | Location: {self.location} | "
f"Confidence: {confidence:.1%} | Duration: {duration_seconds:.1f}s"
)
with open(ALERTS_LOG, "a") as f:
f.write(msg + "\n")
record = {
"timestamp": timestamp,
"location": self.location,
"confidence": round(confidence, 4),
"duration_seconds": round(duration_seconds, 1),
"camera": str(self.location),
}
with open(ALERTS_JSONL, "a") as f:
f.write(json.dumps(record) + "\n")
self.last_alert_record = record
def get_frame(self):
if not self._ensure_camera_open():
return None
ret, frame = self.cap.read()
if not ret:
self.is_open = False
return None
curr_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
if self.prev_gray is not None:
flow = cv2.calcOpticalFlowFarneback(
self.prev_gray, curr_gray, None, 0.5, 3, 15, 3, 5, 1.2, 0
)
mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])
self.motion_mag = float(np.mean(mag))
self.prev_gray = curr_gray.copy()
results = yolo(frame, classes=[0], conf=YOLO_CONFIDENCE, verbose=False)
roi = frame
boxes = results[0].boxes
if boxes is not None and len(boxes) > 0:
xyxy = boxes.xyxy.cpu().numpy()
areas = [(b[2] - b[0]) * (b[3] - b[1]) for b in xyxy]
best = xyxy[np.argmax(areas)].astype(int)
x1, y1, x2, y2 = best
x1, y1 = max(0, x1), max(0, y1)
x2, y2 = min(frame.shape[1], x2), min(frame.shape[0], y2)
if x2 > x1 and y2 > y1:
roi = frame[y1:y2, x1:x2]
try:
rgb = cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)
tensor = transform(rgb)
self.buffer.append(tensor)
except Exception:
return frame
if len(self.buffer) == FRAME_BUFFER_SIZE:
with torch.no_grad():
clip = torch.stack(list(self.buffer)).unsqueeze(0).to(DEVICE)
logits = classifier(clip)
probs = torch.softmax(logits, dim=1)
raw_score = probs[0][1].item()
if self.motion_mag < MOTION_THRESHOLD:
raw_score *= MOTION_SUPPRESS
self.score = raw_score
if self.score > VIOLENCE_THRESHOLD:
if self.violence_start_t is None:
self.violence_start_t = time.time()
duration = time.time() - self.violence_start_t
self._trigger_alert(self.score, duration)
else:
self.violence_start_t = None
annotated = results[0].plot()
is_violent = self.score > VIOLENCE_THRESHOLD
status_color = (0, 60, 220) if is_violent else (20, 180, 60)
status_text = "!! VIOLENCE DETECTED !!" if is_violent else "NORMAL"
suppressed = self.motion_mag < MOTION_THRESHOLD
h_w = annotated.shape[1]
overlay = annotated.copy()
cv2.rectangle(overlay, (0, 0), (h_w, 75), (10, 10, 10), -1)
cv2.addWeighted(overlay, 0.75, annotated, 0.25, 0, annotated)
cv2.putText(
annotated,
f"Status: {status_text}",
(10, 26),
cv2.FONT_HERSHEY_DUPLEX,
0.75,
status_color,
2,
)
cv2.putText(
annotated,
f"Score: {self.score:.2f}{' [suppressed]' if suppressed else ''} | {self.location}",
(10, 52),
cv2.FONT_HERSHEY_SIMPLEX,
0.50,
(255, 255, 255),
1,
)
cv2.putText(
annotated,
f"Motion: {'HIGH' if self.motion_mag >= MOTION_THRESHOLD else 'LOW'} ({self.motion_mag:.2f})",
(10, 72),
cv2.FONT_HERSHEY_SIMPLEX,
0.42,
(100, 200, 100) if self.motion_mag >= MOTION_THRESHOLD else (100, 100, 200),
1,
)
if is_violent and int(time.time() * 2) % 2 == 0:
cv2.rectangle(annotated, (2, 2), (h_w - 2, annotated.shape[0] - 2), (0, 0, 255), 3)
return annotated
def status(self):
duration = 0.0
if self.violence_start_t is not None:
duration = time.time() - self.violence_start_t
return {
"location": self.location,
"camera_id": self.cam_id,
"camera_open": bool(self._ensure_camera_open()),
"score": round(float(self.score), 4),
"is_violent": bool(self.score > VIOLENCE_THRESHOLD),
"motion": round(float(self.motion_mag), 4),
"duration_seconds": round(float(duration), 1),
"threshold": VIOLENCE_THRESHOLD,
}