import cv2
import torch
# Load YOLO model (replace with your chosen YOLO version)
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # Example using yolov5s
#model = torch.hub.load('ultralytics/yolov8', 'yolov8s') # Example using yolov8s
# Define classes to detect (e.g., "weapon", "person")
class_names = ["weapon", "person"] # Replace with your specific classes
# Load an image or video stream
#image_path = "path/to/your/image.jpg"
#img = cv2.imread(image_path)
cap = cv2.VideoCapture(0) # 0 for default camera, or path to video
# Detection loop
while True:
ret, frame = cap.read()
if not ret:
break
# Convert to RGB for YOLO
img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Perform inference
results = model(img_rgb)
# Process results
for result in results.xyxy[0]:
#result is a tensor with [x1, y1, x2, y2, confidence, class_id]
# Extract bounding box coordinates and class ID
x1, y1, x2, y2, confidence, class_id = result.tolist()
# Convert to integer for drawing
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
class_id = int(class_id)
# Draw bounding box and label
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
label = f"{class_names[class_id]}: {confidence:.2f}"
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0),
2)
# Display the image/video
cv2.imshow("Weapon Detection", frame)
# Break loop on 'q' key press
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# Clean up
cap.release()
cv2.destroyAllWindows()
Yolo
import numpy as np
import cv2
import imutils
import datetime
gun_cascade = cv2.CascadeClassifier('cascade.xml')
camera = cv2.VideoCapture(0)
firstFrame = None
gun_exist = False
while True:
ret, frame = camera.read()
if frame is None:
break
frame = imutils.resize(frame, width=500)
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
gun = gun_cascade.detectMultiScale(gray, 1.3, 20, minSize=(100, 100))
if len(gun) > 0:
gun_exist = True
for (x, y, w, h) in gun:
frame = cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
roi_gray = gray[y:y + h, x:x + w]
roi_color = frame[y:y + h, x:x + w]
if firstFrame is None:
firstFrame = gray
continue
cv2.putText(frame, datetime.datetime.now().strftime("%A %d %B %Y %I:%M:%S %p"),
(10, frame.shape[0] - 10),
cv2.FONT_HERSHEY_SIMPLEX,
0.35, (0, 0, 255), 1)
if gun_exist:
print("Guns detected")
plt.imshow(frame)
break
else:
cv2.imshow("Security Feed", frame)
key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
break
camera.release()
cv2.destroyAllWindows()
Jupytyer notebook
import os
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout
from tensorflow.keras.models import Model
from sklearn.metrics import classification_report, confusion_matrix
Setup and import
IMAGE_SIZE = (224, 224)
BATCH_SIZE = 32
train_datagen = ImageDataGenerator(rescale=1./255,
rotation_range=20,
zoom_range=0.15,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.15,
horizontal_flip=True,
fill_mode="nearest")
val_test_datagen = ImageDataGenerator(rescale=1./255)
train_generator = train_datagen.flow_from_directory(
'dataset/train',
target_size=IMAGE_SIZE,
batch_size=BATCH_SIZE,
class_mode='binary'
)
val_generator = val_test_datagen.flow_from_directory(
'dataset/val',
target_size=IMAGE_SIZE,
batch_size=BATCH_SIZE,
class_mode='binary'
)
test_generator = val_test_datagen.flow_from_directory(
'dataset/test',
target_size=IMAGE_SIZE,
batch_size=BATCH_SIZE,
class_mode='binary',
shuffle=False
)
Data preprocessing
base_model = MobileNetV2(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
base_model.trainable = False # Freeze base
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dropout(0.5)(x)
predictions = Dense(1, activation='sigmoid')(x)
model = Model(inputs=base_model.input, outputs=predictions)
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
Build the model
history = model.fit(
train_generator,
validation_data=val_generator,
epochs=10
)
Train the model
# Test accuracy
loss, accuracy = model.evaluate(test_generator)
print(f'Test accuracy: {accuracy:.2f}')
# Predictions and confusion matrix
y_pred = model.predict(test_generator)
y_pred_classes = (y_pred > 0.5).astype("int32").flatten()
y_true = test_generator.classes
print("Classification Report:")
print(classification_report(y_true, y_pred_classes, target_names=['Non-Weapon', 'Weapon']))
print("Confusion Matrix:")
print(confusion_matrix(y_true, y_pred_classes))
Evaluate the model
model.save('weapon_detection_model.h5')
Save the model
