Spaces:
Sleeping
Sleeping

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() | |