| |
| 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" |
|
|
| |
| 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])) |
|
|
| |
| 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 |
| |
| 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, |
| } |
|
|