diff --git a/maps4fs/generator/background.py b/maps4fs/generator/background.py index 80deae06..894c810a 100644 --- a/maps4fs/generator/background.py +++ b/maps4fs/generator/background.py @@ -58,6 +58,10 @@ def preprocess(self) -> None: os.makedirs(self.water_directory, exist_ok=True) self.output_path = os.path.join(self.background_directory, f"{FULL_NAME}.png") + if self.map.custom_background_path: + self.check_custom_background(self.map.custom_background_path) + shutil.copyfile(self.map.custom_background_path, self.output_path) + self.not_substracted_path = os.path.join(self.background_directory, "not_substracted.png") self.not_resized_path = os.path.join(self.background_directory, "not_resized.png") @@ -75,6 +79,28 @@ def preprocess(self) -> None: self.dem.set_output_resolution((self.rotated_size, self.rotated_size)) self.dem.set_dem_path(self.output_path) + def check_custom_background(self, image_path: str) -> None: + """Checks if the custom background image meets the requirements. + + Arguments: + image_path (str): The path to the custom background image. + + Raises: + ValueError: If the custom background image does not meet the requirements. + """ + image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED) # pylint: disable=no-member + if image.shape[0] != image.shape[1]: + raise ValueError("The custom background image must be a square.") + + if image.shape[0] != self.map_size + DEFAULT_DISTANCE * 2: + raise ValueError("The custom background image must have the size of the map + 4096.") + + if len(image.shape) != 2: + raise ValueError("The custom background image must be a grayscale image.") + + if image.dtype != np.uint16: + raise ValueError("The custom background image must be a 16-bit grayscale image.") + def is_preview(self, name: str) -> bool: """Checks if the DEM is a preview. @@ -91,7 +117,9 @@ def process(self) -> None: as a result the DEM files will be saved, then based on them the obj files will be generated.""" self.create_background_textures() - self.dem.process() + + if not self.map.custom_background_path: + self.dem.process() shutil.copyfile(self.dem.dem_path, self.not_substracted_path) self.cutout(self.dem.dem_path, save_path=self.not_resized_path) diff --git a/maps4fs/generator/map.py b/maps4fs/generator/map.py index 4f5e3b03..b501e2a1 100644 --- a/maps4fs/generator/map.py +++ b/maps4fs/generator/map.py @@ -140,6 +140,11 @@ def __init__( # pylint: disable=R0917, R0915 json.dump(self.tree_custom_schema, file, indent=4) self.logger.debug("Tree custom schema saved to %s", save_path) + self.custom_background_path = kwargs.get("custom_background_path", None) + if self.custom_background_path: + save_path = os.path.join(self.map_directory, "custom_background.png") + shutil.copyfile(self.custom_background_path, save_path) + try: shutil.unpack_archive(game.template_path, self.map_directory) self.logger.debug("Map template unpacked to %s", self.map_directory) diff --git a/webui/generator.py b/webui/generator.py index 979edef0..2437689f 100644 --- a/webui/generator.py +++ b/webui/generator.py @@ -340,6 +340,7 @@ def add_left_widgets(self) -> None: on_change=self.map_preview, ) + self.custom_background_path = None self.expert_mode = False self.raw_config = None @@ -422,6 +423,23 @@ def add_left_widgets(self) -> None: label_visibility="collapsed", ) + self.custom_background = st.checkbox( + "Upload custom background", value=False, key="custom_background" + ) + + if self.custom_background: + st.info(Messages.CUSTOM_BACKGROUND_INFO) + + uploaded_file = st.file_uploader("Choose a file", type=["png"]) + if uploaded_file is not None: + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + self.custom_background_path = os.path.join( + config.INPUT_DIRECTORY, f"custom_background_{timestamp}.png" + ) + with open(self.custom_background_path, "wb") as f: + f.write(uploaded_file.read()) + st.success(f"Custom background uploaded: {uploaded_file.name}") + # Add an empty container for status messages. self.status_container = st.empty() @@ -568,6 +586,7 @@ def generate_map(self) -> None: satellite_settings=all_settings["SatelliteSettings"], texture_custom_schema=texture_schema, tree_custom_schema=tree_schema, + custom_background_path=self.custom_background_path, ) if self.public: diff --git a/webui/templates.py b/webui/templates.py index f4ff6311..a78dcff8 100644 --- a/webui/templates.py +++ b/webui/templates.py @@ -61,6 +61,14 @@ class Messages: "https://github.com/iwatkot/maps4fs/blob/main/docs/custom_osm.md). \n" "Note, that incorrect file can lead to errors or completely broken map." ) + CUSTOM_BACKGROUND_INFO = ( + "The uploaded file should be: \n" + "- Single-channel (grayscale) unsigned 16-bit PNG image. \n" + "- The size of the image should be map size + 4096 in each dimension, where the map in " + "the center. \n" + "- If rotation needed, the image should be rotated already. \n \n" + "If any of above conditions are not met, generation will fail." + ) EXPERT_MODE_INFO = ( "In this mode you can edit confuguration of the generation in a raw format. " "Be careful, any incorrect value can lead to errors or completely broken map."