Upload 2 files

app.py

@@ -0,0 +1,269 @@
+import streamlit as st
+from streamlit_option_menu import option_menu
+
+# импортируем библиотеки
+import transformers
+from transformers import pipeline
+
+import torch
+import torch.multiprocessing as mp
+from diffusers import StableDiffusionPipeline
+
+from PIL import Image, ImageFilter
+import requests
+import numpy as np
+# import cairosvg
+from io import BytesIO
+
+import wikipedia
+from wikipedia.exceptions import DisambiguationError
+
+st.set_page_config(
+    page_title="WIKI: Imagination VS reality",
+    page_icon=":vs:",
+    layout="wide"
+)
+
+# +++++++++++++++++++++++++++++++++++++++++++++++++++++
+# +++++++++++++ Обязательные переменные +++++++++++++++
+# +++++++++++++++++++++++++++++++++++++++++++++++++++++
+# определяем доступное ядро
+config_device = "cuda:0" if torch.cuda.is_available() else "cpu"
+# для wiki надо указать в request user-agentа, иначе не открывает картинки
+headers = {'User-Agent': 'My User Agent 1.0'}
+
+# модель генерации картинок
+model_id = "runwayml/stable-diffusion-v1-5"
+
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+if torch.cuda.is_available(): 
+   pipe = pipe.to(config_device)
+   
+# ++++++++++++++++++++++++++++++++++++++++++++++++++++   
+# +++++++++++++ Настройки ++++++++++++++++++++++++++++
+# ++++++++++++++++++++++++++++++++++++++++++++++++++++
+# размеры картинок для вывода 
+GLOBAL_thumb_size = 128, 128
+# количество картинок в ряду коллажа
+GLOBAL_сollage_cols = 4
+# фон картинок если не вписываются в превью, 
+# если не задан в качестве фона используем размытое изображени
+GLOBAL_bg_color = (127, 127, 127)
+#GLOBAL_bg_color = ()
+
+# язык запроса в вики
+GLOBAL_lang = 'en'
+# количество статей в выдачи в вики
+GLOBAL_results = 1
+
+# ++++++++++++++++++++++++++++++++++++++++++++++++++++   
+# +++++++++++++ Функции ++++++++++++++++++++++++++++++
+# ++++++++++++++++++++++++++++++++++++++++++++++++++++
+# для обработки списка любой функцией
+def to_pool(item_list, method):             
+    rezult_list = list(map(method, (i for i in item_list)))
+    
+    return rezult_list
+ 
+# считываем контент с wiki-страницы
+def get_wiki_pages_content(article_name):    
+    rezult = dict()
+    try:
+        wiki_page = wikipedia.page(article_name, auto_suggest=False)
+    except DisambiguationError as e:
+        print("Не удалось прочесть статью:", article_name)
+        
+    else:                                  
+        content = wiki_page.content
+        rezult['title'] = wiki_page.title
+        paragraphs = list(filter(None, content.split("\n")))
+        # убераем служебные параграфы
+        paragraphs = list(filter(lambda x: not('==' in x), paragraphs))
+        rezult['content'] = paragraphs
+        # уберем звуковые файлы из выдачи
+        images = list(filter(lambda x: not(x.endswith(".ogg")), wiki_page.images))
+        
+        rezult['images'] = images
+
+    return rezult
+
+# ищем по запросу, или берем рандомную wiki-страницу
+def wiki_search(query, random = True, results=GLOBAL_results):
+    if random:
+       search_rezult = wikipedia.random(pages=results)       
+    else:
+       search_rezult = wikipedia.search(query, results=results, suggestion=False)
+
+    # для единообразия возвращаем список даже если запрос был на одну страницу 
+    return [search_rezult] if results == 1 else search_rezult
+ 
+# переводчик
+def rus2eng(txt):
+  rezult = translator(txt, max_length=400)
+  return rezult[0]['translation_text'] 
+
+# преведение картинок к заданному размеру для удобства коллажирования
+def resize_img(img, size = GLOBAL_thumb_size, bg_color = GLOBAL_bg_color):  
+    img.thumbnail(size)    
+    current_size = img.size
+    # если картинка не вписывается в квадрат, создаем фон из размытого изображения / или заданного цвета
+    if (current_size[0] < size[0]) | (current_size[1] < size[1]):
+       if bg_color:
+          new_img = Image.new('RGB', size, color = bg_color)
+       else:
+          new_img = img.filter(filter=ImageFilter.GaussianBlur)              
+          new_img = new_img.resize(size)  
+
+       cord_w = (size[0]//2) - current_size[0]//2
+       cord_h = (size[1]//2) - current_size[1]//2  
+       new_img.paste(img, box=(cord_w, cord_h))
+       return new_img
+
+    return img
+ 
+# генерация картинок
+def text2img(prompt, size = (512, 512)):
+    images = pipe(prompt, height=size[1], width=size[0]).images
+    
+    rezult = images[0] if len(images) == 1 else images
+    return rezult
+ 
+# читаем картинки из списка адресов
+def file2img(url):            
+    # if url.endswith(".svg"):
+    #    out = BytesIO()
+    #    cairosvg.svg2png(url=url, write_to=out)
+    #    image = Image.open(out)
+    #    file = out
+    # else:
+    file = requests.get(url, headers=headers, stream=True).raw
+       
+    try:       
+       image = Image.open(file)   
+       return image
+    except OSError:
+       #print("Не получилось конвертировать", url)
+       return None
+    
+# создание коллажа
+def create_collage(img_list, cols = GLOBAL_сollage_cols, size = GLOBAL_thumb_size):
+    thumb_width = size[0]
+    thumb_height = size[1]
+    # если список пустой - создаем пустую картинку заданной ширины
+    if len(img_list) == 0:
+       width = cols*thumb_width
+       height = cols*thumb_height
+
+       new_img = Image.new('RGB', (width, height))
+
+       return new_img
+
+    # определяем высоту и ширину коллажа    
+    # чтобы не подключать math ради одного округления вверх такая странная конструкция
+    rows = len(img_list) // cols if len(img_list) // cols == len(img_list) / cols else (len(img_list) // cols) + 1
+    cols = cols if cols < len(img_list) else len(img_list)
+
+    width = cols*thumb_width
+    height = rows*thumb_height
+
+    new_img = Image.new('RGB', (width, height))
+        
+    i, x, y = 0, 0, 0
+    for row in range(rows):        
+            if i == len(img_list):
+              break
+            for col in range(cols):
+                if i == len(img_list):
+                  break                    
+                new_img.paste(img_list[i], (x, y))
+                i += 1
+                x += thumb_height
+            y += thumb_width
+            x = 0
+
+    return new_img
+
+def main():
+   # поиск в Вики статей (по умолчанию - одной) по запросу или рандом   но
+   search_random = False if input_query else True
+   # установим выбранный язык
+   wiki_lang = input_lang if input_lang else GLOBAL_lang
+   wikipedia.set_lang(wiki_lang)   
+   with st.spinner('Ищем в википедии...'):
+       search_rezult = wiki_search(query = input_query, random=search_random)
+       # из полученной статей считываем картинки и текст
+       pages = to_pool(search_rezult, eval('get_wiki_pages_content'))       
+   st.success('Нашли, теперь повеселимся.')
+   
+   for page in pages:
+       # Получаем интересный нам контент
+       page = pages[0]
+       title = page['title']
+       pharagrafs = page['content']
+       images_urls = page['images']
+       wiki_text = '\n'.join(pharagrafs)
+       
+       # если для поиска использовалась русская вики, то переводим тексты
+       if (input_lang == 'ru') | (GLOBAL_lang == 'ru'):   
+           with st.spinner('Еще минутку - нужен перевод...'):
+                mname = "Helsinki-NLP/opus-mt-ru-en"
+                translator = pipeline("translation", model = mname)            
+                pharagrafs = to_pool(pharagrafs, eval('rus2eng'))       
+           st.success('Готово')
+           
+       # чтение реальных картинок со страницы   
+       with st.spinner('Реальность...'):
+            images_natur = to_pool(images_urls, eval('file2img'))
+            # убираем те, которые не смогли прочесть
+            images_natur = list(filter(None, images_natur))
+            # уменьшаем
+            images_natur_small = to_pool(images_natur, eval('resize_img'))
+            images_natur_collage = create_collage(images_natur_small)       
+       st.success('Готово')
+       
+       # генерация картинок по описанию
+       # можно было бы использовать для генерации размеры GLOBAL_thumb_size
+       # но практика показала, что качество сильно падает
+       # пришлось оставить связку генерация в размере 256*256 + уменьшение размера после
+          
+       with st.spinner('А теперь самое интересное...'):
+            images_gen = text2img(pharagrafs)
+            # уменьшаем
+            images_gen_small = to_pool(images_gen, eval('resize_img'))
+            images_gen_collage = create_collage(images_gen_small)       
+       st.success('Готово')
+       
+       # ++++++++++++++++++++++++++++++++++++++++++++++++++++   
+       # +++++++++++++ Вывод результатов ++++++++++++++++++++
+       # ++++++++++++++++++++++++++++++++++++++++++++++++++++
+       st.subheader(title)
+       
+       col1, col2 = st.columns(2)
+       
+       with col1:
+            st.image(images_gen_collage, caption='Ожидания')
+       with col2:
+            st.image(images_natur_collage, caption='Реальность')       
+            
+       with st.expander("Посмотреть текст"):
+            st.text(wiki_text)
+if __name__ == "__main__":
+# ++++++++++++++++++++++++++++++++++++++++++++++++++++   
+# +++++++++++++ Пользовательские импуты ++++++++++++++
+# ++++++++++++++++++++++++++++++++++++++++++++++++++++  
+   with st.sidebar:       
+        # поскольку переводчик у нас только на два языка надо делать переключателем en/ru
+        # в зависимости от выбора ищем в русскоязычной или англоязычной вики
+        input_lang = st.radio('Какую википедию использовать для поиска?:', ('en', 'ru'), index=0)
+        # запрос на поиск в Вики
+        input_query = st.text_input(
+           label = 'Вы хотите посмотреть на:', 
+           value= ''
+        )              
+   st.header('"WIKI Images: Ожидания vs Реальность"') 
+   st.text('Задача проекта визуализировать, насколько текст статей википедии соответствует иллюстрациям. Для этого на основе текста статьи генерится коллаж изображений (параграф = одно изображение). Второй коллаж формируется из реальных иллюстраций к статье.')
+   
+   st.text('Вы можете сами выбрать запрос (в сайдбаре) или довериться рандому.')
+   
+   if st.button('Начать'):
+      main() 

requirements.txt

@@ -0,0 +1,7 @@
+opencv-python==4.7.0.72
+sacremoses==0.0.53
+torch==2.0.0
+torchvision==0.15.1
+transformers
+sentencepiece
+