Update app.py

app.py

@@ -7,9 +7,8 @@ import time
 from dataclasses import dataclass
 import zipfile
 import logging
-import cv2
-from PIL import Image
-import numpy as np
+import av
+from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, WebRtcMode
 
 # Logging setup
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
@@ -190,6 +189,30 @@ def zip_files(files, zip_name):
             zipf.write(file, os.path.basename(file))
     return zip_name
 
+# Video Processor for WebRTC
+class CameraProcessor(VideoProcessorBase):
+    def __init__(self):
+        self.frame = None
+
+    def recv(self, frame):
+        from PIL import Image
+        img = frame.to_image()
+        self.frame = img
+        return av.VideoFrame.from_image(img)
+
+    def capture_frame(self):
+        from PIL import Image
+        return self.frame
+
+    def capture_video(self):
+        from PIL import Image
+        frames = []
+        start_time = time.time()
+        while time.time() - start_time < 10 and self.frame:
+            frames.append(np.array(self.frame))
+            time.sleep(0.033)  # ~30 FPS
+        return frames
+
 # Main App
 st.title("SFT Tiny Titans 🚀 (Dual Cam Action!)")
 
@@ -204,6 +227,7 @@ def update_gallery():
             for idx, file in enumerate(media_files[:4]):
                 with cols[idx % 2]:
                     if file.endswith(".png"):
+                        from PIL import Image
                         st.image(Image.open(file), caption=file.split('/')[-1], use_container_width=True)
                     elif file.endswith(".mp4"):
                         st.video(file)
@@ -249,82 +273,77 @@ with tab1:
 
 with tab2:
     st.header("Camera Snap 📷 (Dual Live Feed!)")
-    caps = {0: cv2.VideoCapture(0), 1: cv2.VideoCapture(1)}
     cols = st.columns(2)
+    processors = {}
     for i in range(2):
         with cols[i]:
             st.subheader(f"Camera {i}")
-            if caps[i].isOpened():
-                ret, frame = caps[i].read()
-                if ret:
-                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                    st.image(frame_rgb, caption=f"Live Feed Cam {i}", use_container_width=True)
-                else:
-                    st.warning(f"Camera {i} failed to read frame!")
-                    logger.error(f"Failed to read frame from Camera {i}")
-            else:
-                st.warning(f"Camera {i} not detected!")
-                logger.error(f"Camera {i} not opened")
-
+            key = f"camera_{i}"
+            processors[key] = webrtc_streamer(
+                key=key,
+                mode=WebRtcMode.SENDRECV,
+                video_processor_factory=CameraProcessor,
+                frontend_rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
+            )
             if st.button(f"Capture Frame 📸 Cam {i}", key=f"snap_{i}"):
                 logger.info(f"Capturing frame from Camera {i}")
                 try:
-                    ret, frame = caps[i].read()
-                    if ret:
-                        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                        img = Image.fromarray(frame_rgb)
+                    if processors[key].video_processor and processors[key].video_processor.frame:
+                        snapshot = processors[key].video_processor.capture_frame()
                         filename = generate_filename(i)
-                        img.save(filename)
-                        st.image(img, caption=filename, use_container_width=True)
+                        snapshot.save(filename)
+                        st.image(snapshot, caption=filename, use_container_width=True)
                         logger.info(f"Saved snapshot: {filename}")
                         if 'captured_images' not in st.session_state:
                             st.session_state['captured_images'] = []
                         st.session_state['captured_images'].append(filename)
                         update_gallery()
                     else:
-                        st.error("Failed to capture frame!")
+                        st.error("No frame available!")
                         logger.error(f"No frame captured from Camera {i}")
                 except Exception as e:
                     st.error(f"Frame capture failed: {str(e)}")
                     logger.error(f"Error capturing frame: {str(e)}")
-
             if st.button(f"Capture Video 🎥 Cam {i}", key=f"rec_{i}"):
                 logger.info(f"Capturing 10s video from Camera {i}")
                 try:
