Update app.py

app.py

@@ -54,7 +54,7 @@ st.set_page_config(
 
 # Initialize st.session_state
 if 'captured_files' not in st.session_state:
-    st.session_state['captured_files'] = []
+    st.session_state['captured_files'] = {'cam0': None, 'cam1': None}  # Dictionary to store one file per camera
 if 'builder' not in st.session_state:
     st.session_state['builder'] = None
 if 'model_loaded' not in st.session_state:
@@ -430,25 +430,29 @@ async def process_pdf_snapshot(pdf_path, mode="thumbnail"):
     start_time = time.time()
     status = st.empty()
     status.text(f"Processing PDF Snapshot ({mode})... (0s)")
-    images = convert_from_path(pdf_path, dpi=200)
-    output_files = []
-    if mode == "thumbnail":
-        img = images[0].resize((int(images[0].width * 0.5), int(images[0].height * 0.5)), Image.Resampling.LANCZOS)
-        output_file = generate_filename("thumbnail", "png")
-        img.save(output_file)
-        output_files.append(output_file)
-    elif mode == "twopage":
-        for i in range(min(2, len(images))):
-            output_file = generate_filename(f"twopage_{i}", "png")
-            images[i].save(output_file)
+    try:
+        images = convert_from_path(pdf_path, dpi=200)
+        output_files = []
+        if mode == "thumbnail":
+            img = images[0].resize((int(images[0].width * 0.5), int(images[0].height * 0.5)), Image.Resampling.LANCZOS)
+            output_file = generate_filename("thumbnail", "png")
+            img.save(output_file)
             output_files.append(output_file)
-    elapsed = int(time.time() - start_time)
-    status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
-    for file in output_files:
-        if file not in st.session_state['captured_files']:
-            st.session_state['captured_files'].append(file)
-    update_gallery()
-    return output_files
+        elif mode == "twopage":
+            for i in range(min(2, len(images))):
+                output_file = generate_filename(f"twopage_{i}", "png")
+                images[i].save(output_file)
+                output_files.append(output_file)
+        elapsed = int(time.time() - start_time)
+        status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
+        for file in output_files:
+            if file not in st.session_state['captured_files'].values():
+                st.session_state['captured_files'][f"pdf_{len(output_files)}"] = file
+        update_gallery()
+        return output_files
+    except Exception as e:
+        status.error(f"Failed to process PDF: {str(e)}. Install poppler-utils (e.g., 'sudo apt-get install poppler-utils' on Ubuntu) and ensure it's in PATH.")
+        return []
 
 async def process_ocr(image, output_file):
     start_time = time.time()
@@ -461,8 +465,8 @@ async def process_ocr(image, output_file):
     status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
     async with aiofiles.open(output_file, "w") as f:
         await f.write(result)
-    if output_file not in st.session_state['captured_files']:
-        st.session_state['captured_files'].append(output_file)
+    if output_file not in st.session_state['captured_files'].values():
+        st.session_state['captured_files']['ocr'] = output_file
     update_gallery()
     return result
 
@@ -475,8 +479,8 @@ async def process_image_gen(prompt, output_file):
     elapsed = int(time.time() - start_time)
     status.text(f"Image Gen completed in {elapsed}s!")
     gen_image.save(output_file)
-    if output_file not in st.session_state['captured_files']:
-        st.session_state['captured_files'].append(output_file)
+    if output_file not in st.session_state['captured_files'].values():
+        st.session_state['captured_files']['gen'] = output_file
     update_gallery()
     return gen_image
 
@@ -492,8 +496,8 @@ async def process_custom_diffusion(images, output_file, model_name):
     elapsed = int(time.time() - start_time)
     status.text(f"{model_name} completed in {elapsed}s!")
     upscaled_image.save(output_file)
-    if output_file not in st.session_state['captured_files']:
-        st.session_state['captured_files'].append(output_file)
+    if output_file not in st.session_state['captured_files'].values():
+        st.session_state['captured_files']['diffusion'] = output_file
     update_gallery()
     return upscaled_image
 
@@ -504,17 +508,18 @@ st.title("AI Vision & SFT Titans 🚀")
 st.sidebar.header("Captured Files 📜")
 gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 4)
 def update_gallery():
-    media_files = get_gallery_files(["png", "txt"])
+    media_files = list(st.session_state['captured_files'].values())
     if media_files:
         cols = st.sidebar.columns(2)
         for idx, file in enumerate(media_files[:gallery_size * 2]):
-            with cols[idx % 2]:
-                if file.endswith(".png"):
-                    st.image(Image.open(file), caption=file, use_container_width=True)
-                elif file.endswith(".txt"):
-                    with open(file, "r") as f:
-                        content = f.read()
-                        st.text(content[:50] + "..." if len(content) > 50 else content, help=file)
+            if file and os.path.exists(file):  # Check if file exists
+                with cols[idx % 2]:
+                    if file.endswith(".png"):
+                        st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
+                    elif file.endswith(".txt"):
+                        with open(file, "r") as f:
+                            content = f.read()
+                            st.text(content[:50] + "..." if len(content) > 50 else content)
 update_gallery()
 
 st.sidebar.subheader("Model Management 🗂️")
