ishworrsubedii's picture
feat: add complete pipeline and Streamlit code This commit introduces a complete pipeline for both single and real-time inferences using cameras. It includes the implementation of Streamlit code to facilitate the process.
c640bc9 verified
"""
Author: ishwor subedi
Date: 2023-12-28
Project Name:gun-detection
"""
import os
from PIL import Image
from datetime import datetime
import threading
import imagehash
import cv2 as cv
from utils.settings import get_frame_save_dir
class FrameSaver:
def __init__(self, source, image_path_to_save, image_hash_threshold):
self.source = source
self.image_path_to_save = image_path_to_save
if self.image_path_to_save is None or not os.path.exists(self.image_path_to_save):
self.image_path_to_save = get_frame_save_dir()
self.thread_running = False
self.frame_save_thread = None
self.cap = cv.VideoCapture(self.source)
self.image_hash_threshold = image_hash_threshold
def start_stream(self):
self.frame_save_thread = threading.Thread(target=self.video_webcam_frame_extraction)
self.frame_save_thread.start()
self.thread_running = True
def stop_stream(self):
if self.thread_running:
self.thread_running = False
self.frame_save_thread.join()
self.cap.release()
def hashing_diff(self, prev_frame, current_frame):
if prev_frame is None:
return None
else:
hash1 = imagehash.average_hash(Image.fromarray(prev_frame))
hash2 = imagehash.average_hash(Image.fromarray(current_frame))
return hash2 - hash1
def video_webcam_frame_extraction(self):
"""
:return:
"""
try:
previous_frame = None
while self.cap.isOpened():
success, frame = self.cap.read()
if success:
hash_diff = self.hashing_diff(previous_frame, frame)
if hash_diff is None or hash_diff > self.image_hash_threshold:
current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
name = os.path.join(self.image_path_to_save, f"{current_time}.jpg")
success = cv.imwrite(name, frame)
if success:
print(f"Saved {name}...")
else:
print(f"Failed to save image: {name}")
previous_frame = frame
except Exception as e:
print(f"Exception occurred: {e}")
try:
self.stop_stream()
except Exception as e:
print(f"Failed to stop stream: {e}")
if __name__ == '__main__':
source = 'rtsp://ishwor:subedi@192.168.1.106:5555/h264_opus.sdp'
image_path_to_save = "images/cam_images"
image_hash_threshold = 5
image_capture_service = FrameSaver(source, image_path_to_save, image_hash_threshold)
image_capture_service.video_webcam_frame_extraction()