graxpert 3.2.0a0.dev4__py3-none-any.whl

graxpert/AstroImageRepository.py

@@ -0,0 +1,114 @@
+from enum import StrEnum
+from typing import Dict
+
+from graxpert.astroimage import AstroImage
+from graxpert.stretch import StretchParameters, calculate_mtf_stretch_parameters_for_image, stretch_all
+
+
+class ImageTypes(StrEnum):
+    Original = "Original"
+    Gradient_Corrected = "Gradient-Corrected"
+    Background = "Background"
+    Deconvolved_Object_only = "Deconvolved Object-only"
+    Deconvolved_Stars_only = "Deconvolved Stars-only"
+    Denoised = "Denoised"
+
+
+class AstroImageRepository:
+
+    images: Dict = {
+        ImageTypes.Original: None,
+        ImageTypes.Gradient_Corrected: None,
+        ImageTypes.Background: None,
+        ImageTypes.Deconvolved_Object_only: None,
+        ImageTypes.Deconvolved_Stars_only: None,
+        ImageTypes.Denoised: None,
+    }
+
+    def set(self, type: ImageTypes, image: AstroImage):
+        self.images[type] = image
+
+    def get(self, type: ImageTypes):
+        return self.images[type]
+
+    def stretch_all(self, stretch_params: StretchParameters, saturation: float):
+
+        if self.get(ImageTypes.Original) is None:
+            return
+
+        stretches = []
+
+        if not stretch_params.do_stretch:
+            for key, image in self.images.items():
+                if image is not None:
+                    stretches.append(image.img_array)
+
+        else:
+
+            all_image_arrays = []
+            all_mtf_stretch_params = []
+
+            all_image_arrays.append(self.get(ImageTypes.Original).img_array)
+            all_mtf_stretch_params.append(calculate_mtf_stretch_parameters_for_image(stretch_params, self.get(ImageTypes.Original).img_array))
+
+            if self.get(ImageTypes.Gradient_Corrected) is not None and self.get(ImageTypes.Background) is not None:
+                all_image_arrays.append(self.get(ImageTypes.Gradient_Corrected).img_array)
+                all_mtf_stretch_params.append(calculate_mtf_stretch_parameters_for_image(stretch_params, self.get(ImageTypes.Gradient_Corrected).img_array))
+
+                all_image_arrays.append(self.get(ImageTypes.Background).img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[0])
+
+            if self.get(ImageTypes.Deconvolved_Object_only) is not None and self.get(ImageTypes.Gradient_Corrected) is None:
+                all_image_arrays.append(self.get(ImageTypes.Deconvolved_Object_only).img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[0])
+
+            elif self.get(ImageTypes.Deconvolved_Object_only) is not None and self.get(ImageTypes.Gradient_Corrected) is not None:
+                all_image_arrays.append(self.get(ImageTypes.Deconvolved_Object_only).img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[1])
+
+            if self.get(ImageTypes.Deconvolved_Stars_only) is not None and self.get(ImageTypes.Gradient_Corrected) is None:
+                all_image_arrays.append(self.get(ImageTypes.Deconvolved_Stars_only).img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[0])
+
+            elif self.get(ImageTypes.Deconvolved_Stars_only) is not None and self.get(ImageTypes.Gradient_Corrected) is not None:
+                all_image_arrays.append(self.get(ImageTypes.Deconvolved_Stars_only).img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[1])
+
+            if self.get(ImageTypes.Denoised) is not None and self.get(ImageTypes.Gradient_Corrected) is None:
+                all_image_arrays.append(self.get(ImageTypes.Denoised).img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[0])
+
+            elif self.get(ImageTypes.Denoised) is not None and self.get(ImageTypes.Gradient_Corrected) is not None:
+                all_image_arrays.append(self.get(ImageTypes.Denoised).img_array)
+                all_mtf_stretch_params.append(all_mtf_stretch_params[1])
+
+            stretches = stretch_all(all_image_arrays, all_mtf_stretch_params)
+
+        i = 0
+        for key, image in self.images.items():
+            if image is not None:
+                image.update_display_from_array(stretches[i], saturation)
+                i = i + 1
+
+    def crop_all(self, start_x: float, end_x: float, start_y: float, end_y: float):
+        for key, astroimg in self.images.items():
+            if astroimg is not None:
+                astroimg.crop(start_x, end_x, start_y, end_y)
+
+    def update_saturation(self, saturation):
+        for key, value in self.images.items():
+            if value is not None:
+                value.update_saturation(saturation)
+
+    def reset(self):
+        for key, value in self.images.items():
+            self.images[key] = None
+
+    def display_options(self):
+        display_options = []
+
+        for key, value in self.images.items():
+            if self.images[key] is not None:
+                display_options.append(key)
+
+        return display_options

