Initial commit
This commit is contained in:
@@ -0,0 +1,4 @@
|
||||
venv/
|
||||
test/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
@@ -0,0 +1,37 @@
|
||||
# SUV-Helper
|
||||
|
||||
## What is it
|
||||
|
||||
This is a project to get security camera stream and apply smart assistant to it.
|
||||
|
||||
## Design
|
||||
|
||||
Producer (multiple) -> Server -> Processor (multiple)
|
||||
|
||||
## Current status
|
||||
|
||||
Scuffolded.
|
||||
|
||||
Tested on 145 devices.
|
||||
|
||||
## How to play it
|
||||
```
|
||||
$ apt install ffmpeg
|
||||
|
||||
# Create venv
|
||||
$ virtualenv venv
|
||||
$ source venv/bin/activate
|
||||
|
||||
# Install requirements
|
||||
$ pip install -r requirements.txt
|
||||
|
||||
# Start server
|
||||
$ python src/server/server.py
|
||||
|
||||
# Start processor
|
||||
$ python src/processor/processor.py
|
||||
|
||||
# Start producer
|
||||
RTSP_URL="rtsp://user:password@ip_address:554/"
|
||||
$ python src/producer/producer.py --rtsp_url "${RTSP_URL}" --output_dir test --capture_interval 0.1
|
||||
```
|
||||
@@ -0,0 +1,4 @@
|
||||
numpy==2.2.1
|
||||
opencv-python==4.10.0.84
|
||||
pillow==11.1.0
|
||||
rtsp==1.1.12
|
||||
@@ -0,0 +1,140 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
processor.py
|
||||
|
||||
A processor client that connects to the central server, receives image processing tasks,
|
||||
processes the images using OpenCV, and notifies the server upon completion.
|
||||
"""
|
||||
|
||||
import socket
|
||||
import time
|
||||
import cv2
|
||||
import logging
|
||||
import os
|
||||
|
||||
# Configure logging format and level
|
||||
logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
|
||||
|
||||
|
||||
class ProcessorClient:
|
||||
"""
|
||||
ProcessorClient connects to the central server as a processor.
|
||||
It listens for tasks (i.e. image paths), processes images using OpenCV,
|
||||
and sends a "DONE" message to the server once processing is complete.
|
||||
"""
|
||||
|
||||
def __init__(self, server_ip: str, server_port: int, buffer_size: int = 1024):
|
||||
"""
|
||||
Initialize the ProcessorClient with server connection details.
|
||||
"""
|
||||
self.server_ip = server_ip
|
||||
self.server_port = server_port
|
||||
self.buffer_size = buffer_size
|
||||
self.socket = None
|
||||
|
||||
def connect(self) -> None:
|
||||
"""
|
||||
Establish a connection with the server and register as a processor.
|
||||
"""
|
||||
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
try:
|
||||
self.socket.connect((self.server_ip, self.server_port))
|
||||
logging.info(f"Connected to server at {self.server_ip}:{self.server_port}")
|
||||
# Send registration message
|
||||
self.socket.sendall("ROLE:processor".encode())
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to connect or register with the server: {e}")
|
||||
raise
|
||||
|
||||
def process_image(self, image_path: str) -> None:
|
||||
"""
|
||||
Process the image using OpenCV.
|
||||
For demonstration, this method converts the image to grayscale and saves it.
|
||||
"""
|
||||
logging.info(f"Processing image: {image_path}")
|
||||
if not os.path.exists(image_path):
|
||||
logging.error(f"Image file not found: {image_path}")
|
||||
return
|
||||
|
||||
# Load the image using OpenCV
|
||||
image = cv2.imread(image_path)
|
||||
if image is None:
|
||||
logging.error(f"Failed to load image: {image_path}")
|
||||
return
|
||||
|
||||
# Convert to grayscale as an example processing step
|
||||
processed_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# Generate the processed image path and save the processed image
|
||||
processed_path = self._get_processed_path(image_path)
|
||||
cv2.imwrite(processed_path, processed_image)
|
||||
logging.info(f"Processed image saved to: {processed_path}")
|
||||
|
||||
# (Optional) Simulate additional processing delay
|
||||
time.sleep(1)
|
||||
|
||||
def _get_processed_path(self, image_path: str) -> str:
|
||||
"""
|
||||
Generate a processed image path based on the original image path.
|
||||
For example, '/path/to/image.jpg' becomes '/path/to/processed_image.jpg'.
|
||||
"""
|
||||
directory, filename = os.path.split(image_path)
|
||||
processed_filename = f"processed_{filename}"
|
||||
return os.path.join(directory, processed_filename)
|
||||
|
||||
def listen_for_tasks(self) -> None:
|
||||
"""
|
||||
Listen for image processing tasks from the server.
|
||||
Each received message is assumed to be an image file path.
|
||||
"""
|
||||
try:
|
||||
while True:
|
||||
data = self.socket.recv(self.buffer_size).decode().strip()
|
||||
if not data:
|
||||
logging.info("No data received; the server may have closed the connection.")
|
||||
break
|
||||
|
||||
# For this design, we assume 'data' is the image path.
|
||||
image_path = data
|
||||
logging.info(f"Received task: {image_path}")
|
||||
|
||||
# Process the image
|
||||
self.process_image(image_path)
|
||||
|
||||
# Notify the server that processing is complete
|
||||
self.socket.sendall("DONE".encode())
|
||||
logging.info("Notified server: DONE")
|
||||
except Exception as e:
|
||||
logging.error(f"Error during task processing: {e}")
|
||||
finally:
|
||||
self.disconnect()
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""
|
||||
Close the connection with the server.
|
||||
"""
|
||||
if self.socket:
|
||||
self.socket.close()
|
||||
logging.info("Disconnected from server.")
|
||||
|
||||
def run(self) -> None:
|
||||
"""
|
||||
Connect to the server and start listening for tasks.
|
||||
"""
|
||||
try:
|
||||
self.connect()
|
||||
self.listen_for_tasks()
|
||||
except Exception as e:
|
||||
logging.error(f"Processor client encountered an error: {e}")
|
||||
self.disconnect()
|
||||
|
||||
|
||||
def main():
|
||||
SERVER_IP = "127.0.0.1"
|
||||
SERVER_PORT = 5000
|
||||
client = ProcessorClient(SERVER_IP, SERVER_PORT)
|
||||
client.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,193 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
producer.py
|
||||
|
||||
A ProducerClient that connects to a central server, reads frames from an RTSP stream using OpenCV,
|
||||
saves them to disk, and notifies the server when a new image is produced.
|
||||
|
||||
Sensitive data (e.g. the RTSP URL) is passed via command-line arguments.
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import socket
|
||||
import time
|
||||
import logging
|
||||
import os
|
||||
import argparse
|
||||
|
||||
|
||||
# Configure logging for the module.
|
||||
logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
|
||||
|
||||
|
||||
class ProducerClient:
|
||||
"""
|
||||
ProducerClient connects to the central server as a producer.
|
||||
It reads frames from an RTSP stream, saves them as images, and notifies the server
|
||||
of newly produced images.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rtsp_url: str,
|
||||
server_ip: str,
|
||||
server_port: int,
|
||||
num_frames: int = 10,
|
||||
output_dir: str = "./",
|
||||
capture_interval: float = 2.0,
|
||||
):
|
||||
"""
|
||||
Initialize the ProducerClient.
|
||||
|
||||
Args:
|
||||
rtsp_url (str): The RTSP stream URL (sensitive data).
|
||||
server_ip (str): The IP address of the central server.
|
||||
server_port (int): The port number of the central server.
|
||||
num_frames (int): The number of frames to capture.
|
||||
output_dir (str): Directory to save the captured images.
|
||||
capture_interval (float): Time (in seconds) to wait between captures.
|
||||
"""
|
||||
self.rtsp_url = rtsp_url
|
||||
self.server_ip = server_ip
|
||||
self.server_port = server_port
|
||||
self.num_frames = num_frames
|
||||
self.output_dir = output_dir
|
||||
self.capture_interval = capture_interval
|
||||
self.socket = None
|
||||
|
||||
def connect(self) -> None:
|
||||
"""
|
||||
Establish a connection to the central server and register as a producer.
|
||||
"""
|
||||
try:
|
||||
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
self.socket.connect((self.server_ip, self.server_port))
|
||||
logging.info(f"Connected to server at {self.server_ip}:{self.server_port}")
|
||||
# Register with the server as a producer.
|
||||
self.socket.sendall("ROLE:producer".encode())
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to connect or register with the server: {e}")
|
||||
raise
|
||||
|
||||
def send_message(self, message: str) -> None:
|
||||
"""
|
||||
Send a message to the server over the established socket connection.
|
||||
|
||||
Args:
|
||||
message (str): The message to send.
|
||||
"""
|
||||
try:
|
||||
self.socket.sendall(message.encode())
|
||||
logging.info(f"Sent message to server: {message}")
|
||||
except Exception as e:
|
||||
logging.error(f"Error sending message to server: {e}")
|
||||
|
||||
def produce_frames(self) -> None:
|
||||
"""
|
||||
Capture frames from the RTSP stream, save them as images, and notify the server.
|
||||
"""
|
||||
cap = cv2.VideoCapture(self.rtsp_url)
|
||||
if not cap.isOpened():
|
||||
logging.error(f"Error: Could not open RTSP stream: {self.rtsp_url}")
|
||||
return
|
||||
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
frame_count = 0
|
||||
|
||||
try:
|
||||
while frame_count < self.num_frames:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
logging.error("Error: Could not read frame from stream.")
|
||||
break
|
||||
|
||||
filename = f"frame_{frame_count:03d}.jpg"
|
||||
filepath = os.path.join(self.output_dir, filename)
|
||||
|
||||
if cv2.imwrite(filepath, frame):
|
||||
logging.info(f"Saved frame {frame_count} to {filepath}")
|
||||
else:
|
||||
logging.error(f"Failed to save frame {frame_count} to {filepath}")
|
||||
continue # Optionally, decide if you want to break or try again
|
||||
|
||||
# Send a message to the server notifying a new image is produced.
|
||||
message = f"{filepath}"
|
||||
self.send_message(message)
|
||||
|
||||
frame_count += 1
|
||||
time.sleep(self.capture_interval)
|
||||
|
||||
logging.info(f"Finished capturing. Total frames saved: {frame_count}")
|
||||
except Exception as e:
|
||||
logging.error(f"An error occurred during frame production: {e}")
|
||||
finally:
|
||||
cap.release()
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""
|
||||
Disconnect from the central server.
|
||||
"""
|
||||
if self.socket:
|
||||
self.socket.close()
|
||||
logging.info("Disconnected from server.")
|
||||
|
||||
def run(self) -> None:
|
||||
"""
|
||||
Connect to the server, capture frames, and then disconnect.
|
||||
"""
|
||||
try:
|
||||
self.connect()
|
||||
self.produce_frames()
|
||||
except Exception as e:
|
||||
logging.error(f"Error in ProducerClient: {e}")
|
||||
finally:
|
||||
self.disconnect()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Producer Client for capturing RTSP stream frames and notifying a central server."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rtsp_url",
|
||||
type=str,
|
||||
required=True,
|
||||
help="RTSP stream URL (sensitive data; do not hard code)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--server_ip", type=str, default="127.0.0.1", help="Central server IP address"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--server_port", type=int, default=5000, help="Central server port number"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_frames", type=int, default=10, help="Number of frames to capture"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default="./",
|
||||
help="Directory to save captured frames",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--capture_interval",
|
||||
type=float,
|
||||
default=2.0,
|
||||
help="Interval (in seconds) between frame captures",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
producer = ProducerClient(
|
||||
rtsp_url=args.rtsp_url,
|
||||
server_ip=args.server_ip,
|
||||
server_port=args.server_port,
|
||||
num_frames=args.num_frames,
|
||||
output_dir=args.output_dir,
|
||||
capture_interval=args.capture_interval,
|
||||
)
|
||||
producer.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,144 @@
|
||||
# server.py
|
||||
import socket
|
||||
import threading
|
||||
import queue
|
||||
import time
|
||||
|
||||
class ImageServer:
|
||||
def __init__(self, host='127.0.0.1', port=5000):
|
||||
self.server_ip = host
|
||||
self.server_port = port
|
||||
self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
self.server_socket.bind((self.server_ip, self.server_port))
|
||||
self.server_socket.listen()
|
||||
print(f"Server listening on {self.server_ip}:{self.server_port}")
|
||||
|
||||
# Queue for new image tasks (the image paths)
|
||||
self.image_queue = queue.Queue()
|
||||
|
||||
# Registry for connected clients.
|
||||
# Processors are stored as: {client_socket: {'address': addr, 'status': 'idle'/'busy'}}
|
||||
self.processors = {}
|
||||
# (Optional) keep track of producers if needed.
|
||||
self.producers = []
|
||||
|
||||
# A lock to protect shared data structures
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def start(self):
|
||||
# Start a thread that continuously dispatches image tasks to idle processors.
|
||||
dispatcher_thread = threading.Thread(target=self.dispatcher, daemon=True)
|
||||
dispatcher_thread.start()
|
||||
|
||||
while True:
|
||||
client_socket, addr = self.server_socket.accept()
|
||||
print(f"[Server] New connection from {addr}")
|
||||
threading.Thread(target=self.handle_client, args=(client_socket, addr), daemon=True).start()
|
||||
|
||||
def handle_client(self, client_socket, addr):
|
||||
"""
|
||||
First message from the client must be its registration,
|
||||
for example: "ROLE:producer" or "ROLE:processor"
|
||||
"""
|
||||
try:
|
||||
reg_data = client_socket.recv(1024).decode().strip()
|
||||
if reg_data.startswith("ROLE:"):
|
||||
role = reg_data.split(":", 1)[1].strip()
|
||||
if role == "producer":
|
||||
with self.lock:
|
||||
self.producers.append(client_socket)
|
||||
print(f"[Server] Registered producer from {addr}")
|
||||
self.handle_producer(client_socket, addr)
|
||||
elif role == "processor":
|
||||
with self.lock:
|
||||
self.processors[client_socket] = {'address': addr, 'status': 'idle'}
|
||||
print(f"[Server] Registered processor from {addr}")
|
||||
self.handle_processor(client_socket, addr)
|
||||
else:
|
||||
print(f"[Server] Unknown role '{role}' from {addr}. Disconnecting.")
|
||||
client_socket.close()
|
||||
else:
|
||||
print(f"[Server] No valid registration from {addr}. Disconnecting.")
|
||||
client_socket.close()
|
||||
except Exception as e:
|
||||
print(f"[Server] Error during registration from {addr}: {e}")
|
||||
client_socket.close()
|
||||
|
||||
def handle_producer(self, client_socket, addr):
|
||||
"""
|
||||
Receives image paths from the producer and enqueues them.
|
||||
"""
|
||||
try:
|
||||
while True:
|
||||
data = client_socket.recv(1024).decode().strip()
|
||||
if not data:
|
||||
break # connection closed
|
||||
# Expect data to be an image path (or any image-notification)
|
||||
print(f"[Server] Received image path from producer {addr}: {data}")
|
||||
self.image_queue.put(data)
|
||||
except Exception as e:
|
||||
print(f"[Server] Producer {addr} error: {e}")
|
||||
finally:
|
||||
print(f"[Server] Producer {addr} disconnected.")
|
||||
with self.lock:
|
||||
if client_socket in self.producers:
|
||||
self.producers.remove(client_socket)
|
||||
client_socket.close()
|
||||
|
||||
def handle_processor(self, client_socket, addr):
|
||||
"""
|
||||
Listens for processor messages (e.g. a "DONE" notification after processing).
|
||||
"""
|
||||
try:
|
||||
while True:
|
||||
data = client_socket.recv(1024).decode().strip()
|
||||
if not data:
|
||||
break
|
||||
print(f"[Server] Received from processor {addr}: {data}")
|
||||
if data.upper() == "DONE":
|
||||
with self.lock:
|
||||
if client_socket in self.processors:
|
||||
self.processors[client_socket]['status'] = 'idle'
|
||||
print(f"[Server] Processor {addr} set to idle.")
|
||||
except Exception as e:
|
||||
print(f"[Server] Processor {addr} error: {e}")
|
||||
finally:
|
||||
print(f"[Server] Processor {addr} disconnected.")
|
||||
with self.lock:
|
||||
if client_socket in self.processors:
|
||||
del self.processors[client_socket]
|
||||
client_socket.close()
|
||||
|
||||
def dispatcher(self):
|
||||
"""
|
||||
Waits for new image tasks and assigns them to an idle processor.
|
||||
If no idle processor is available, waits (polling every 0.5 sec).
|
||||
"""
|
||||
while True:
|
||||
image_path = self.image_queue.get() # Wait for a new image task
|
||||
assigned = False
|
||||
while not assigned:
|
||||
with self.lock:
|
||||
# Find any idle processor
|
||||
idle_processors = [
|
||||
sock for sock, info in self.processors.items()
|
||||
if info['status'] == 'idle'
|
||||
]
|
||||
if idle_processors:
|
||||
processor_socket = idle_processors[0]
|
||||
try:
|
||||
processor_socket.sendall(image_path.encode())
|
||||
self.processors[processor_socket]['status'] = 'busy'
|
||||
print(f"[Server] Dispatched image '{image_path}' to processor {self.processors[processor_socket]['address']}")
|
||||
assigned = True
|
||||
except Exception as e:
|
||||
print(f"[Server] Error sending image to processor {self.processors[processor_socket]['address']}: {e}")
|
||||
# Remove this processor if there's an error
|
||||
del self.processors[processor_socket]
|
||||
# If no processor is idle, we simply wait a little before trying again.
|
||||
if not assigned:
|
||||
time.sleep(0.5)
|
||||
|
||||
if __name__ == "__main__":
|
||||
server = ImageServer(host="127.0.0.1", port=5000)
|
||||
server.start()
|
||||
Reference in New Issue
Block a user