SmartSurveillanceAI / detect.py
asad231's picture
Create detect.py
845aa73 verified
from ultralytics import YOLO
from PIL import Image
import numpy as np
import cv2
import time
# Load the YOLOv8 model
model = YOLO("yolov8n.pt") # You can use yolov8s.pt for better accuracy
# Detect theft from an image
def detect_theft(pil_image):
frame = np.array(pil_image)
results = model(frame)
# Save thief photo with timestamp
timestamp = time.strftime("%Y%m%d-%H%M%S")
thief_path = f"thief.jpg"
results[0].save(filename=thief_path)
# Annotated image
annotated = results[0].plot()
return Image.fromarray(annotated), thief_path