graxpert/__init__.py

@@ -0,0 +1 @@
+

graxpert/ai_model_handling.py

@@ -0,0 +1,215 @@
+import logging
+import os
+import sys
+import re
+import shutil
+import zipfile
+
+from appdirs import user_data_dir
+from minio import Minio
+from packaging import version
+
+try:
+    from graxpert.s3_secrets import endpoint, ro_access_key, ro_secret_key
+
+    client = Minio(endpoint, ro_access_key, ro_secret_key)
+except Exception as e:
+    logging.exception(e)
+    client = None
+
+from graxpert.ui.loadingframe import DynamicProgressThread
+
+ai_models_dir = os.path.join(user_data_dir(appname="GraXpert"), "ai-models")
+bge_ai_models_dir = os.path.join(user_data_dir(appname="GraXpert"), "bge-ai-models")
+
+# old ai-models folder exists, rename to 'bge-ai-models'
+if os.path.exists(ai_models_dir):
+    logging.warning(f"Older 'ai_models_dir' {ai_models_dir} exists. Renaming to {bge_ai_models_dir} due to introduction of new denoising models in GraXpert 3.")
+    try:
+        os.rename(ai_models_dir, bge_ai_models_dir)
+    except Exception as e:
+        logging.error(f"Renaming {ai_models_dir} to {bge_ai_models_dir} failed. {bge_ai_models_dir} will be newly created. Consider deleting obsolete {ai_models_dir} manually.")
+
+os.makedirs(bge_ai_models_dir, exist_ok=True)
+
+deconvolution_object_ai_models_dir = os.path.join(user_data_dir(appname="GraXpert"), "deconvolution-object-ai-models")
+os.makedirs(deconvolution_object_ai_models_dir, exist_ok=True)
+deconvolution_stars_ai_models_dir = os.path.join(user_data_dir(appname="GraXpert"), "deconvolution-stars-ai-models")
+os.makedirs(deconvolution_stars_ai_models_dir, exist_ok=True)
+denoise_ai_models_dir = os.path.join(user_data_dir(appname="GraXpert"), "denoise-ai-models")
+os.makedirs(denoise_ai_models_dir, exist_ok=True)
+
+
+# ui operations
+def list_remote_versions(bucket_name):
+    if client is None:
+        return []
+    try:
+        objects = client.list_objects(bucket_name)
+        versions = []
+
+        for o in objects:
+            tags = client.get_object_tags(o.bucket_name, o.object_name)
+            if tags is not None and "ai-version" in tags:
+                versions.append(
+                    {
+                        "bucket": o.bucket_name,
+                        "object": o.object_name,
+                        "version": tags["ai-version"],
+                    }
+                )
+        return versions
+
+    except Exception as e:
+        logging.exception(e)
+    finally:
+        return versions
+
+
+def list_local_versions(ai_models_dir):
+    try:
+        model_dirs = [
+            {"path": os.path.join(ai_models_dir, f), "version": f}
+            for f in os.listdir(ai_models_dir)
+            if re.search(r"\d\.\d\.\d", f) and len(os.listdir(os.path.join(ai_models_dir, f))) > 0  # match semantic version
+        ]
+        return model_dirs
+    except Exception as e:
+        logging.exception(e)
+        return None
+
+
+def latest_version(ai_models_dir, bucket_name):
+    try:
+        remote_versions = list_remote_versions(bucket_name)
+    except Exception as e:
+        remote_versions = []
+        logging.exception(e)
+    try:
+        local_versions = list_local_versions(ai_models_dir)
+    except Exception as e:
+        local_versions = []
+        logging.exception(e)
+    ai_options = set([])
+    ai_options.update([rv["version"] for rv in remote_versions])
+    ai_options.update(set([lv["version"] for lv in local_versions]))
+    ai_options = sorted(ai_options, key=lambda k: version.parse(k), reverse=True)
+    return ai_options[0]
+
+
+def ai_model_path_from_version(ai_models_dir, local_version):
+    if local_version is None:
+        return None
+
+    return os.path.join(ai_models_dir, local_version, "model.onnx")
+
+
+def compute_orphaned_local_versions(ai_models_dir):
+    remote_versions = list_remote_versions(ai_models_dir)
+
+    if remote_versions is None:
+        logging.warning("Could not fetch remote versions. Thus, aborting cleaning of local versions in {}. Consider manual cleaning".format(ai_models_dir))
+        return
+
+    local_versions = list_local_versions()
+
+    if local_versions is None:
+        logging.warning("Could not read local versions in {}. Thus, aborting cleaning. Consider manual cleaning".format(ai_models_dir))
+        return
+
+    orphaned_local_versions = [{"path": lv["path"], "version": lv["version"]} for lv in local_versions if lv["version"] not in [rv["version"] for rv in remote_versions]]
+
+    return orphaned_local_versions
+
+
+def cleanup_orphaned_local_versions(orphaned_local_versions):
+    for olv in orphaned_local_versions:
+        try:
+            shutil.rmtree(olv["path"])
+        except Exception as e:
+            logging.exception(e)
+
+
+def download_version(ai_models_dir, bucket_name, target_version, progress=None):
+    try:
+        remote_versions = list_remote_versions(bucket_name)
+        for r in remote_versions:
+            if target_version == r["version"]:
+                remote_version = r
+                break
+
+        ai_model_dir = os.path.join(ai_models_dir, "{}".format(remote_version["version"]))
+        os.makedirs(ai_model_dir, exist_ok=True)
+
+        ai_model_file = os.path.join(ai_model_dir, "model.onnx")
+        ai_model_zip = os.path.join(ai_model_dir, "model.zip")
+        client.fget_object(
+            remote_version["bucket"],
+            remote_version["object"],
+            ai_model_zip,
+            progress=DynamicProgressThread(callback=progress),
+        )
+
+        with zipfile.ZipFile(ai_model_zip, "r") as zip_ref:
+            zip_ref.extractall(ai_model_dir)
+
+        if not os.path.isfile(ai_model_file):
+            raise ValueError(f"Could not find ai 'model.onnx' file after extracting {ai_model_zip}")
+        os.remove(ai_model_zip)
+
+    except Exception as e:
+        # try to delete (rollback) ai_model_dir in case of errors
+        logging.exception(e)
+        try:
+            shutil.rmtree(ai_model_dir)
+        except Exception as e2:
+            logging.exception(e2)
+
+
+def validate_local_version(ai_models_dir, local_version):
+    return os.path.isfile(os.path.join(ai_models_dir, local_version, "model.onnx"))
+
+
+def get_execution_providers_ordered(gpu_acceleration=True):
+    if gpu_acceleration:
+        supported_providers = [
+            (
+                "OpenVINOExecutionProvider",
+                {
+                    # per https://onnxruntime.ai/docs/execution-providers/OpenVINO-ExecutionProvider.html#summary-of-options
+                    # "device_type": "HETERO:GPU,CPU,NPU AUTO:GPU,CPU,NPU MULTI:GPU,CPU,NPU", # Will prefer dGPU, fallback to iGPU, NPU or CPU with extra Intel specific optimizations
+                    "device_type": "AUTO:GPU,CPU", # Will prefer dGPU, fallback to iGPU, NPU or CPU with extra Intel specific optimizations
+                }
+            ),
+            "ROCMExecutionProvider",
+            "DmlExecutionProvider",
+            (
+                "CoreMLExecutionProvider",
+                {
+                    "flags": "COREML_FLAG_CREATE_MLPROGRAM",
+                },
+            ),
+            "CUDAExecutionProvider",
+            "CPUExecutionProvider",
+        ]
+    else:
+        supported_providers = ["CPUExecutionProvider"]
+
+    result = []
+    try:
+        import onnxruntime as ort    
+        available = ort.get_available_providers()
+        for provider in supported_providers:
+            if isinstance(provider, tuple):
+                if provider[0] in available:
+                    result.append(provider)  # Append the entire tuple
+            else:
+                if provider in available:
+                    result.append(provider)
+        return result
+    except Exception as e:
+        logging.error("Critical error!  The required ONNX Runtime (AI library) package is misconfigured.\n" \
+        "Please read the README.md to confirm that you've selected the correct build for your hardware.\n" \
+        "If you are using one of our prebuilt executables, please file a bug with the following information:\n"
+        "{}".format(e))
+        sys.exit(1)

graxpert/app_state.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass, field
+from typing import List
+
+
+@dataclass
+class AppState:
+    background_points: List = field(default_factory=list)
+
+
+INITIAL_STATE = AppState()