Update app.py

app.py

@@ -107,7 +107,10 @@ def generate_charts(ner_output_bin: dict) -> Tuple[go.Figure, np.ndarray]:
     
     return fig2, wordcloud_image
 
-def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str]) -> np.ndarray:
+def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str], image_path: str) -> np.ndarray:
+
+    mask_image = np.array(Image.open(image_path))
+    
     token_texts = []
     token_scores = []
     token_types = []
@@ -128,7 +131,7 @@ def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str]) -> np.nd
         entity_type = next((t for t, w in zip(token_types, token_texts) if w == word), None)
         return color_map.get(entity_type, "#FFFFFF")
 
-    wordcloud = WordCloud(width=800, height=400, background_color='#121212', color_func=color_func).generate_from_frequencies(word_freq)
+    wordcloud = WordCloud(width=800, height=400, background_color='#121212', mask=mask_image, color_func=color_func).generate_from_frequencies(word_freq)
 
     # Convert to image array
     plt.figure(figsize=(10, 5))