graxpert 3.2.0a0.dev4__py3-none-any.whl

graxpert/commands.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from copy import deepcopy
+from typing import Dict
+
+import numpy as np
+
+from graxpert.app_state import INITIAL_STATE, AppState
+from graxpert.background_grid_selection import background_grid_selection
+from graxpert.background_flood_selection import background_flood_selection
+
+
+class ICommandHandler(ABC):
+    @abstractmethod
+    def execute(self, app_state: AppState, cmd_args: Dict) -> AppState:
+        pass
+
+    @abstractmethod
+    def undo(
+        self, cur_state: AppState, prev_state: AppState, cmd_args: Dict
+    ) -> AppState:
+        pass
+
+    @abstractmethod
+    def redo(
+        self, cur_state: AppState, next_state: AppState, cmd_args: Dict
+    ) -> AppState:
+        pass
+
+    @abstractmethod
+    def progress(self) -> float:
+        pass
+
+
+class Command:
+
+    def __init__(
+        self, handler: ICommandHandler, prev: Command = None, **kwargs
+    ) -> None:
+        self.prev: Command = prev
+        self.next: Command = None
+        self.handler: ICommandHandler = handler
+        self.app_state: AppState = None
+        self.cmd_args = kwargs
+
+    def execute(self) -> AppState:
+        if self.prev is None:
+            app_state = INITIAL_STATE
+        else:
+            app_state = self.prev.app_state
+        self.app_state = self.handler.execute(app_state, self.cmd_args)
+        if self.prev is not None:
+            self.prev.next = self
+        return self.app_state
+
+    def undo(self) -> Command:
+        if self.prev is None:
+            prev_state = INITIAL_STATE
+            return self
+        else:
+            prev_state = self.prev.app_state
+            self.prev.app_state = self.handler.undo(self.app_state, prev_state, self.cmd_args)
+            return self.prev
+        
+
+    def redo(self) -> Command:
+        assert self.next is not None
+        next_state = self.next.app_state
+        self.next.app_state = self.handler.redo(self.app_state, next_state, self.cmd_args)
+        return self.next
+
+
+class InitHandler(ICommandHandler):
+
+    def execute(self, app_state: AppState, cmd_args: Dict) -> AppState:
+        state = INITIAL_STATE
+        state.background_points = cmd_args["background_points"]
+        return state
+
+    def undo(
+        self, cur_state: AppState, prev_state: AppState, cmd_args: Dict
+    ) -> AppState:
+        state = INITIAL_STATE
+        state.background_points = cmd_args["background_points"]
+        return state
+
+    def redo(
+        self, cur_state: AppState, next_state: AppState, cmd_args: Dict
+    ) -> AppState:
+        return deepcopy(next_state)
+
+    def progress(self) -> float:
+        return 1.0
+
+class PointHandler(ICommandHandler):
+    def undo(
+        self, cur_state: AppState, prev_state: AppState, cmd_args: Dict
+    ) -> AppState:
+        app_state_copy = deepcopy(cur_state)
+        prev_background_points = deepcopy(prev_state.background_points)
+        app_state_copy.background_points = prev_background_points
+        return app_state_copy
+
+    def redo(
+        self, cur_state: AppState, next_state: AppState, cmd_args: Dict
+    ) -> AppState:
+        app_state_copy = deepcopy(cur_state)
+        next_background_points = deepcopy(next_state.background_points)
+        app_state_copy.background_points = next_background_points
+        return app_state_copy
+
+
+class AddPointHandler(PointHandler):
+    def execute(self, app_state: AppState, cmd_args: Dict) -> AppState:
+        app_state_copy = deepcopy(app_state)
+        point = cmd_args["point"]
+        app_state_copy.background_points.append(point)
+        return app_state_copy
+
+    def progress(self) -> float:
+        return 1.0
+
+
+class AddPointsHandler(PointHandler):
+    def execute(self, app_state: AppState, cmd_args: Dict) -> AppState:
+        app_state_copy = deepcopy(app_state)
+        point = cmd_args["point"]
+        background_points = app_state_copy.background_points
+        tol = cmd_args["tol"]
+        bg_pts = cmd_args["bg_pts"]
+        sample_size = cmd_args["sample_size"]
+        image = cmd_args["image"]
+        new_points = background_flood_selection(point, background_points, tol, bg_pts, sample_size, image)
+        app_state_copy.background_points.extend(new_points)
+        return app_state_copy
+
+    def progress(self) -> float:
+        return 1.0
+
+
+class RemovePointHandler(PointHandler):
+    def execute(self, app_state: AppState, cmd_args: Dict) -> AppState:
+        app_state_copy = deepcopy(app_state)
+        idx = cmd_args["idx"]
+        app_state_copy.background_points.pop(idx)
+        return app_state_copy
+
+    def progress(self) -> float:
+        return 1.0
+
+
+class MovePointHandler(PointHandler):
+    def execute(self, app_state: AppState, cmd_args: Dict) -> AppState:
+        app_state_copy = deepcopy(app_state)
+        idx = cmd_args["idx"]
+        new_point = cmd_args["new_point"]
+        
+        if len(new_point) == 0:
+            app_state_copy.background_points.pop(idx)
+        else:
+            app_state_copy.background_points[idx] = new_point
+        
+        return app_state_copy
+        
+    def progress(self) -> float:
+        return 1.0
+    
+    
+class SelectPointsHandler(PointHandler):
+    def execute(self, app_state: AppState, cmd_args: Dict) -> AppState:
+        app_state_copy = deepcopy(app_state)
+        data = cmd_args["data"]
+        num_pts = cmd_args["num_pts"]
+        tol = cmd_args["tol"]
+        sample_size = cmd_args["sample_size"]
+        automatic_points = background_grid_selection(data, num_pts, tol, sample_size)
+        app_state_copy.background_points = automatic_points
+        return app_state_copy
+
+    def progress(self) -> float:
+        return 1.0
+
+class ResetPointsHandler(PointHandler):
+    def execute(self, app_state: AppState, cmd_args: Dict) -> AppState:
+        app_state_copy = deepcopy(app_state)
+        app_state_copy.background_points.clear()
+        return app_state_copy
+
+    def progress(self) -> float:
+        return 1.0
+
+
+INIT_HANDLER = InitHandler()
+ADD_POINT_HANDLER = AddPointHandler()
+ADD_POINTS_HANDLER = AddPointsHandler()
+RM_POINT_HANDLER = RemovePointHandler()
+MOVE_POINT_HANDLER = MovePointHandler()
+SEL_POINTS_HANDLER = SelectPointsHandler()
+RESET_POINTS_HANDLER = ResetPointsHandler()