@@ -554,50 +559,25 @@ with tab1:
     with cols[0]:
         cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
         if cam0_img:
-            filename = generate_filename(0)
-            if filename not in st.session_state['captured_files']:
-                with open(filename, "wb") as f:
-                    f.write(cam0_img.getvalue())
-                st.image(Image.open(filename), caption=filename, use_container_width=True)
-                logger.info(f"Saved snapshot from Camera 0: {filename}")
-                st.session_state['captured_files'].append(filename)
-                st.session_state['history'].append(f"Snapshot from Cam 0: {filename}")
-                update_gallery()
+            filename = generate_filename("cam0")
+            with open(filename, "wb") as f:
+                f.write(cam0_img.getvalue())
+            st.session_state['captured_files']['cam0'] = filename
+            st.image(Image.open(filename), caption="Camera 0", use_container_width=True)
+            logger.info(f"Saved snapshot from Camera 0: {filename}")
+            st.session_state['history'].append(f"Snapshot from Cam 0: {filename}")
+            update_gallery()
     with cols[1]:
         cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
         if cam1_img:
-            filename = generate_filename(1)
-            if filename not in st.session_state['captured_files']:
-                with open(filename, "wb") as f:
-                    f.write(cam1_img.getvalue())
-                st.image(Image.open(filename), caption=filename, use_container_width=True)
-                logger.info(f"Saved snapshot from Camera 1: {filename}")
-                st.session_state['captured_files'].append(filename)
-                st.session_state['history'].append(f"Snapshot from Cam 1: {filename}")
-                update_gallery()
-
-    st.subheader("Burst Capture")
-    slice_count = st.number_input("Number of Frames", min_value=1, max_value=20, value=10, key="burst_count")
-    if st.button("Start Burst Capture 📸"):
-        st.session_state['burst_frames'] = []
-        placeholder = st.empty()
-        for i in range(slice_count):
-            with placeholder.container():
-                st.write(f"Capturing frame {i+1}/{slice_count}...")
-                img = st.camera_input(f"Frame {i}", key=f"burst_{i}_{time.time()}")
-                if img:
-                    filename = generate_filename(f"burst_{i}")
-                    if filename not in st.session_state['captured_files']:
-                        with open(filename, "wb") as f:
-                            f.write(img.getvalue())
-                        st.session_state['burst_frames'].append(filename)
-                        logger.info(f"Saved burst frame {i}: {filename}")
-                        st.session_state['history'].append(f"Burst frame {i}: {filename}")
-                        st.image(Image.open(filename), caption=filename, use_container_width=True)
-                    time.sleep(0.5)
-        st.session_state['captured_files'].extend([f for f in st.session_state['burst_frames'] if f not in st.session_state['captured_files']])
-        update_gallery()
-        placeholder.success(f"Captured {len(st.session_state['burst_frames'])} frames!")
+            filename = generate_filename("cam1")
+            with open(filename, "wb") as f:
+                f.write(cam1_img.getvalue())
+            st.session_state['captured_files']['cam1'] = filename
+            st.image(Image.open(filename), caption="Camera 1", use_container_width=True)
+            logger.info(f"Saved snapshot from Camera 1: {filename}")
+            st.session_state['history'].append(f"Snapshot from Cam 1: {filename}")
+            update_gallery()
 
 with tab2:
     st.header("Download PDFs 📥")
@@ -671,9 +651,9 @@ with tab4:
                 st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Titan"), unsafe_allow_html=True)
                 st.rerun()
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
-            captured_files = get_gallery_files(["png"])
+            captured_files = list(st.session_state['captured_files'].values())
             if len(captured_files) >= 2:
-                demo_data = [{"image": img, "text": f"Superhero {os.path.basename(img).split('.')[0]}"} for img in captured_files[:min(len(captured_files), slice_count)]]
+                demo_data = [{"image": img, "text": f"Superhero {os.path.basename(img).split('.')[0]}"} for img in captured_files if img]
                 edited_data = st.data_editor(pd.DataFrame(demo_data), num_rows="dynamic")
                 if st.button("Fine-Tune with Dataset 🔄"):
                     images = [Image.open(row["image"]) for _, row in edited_data.iterrows()]
@@ -726,7 +706,7 @@ with tab5:
                 image = st.session_state['builder'].generate(test_prompt)
                 output_file = generate_filename("diffusion_test", "png")
                 image.save(output_file)
-                st.session_state['captured_files'].append(output_file)
+                st.session_state['captured_files']['diffusion_test'] = output_file
                 st.session_state['history'].append(f"Diffusion Test: {test_prompt} -> {output_file}")
                 st.image(image, caption="Generated Image")
                 update_gallery()
@@ -754,54 +734,56 @@ with tab6:
                     image = agent.generate(row["Image Idea"])
                     output_file = generate_filename(f"cv_rag_{row['Theme'].lower()}", "png")
                     image.save(output_file)
