Update sozo_gen.py

sozo_gen.py

@@ -107,19 +107,19 @@ def detect_dataset_domain(df: pd.DataFrame) -> str:
 # NEW: Keyword extraction for better Pexels searches
 def extract_keywords_for_query(text: str, llm) -> str:
     prompt = f"""
-    Extract 2-4 key nouns and verbs from the following text to use as a search query for a stock video. 
+    Extract a maximum of 3 key nouns or verbs from the following text to use as a search query for a stock video. 
     Focus on concrete actions and subjects.
-    Example: 'Our analysis shows a significant growth in quarterly revenue and strong partnerships.' -> 'data analysis growth chart business'
+    Example: 'Our analysis shows a significant growth in quarterly revenue and strong partnerships.' -> 'data analysis growth'
     Output only the search query keywords, separated by spaces.
 
     Text: "{text}"
     """
     try:
         response = llm.invoke(prompt).content.strip()
-        return response
+        return response if response else text
     except Exception as e:
         logging.error(f"Keyword extraction failed: {e}. Using original text.")
-        return text # Fallback to the original text if LLM fails
+        return text
 
 # UPDATED: Pexels search now loops short videos
 def search_and_download_pexels_video(query: str, duration: float, out_path: Path) -> str:
@@ -265,7 +265,7 @@ def prepare_plot_data(spec: ChartSpecification, df: pd.DataFrame):
 
 def animate_chart(spec: ChartSpecification, df: pd.DataFrame, dur: float, out: Path, fps: int = FPS) -> str:
     plot_data = prepare_plot_data(spec, df)
-    frames = math.ceil(dur * fps) # Use math.ceil to always round up frames
+    frames = math.ceil(dur * fps)
     fig, ax = plt.subplots(figsize=(WIDTH / 100, HEIGHT / 100), dpi=100)
     plt.tight_layout(pad=3.0)
     ctype = spec.chart_type
