Update app.py

app.py

@@ -1135,7 +1135,7 @@ class VideoGeneratorUI:
                 st.rerun()
 
     def show_image_selection(self):
-        st.subheader("Select Images for Your Video")
+        st.subheader("AI-Selected Images for Your Video")
         
         # Display back button
         if st.button("← Back to Script"):
@@ -1154,20 +1154,21 @@ class VideoGeneratorUI:
                     col1, col2 = st.columns([3, 1])
                     
                     with col1:
-                        # Display images with selection
-                        selected_count = 0
-                        st.session_state.selected_images = []  # Reset selected images
+                        selected_images = []
                         
                         for category, images in image_categories.items():
                             if images:
                                 st.subheader(f"{category.title()} Images")
                                 selected = self.display_image_grid(images)
-                                selected_count += len(selected)
-                                st.session_state.selected_images.extend(selected)
+                                selected_images.extend(selected)
+                        
+                        # Update session state with selected images
+                        st.session_state.selected_images = selected_images
                     
                     with col2:
                         # Show selection summary
                         st.subheader("Selection Summary")
+                        selected_count = len(st.session_state.selected_images)
                         st.write(f"Selected: {selected_count} images")
                         
                         if selected_count > 0:
@@ -1183,12 +1184,16 @@ class VideoGeneratorUI:
         except Exception as e:
             st.error(f"An error occurred: {str(e)}")
             print(f"Error in image selection: {str(e)}")
-
-
+    
+    
     def display_image_grid(self, images: List[Dict[str, str]], cols: int = 3) -> List[str]:
         """Display images in a grid and return selected image URLs"""
         selected_urls = []
         
+        # Initialize session state for image selection if not exists
+        if 'image_selections' not in st.session_state:
+            st.session_state.image_selections = {}
+        
         n_images = len(images)
         n_rows = (n_images + cols - 1) // cols
         
@@ -1204,18 +1209,27 @@ class VideoGeneratorUI:
                                 st.image(img['url'], use_container_width=True)
                                 
                                 # Create unique key for each checkbox
-                                checkbox_key = f"img_select_{row}_{col}_{img['url'][-10:]}"
+                                checkbox_key = f"img_select_{row}_{col}_{hash(img['url'])}"
+                                
+                                # Pre-select images based on AI confidence score
+                                default_selected = img.get('relevance_score', 0) > 0.6
+                                
+                                # Initialize checkbox state if not exists
+                                if checkbox_key not in st.session_state.image_selections:
+                                    st.session_state.image_selections[checkbox_key] = default_selected
                                 
+                                # Create checkbox with persistent state
                                 if st.checkbox(
-                                    f"Select (Confidence: {img.get('relevance_score', 0)*100:.1f}%)",
-                                    key=checkbox_key
+                                    f"Select (AI Confidence: {img.get('relevance_score', 0)*100:.1f}%)",
+                                    key=checkbox_key,
+                                    value=st.session_state.image_selections[checkbox_key]
                                 ):
                                     selected_urls.append(img['url'])
                                 
                                 # Show image metadata
                                 st.markdown(
                                     f"""<div style='font-size: 0.8em;'>
-                                    Keywords: {img['keyword']}<br>
+                                    Keywords: {img.get('keyword', 'N/A')}<br>
                                     Category: {img.get('category', 'General')}
                                     </div>""",
                                     unsafe_allow_html=True
@@ -1224,8 +1238,8 @@ class VideoGeneratorUI:
                                 st.error(f"Error displaying image: {str(e)}")
         
         return selected_urls
-
     
+        
     def show_video_generation(self):
         st.subheader("Video Generation Settings")