Add working WebRTC code with image capturing!

.gitignore

@@ -91,4 +91,6 @@ temp/
 
 logs/
 
-aura/cnn/checkpoints/
+aura/cnn/checkpoints/
+
+*.pth

aura/webrtc/signaling.py

@@ -4,7 +4,7 @@
 import socket
 import signal
 import time
-from camera import ProcessingPipeline, FaceNotFoundException
+from aura.camera import ProcessingPipeline, FaceNotFoundException
 import os
 import numpy as np
 import cv2
@@ -14,57 +14,57 @@
 def get_free_port():
     """Get an unused TCP port."""
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(('', 0))
+        s.bind(("", 0))
         return s.getsockname()[1]
 
+
 def signal_handler(sig, frame):
     """Handle Ctrl+C gracefully"""
     print("\nShutting down signaling server...")
     sys.exit(0)
 
+
 def capture_images(server, output_dir="../logs", verbose=2):
     """Capture, process, and save images with face detection"""
     os.makedirs(output_dir, exist_ok=True)
-
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") 
-
-    pipeline = ProcessingPipeline(
-        log_path=output_dir,
-        verbose=verbose
-    )
-
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    pipeline = ProcessingPipeline(log_path=output_dir, verbose=verbose)
+
     server.capture()
     time.sleep(2)
-    
+
     image_bytes = server.get_capture()
     if image_bytes:
         nparr = np.frombuffer(image_bytes, np.uint8)
         image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-        
-        try:        
+
+        try:
             annotated_image = pipeline.annotate_face(image)
             processed_face = pipeline.process_image(image)
-            
-            # TODO: Send processed face to the model 
-            
+
+            # TODO: Send processed face to the model
+
         except FaceNotFoundException:
             print(f"No face detected in captured image.")
         except Exception as e:
             print(f"Error processing image: {str(e)}")
     else:
         print("No image captured")
 
+
 def main():
     signal.signal(signal.SIGINT, signal_handler)
 
-    port = 8765
-    
+    port = 3030
+
     server = SignalingServer(port=port)
     server.start()
-    
+
     print(f"Signaling server started on port {port}")
     print("Press Ctrl+C to stop the server")
-    
+
     try:
         while True:
             time.sleep(10)
@@ -73,7 +73,6 @@ def main():
     except KeyboardInterrupt:
         print("\nShutting down signaling server...")
 
+
 if __name__ == "__main__":
-    help(SignalingServer)
-    help(VideoStreamer)
-
+    main()

aura/webrtc/streamer.py

@@ -4,31 +4,33 @@
 import time
 from aura import VideoStreamer
 
+
 def signal_handler(sig, frame):
     """Handle Ctrl+C gracefully"""
     print("\nShutting down video streamer...")
     sys.exit(0)
 
+
 def main():
     signal.signal(signal.SIGINT, signal_handler)
 
     ws_ip = "127.0.0.1"  # WebSocket IP address
-    ws_port = 8765       # WebSocket port
-    ivf_dir = "../ivf_dir"    # Directory to watch for IVF files
+    ws_port = 3030  # WebSocket port
+    ivf_dir = "./ivf_files"  # Match RustWebRTC directory structure
 
     streamer = VideoStreamer(ws_ip, ws_port, ivf_dir)
     streamer.start_streaming()
-    
+
     print(f"Video streamer started - WebSocket server at ws://{ws_ip}:{ws_port}")
     print(f"Watching directory: {ivf_dir}")
     print("Press Ctrl+C to stop the streamer")
-    
+
     try:
         while True:
             time.sleep(1)
     except KeyboardInterrupt:
         print("\nShutting down video streamer...")
 
+
 if __name__ == "__main__":
     main()
-

scripts/convert_to_ivf.sh

@@ -1,5 +1,6 @@
 #!/bin/bash
 
+# Check if input file is provided
 if [ $# -ne 1 ]; then
     echo "Usage: $0 input.mp4"
     exit 1
@@ -8,11 +9,13 @@ fi
 input_file="$1"
 output_file="${input_file%.*}.ivf"
 
+# Check if input file exists
 if [ ! -f "$input_file" ]; then
     echo "Error: Input file '$input_file' not found"
     exit 1
 fi
 
+# Convert to IVF using VP8 codec
 ffmpeg -i "$input_file" -c:v libvpx -an -f ivf "$output_file"
 
 if [ $? -eq 0 ]; then

src/server.rs

@@ -100,6 +100,7 @@ impl SignalingServer {
 
         Ok(())
     }
+
     #[pyo3(text_signature = "(self, client_id: str, message: str) -> bool")]
     pub fn send_to_client(&self, client_id: String, message: String) -> PyResult<bool> {
         let peers = self.peers.clone();
@@ -243,9 +244,14 @@ async fn handle_connection(
         match result {
             Ok(msg) => {
                 if let Ok(text) = msg.to_str() {
+                    println!("Received message from {}: {}", client_id, text);
+
+                    // Attempt to parse the message
                     let signaling_message: Result<SignalingMessage, _> = serde_json::from_str(text);
                     match signaling_message {
                         Ok(SignalingMessage::Image { data }) => {
+                            // Handle image message
+                            println!("Handling image message from client {}", client_id);
                             handle_image_message(data.clone()).await;
 
                             if let Some(base64_data) = data.split(',').nth(1) {
@@ -258,16 +264,21 @@ async fn handle_connection(
                             }
                         }
                         Ok(message) => {
+                            // Handle other signaling messages
+                            println!("Parsed signaling message: {:?}", message);
                             forward_message(&client_id, &message, &peers).await;
                         }
                         Err(e) => {
-                            eprintln!("Error parsing message: {:?}", e);
+                            eprintln!(
+                                "Error parsing message from client {}: {} - Error: {:?}",
+                                client_id, text, e
+                            );
                         }
                     }
                 }
             }
             Err(e) => {
-                eprintln!("Error receiving message: {}", e);
+                eprintln!("Error receiving message for client {}: {}", client_id, e);
                 break;
             }
         }
@@ -302,6 +313,16 @@ async fn handle_image_message(data: String) {
     }
 }
 
+async fn trigger_image_capture(
+    sender: Arc<Mutex<SplitSink<WebSocket, Message>>>,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let message = serde_json::to_string(&SignalingMessage::TriggerImageCapture)?;
+    let mut sender = sender.lock().await;
+    sender.send(Message::text(message)).await?;
+    println!("Sent image capture trigger to client.");
+    Ok(())
+}
+
 async fn forward_message(sender_id: &str, message: &SignalingMessage, peers: &Peers) {
     let serialized_message = match serde_json::to_string(message) {
         Ok(json) => json,
@@ -311,7 +332,7 @@ async fn forward_message(sender_id: &str, message: &SignalingMessage, peers: &Pe
         }
     };
 
-    let peers = peers.lock().await;
+    let peers = peers.lock().await; // Await the async Mutex lock
     for (client_id, client) in peers.iter() {
         if client_id != sender_id {
             let mut client = client.lock().await; // Await the async Mutex lock
@@ -320,14 +341,4 @@ async fn forward_message(sender_id: &str, message: &SignalingMessage, peers: &Pe
             }
         }
     }
-}
-
-async fn trigger_image_capture(
-    sender: Arc<Mutex<SplitSink<WebSocket, Message>>>,
-) -> Result<(), Box<dyn std::error::Error>> {
-    let message = serde_json::to_string(&SignalingMessage::TriggerImageCapture)?;
-    let mut sender = sender.lock().await;
-    sender.send(Message::text(message)).await?;
-    println!("Sent image capture trigger to client.");
-    Ok(())
-}
+}