-                    st.session_state['captured_files'].append(output_file)
+                    st.session_state['captured_files'][f"cv_rag_{row['Theme'].lower()}"] = output_file
                     st.image(image, caption=f"{row['Theme']} - {row['Image Idea']}")
                 update_gallery()
 
 with tab7:
     st.header("Test OCR 🔍")
-    captured_files = get_gallery_files(["png"])
+    captured_files = list(st.session_state['captured_files'].values())
     if captured_files:
-        selected_file = st.selectbox("Select Image", captured_files, key="ocr_select")
-        image = Image.open(selected_file)
-        st.image(image, caption="Input Image", use_container_width=True)
-        if st.button("Run OCR 🚀", key="ocr_run"):
-            output_file = generate_filename("ocr_output", "txt")
-            st.session_state['processing']['ocr'] = True
-            result = asyncio.run(process_ocr(image, output_file))
-            st.session_state['history'].append(f"OCR Test: {selected_file} -> {output_file}")
-            st.text_area("OCR Result", result, height=200, key="ocr_result")
-            st.success(f"OCR output saved to {output_file}")
-            st.session_state['processing']['ocr'] = False
+        selected_file = st.selectbox("Select Image", [f for f in captured_files if f and f.endswith(".png")], key="ocr_select")
+        if selected_file:
+            image = Image.open(selected_file)
+            st.image(image, caption="Input Image", use_container_width=True)
+            if st.button("Run OCR 🚀", key="ocr_run"):
+                output_file = generate_filename("ocr_output", "txt")
+                st.session_state['processing']['ocr'] = True
+                result = asyncio.run(process_ocr(image, output_file))
+                st.session_state['history'].append(f"OCR Test: {selected_file} -> {output_file}")
+                st.text_area("OCR Result", result, height=200, key="ocr_result")
+                st.success(f"OCR output saved to {output_file}")
+                st.session_state['processing']['ocr'] = False
     else:
         st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")
 
 with tab8:
     st.header("Test Image Gen 🎨")
-    captured_files = get_gallery_files(["png"])
+    captured_files = list(st.session_state['captured_files'].values())
     if captured_files:
-        selected_file = st.selectbox("Select Image", captured_files, key="gen_select")
-        image = Image.open(selected_file)
-        st.image(image, caption="Reference Image", use_container_width=True)
-        prompt = st.text_area("Prompt", "Generate a similar superhero image", key="gen_prompt")
-        if st.button("Run Image Gen 🚀", key="gen_run"):
-            output_file = generate_filename("gen_output", "png")
-            st.session_state['processing']['gen'] = True
-            result = asyncio.run(process_image_gen(prompt, output_file))
-            st.session_state['history'].append(f"Image Gen Test: {prompt} -> {output_file}")
-            st.image(result, caption="Generated Image", use_container_width=True)
-            st.success(f"Image saved to {output_file}")
-            st.session_state['processing']['gen'] = False
+        selected_file = st.selectbox("Select Image", [f for f in captured_files if f and f.endswith(".png")], key="gen_select")
+        if selected_file:
+            image = Image.open(selected_file)
+            st.image(image, caption="Reference Image", use_container_width=True)
+            prompt = st.text_area("Prompt", "Generate a similar superhero image", key="gen_prompt")
+            if st.button("Run Image Gen 🚀", key="gen_run"):
+                output_file = generate_filename("gen_output", "png")
+                st.session_state['processing']['gen'] = True
+                result = asyncio.run(process_image_gen(prompt, output_file))
+                st.session_state['history'].append(f"Image Gen Test: {prompt} -> {output_file}")
+                st.image(result, caption="Generated Image", use_container_width=True)
+                st.success(f"Image saved to {output_file}")
+                st.session_state['processing']['gen'] = False
     else:
         st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")
 
 with tab9:
     st.header("Custom Diffusion 🎨🤓")
     st.write("Unleash your inner artist with our tiny diffusion models!")
-    captured_files = get_gallery_files(["png"])
+    captured_files = list(st.session_state['captured_files'].values())
     if captured_files:
         st.subheader("Select Images to Train")
-        selected_files = st.multiselect("Pick Images", captured_files, key="diffusion_select")
+        selected_files = st.multiselect("Pick Images", [f for f in captured_files if f and f.endswith(".png")], key="diffusion_select")
         images = [Image.open(file) for file in selected_files]
         
         model_options = [
@@ -822,7 +804,7 @@ with tab9:
                 builder.load_model(model_name)
                 result = builder.generate("A superhero scene inspired by captured images")
                 result.save(output_file)
-                st.session_state['captured_files'].append(output_file)
+                st.session_state['captured_files']['diffusion'] = output_file
             st.session_state['history'].append(f"Custom Diffusion: {model_choice} -> {output_file}")
             st.image(result, caption=f"{model_choice} Masterpiece", use_container_width=True)
             st.success(f"Image saved to {output_file}")