Spaces:
Runtime error
Runtime error
julianna-fil
commited on
Commit
•
b9c2b0f
1
Parent(s):
0cced28
Upload 2 files
Browse files- app.py +269 -0
- 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 |
+
|