File size: 4,089 Bytes
545f568 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import io
import cv2
import numpy as np
from PIL import Image
import streamlit as st
# ===================== Настройки страницы ======================
st.set_page_config(
layout='wide',
initial_sidebar_state='auto',
page_title='Face Emotion Recognition',
page_icon='👻',
)
st.write("#### Детекция лиц, эмоций, пола, расы и возраста на изображении")
# отображение картинки с примерами детекции
@st.cache_data
def load_main_image(image_path: str) -> Image.Image:
main_pil_image = Image.open(image_path)
return main_pil_image
MAIN_IMAGE_PATH = './media/emotions_detect.png'
main_pil_image = load_main_image(MAIN_IMAGE_PATH)
st.image(main_pil_image, width=800)
# ===================== Боковое меню настроек ======================
st.sidebar.header('Настройки')
st.sidebar.write('---')
face_conf_threshold = st.sidebar.slider(
label='Порог уверенности для детекции лиц',
min_value=0.0,
max_value=1.0,
value=0.7,
step=0.01,
)
st.sidebar.write('---')
# какие действия нужно детектить (пока не реализовано переключение)
actions = ['age', 'gender', 'race', 'emotion']
# применять ли дополнительное выравнивание
st.sidebar.write('Применять ли дополнительное выравнивание')
align = st.sidebar.checkbox(label='Align', value=False)
# загрузка и инициализация моделей
with st.spinner('Инициализация/загрузка моделей...'):
from detector import detector_model
# ==================== Загрузка изображения и распознавание ============
st_image = st.file_uploader(label='Выберите изображение')
st.session_state['detect_image_ready'] = False
if st_image:
pil_image = Image.open(st_image)
np_image_rgb = np.array(pil_image)
np_image_bgr = cv2.cvtColor(np_image_rgb, cv2.COLOR_RGB2BGR)
st.image(np_image_bgr, width=400, channels='BGR')
# ================ Кнопка распознать ==================
if st.button('Распознать'):
if detector_model.is_first_run:
spinner_text = 'Первоначальная инициализация моделей и распознавание фото...'
else:
spinner_text = 'Распознавание фото ...'
with st.spinner(spinner_text):
detections = detector_model.detect_image(
np_image_bgr,
actions=actions,
align=align,
)
detector_model.is_first_run = False
with st.spinner('Отрисовка результата'):
result_np_image = detector_model.draw_detections(
np_image_rgb=np_image_rgb,
detections=detections,
face_conf_threshold=face_conf_threshold,
)
st.session_state['detect_image_ready'] = True
# отображение результата детекции
if st.session_state['detect_image_ready']:
st.image(result_np_image)
# сохранение картинки с результатом чтобы ее скачать
image_name = f"{st_image.name.split('.')[0]}_detect.png"
file_buffer = io.BytesIO()
result_pil_image = Image.fromarray(result_np_image)
result_pil_image.save(file_buffer, format='PNG')
image_bytes = file_buffer.getvalue()
# ================ Кнопка скачать ==================
st.download_button(
label='Скачать изображение',
data=image_bytes,
file_name=image_name,
mime='image/png',
)
st.session_state['detect_image_ready'] = False
|