@@ -275,20 +275,33 @@ def animate_chart(spec: ChartSpecification, df: pd.DataFrame, dur: float, out: P
     if ctype == "line":
         plot_data = plot_data.sort_index()
         x_full, y_full = plot_data.index, plot_data.values
+        
+        ax.set_xlim(x_full.min(), x_full.max())
+        ax.set_ylim(y_full.min() * 0.9, y_full.max() * 1.1)
+        ax.set_title(spec.title); ax.grid(alpha=.3); ax.set_xlabel(spec.x_col); ax.set_ylabel(spec.y_col)
+        
+        line, = ax.plot([], [], lw=2, color='#A23B72')
+        markers, = ax.plot([], [], 'o', color='#A23B72', markersize=5)
+
         def init():
-            ax.set_xlim(x_full.min(), x_full.max())
-            ax.set_ylim(y_full.min() * 0.9, y_full.max() * 1.1)
-            return []
+            line.set_data([], [])
+            markers.set_data([], [])
+            return line, markers
         def update(i):
-            ax.clear()
-            ax.set_xlim(x_full.min(), x_full.max())
-            ax.set_ylim(y_full.min() * 0.9, y_full.max() * 1.1)
-            ax.set_title(spec.title); ax.grid(alpha=.3); ax.set_xlabel(spec.x_col); ax.set_ylabel(spec.y_col)
             k = max(2, int(len(x_full) * (i / (frames - 1))))
-            ax.plot(x_full[:k], y_full[:k], lw=2, color='#A23B72', marker='o', markersize=5)
-            return []
+            line.set_data(x_full[:k], y_full[:k])
+            markers.set_data(x_full[:k], y_full[:k])
+            return line, markers
         init_func, update_func = init, update
-    elif ctype == "bar":
+        
+        anim = FuncAnimation(fig, update, init_func=init, frames=frames, blit=True, interval=1000 / fps)
+        anim.save(str(out), writer=FFMpegWriter(fps=fps), dpi=144)
+        plt.close(fig)
+        return str(out)
+
+    # Fallback to the slightly slower but reliable blit=False for other types
+    # This ensures stability across all chart types while the line chart is optimized
+    if ctype == "bar":
         bars = ax.bar(plot_data.index.astype(str), np.zeros_like(plot_data.values, dtype=float), color="#1f77b4")
         ax.set_ylim(0, plot_data.max() * 1.1 if not pd.isna(plot_data.max()) and plot_data.max() > 0 else 1)
         ax.set_title(spec.title); plt.xticks(rotation=45, ha="right")
@@ -308,7 +321,7 @@ def animate_chart(spec: ChartSpecification, df: pd.DataFrame, dur: float, out: P
         ax.set_title(spec.title); ax.grid(alpha=.3); ax.set_xlabel(spec.x_col); ax.set_ylabel(spec.y_col)
         def init():
             scat.set_offsets(np.empty((0, 2))); line.set_data([], [])
-            return [scat, line]
+            return []
         def update(i):
             point_frames = int(frames * 0.7)
             if i <= point_frames:
@@ -318,25 +331,25 @@ def animate_chart(spec: ChartSpecification, df: pd.DataFrame, dur: float, out: P
                 line_frame = i - point_frames; line_total_frames = frames - point_frames
                 current_x = reg_line_x[0] + (reg_line_x[1] - reg_line_x[0]) * (line_frame / line_total_frames)
                 line.set_data([reg_line_x[0], current_x], [reg_line_y[0], slope * current_x + intercept])
-            return [scat, line]
+            return []
         init_func, update_func = init, update
     elif ctype == "pie":
         wedges, _, _ = ax.pie(plot_data, labels=plot_data.index, startangle=90, autopct='%1.1f%%')
         ax.set_title(spec.title); ax.axis('equal')
-        def init(): [w.set_alpha(0) for w in wedges]; return wedges
-        def update(i): [w.set_alpha(i / (frames - 1)) for w in wedges]; return wedges
+        def init(): [w.set_alpha(0) for w in wedges]; return []
+        def update(i): [w.set_alpha(i / (frames - 1)) for w in wedges]; return []
         init_func, update_func = init, update
     elif ctype == "hist":
         _, _, patches = ax.hist(plot_data, bins=20, alpha=0)
         ax.set_title(spec.title); ax.set_xlabel(spec.x_col); ax.set_ylabel("Frequency")
-        def init(): [p.set_alpha(0) for p in patches]; return patches
-        def update(i): [p.set_alpha((i / (frames - 1)) * 0.7) for p in patches]; return patches
+        def init(): [p.set_alpha(0) for p in patches]; return []
+        def update(i): [p.set_alpha((i / (frames - 1)) * 0.7) for p in patches]; return []
         init_func, update_func = init, update
     elif ctype == "heatmap":
         sns.heatmap(plot_data, annot=True, cmap="viridis", ax=ax, alpha=0)
         ax.set_title(spec.title)
-        def init(): ax.collections[0].set_alpha(0); return [ax.collections[0]]
-        def update(i): ax.collections[0].set_alpha(i / (frames - 1)); return [ax.collections[0]]
+        def init(): ax.collections[0].set_alpha(0); return []
+        def update(i): ax.collections[0].set_alpha(i / (frames - 1)); return []
         init_func, update_func = init, update
     else:
         ax.text(0.5, 0.5, f"'{ctype}' animation not implemented", ha='center', va='center')
@@ -348,16 +361,6 @@ def animate_chart(spec: ChartSpecification, df: pd.DataFrame, dur: float, out: P
     anim.save(str(out), writer=FFMpegWriter(fps=fps), dpi=144)
     plt.close(fig)
     return str(out)
-    
-def animate_image_fade(img: np.ndarray, dur: float, out: Path, fps: int = 24) -> str:
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v'); video_writer = cv2.VideoWriter(str(out), fourcc, fps, (WIDTH, HEIGHT))
-    total_frames = max(1, int(dur * fps))
-    for i in range(total_frames):
-        alpha = i / (total_frames - 1) if total_frames > 1 else 1.0
-        frame = cv2.addWeighted(img, alpha, np.zeros_like(img), 1 - alpha, 0)
-        video_writer.write(frame)
-    video_writer.release()
-    return str(out)
 
 def safe_chart(desc: str, df: pd.DataFrame, dur: float, out: Path, context: Dict) -> str:
     try:
@@ -608,38 +611,32 @@ def generate_video_from_project(df: pd.DataFrame, raw_md: str, data_context: Dic
         video_dur = audio_dur + 1.5 
         
         try:
-            primary_query = None
-            narration_lower = narrative.lower()
-            is_conclusion_scene = any(k in narration_lower for k in ["conclusion", "summary", "in closing", "final thoughts"])
-
-            if any(k in narration_lower for k in ["introduction", "welcome", "let's begin"]):
-                primary_query = f"abstract technology background {domain}"
-            elif is_conclusion_scene:
-                primary_query = f"future strategy business meeting {domain}"
-
-            if primary_query:
-                logging.info(f"Scene {i+1}: Pre-emptive guard triggered. Query: '{primary_query}'")
-                video_path = search_and_download_pexels_video(primary_query, video_dur, mp4)
-                if not video_path: raise ValueError("Pexels search failed for guarded query.")
+            chart_descs = extract_chart_tags(sc)
+            pexels_descs = extract_pexels_tags(sc)
+            is_conclusion_scene = any(k in narrative.lower() for k in ["conclusion", "summary", "in closing", "final thoughts"])
+
+            if pexels_descs:
+                logging.info(f"Scene {i+1}: Processing Pexels scene.")
+                base_keywords = extract_keywords_for_query(narrative, llm)
+                final_query = f"{base_keywords} {domain}"
+                video_path = search_and_download_pexels_video(final_query, video_dur, mp4)
+                if not video_path: raise ValueError("Pexels search returned no results for chained query.")
                 video_parts.append(video_path)
                 if is_conclusion_scene:
                     conclusion_video_path = video_path
+            elif chart_descs:
+                logging.info(f"Scene {i+1}: Primary attempt with animated chart.")
+                safe_chart(chart_descs[0], df, video_dur, mp4, data_context)
+                video_parts.append(str(mp4))
             else:
-                chart_descs = extract_chart_tags(sc)
-                if chart_descs:
-                    logging.info(f"Scene {i+1}: Primary attempt with animated chart.")
-                    safe_chart(chart_descs[0], df, video_dur, mp4, data_context)
-                    video_parts.append(str(mp4))
-                else:
-                    raise ValueError("No chart tag found in a middle scene.")
+                raise ValueError("No visual tag found in scene.")
         except Exception as e:
             logging.warning(f"Scene {i+1}: Primary visual failed ({e}). Marking for fallback.")
             video_parts.append("FALLBACK_NEEDED")
 
         temps.append(mp4)
 
-    # Post-processing loop to apply the conclusion video as a fallback
-    if not conclusion_video_path: # Failsafe if conclusion scene itself failed
+    if not conclusion_video_path:
         logging.warning("No conclusion video was generated; creating a generic one for fallbacks.")
         fallback_mp4 = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.mp4"
         conclusion_video_path = search_and_download_pexels_video(f"data visualization abstract {domain}", 5.0, fallback_mp4)