Update app.py

app.py

@@ -7,11 +7,13 @@ import time
 from dataclasses import dataclass
 import zipfile
 import logging
+import av
+from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, WebRtcMode
 
-# Logging setup with custom log storage
+# Logging setup
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
-log_records = []  # Custom list to store logs
+log_records = []
 
 class LogCaptureHandler(logging.Handler):
     def emit(self, record):
@@ -170,12 +172,12 @@ def get_download_link(file_path, mime_type="text/plain", label="Download"):
     b64 = base64.b64encode(data).decode()
     return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label} 📥</a>'
 
-def generate_filename(sequence):
+def generate_filename(sequence, ext="png"):
     from datetime import datetime
     import pytz
     central = pytz.timezone('US/Central')
     timestamp = datetime.now(central).strftime("%d%m%Y%H%M%S%p")
-    return f"{sequence}{timestamp}.png"
+    return f"{sequence}{timestamp}.{ext}"
 
 def get_gallery_files(file_types):
     import glob
@@ -188,29 +190,52 @@ def zip_files(files, zip_name):
     return zip_name
 
 # Video Processor for WebRTC
-class VideoSnapshot:
+class CameraProcessor(VideoProcessorBase):
     def __init__(self):
         self.snapshot = None
+        self.recording = False
+        self.frames = []
+        self.start_time = None
+
     def recv(self, frame):
         from PIL import Image
         img = frame.to_image()
         self.snapshot = img
-        return frame
+        if self.recording and time.time() - self.start_time < 10:
+            self.frames.append(frame.to_ndarray(format="bgr24"))
+        return av.VideoFrame.from_image(img)
+
     def take_snapshot(self):
+        from PIL import Image
         return self.snapshot
 
+    def start_recording(self):
+        self.recording = True
+        self.frames = []
+        self.start_time = time.time()
+
+    def stop_recording(self):
+        self.recording = False
+        return self.frames
+
 # Main App
-st.title("SFT Tiny Titans 🚀 (Capture & Tune!)")
+st.title("SFT Tiny Titans 🚀 (Dual Cam Action!)")
 
 # Sidebar Galleries
-st.sidebar.header("Captured Images 🎨")
-image_files = get_gallery_files(["png"])
-if image_files:
-    cols = st.sidebar.columns(2)
-    for idx, file in enumerate(image_files[:4]):
-        with cols[idx % 2]:
-            from PIL import Image
-            st.image(Image.open(file), caption=file.split('/')[-1], use_container_width=True)
+st.sidebar.header("Captured Media 🎨")
+gallery_container = st.sidebar.empty()
+def update_gallery():
+    media_files = get_gallery_files(["png", "mp4"])
+    with gallery_container:
+        if media_files:
+            cols = st.columns(2)
+            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)
 
 # Sidebar Model Management
 st.sidebar.subheader("Model Hub 🗂️")
@@ -252,63 +277,83 @@ with tab1:
                 st.error(f"Download failed: {str(e)}")
 
 with tab2:
-    st.header("Camera Snap 📷 (Sequence Shots!)")
-    from streamlit_webrtc import webrtc_streamer
-    ctx = webrtc_streamer(
-        key="camera",
-        video_processor_factory=VideoSnapshot,
-        frontend_rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
-    )
-    if ctx.video_processor:
-        delay = st.slider("Delay between captures (seconds)", 0, 10, 2)
-        if st.button("Capture 6 Frames 📸"):
-            logger.info("Starting 6-frame capture")
-            captured_images = []
-            try:
-                for i in range(6):
-                    snapshot = ctx.video_processor.take_snapshot()
-                    if snapshot:
-                        filename = generate_filename(i)
-                        snapshot.save(filename)
-                        st.image(snapshot, caption=filename, use_container_width=True)
-                        captured_images.append(filename)
-                        logger.info(f"Captured frame {i}: {filename}")
-                    time.sleep(delay)
-                st.success("6 frames captured! 🎉")
-                st.session_state['captured_images'] = captured_images
-            except Exception as e:
-                st.error(f"Capture failed: {str(e)}")
-                logger.error(f"Error during capture: {str(e)}")
-
-    if 'captured_images' in st.session_state and len(st.session_state['captured_images']) >= 2:
-        st.subheader("Diffusion SFT Dataset 🎨")
-        sample_texts = ["Neon Hero", "Glowing Cape", "Spark Flyer", "Dark Knight", "Iron Shine", "Thunder Bolt"]
-        dataset = list(zip(st.session_state['captured_images'], sample_texts[:len(st.session_state['captured_images'])]))
-        st.code("\n".join([f"{i+1}. {text} -> {img}" for i, (img, text) in enumerate(dataset)]), language="text")
-        if st.button("Download Dataset CSV 📝"):
-            logger.info("Generating dataset CSV")
-            try:
-                csv_path = f"diffusion_sft_{int(time.time())}.csv"
-                with open(csv_path, "w", newline="") as f:
-                    writer = csv.writer(f)
-                    writer.writerow(["image", "text"])
-                    for img, text in dataset:
-                        writer.writerow([img, text])
-                st.markdown(get_download_link(csv_path, "text/csv", "Download Dataset CSV"), unsafe_allow_html=True)
-                logger.info("Dataset CSV generated")
-            except Exception as e:
-                st.error(f"CSV generation failed: {str(e)}")
-                logger.error(f"Error generating CSV: {str(e)}")
-        if st.button("Download Images ZIP 📦"):
-            logger.info("Generating images ZIP")
-            try:
-                zip_path = f"captured_images_{int(time.time())}.zip"
-                zip_files(st.session_state['captured_images'], zip_path)
-                st.markdown(get_download_link(zip_path, "application/zip", "Download Images ZIP"), unsafe_allow_html=True)
-                logger.info("Images ZIP generated")
-            except Exception as e:
-                st.error(f"ZIP generation failed: {str(e)}")
-                logger.error(f"Error generating ZIP: {str(e)}")
+    st.header("Camera Snap 📷 (Dual Live Feed!)")
+    cols = st.columns(2)
+    processors = {}
+    for i in range(2):
+        with cols[i]:
+            st.subheader(f"Camera {i}")
+            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 processors[key].video_processor:
+                if st.button(f"Capture 📸 Cam {i}", key=f"snap_{i}"):
+                    logger.info(f"Capturing snapshot from Camera {i}")
+                    try:
+                        snapshot = processors[key].video_processor.take_snapshot()
+                        if snapshot:
+                            filename = generate_filename(i)
+                            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()
+                    except Exception as e:
+                        st.error(f"Snapshot failed: {str(e)}")
+                        logger.error(f"Error capturing snapshot: {str(e)}")
+                record_key = f"record_{i}"
+                if record_key not in st.session_state:
+                    st.session_state[record_key] = False
+                if st.button(f"{'Stop' if st.session_state[record_key] else 'Record'} 🎥 Cam {i}", key=f"rec_{i}"):
+                    if not st.session_state[record_key]:
+                        logger.info(f"Starting recording from Camera {i}")
+                        try:
+                            processors[key].video_processor.start_recording()
+                            st.session_state[record_key] = True
+                        except Exception as e:
+                            st.error(f"Start recording failed: {str(e)}")
+                            logger.error(f"Error starting recording: {str(e)}")
+                    else:
+                        logger.info(f"Stopping recording from Camera {i}")
+                        try:
+                            frames = processors[key].video_processor.stop_recording()
+                            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="bgr24")
+                                        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}")
+                                # 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:
+                                        img = Image.fromarray(frames[j][:, :, ::-1])  # BGR to 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()
+                            st.session_state[record_key] = False
+                        except Exception as e:
+                            st.error(f"Stop recording failed: {str(e)}")
+                            logger.error(f"Error stopping recording: {str(e)}")
 
 with tab3:
     st.header("Fine-Tune Titans 🔧 (Tune Fast!)")
@@ -345,7 +390,7 @@ with tab3:
                 st.warning("Capture at least 2 images first! ⚠️")
 
 with tab4:
-    st.header("Test Titans 🧪 (Quick Check!)")
+    st.header("Test Titans 🧪 (Image Agent Demo!)")
     if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
         st.warning("Load a Titan first! ⚠️")
     else:
@@ -360,20 +405,43 @@ with tab4:
                 except Exception as e:
                     st.error(f"NLP test failed: {str(e)}")
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
-            st.subheader("CV Test 🎨")
-            prompt = st.text_area("Prompt", "Neon Batman", key="cv_test")
-            if st.button("Test CV ▶️"):
-                logger.info("Running CV test")
-                try:
-                    with st.spinner("Generating... ⏳"):
-                        img = st.session_state['builder'].generate(prompt)
-                    st.image(img, caption="Generated Art", use_container_width=True)
-                except Exception as e:
-                    st.error(f"CV test failed: {str(e)}")
+            st.subheader("CV Test 🎨 (Image Set Demo)")
+            captured_images = get_gallery_files(["png"])
+            if len(captured_images) >= 2:
+                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 = []
+                        for prompt in prompts:
+                            img = st.session_state['builder'].generate(prompt)
+                            generated_images.append(img)
+                        cols = st.columns(2)
+                        for idx, (orig, gen) in enumerate(zip(images, generated_images)):
+                            with cols[idx % 2]:
+                                st.image(orig, caption=f"Original: {captured_images[idx]}", use_container_width=True)
+                                st.image(gen, caption=f"Generated: {prompts[idx]}", use_container_width=True)
+                        md_content = "# Image Set Demo\n\nScript of filenames and descriptions:\n"
+                        for i, (img, prompt) in enumerate(zip(captured_images[:10], prompts)):
+                            md_content += f"{i+1}. `{img}` - {prompt}\n"
+                        md_filename = f"demo_metadata_{int(time.time())}.md"
+                        with open(md_filename, "w") as f:
+                            f.write(md_content)
+                        st.markdown(get_download_link(md_filename, "text/markdown", "Download Metadata .md"), unsafe_allow_html=True)
+                        logger.info("CV demo completed with metadata")
+                    except Exception as e:
+                        st.error(f"CV demo failed: {str(e)}")
+                        logger.error(f"Error in CV demo: {str(e)}")
+            else:
+                st.warning("Capture at least 2 images first! ⚠️")
 
 # Display Logs
 st.sidebar.subheader("Action Logs 📜")
 log_container = st.sidebar.empty()
 with log_container:
     for record in log_records:
-        st.write(f"{record.asctime} - {record.levelname} - {record.message}")
+        st.write(f"{record.asctime} - {record.levelname} - {record.message}")
+
+update_gallery()  # Initial gallery update