graxpert/deconvolution.py

@@ -0,0 +1,176 @@
+import copy
+import logging
+
+import numpy as np
+
+from graxpert.ai_model_handling import get_execution_providers_ordered
+from graxpert.application.app_events import AppEvents
+from graxpert.application.eventbus import eventbus
+
+
+def deconvolve(image, ai_path, strength, psfsize, batch_size=4, window_size=512, stride=448, progress=None, ai_gpu_acceleration=True):
+    logging.info("Starting deconvolution")
+    if "stars" in ai_path:
+        type = "Stellar"
+    elif "obj" in ai_path:
+        type = "Obj"
+    strength = 0.95 * strength  # TODO : strenght of exactly 1.0 brings no results, to fix
+
+    if type == "Stellar":
+        psfsize = np.clip((psfsize / 2.355 - 1.5) / 3.0, 0.05, 0.95)  # Stellar
+    else:
+        if "1.0.0" in ai_path:
+            psfsize = np.clip((psfsize / 2.355 - 1.0) / 5.0, 0.05, 0.95)  # Object v1.0.0
+        else:
+            psfsize = np.clip((psfsize / 2.355 - 0.5) / 5.5, 0.05, 0.95)  # Object v1.0.1
+
+    logging.info(f"Calculated normalized PSFsize value: {psfsize}")
+
+    if batch_size < 1:
+        logging.info(f"mapping batch_size of {batch_size} to 1")
+        batch_size = 1
+    elif batch_size > 32:
+        logging.info(f"mapping batch_size of {batch_size} to 32")
+        batch_size = 32
+    elif not (batch_size & (batch_size - 1) == 0):  # check if batch_size is power of two
+        logging.info(f"mapping batch_size of {batch_size} to {2 ** (batch_size).bit_length() // 2}")
+        batch_size = 2 ** (batch_size).bit_length() // 2  # map batch_size to power of two
+
+    if batch_size >= 4 and image.shape[-1] == 3:
+        batch_size = batch_size // 4
+
+    num_colors = image.shape[-1]
+
+    H, W, _ = image.shape
+    offset = int((window_size - stride) / 2)
+
+    h, w, _ = image.shape
+
+    ith = int(h / stride) + 1
+    itw = int(w / stride) + 1
+
+    dh = ith * stride - h
+    dw = itw * stride - w
+
+    image = np.concatenate((image, image[(h - dh) :, :, :]), axis=0)
+    image = np.concatenate((image, image[:, (w - dw) :, :]), axis=1)
+
+    h, w, _ = image.shape
+    image = np.concatenate((image, image[(h - offset) :, :, :]), axis=0)
+    image = np.concatenate((image[:offset, :, :], image), axis=0)
+    image = np.concatenate((image, image[:, (w - offset) :, :]), axis=1)
+    image = np.concatenate((image[:, :offset, :], image), axis=1)
+
+    output = copy.deepcopy(image)
+
+    providers = get_execution_providers_ordered(ai_gpu_acceleration)
+    import onnxruntime as ort # Must be after get_execution_providers_ordered
+    session = ort.InferenceSession(ai_path, providers=providers)
+
+    logging.info(f"Available inference providers : {providers}")
+    logging.info(f"Used inference providers : {session.get_providers()}")
+
+    cancel_flag = False
+
+    def cancel_listener(event):
+        nonlocal cancel_flag
+        cancel_flag = True
+
+    eventbus.add_listener(AppEvents.CANCEL_PROCESSING, cancel_listener)
+
+    last_progress = 0
+    for b in range(0, ith * itw + batch_size, batch_size):
+
+        if cancel_flag:
+            logging.info("Deconvolution cancelled")
+            eventbus.remove_listener(AppEvents.CANCEL_PROCESSING, cancel_listener)
+            return None
+
+        input_tiles = []
+        input_tile_copies = []
+        params = []
+        for t_idx in range(0, batch_size):
+
+            index = b + t_idx
+            i = index % ith
+            j = index // ith
+
+            if i >= ith or j >= itw:
+                break
+
+            x = stride * i
+            y = stride * j
+
+            tile = image[x : x + window_size, y : y + window_size, :]
+
+            _min = np.min(tile, axis=(0, 1))
+            tile = tile - _min + 1e-5
+            tile = np.log(tile)
+
+            _mean = tile.mean()
+            _std = tile.std()
+            _mean, _std = _mean.astype(np.float32), _std.astype(np.float32)
+            tile = (tile - _mean) / _std * 0.1
+            params.append([_mean, _std, _min])
+
+            input_tile_copies.append(np.copy(tile))
+
+            input_tiles.append(tile)
+
+        if not input_tiles:
+            continue
+
+        input_tiles = np.array(input_tiles)
+        input_tiles = np.moveaxis(input_tiles, -1, 1)
+        input_tiles = np.reshape(input_tiles, [input_tiles.shape[0] * num_colors, 1, window_size, window_size])
+
+        output_tiles = []
+        sigma = np.full(shape=(input_tiles.shape[0], 1), fill_value=psfsize, dtype=np.float32)
+        strenght_p = np.full(shape=(input_tiles.shape[0], 1), fill_value=strength, dtype=np.float32)
+        conds = np.concatenate([sigma, strenght_p], axis=-1)
+        if type == "Obj" and "1.0.0" in ai_path:
+            session_result = session.run(None, {"gen_input_image": input_tiles, "sigma": sigma, "strenght": strenght_p})[0]
+        else:
+            session_result = session.run(None, {"gen_input_image": input_tiles, "params": conds})[0]
+        for e in session_result:
+            output_tiles.append(e)
+
+        output_tiles = np.array(output_tiles)
+        output_tiles = input_tiles - output_tiles
+        output_tiles = np.reshape(output_tiles, [output_tiles.shape[0] // num_colors, num_colors, window_size, window_size])
+        output_tiles = np.moveaxis(output_tiles, 1, -1)
+
+        for idx in range(len(params)):
+            output_tiles[idx] = output_tiles[idx] * params[idx][1] / 0.1 + params[idx][0]
+            output_tiles[idx] = np.exp(output_tiles[idx])
+            output_tiles[idx] = output_tiles[idx] + params[idx][2] - 1e-5
+
+        for t_idx, tile in enumerate(output_tiles):
+
+            index = b + t_idx
+            i = index % ith
+            j = index // ith
+
+            if i >= ith or j >= itw:
+                break
+
+            x = stride * i
+            y = stride * j
+            tile = tile[offset : offset + stride, offset : offset + stride, :]
+            output[x + offset : stride * (i + 1) + offset, y + offset : stride * (j + 1) + offset, :] = tile
+
+        p = int(b / (ith * itw + batch_size) * 100)
+        if p > last_progress:
+            if progress is not None:
+                progress.update(p - last_progress)
+            else:
+                logging.info(f"Progress: {p}%")
+            last_progress = p
+
+    output = output[offset : H + offset, offset : W + offset, :]
+    output = np.clip(output, 0.0, 1.0)
+
+    eventbus.remove_listener(AppEvents.CANCEL_PROCESSING, cancel_listener)
+    logging.info("Finished deconvolution")
+
+    return output

graxpert/denoising.py

@@ -0,0 +1,178 @@
+import copy
+import logging
+import time
+
+import numpy as np
+
+from graxpert.ai_model_handling import get_execution_providers_ordered
+from graxpert.application.app_events import AppEvents
+from graxpert.application.eventbus import eventbus
+from graxpert.ui.ui_events import UiEvents
+
+
+def denoise(image, ai_path, strength, batch_size=4, window_size=256, stride=128, progress=None, ai_gpu_acceleration=True):
+    logging.info("Starting denoising")
+
+    if batch_size < 1:
+        logging.info(f"mapping batch_size of {batch_size} to 1")
+        batch_size = 1
+    elif batch_size > 32:
+        logging.info(f"mapping batch_size of {batch_size} to 32")
+        batch_size = 32
+    elif not (batch_size & (batch_size - 1) == 0):  # check if batch_size is power of two
+        logging.info(f"mapping batch_size of {batch_size} to {2 ** (batch_size).bit_length() // 2}")
+        batch_size = 2 ** (batch_size).bit_length() // 2  # map batch_size to power of two
+
+    input = copy.deepcopy(image)
+
+    median = np.median(image[::4, ::4, :], axis=[0, 1])
+    mad = np.median(np.abs(image[::4, ::4, :] - median), axis=[0, 1])
+    
+    if "1.0.0" in ai_path or "1.1.0" in ai_path:
+        model_threshold = 1.0
+    else:
+        model_threshold = 10.0
+
+    global cached_denoised_image
+    if cached_denoised_image is not None:
+        return blend_images(input, cached_denoised_image, strength, model_threshold, median, mad)
+
+    num_colors = image.shape[-1]
+    if num_colors == 1:
+        image = np.array([image[:, :, 0], image[:, :, 0], image[:, :, 0]])
+        image = np.moveaxis(image, 0, -1)
+
+    H, W, _ = image.shape
+    offset = int((window_size - stride) / 2)
+
+    h, w, _ = image.shape
+
+    ith = int(h / stride) + 1
+    itw = int(w / stride) + 1
+
+    dh = ith * stride - h
+    dw = itw * stride - w
+
+    image = np.concatenate((image, image[(h - dh) :, :, :]), axis=0)
+    image = np.concatenate((image, image[:, (w - dw) :, :]), axis=1)
+
+    h, w, _ = image.shape
+    image = np.concatenate((image, image[(h - offset) :, :, :]), axis=0)
+    image = np.concatenate((image[:offset, :, :], image), axis=0)
+    image = np.concatenate((image, image[:, (w - offset) :, :]), axis=1)
+    image = np.concatenate((image[:, :offset, :], image), axis=1)
+
+    output = copy.deepcopy(image)
+
+    providers = get_execution_providers_ordered(ai_gpu_acceleration)
+    import onnxruntime as ort # Must be after get_execution_providers_ordered
+    session = ort.InferenceSession(ai_path, providers=providers)
+
+    logging.info(f"Available inference providers : {providers}")
+    logging.info(f"Used inference providers : {session.get_providers()}")
+
+    cancel_flag = False
+
+    def cancel_listener(event):
+        nonlocal cancel_flag
+        cancel_flag = True
+
+    eventbus.add_listener(AppEvents.CANCEL_PROCESSING, cancel_listener)
+
+    last_progress = 0
+    for b in range(0, ith * itw + batch_size, batch_size):
+
+        if cancel_flag:
+            logging.info("Denoising cancelled")
+            eventbus.remove_listener(AppEvents.CANCEL_PROCESSING, cancel_listener)
+            return None
+
+        input_tiles = []
+        input_tile_copies = []
+        for t_idx in range(0, batch_size):
+
+            index = b + t_idx
+            i = index % ith
+            j = index // ith
+
+            if i >= ith or j >= itw:
+                break
+
+            x = stride * i
+            y = stride * j
+
+            tile = image[x : x + window_size, y : y + window_size, :]
+            tile = (tile - median) / mad * 0.04
+            input_tile_copies.append(np.copy(tile))
+            tile = np.clip(tile, -model_threshold, model_threshold)
+
+            input_tiles.append(tile)
+
+        if not input_tiles:
+            continue
+
+        input_tiles = np.array(input_tiles)
+
+        output_tiles = []
+        session_result = session.run(None, {"gen_input_image": input_tiles})[0]
+        for e in session_result:
+            output_tiles.append(e)
+
+        output_tiles = np.array(output_tiles)
+
+        for t_idx, tile in enumerate(output_tiles):
+
+            index = b + t_idx
+            i = index % ith
+            j = index // ith
+
+            if i >= ith or j >= itw:
+                break
+
+            x = stride * i
+            y = stride * j
+            tile = np.where(input_tile_copies[t_idx] < model_threshold, tile, input_tile_copies[t_idx])
+            tile = tile / 0.04 * mad + median
+            tile = tile[offset : offset + stride, offset : offset + stride, :]
+            output[x + offset : stride * (i + 1) + offset, y + offset : stride * (j + 1) + offset, :] = tile
+
+        p = int(b / (ith * itw + batch_size) * 100)
+        if p > last_progress:
+            if progress is not None:
+                progress.update(p - last_progress)
+            else:
+                logging.info(f"Progress: {p}%")
+            last_progress = p
+
+    output = output[offset : H + offset, offset : W + offset, :]
+
+    if num_colors == 1:
+        output = np.array([output[:, :, 0]])
+        output = np.moveaxis(output, 0, -1)
+
+    cached_denoised_image = output
+    output = blend_images(input, output, strength, model_threshold, median, mad)
+
+    eventbus.remove_listener(AppEvents.CANCEL_PROCESSING, cancel_listener)
+    logging.info("Finished denoising")
+
+    return output
+
+
+def blend_images(original_image, denoised_image, strength, threshold, median, mad):
+    threshold = threshold / 0.04 * mad + median
+    blend = np.where(original_image < threshold, denoised_image, original_image)
+    blend = blend * strength + original_image * (1 - strength)
+    return np.clip(blend, 0, 1)
+
+
+def reset_cached_denoised_image(event):
+    global cached_denoised_image
+    cached_denoised_image = None
+
+
+cached_denoised_image = None
+eventbus.add_listener(AppEvents.LOAD_IMAGE_REQUEST, reset_cached_denoised_image)
+eventbus.add_listener(AppEvents.CALCULATE_REQUEST, reset_cached_denoised_image)
+eventbus.add_listener(UiEvents.APPLY_CROP_REQUEST, reset_cached_denoised_image)
+eventbus.add_listener(AppEvents.DENOISE_AI_VERSION_CHANGED, reset_cached_denoised_image)

graxpert/grid_utils.py

@@ -0,0 +1,38 @@
+import logging
+import math
+
+import numpy as np
+
+
+def find_darkest_quadrant(x, y, data_padded, sample_size):
+
+    cords = [
+        [x + 0, y + 0],
+        [x + sample_size, y + sample_size],
+        [x - sample_size, y + sample_size],
+        [x + sample_size, y - sample_size],
+        [x - sample_size, y - sample_size],
+    ]
+    median = []
+    for point in cords:
+        if (
+            point[0] < 0
+            or point[0] > (data_padded.shape[1] - 2 * sample_size)
+            or point[1] < 0
+            or point[1] > (data_padded.shape[0] - 2 * sample_size)
+        ):
+            median.append(2)
+        else:
+            m = np.median(
+                data_padded[
+                    point[1] : point[1] + 2 * sample_size,
+                    point[0] : point[0] + 2 * sample_size,
+                ]
+            )
+            if math.isnan(m):
+                logging.error("computed median is NaN", stack_info=True)
+            median.append(m)
+
+    min_idx = np.argmin(median)
+
+    return cords[min_idx], median[min_idx]