julianna-fil commited on
Commit
b9c2b0f
1 Parent(s): 0cced28

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +269 -0
  2. requirements.txt +7 -0
app.py ADDED
@@ -0,0 +1,269 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from streamlit_option_menu import option_menu
3
+
4
+ # импортируем библиотеки
5
+ import transformers
6
+ from transformers import pipeline
7
+
8
+ import torch
9
+ import torch.multiprocessing as mp
10
+ from diffusers import StableDiffusionPipeline
11
+
12
+ from PIL import Image, ImageFilter
13
+ import requests
14
+ import numpy as np
15
+ # import cairosvg
16
+ from io import BytesIO
17
+
18
+ import wikipedia
19
+ from wikipedia.exceptions import DisambiguationError
20
+
21
+ st.set_page_config(
22
+ page_title="WIKI: Imagination VS reality",
23
+ page_icon=":vs:",
24
+ layout="wide"
25
+ )
26
+
27
+ # +++++++++++++++++++++++++++++++++++++++++++++++++++++
28
+ # +++++++++++++ Обязательные переменные +++++++++++++++
29
+ # +++++++++++++++++++++++++++++++++++++++++++++++++++++
30
+ # определяем доступное ядро
31
+ config_device = "cuda:0" if torch.cuda.is_available() else "cpu"
32
+ # для wiki надо указать в request user-agentа, иначе не открывает картинки
33
+ headers = {'User-Agent': 'My User Agent 1.0'}
34
+
35
+ # модель генерации картинок
36
+ model_id = "runwayml/stable-diffusion-v1-5"
37
+
38
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
39
+ if torch.cuda.is_available():
40
+ pipe = pipe.to(config_device)
41
+
42
+ # ++++++++++++++++++++++++++++++++++++++++++++++++++++
43
+ # +++++++++++++ Настройки ++++++++++++++++++++++++++++
44
+ # ++++++++++++++++++++++++++++++++++++++++++++++++++++
45
+ # размеры картинок для вывода
46
+ GLOBAL_thumb_size = 128, 128
47
+ # количество картинок в ряду коллажа
48
+ GLOBAL_сollage_cols = 4
49
+ # фон картинок если не вписываются в превью,
50
+ # если не задан в качестве фона используем размытое изображени
51
+ GLOBAL_bg_color = (127, 127, 127)
52
+ #GLOBAL_bg_color = ()
53
+
54
+ # язык запроса в вики
55
+ GLOBAL_lang = 'en'
56
+ # количество статей в выдачи в вики
57
+ GLOBAL_results = 1
58
+
59
+ # ++++++++++++++++++++++++++++++++++++++++++++++++++++
60
+ # +++++++++++++ Функции ++++++++++++++++++++++++++++++
61
+ # ++++++++++++++++++++++++++++++++++++++++++++++++++++
62
+ # для обработки списка любой функцией
63
+ def to_pool(item_list, method):
64
+ rezult_list = list(map(method, (i for i in item_list)))
65
+
66
+ return rezult_list
67
+
68
+ # считываем контент с wiki-страницы
69
+ def get_wiki_pages_content(article_name):
70
+ rezult = dict()
71
+ try:
72
+ wiki_page = wikipedia.page(article_name, auto_suggest=False)
73
+ except DisambiguationError as e:
74
+ print("Не удалось прочесть статью:", article_name)
75
+
76
+ else:
77
+ content = wiki_page.content
78
+ rezult['title'] = wiki_page.title
79
+ paragraphs = list(filter(None, content.split("\n")))
80
+ # убераем служебные параграфы
81
+ paragraphs = list(filter(lambda x: not('==' in x), paragraphs))
82
+ rezult['content'] = paragraphs
83
+ # уберем звуковые файлы из выдачи
84
+ images = list(filter(lambda x: not(x.endswith(".ogg")), wiki_page.images))
85
+
86
+ rezult['images'] = images
87
+
88
+ return rezult
89
+
90
+ # ищем по запросу, или берем рандомную wiki-страницу
91
+ def wiki_search(query, random = True, results=GLOBAL_results):
92
+ if random:
93
+ search_rezult = wikipedia.random(pages=results)
94
+ else:
95
+ search_rezult = wikipedia.search(query, results=results, suggestion=False)
96
+
97
+ # для единообразия возвращаем список даже если запрос был на одну страницу
98
+ return [search_rezult] if results == 1 else search_rezult
99
+
100
+ # переводчик
101
+ def rus2eng(txt):
102
+ rezult = translator(txt, max_length=400)
103
+ return rezult[0]['translation_text']
104
+
105
+ # преведение картинок к заданному размеру для удобства коллажирования
106
+ def resize_img(img, size = GLOBAL_thumb_size, bg_color = GLOBAL_bg_color):
107
+ img.thumbnail(size)
108
+ current_size = img.size
109
+ # если картинка не вписывается в квадрат, создаем фон из размытого изображения / или заданного цвета
110
+ if (current_size[0] < size[0]) | (current_size[1] < size[1]):
111
+ if bg_color:
112
+ new_img = Image.new('RGB', size, color = bg_color)
113
+ else:
114
+ new_img = img.filter(filter=ImageFilter.GaussianBlur)
115
+ new_img = new_img.resize(size)
116
+
117
+ cord_w = (size[0]//2) - current_size[0]//2
118
+ cord_h = (size[1]//2) - current_size[1]//2
119
+ new_img.paste(img, box=(cord_w, cord_h))
120
+ return new_img
121
+
122
+ return img
123
+
124
+ # генерация картинок
125
+ def text2img(prompt, size = (512, 512)):
126
+ images = pipe(prompt, height=size[1], width=size[0]).images
127
+
128
+ rezult = images[0] if len(images) == 1 else images
129
+ return rezult
130
+
131
+ # читаем картинки из списка адресов
132
+ def file2img(url):
133
+ # if url.endswith(".svg"):
134
+ # out = BytesIO()
135
+ # cairosvg.svg2png(url=url, write_to=out)
136
+ # image = Image.open(out)
137
+ # file = out
138
+ # else:
139
+ file = requests.get(url, headers=headers, stream=True).raw
140
+
141
+ try:
142
+ image = Image.open(file)
143
+ return image
144
+ except OSError:
145
+ #print("Не получилось конвертировать", url)
146
+ return None
147
+
148
+ # создание коллажа
149
+ def create_collage(img_list, cols = GLOBAL_сollage_cols, size = GLOBAL_thumb_size):
150
+ thumb_width = size[0]
151
+ thumb_height = size[1]
152
+ # если список пустой - создаем пустую картинку заданной ширины
153
+ if len(img_list) == 0:
154
+ width = cols*thumb_width
155
+ height = cols*thumb_height
156
+
157
+ new_img = Image.new('RGB', (width, height))
158
+
159
+ return new_img
160
+
161
+ # определяем высоту и ширину коллажа
162
+ # чтобы не подключать math ради одного округления вверх такая странная конструкция
163
+ rows = len(img_list) // cols if len(img_list) // cols == len(img_list) / cols else (len(img_list) // cols) + 1
164
+ cols = cols if cols < len(img_list) else len(img_list)
165
+
166
+ width = cols*thumb_width
167
+ height = rows*thumb_height
168
+
169
+ new_img = Image.new('RGB', (width, height))
170
+
171
+ i, x, y = 0, 0, 0
172
+ for row in range(rows):
173
+ if i == len(img_list):
174
+ break
175
+ for col in range(cols):
176
+ if i == len(img_list):
177
+ break
178
+ new_img.paste(img_list[i], (x, y))
179
+ i += 1
180
+ x += thumb_height
181
+ y += thumb_width
182
+ x = 0
183
+
184
+ return new_img
185
+
186
+ def main():
187
+ # поиск в Вики статей (по умолчанию - одной) по запросу или рандом но
188
+ search_random = False if input_query else True
189
+ # установим выбранный язык
190
+ wiki_lang = input_lang if input_lang else GLOBAL_lang
191
+ wikipedia.set_lang(wiki_lang)
192
+ with st.spinner('Ищем в википедии...'):
193
+ search_rezult = wiki_search(query = input_query, random=search_random)
194
+ # из полученной статей считываем картинки и текст
195
+ pages = to_pool(search_rezult, eval('get_wiki_pages_content'))
196
+ st.success('Нашли, теперь повеселимся.')
197
+
198
+ for page in pages:
199
+ # Получаем интересный нам контент
200
+ page = pages[0]
201
+ title = page['title']
202
+ pharagrafs = page['content']
203
+ images_urls = page['images']
204
+ wiki_text = '\n'.join(pharagrafs)
205
+
206
+ # если для поиска использовалась русская вики, то переводим тексты
207
+ if (input_lang == 'ru') | (GLOBAL_lang == 'ru'):
208
+ with st.spinner('Еще минутку - нужен перевод...'):
209
+ mname = "Helsinki-NLP/opus-mt-ru-en"
210
+ translator = pipeline("translation", model = mname)
211
+ pharagrafs = to_pool(pharagrafs, eval('rus2eng'))
212
+ st.success('Готово')
213
+
214
+ # чтение реальных картинок со страницы
215
+ with st.spinner('Реальность...'):
216
+ images_natur = to_pool(images_urls, eval('file2img'))
217
+ # убираем те, которые не смогли прочесть
218
+ images_natur = list(filter(None, images_natur))
219
+ # уменьшаем
220
+ images_natur_small = to_pool(images_natur, eval('resize_img'))
221
+ images_natur_collage = create_collage(images_natur_small)
222
+ st.success('Готово')
223
+
224
+ # генерация картинок по описанию
225
+ # можно было бы использовать для генерации размеры GLOBAL_thumb_size
226
+ # но практика показала, что качество сильно падает
227
+ # пришлось оставить связку генерация в размере 256*256 + уменьшение размера после
228
+
229
+ with st.spinner('А теперь самое интересное...'):
230
+ images_gen = text2img(pharagrafs)
231
+ # уменьшаем
232
+ images_gen_small = to_pool(images_gen, eval('resize_img'))
233
+ images_gen_collage = create_collage(images_gen_small)
234
+ st.success('Готово')
235
+
236
+ # ++++++++++++++++++++++++++++++++++++++++++++++++++++
237
+ # +++++++++++++ Вывод результатов ++++++++++++++++++++
238
+ # ++++++++++++++++++++++++++++++++++++++++++++++++++++
239
+ st.subheader(title)
240
+
241
+ col1, col2 = st.columns(2)
242
+
243
+ with col1:
244
+ st.image(images_gen_collage, caption='Ожидания')
245
+ with col2:
246
+ st.image(images_natur_collage, caption='Реальность')
247
+
248
+ with st.expander("Посмотреть текст"):
249
+ st.text(wiki_text)
250
+ if __name__ == "__main__":
251
+ # ++++++++++++++++++++++++++++++++++++++++++++++++++++
252
+ # +++++++++++++ Пользовательские импуты ++++++++++++++
253
+ # ++++++++++++++++++++++++++++++++++++++++++++++++++++
254
+ with st.sidebar:
255
+ # поскольку переводчик у нас только на два языка надо делать переключателем en/ru
256
+ # в зависимости от выбора ищем в русскоязычной или англоязычной вики
257
+ input_lang = st.radio('Какую википедию использовать для поиска?:', ('en', 'ru'), index=0)
258
+ # запрос на поиск в Вики
259
+ input_query = st.text_input(
260
+ label = 'Вы хотите посмотреть на:',
261
+ value= ''
262
+ )
263
+ st.header('"WIKI Images: Ожидания vs Реальность"')
264
+ st.text('Задача проекта визуализировать, насколько текст статей википедии соответствует иллюстрациям. Для этого на основе текста статьи генерится коллаж изображений (параграф = одно изображение). Второй коллаж формируется из реальных иллюстраций к статье.')
265
+
266
+ st.text('Вы можете сами выбрать запрос (в сайдбаре) или довериться рандому.')
267
+
268
+ if st.button('Начать'):
269
+ main()
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ opencv-python==4.7.0.72
2
+ sacremoses==0.0.53
3
+ torch==2.0.0
4
+ torchvision==0.15.1
5
+ transformers
6
+ sentencepiece
7
+