-                    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-                    mp4_filename = generate_filename(i, "mp4")
-                    out = cv2.VideoWriter(mp4_filename, fourcc, 30.0, (int(caps[i].get(3)), int(caps[i].get(4))))
-                    frames = []
-                    start_time = time.time()
-                    while time.time() - start_time < 10:
-                        ret, frame = caps[i].read()
-                        if ret:
-                            frames.append(frame)
-                            out.write(frame)
-                        time.sleep(0.033)  # ~30 FPS
-                    out.release()
-                    st.video(mp4_filename)
-                    logger.info(f"Saved video: {mp4_filename}")
-                    # Slice into 10 frames
-                    sliced_images = []
-                    step = max(1, len(frames) // 10)
-                    for j in range(0, len(frames), step):
-                        if len(sliced_images) < 10:
-                            frame_rgb = cv2.cvtColor(frames[j], cv2.COLOR_BGR2RGB)
-                            img = Image.fromarray(frame_rgb)
-                            img_filename = generate_filename(f"{i}_{len(sliced_images)}")
-                            img.save(img_filename)
-                            sliced_images.append(img_filename)
-                            st.image(img, caption=img_filename, use_container_width=True)
-                    st.session_state['captured_images'] = st.session_state.get('captured_images', []) + sliced_images
-                    logger.info(f"Sliced video into {len(sliced_images)} images")
-                    update_gallery()
+                    if processors[key].video_processor:
+                        frames = processors[key].video_processor.capture_video()
+                        if frames:
+                            mp4_filename = generate_filename(i, "mp4")
+                            with av.open(mp4_filename, "w") as container:
+                                stream = container.add_stream("h264", rate=30)
+                                stream.width = frames[0].shape[1]
+                                stream.height = frames[0].shape[0]
+                                for frame in frames:
+                                    av_frame = av.VideoFrame.from_ndarray(frame, format="rgb24")
+                                    for packet in stream.encode(av_frame):
+                                        container.mux(packet)
+                                for packet in stream.encode():
+                                    container.mux(packet)
+                            st.video(mp4_filename)
+                            logger.info(f"Saved video: {mp4_filename}")
+                            sliced_images = []
+                            step = max(1, len(frames) // 10)
+                            for j in range(0, len(frames), step):
+                                if len(sliced_images) < 10:
+                                    img = Image.fromarray(frames[j])
+                                    img_filename = generate_filename(f"{i}_{len(sliced_images)}")
+                                    img.save(img_filename)
+                                    sliced_images.append(img_filename)
+                                    st.image(img, caption=img_filename, use_container_width=True)
+                            st.session_state['captured_images'] = st.session_state.get('captured_images', []) + sliced_images
+                            logger.info(f"Sliced video into {len(sliced_images)} images")
+                            update_gallery()
+                        else:
+                            st.error("No frames recorded!")
+                            logger.error("No frames captured during video recording")
+                    else:
+                        st.error("Camera processor not initialized!")
+                        logger.error(f"Processor not ready for Camera {i}")
                 except Exception as e:
                     st.error(f"Video capture failed: {str(e)}")
                     logger.error(f"Error capturing video: {str(e)}")
-    # Release cameras after use
-    for cap in caps.values():
-        cap.release()
 
 with tab3:
     st.header("Fine-Tune Titans 🔧 (Tune Fast!)")
@@ -351,6 +370,7 @@ with tab3:
                 if st.button("Tune CV 🔄"):
                     logger.info("Initiating CV fine-tune")
                     try:
+                        from PIL import Image
                         images = [Image.open(img) for img in captured_images]
                         st.session_state['builder'].fine_tune(images, texts)
                         st.success("CV polished! 🎉")
@@ -381,6 +401,7 @@ with tab4:
                 if st.button("Run CV Demo ▶️"):
                     logger.info("Running CV image set demo")
                     try:
+                        from PIL import Image
                         images = [Image.open(img) for img in captured_images[:10]]
                         prompts = ["Neon " + os.path.basename(img).split('.')[0] for img in captured_images[:10]]
                         generated_images = []