| import streamlit as st |
| import moduls.Texture as texture |
| from PIL import Image |
| import mmcv |
| import moduls.model |
| from moduls.texture_for_point import texture_transformation_1 |
| from main import model |
| import numpy as np |
| import cv2 |
| import os |
| from streamlit_image_select import image_select |
| from streamlit_image_annotation import pointdet |
| from moduls.conturs import conturs_on_image |
| from streamlit_image_comparison import image_comparison |
| from moduls.settings import path_for_streamlit |
| st.set_page_config(layout="wide") |
|
|
|
|
| @st.cache_data |
| def load_image(image_file): |
| img = Image.open(image_file) |
| img = mmcv.imread(np.array(img)) |
| return img |
| def save_file_st(path_file): |
| file_name = path_file.name |
| img = load_image(path_file) |
| im = Image.fromarray(img) |
| output_filename = 'processing.jpg' |
| im.save(output_filename) |
| return output_filename |
| @st.dialog("Выбор текстуры пола", width = 'large') |
| def select_floor(img_list_1, image_1, img_list): |
| tex_floor = image_select("Выбор текстуры пола", img_list_1, return_value= 'index') |
| if st.button("Выбрать"): |
| st.session_state.json_list[0] = img_list[tex_floor] |
| modificate_image('floor') |
| st.session_state.select_floor = [tex_floor] |
| st.rerun() |
| @st.dialog("Выбор текстуры стен", width = 'large') |
| def select_wall(img_list_1, image_1, img_list): |
| tex_wall = image_select("Выбор текстуры стен", img_list_1, return_value= 'index') |
| if st.button("Выбрать"): |
| st.session_state.json_list[1] = img_list[tex_wall] |
| modificate_image('wall') |
| st.session_state.select_wall = [tex_wall] |
| st.rerun() |
| @st.dialog("Выбор текстуры потолка", width = 'large') |
| def select_ceiling(img_list_1, image_1, img_list): |
| tex_ceiling = image_select("Выбор текстуры потолка", img_list_1, return_value= 'index') |
| if st.button("Выбрать"): |
| st.session_state.json_list[2] = img_list[tex_ceiling] |
| modificate_image('ceiling') |
| st.session_state.select_ceiling = [tex_ceiling] |
| st.rerun() |
| def points_sorter(points): |
| messages_list = ['Недостаточно точек. Добавьте точек!', |
| 'Точек не четное количество. Добавьте или уберите точку!', |
| 'Точек слишком много. Уберите лишние точки!', |
| 'Точки успешно отсортированы!', |
| 'Пол или потолок имеют обратную перспективу. Скорректируйте точки!'] |
| print('len(points)', len(points)) |
| print(points) |
| if (num_points := len(points)) < 6: |
| return [], [] |
| elif num_points%2 != 0: |
| return [], [] |
| elif num_points > 8: |
| return [], [] |
| else: |
| |
| sorted_list = sorted(points, key=lambda k: [k[1], k[0]]) |
| print('sorted_list', sorted_list) |
| |
| ceiling_points = sorted(sorted_list[:num_points//2], key=lambda k: [k[0], k[1]]) |
| floor_points = sorted(sorted_list[num_points//2:], key=lambda k: [k[0], k[1]]) |
| bad_points = [] |
| print('ceiling_points, floor_points', ceiling_points, floor_points) |
| |
| if ceiling_points[0][1] > ceiling_points[1][1]: |
| bad_points.append(ceiling_points[0]) |
| bad_points.append(ceiling_points[1]) |
| if ceiling_points[-1][1] > ceiling_points[-2][1]: |
| bad_points.append(ceiling_points[-1]) |
| bad_points.append(ceiling_points[-2]) |
| |
| if floor_points[0][1] < floor_points[1][1]: |
| bad_points.append(floor_points[0]) |
| bad_points.append(floor_points[1]) |
| if floor_points[-1][1] < floor_points[-2][1]: |
| bad_points.append(floor_points[-1]) |
| bad_points.append(floor_points[-2]) |
| c_f_point = [ceiling_points, floor_points] |
| return c_f_point |
| |
| |
| |
| |
| |
| def corner_change_2(image_path): |
| list_1 = ['Угол'] |
| list_label_1 = [] |
| list_label_2 = [] |
| list_point_1 = [] |
| list_point_2 = [] |
| if 'ceiling_point' in st.session_state and 'floor_point' in st.session_state and st.session_state.ceiling_point is not None and st.session_state.floor_point is not None: |
| list_point_1 = st.session_state.ceiling_point |
| print('list_point_1', list_point_1) |
| for i in range(len(st.session_state.ceiling_point)): |
| list_label_1.append(0) |
| list_point_2 = st.session_state.floor_point |
| for i in range(len(st.session_state.floor_point)): |
| list_label_2.append(0) |
| else: |
| list_point_1 = [] |
| list_point_2 = [] |
| list_point = list_point_1 + list_point_2 |
| list_label = list_label_1 + list_label_2 |
| print('list_point', list_point) |
| print('list_label', list_label) |
| point = pointdet(image_path , list_1, |
| points=list_point, |
| |
| labels=list_label, |
| |
| use_space=True) |
| if point is not None: |
| st.session_state.point = point |
| st.session_state.floor_point = [] |
| st.session_state.ceiling_point = [] |
| |
| st.session_state.points = [] |
| if point is not None: |
| for i in st.session_state.point: |
| if i['label_id'] == 0: |
| |
| st.session_state.points.append([int(i['point'][0]), int(i['point'][1])]) |
| else: |
| |
| st.session_state.points.append([int(i['point'][0]), int(i['point'][1])]) |
| print('st.session_state.points', st.session_state.points) |
| c_f_point = points_sorter(st.session_state.points) |
| st.session_state.ceiling_point = c_f_point[0] |
| st.session_state.floor_point = c_f_point[1] |
| |
| |
| print(st.session_state.ceiling_point) |
| |
| |
| |
| |
| |
| |
| def sort_list_ceiling(no_change_list): |
| if len(no_change_list) >= 4: |
| no_change_list.sort() |
| if no_change_list[0][1] > no_change_list[1][1]: |
| no_change_list[0], no_change_list[1] = no_change_list[1], no_change_list[0] |
| if no_change_list[-1][1] > no_change_list[-2][1]: |
| no_change_list[-1], no_change_list[-2] = no_change_list[-2], no_change_list[-1] |
| if len(no_change_list) == 3: |
| no_change_list.sort() |
| if no_change_list[0][1] > no_change_list[1][1]: |
| no_change_list[0], no_change_list[1] = no_change_list[1], no_change_list[0] |
| return no_change_list |
| def sort_list_floor(no_change_list): |
| if len(no_change_list) >= 4: |
| no_change_list.sort() |
| if no_change_list[0][1] < no_change_list[1][1]: |
| no_change_list[0], no_change_list[1] = no_change_list[1], no_change_list[0] |
| if no_change_list[-1][1] < no_change_list[-2][1]: |
| no_change_list[-1], no_change_list[-2] = no_change_list[-2], no_change_list[-1] |
| if len(no_change_list) == 3: |
| no_change_list.sort() |
| if no_change_list[0][1] < no_change_list[1][1]: |
| no_change_list[0], no_change_list[1] = no_change_list[1], no_change_list[0] |
| return no_change_list |
| |
| def corner_change(image_path): |
| list_1 = ['Угол потолка', 'Угол пола'] |
| list_label_1 = [] |
| list_label_2 = [] |
| list_point_1 = [] |
| list_point_2 = [] |
| if 'ceiling_point' in st.session_state and 'floor_point' in st.session_state and st.session_state.ceiling_point is not None and st.session_state.floor_point is not None: |
| list_point_1 = st.session_state.ceiling_point |
| for i in range(len(st.session_state.ceiling_point)): |
| list_label_1.append(0) |
| list_point_2 = st.session_state.floor_point |
| for i in range(len(st.session_state.floor_point)): |
| list_label_2.append(1) |
| else: |
| list_point_1 = [] |
| list_point_2 = [] |
| list_point = list_point_1 + list_point_2 |
| list_label = list_label_1 + list_label_2 |
| print('list_point', list_point) |
| print('list_label', list_label) |
| point = pointdet(image_path , list_1, |
| points=list_point, |
| |
| labels=list_label, |
| |
| use_space=True) |
| if point is not None: |
| st.session_state.point = point |
| st.session_state.floor_point = [] |
| st.session_state.ceiling_point = [] |
| |
| if point is not None: |
| for i in st.session_state.point: |
| if i['label_id'] == 0: |
|
|
| st.session_state.ceiling_point.append([int(i['point'][0]), int(i['point'][1])]) |
| else: |
| st.session_state.floor_point.append([int(i['point'][0]), int(i['point'][1])]) |
| st.session_state.ceiling_point = sort_list_ceiling(st.session_state.ceiling_point) |
| print(st.session_state.ceiling_point) |
| st.session_state.floor_point = sort_list_floor(st.session_state.floor_point) |
| |
| |
| print('st.session_state.ceiling_point', st.session_state.ceiling_point) |
| print('st.session_state.floor_point', st.session_state.floor_point) |
| |
| return point |
| @st.dialog("modification_photo") |
| def modification_photo(image_1): |
| st.session_state.modification_photo = [image_1] |
| st.rerun() |
| def a(path_photo): |
| st.session_state.a = path_photo |
| def predict(path_photo, model): |
| path_file =save_file_st(path_photo) |
| _, seg_map, seg_all, image = moduls.model.predict_and_visualize(model, path_file) |
| print(seg_map.shape) |
| st.session_state.predict =[seg_all, image] |
| contur_image, simplified_contour = conturs_on_image(image, seg_map) |
| st.session_state.contur_image =contur_image |
| st.session_state.simplified_contour = simplified_contour |
| return image |
| def modificate_image(surface): |
| original_img = cv2.cvtColor(st.session_state.predict[1], cv2.COLOR_BGR2RGB) |
| if st.session_state.f_h !=0: |
| f_h = st.session_state.f_h |
| else: |
| f_h = None |
| if st.session_state.w_h !=0: |
| w_h = st.session_state.w_h |
| else: |
| w_h = None |
| if st.session_state.c_h !=0: |
| c_h = st.session_state.c_h |
| else: |
| c_h = None |
| image_1 = texture_transformation_1(original_img, st.session_state.predict[0], |
| st.session_state.room_l*100, st.session_state.room_w*100, st.session_state.room_h*100, |
| st.session_state.ceiling_point, st.session_state.floor_point, |
| st.session_state.floor_angel, st.session_state.wall_angel, st.session_state.ceiling_angel, |
| st.session_state.json_list, surface = st.session_state.surface, |
| h_f= f_h, h_w = w_h, h_c = c_h, |
| rescale = 1.) |
| st.session_state.image_1 = image_1 |
| st.session_state.modification_photo = [image_1] |
| def save_angel(): |
| st.session_state.floor_angel = floor_angel |
| st.session_state.ceiling_angel = ceiling_angel |
| st.session_state.wall_angel = wall_angel |
| path_icon, path_json, path_logo = path_for_streamlit() |
| if 'json_list' not in st.session_state: |
| st.session_state.json_list = [None, None, None] |
| img_list_1 = [] |
| img_list = [] |
| for path_file in os.listdir(path_icon): |
| img_list_1.append(path_icon + f'\{path_file}') |
| for path_json_file in os.listdir(path_json): |
| if path_json_file.endswith(('.json')): |
| img_list.append(path_json + f'\{path_json_file}') |
| img_list_1.sort() |
| img_list.sort() |
| title_col1, title_col2 = st.columns([1,3]) |
| with title_col1: |
| st.image(path_logo) |
| with title_col2: |
| st.title(":orange[Замена текстур на фотографии вашей комнаты]",) |
|
|
| change_col_1, change_col_2, change_col_3 = st.columns([1, 4, 1]) |
| |
| with change_col_1: |
| path_photo = st.file_uploader("Загрузка фото", type=['png','jpeg','jpg'], accept_multiple_files=False, key=None, |
| help=None, on_change=None, args=None, kwargs=None, |
| disabled=False, label_visibility = 'collapsed') |
| if path_photo == None: |
| st.session_state.a = None |
| st.session_state.pp = None |
| if path_photo is not None: |
| if st.session_state.a != path_photo: |
| st.session_state.ceiling_point = None |
| st.session_state.floor_point = None |
| st.session_state.pp = None |
| st.session_state.a = path_photo |
| predict(st.session_state.a, model) |
| st.session_state.modification_photo = [cv2.cvtColor(st.session_state.predict[1], cv2.COLOR_BGR2RGB)] |
| |
| with change_col_2: |
| if path_photo is not None: |
| cv2.imwrite('corners.png',st.session_state.contur_image) |
| |
| st.session_state.corners = corner_change_2('corners.png') |
| |
| with change_col_3: |
| if path_photo is not None: |
| st.session_state.room_l = st.number_input('Глубина комнаты', min_value=0., max_value=6., value=4., step=0.1, format=None, key='room_l_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible") |
| st.session_state.room_w = st.number_input('Ширина комнаты', min_value=0., max_value=6., value=3., step=0.1, format=None, key='room_w_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible") |
| st.session_state.room_h = st.number_input('Высота потолка', min_value=0., max_value=4., value=3., step=0.1, format=None, key='room_h_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible") |
| |
| if 'corners' in st.session_state and st.session_state.corners != None: |
| st.session_state.pp = 1 |
| col1, col2, col3 = st.columns(3) |
| |
| with col1: |
| if path_photo is not None and st.session_state.pp is not None: |
| floor = st.checkbox('Текстура пола', value=False, key=None, help=None, on_change=None, args=None, kwargs=None, disabled=False, label_visibility="visible") |
| if floor: |
| floor_select = st.button("Выбор текстуры пола", key = '1') |
| if floor_select: |
| select_floor(img_list_1, cv2.cvtColor(st.session_state.modification_photo[0], cv2.COLOR_BGR2RGB), img_list) |
| try: |
| st.image(img_list_1[st.session_state.select_floor[0]], caption=None, width=200, use_column_width=None, clamp=False, channels='RGB', output_format="auto") |
| st.session_state.json_list[0] = img_list[st.session_state.select_floor[0]] |
| except: |
| b=1 |
| else: |
| st.session_state.json_list[0] = None |
| |
| with col2: |
| if path_photo is not None and st.session_state.pp is not None: |
| wall = st.checkbox('Текстура стен', value=False, key=None, help=None, on_change=None, args=None, kwargs=None, disabled=False, label_visibility="visible") |
| if wall: |
| wall_select = st.button("Выбор текстуры стен", key = '2') |
| if wall_select: |
| select_wall(img_list_1, cv2.cvtColor(st.session_state.modification_photo[0], cv2.COLOR_BGR2RGB), img_list) |
| try: |
| st.image(img_list_1[st.session_state.select_wall[0]], caption=None, width=200, use_column_width=None, clamp=False, channels='RGB', output_format="auto") |
| st.session_state.json_list[1] = img_list[st.session_state.select_wall[0]] |
| except: |
| b=1 |
| else: |
| st.session_state.json_list[1] = None |
| |
| with col3: |
| if path_photo is not None and st.session_state.pp is not None: |
| ceiling = st.checkbox('Текстура потолка', value=False, key=None, help=None, on_change=None, args=None, kwargs=None, disabled=False, label_visibility="visible") |
| if ceiling: |
| ceiling_select = st.button("Выбор текстуры потолка", key = '3') |
| if ceiling_select: |
| select_ceiling(img_list_1, cv2.cvtColor(st.session_state.modification_photo[0], cv2.COLOR_BGR2RGB), img_list) |
| try: |
| st.image(img_list_1[st.session_state.select_ceiling[0]], caption=None, width=200, use_column_width=None, clamp=False, channels='RGB', output_format="auto") |
| st.session_state.json_list[2] = img_list[st.session_state.select_ceiling[0]] |
| except: |
| b=1 |
| else: |
| st.session_state.json_list[2] = None |
| if path_photo is not None and st.session_state.pp is not None: |
| st.session_state.surface = [floor, wall, ceiling] |
|
|
| |
| col1_1, col2_1 = st.columns([3,1]) |
| try: |
| with col2_1: |
| if path_photo is not None and st.session_state.pp is not None: |
| floor_angel = st.slider(label = 'Угол поворота текстуры пола',min_value=0, max_value=360, value=0, step=1, format=None, key='floor_angle_key', |
| help=None, on_change=save_angel, args=None, kwargs=None, disabled=False, label_visibility= 'visible') |
| wall_angel = st.slider(label = 'Угол поворота текстуры стен',min_value=0, max_value=360, value=0, step=1, format=None, key='wall_angle_key', |
| help=None, on_change=save_angel, args=None, kwargs=None, disabled=False, label_visibility= 'visible') |
| ceiling_angel = st.slider(label = 'Угол поворота текстуры потолка',min_value=0, max_value=360, value=0, step=1, format=None, key='ceiling_angle_key', |
| help=None, on_change=save_angel, args=None, kwargs=None, disabled=False, label_visibility= 'visible') |
| st.write(":orange[Выбор способа набора текстур]") |
| st.write(":orange[Если больше 1 то набор ёлочкой]") |
| st.session_state.f_h = st.number_input('Пол', min_value=0, max_value=3, value=0, step=1, format=None, key='f_h_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible") |
| st.session_state.w_h = st.number_input('Стены', min_value=0, max_value=3, value=0, step=1, format=None, key='w_h_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible") |
| st.session_state.c_h = st.number_input('Потолок', min_value=0, max_value=3, value=0, step=1, format=None, key='c_h_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible") |
|
|
| if 'floor_angel' not in st.session_state: |
| st.session_state.floor_angel = 0 |
| if 'wall_angel' not in st.session_state: |
| st.session_state.wall_angel = 0 |
| if 'ceiling_angel' not in st.session_state: |
| st.session_state.ceiling_angel = 0 |
| try: |
| if st.session_state.floor_angel != floor_angel: |
| modificate_image('floor') |
| except: |
| print('no_change_angel_floor') |
| try: |
| if st.session_state.wall_angel != wall_angel: |
| modificate_image('floor') |
| except: |
| print('no_change_angel_wall') |
| try: |
| if st.session_state.ceiling_angel != ceiling_angel: |
| modificate_image('ceiling') |
| except: |
| print('no_change_angel_ceiling') |
|
|
| |
| with col1_1: |
| |
| |
| if path_photo is not None and st.session_state.pp is not None: |
| if st.session_state.modification_photo is not None: |
| image_1 = st.session_state.modification_photo[0] |
| else: |
| image_1 = cv2.cvtColor(st.session_state.predict[1], cv2.COLOR_BGR2RGB) |
| image = cv2.cvtColor(st.session_state.predict[1], cv2.COLOR_BGR2RGB) |
| |
| image_comparison( |
| img1=image, |
| img2=image_1, |
| width = 1080, |
| starting_position = 1 |
| ) |
| |
| |
| |
| |
| |
| cv2.imwrite('Change_photo.jpg', cv2.cvtColor(st.session_state.modification_photo[0], cv2.COLOR_BGR2RGB)) |
| with open('Change_photo.jpg', "rb") as file: |
| print(type(file)) |
| btn = st.download_button( |
| label="Скачать изображение с наложенной текстурой", |
| data=file, |
| file_name="Change_photo.png", |
| mime="image/png" |
| ) |
| except: |
| b=1 |
|
|
|
|
