CelestialSurveyor

1
import datetime
2
import json
3
import numpy as np
4
import os.path
5
import threading
6
import wx
7
import wx.lib.scrolledpanel as scrolled
8

9
from dataclasses import dataclass
10
from decimal import Decimal
11
from ObjectListView import ObjectListView, ColumnDefn
12
from PIL import Image
13
from threading import Event
14
from typing import Optional, Callable, Any
15

16
from backend.find_asteroids import predict_asteroids, save_results, annotate_results
17
from backend.progress_bar import ProgressBarFactory
18
from backend.source_data_v2 import SourceDataV2
19
from logger.logger import get_logger
20

21

22
logger = get_logger()
23

24

25
@dataclass(frozen=True)
26
class MyDataObject:
27
    """
28
    Represents an object within file list.
29

30
    Attributes:
31
        file_path (str): The path to the file.
32
        timestamp (datetime.datetime): The timestamp of the object.
33
        exposure (Decimal): The exposure value of the object.
34
        checked (bool): The status of the object if it is checked or not in the list, default is False.
35
    """
36
    file_path: str
37
    timestamp: datetime.datetime
38
    exposure: Decimal
39
    checked: bool = False
40

41

42
class ImagePanel(scrolled.ScrolledPanel):
43
    """
44
    A panel for displaying images with zoom and scrolling capability.
45
    Note: Scrolling doesn't work for now.
46
    """
47
    def __init__(self, parent: wx.Window, image_array: Optional[np.ndarray] = None):
48
        super().__init__(parent, style=wx.BORDER_SIMPLE)
49
        self.image_array = image_array
50
        self.Bind(wx.EVT_PAINT, self.on_paint)
51
        self.Bind(wx.EVT_MOUSEWHEEL, self.on_scroll)
52
        self.scale_factor = None
53
        self.default_scale_factor = 1.0
54
        self.scroll_x = 0
55
        self.scroll_y = 0
56
        self.SetupScrolling(scroll_x=True, scroll_y=True, scrollToTop=False)
57

58
    def on_paint(self, event: Event) -> None:
59
        """
60
        Callback function for when the panel needs to paint its content.
61

62
        Parameters:
63
            event (wx.Event): A wxPython event object.
64

65
        Returns:
66
            None
67
        """
68
        if self.image_array is not None:
69
            display_width, display_height = self.GetSize()
70
            image_height, image_width = self.image_array.shape[:2]
71
            dc = wx.PaintDC(self)
72
            if self.scale_factor is None:
73
                self.scale_factor = display_width / image_width
74
                self.default_scale_factor = self.scale_factor
75

76
            dc.SetUserScale(self.scale_factor, self.scale_factor)
77
            image = self.convert_array_to_bitmap(self.image_array)
78
            dc.DrawBitmap(image, 0, 0)
79

80
    def on_scroll(self, event: Event) -> None:
81
        """
82
        Callback function for when the mouse wheel is scrolled.
83
        It's zooming.
84

85
        Parameters:
86
            event (wx.Event): A wxPython event object.
87

88
        Returns:
89
            None
90
        """
91

92
        rotation = event.GetWheelRotation()
93

94
        # Calculate the new scale factor based on the mouse wheel rotation
95
        if rotation > 0:
96
            self.scale_factor *= 1.1
97
        else:
98
            self.scale_factor /= 1.1
99

100
        self.Refresh(eraseBackground=False)
101

102
    @staticmethod
103
    def convert_array_to_bitmap(array: np.ndarray, display_width: int = None, display_height: int = None) -> wx.Bitmap:
104
        """
105
        Convert a NumPy array to a wx.Bitmap to be displayed on the image panel.
106

107
        Parameters:
108
            array (np.ndarray): The input NumPy array.
109
            display_width (int): The width for displaying the bitmap.
110
            display_height (int): The height for displaying the bitmap.
111

112
        Returns:
113
            wx.Bitmap: The converted wx.Bitmap object.
114
        """
115
        # Create a PIL Image from the NumPy array
116
        array = array.reshape(array.shape[:2])
117
        image = Image.fromarray(array)
118

119
        # Convert PIL Image to wx.Image
120
        pil_image = image.convert("RGB")
121
        if display_width is not None and display_height is not None:
122
            pil_image = pil_image.resize((display_width, display_height))
123

124
        width, height = pil_image.size
125
        image = wx.Image(width, height)
126
        image.SetData(pil_image.tobytes())
127

128
        # Convert wx.Image to wx.Bitmap
129
        bitmap = image.ConvertToBitmap()
130

131
        return bitmap
132

133

134
class ProgressFrame(wx.Dialog):
135
    """Progress dialog to display progress of the current operation"""
136

137
    def __init__(self, parent: wx.Frame, title: str, stop_event: Event):
138
        super().__init__(parent=parent, title=title, size=(500, 200))
139
        self.parent = parent
140
        self.stop_event = stop_event
141
        panel = wx.Panel(self)
142
        self.label = wx.StaticText(panel, label="Working...")
143
        self.progress = wx.Gauge(panel, range=100, size=(450, 30))
144
        self.cancel_button = wx.Button(panel, label='Cancel')
145
        self.Bind(wx.EVT_BUTTON, self.on_cancel, self.cancel_button)
146
        sizer = wx.BoxSizer(wx.VERTICAL)
147
        sizer.Add(self.label, 0, wx.ALL | wx.LEFT, 5)
148
        sizer.Add(self.progress, 0, wx.ALL | wx.CENTER, 5)
149
        sizer.Add(self.cancel_button, 0, wx.ALL | wx.CENTER, 5)
150
        panel.SetSizer(sizer)
151

152
        self.failed = False
153

154
    def set_failed(self, value: bool) -> None:
155
        """
156
        Sets the 'failed' attribute of the ProgressFrame instance.
157
        Used to indicate that the operation has failed and display an error message.
158

159
        Args:
160
            value (bool): The value to set the 'failed' attribute to.
161

162
        Returns:
163
            None
164
        """
165
        self.failed = value
166

167
    def on_cancel(self, event: Event) -> None:
168
        """
169
        Callback function for when the cancel button is clicked.
170
        It's stopping the current operation.
171

172
        Parameters:
173
            event (wx.Event): A wxPython event object.
174

175
        Returns:
176
            None
177
        """
178
        if self.failed:
179
            self.Close(force=True)
180
            self.set_failed(False)
181
        else:
182
            self.label.SetLabel("Stopping...")
183
            self.stop_event.set()
184

185

186
class MyFrame(wx.Frame):
187
    """Main frame of the application"""
188

189
    def __init__(self, *args, **kw):
190
        super(MyFrame, self).__init__(*args, **kw)
191

192
        self.source_data: SourceDataV2 = None
193

194
        # Create the main panel
195
        self.panel = wx.Panel(self)
196
        panel_sizer = wx.BoxSizer(wx.VERTICAL)
197

198
        # Create a horizontal box sizer to organize the elements
199
        hbox = wx.BoxSizer(wx.HORIZONTAL)
200

201
        # Create the controls area (1/6 of window width)
202
        controls_sizer = wx.BoxSizer(wx.VERTICAL)
203

204
        controls_label = wx.StaticText(self.panel)
205
        flat_label = wx.StaticText(self.panel, label="Select folder with flat frames")
206
        self.flat_path_picker = wx.DirPickerCtrl(self.panel, style=wx.DIRP_USE_TEXTCTRL)
207
        dark_flat_label = wx.StaticText(self.panel, label="Select folder with dark flat frames")
208
        self.dark_flat_path_picker = wx.DirPickerCtrl(self.panel, style=wx.DIRP_USE_TEXTCTRL)
209
        dark_label = wx.StaticText(self.panel, label="Select folder with dark frames")
210
        self.dark_path_picker = wx.DirPickerCtrl(self.panel, style=wx.DIRP_USE_TEXTCTRL)
211
        self.btn_add_files = wx.Button(self.panel, label="Add files")
212
        self.Bind(wx.EVT_BUTTON, self.on_add_files, self.btn_add_files)
213
        self.chk_to_align = wx.CheckBox(self.panel, label="Align images")
214
        self.chk_to_align.SetValue(True)
215
        self.chk_non_linear = wx.CheckBox(self.panel, label="Non-linear")
216
        self.chk_non_linear.SetValue(False)
217
        self.to_debayer = wx.CheckBox(self.panel, label="Debayer")
218
        self.to_debayer.SetValue(True)
219
        self.btn_load_files = wx.Button(self.panel, label="Load files")
220
        self.progress_frame = None
221
        self.process_thread = None
222
        self.stop_event = None
223

224
        self.Bind(wx.EVT_BUTTON, self.on_load_files, self.btn_load_files)
225
        self.results_label = wx.StaticText(self.panel, label="Select folder to store results")
226
        self.results_path_picker = wx.DirPickerCtrl(self.panel, style=wx.DIRP_USE_TEXTCTRL)
227
        self.magnitude_label = wx.StaticText(self.panel, label="Annotation magnitude limit")
228
        self.magnitude_input = wx.TextCtrl(self.panel)
229
        self.magnitude_input.SetValue("18.0")
230
        self.magnitude_input.Bind(wx.EVT_TEXT, self.on_text_change)
231
        self.btn_process = wx.Button(self.panel, label="Process")
232
        self.Bind(wx.EVT_BUTTON, self.on_process, self.btn_process)
233
        self.btn_start_again = wx.Button(self.panel, label="Start again")
234
        self.Bind(wx.EVT_BUTTON, self.on_start_again, self.btn_start_again)
235

236
        controls_sizer.Add(controls_label, 0, wx.EXPAND | wx.ALL, 5)
237
        controls_sizer.Add(flat_label, 0, wx.EXPAND | wx.ALL, 5)
238
        controls_sizer.Add(self.flat_path_picker, 0, wx.EXPAND | wx.ALL, 5)
239
        controls_sizer.Add(dark_flat_label, 0, wx.EXPAND | wx.ALL, 5)
240
        controls_sizer.Add(self.dark_flat_path_picker, 0, wx.EXPAND | wx.ALL, 5)
241
        controls_sizer.Add(dark_label, 0, wx.EXPAND | wx.ALL, 5)
242
        controls_sizer.Add(self.dark_path_picker, 0, wx.EXPAND | wx.ALL, 5)
243
        controls_sizer.Add(self.btn_add_files, 0, wx.EXPAND | wx.ALL, 5)
244
        controls_sizer.Add(self.to_debayer, 0, wx.EXPAND | wx.ALL, 5)
245
        controls_sizer.Add(self.chk_to_align, 0, wx.EXPAND | wx.ALL, 5)
246
        controls_sizer.Add(self.chk_non_linear, 0, wx.EXPAND | wx.ALL, 5)
247
        controls_sizer.Add(self.btn_load_files, 0, wx.EXPAND | wx.ALL, 5)
248
        controls_sizer.Add(self.results_label, 0, wx.EXPAND | wx.ALL, 5)
249
        controls_sizer.Add(self.results_path_picker, 0, wx.EXPAND | wx.ALL, 5)
250
        controls_sizer.Add(self.magnitude_label, 0, wx.EXPAND | wx.ALL, 5)
251
        controls_sizer.Add(self.magnitude_input, 0, wx.EXPAND | wx.ALL, 5)
252
        controls_sizer.Add(self.btn_process, 0, wx.EXPAND | wx.ALL, 5)
253
        controls_sizer.Add(self.btn_start_again, 0, wx.EXPAND | wx.ALL, 5)
254

255
        hbox.Add(controls_sizer, 1, wx.EXPAND | wx.ALL, 5)
256

257
        # Create the checkbox list area (2/6 of window width) using ObjectListView
258
        checkbox_label = wx.StaticText(self.panel)
259
        self.checkbox_list = ObjectListView(self.panel, style=wx.LC_REPORT | wx.SUNKEN_BORDER)
260
        self.checkbox_list.Bind(wx.EVT_LIST_ITEM_SELECTED, self.on_item_selected)
261
        columns = [
262
            ColumnDefn("Item", "left", 350, "file_path"),
263
            ColumnDefn("Timestamp", "left", 130, "timestamp"),
264
            ColumnDefn("Exposure", "left", 60, "exposure"),
265
        ]
266

267
        self.checkbox_list.SetColumns(columns)
268
        self.checkbox_list.CreateCheckStateColumn()
269
        hbox.Add(checkbox_label, 0, wx.EXPAND | wx.ALL, 5)
270
        hbox.Add(self.checkbox_list, 2, wx.EXPAND | wx.ALL, 5)
271

272
        # Create the picture representation area (3/6 of window width)
273
        self.draw_panel = ImagePanel(self.panel)
274
        hbox.Add(self.draw_panel, 3, wx.EXPAND | wx.ALL, 5)
275

276
        self.log_text = wx.TextCtrl(self.panel, style=wx.TE_MULTILINE | wx.TE_READONLY | wx.HSCROLL)
277
        logger.text_ctrl = self.log_text
278
        panel_sizer.Add(hbox, 3, wx.EXPAND | wx.ALL, 5)
279
        panel_sizer.Add(self.log_text, 1, wx.EXPAND | wx.ALL, 5)
280

281
        # Set the sizer for the main panel
282
        self.panel.SetSizer(panel_sizer)
283

284
        # Create the menu bar
285
        menubar = wx.MenuBar()
286

287
        # Create a File menu
288
        file_menu = wx.Menu()
289
        file_menu.Append(wx.ID_EXIT, "Exit", "Exit the application")
290
        menubar.Append(file_menu, "&File")
291
        self.set_startup_states()
292

293
        # Set the menu bar for the frame
294
        self.SetMenuBar(menubar)
295
        # Set the frame properties
296
        self.SetSize((1900, 1000))
297
        self.SetTitle("CelestialSurveyor")
298
        self.Centre()
299
        self.Maximize(True)
300

301
        # Bind events
302
        self.Bind(wx.EVT_MENU, self.on_exit, id=wx.ID_EXIT)
303
        logger.log.info("App initialized")
304

305
    def on_text_change(self, event: Event) -> None:
306
        """
307
        Event handler for text change in the magnitude limit input field.
308
        If the input value is greater than 25, an error message is shown.
309

310
        Args:
311
            event: The event object that triggered this function.
312

313
        Returns:
314
            None
315
        """
316
        input_value = self.magnitude_input.GetValue()
317
        try:
318
            float_value = float(input_value)
319
            if float_value > 25:
320
                wx.MessageBox("Input value must be less than 25", "Error", wx.OK | wx.ICON_ERROR)
321
        except ValueError:
322
            wx.MessageBox("Input value must be a float", "Error", wx.OK | wx.ICON_ERROR)
323

324
    @staticmethod
325
    def __get_previous_settings() -> dict:
326
        """
327
        Retrieves the previous settings from the 'settings.json' file.
328
        Used to populate calibration folder paths. It's too annoying to type them every time.
329

330
        Returns:
331
            dict: The previous settings loaded from the file, or an empty dictionary if the file does not exist.
332
        """
333
        settings_file = "settings.json"
334
        if os.path.exists(settings_file):
335
            with open(settings_file, 'r') as fileo:
336
                return json.load(fileo)
337
        else:
338
            return {}
339

340
    @staticmethod
341
    def __write_previous_settings(settings: dict) -> None:
342
        """Writes the previous settings to the 'settings.json' file.
343

344
        Args:
345
            settings (dict): The settings to be written to the file.
346

347
        Returns:
348
            None
349
        """
350
        with open("settings.json", 'w') as fileo:
351
            json.dump(settings, fileo)
352

353
    def set_startup_states(self):
354
        """
355
        This function sets the startup states of the UI elements taking in account the previous settings.
356
        """
357
        prev_settings = self.__get_previous_settings()
358
        dark_path = prev_settings.get('dark_path', '')
359
        dark_path = "" if dark_path is None else dark_path
360
        self.dark_path_picker.SetPath(dark_path)
361
        flat_path = prev_settings.get('flat_path', '')
362
        flat_path = "" if flat_path is None else flat_path
363
        self.flat_path_picker.SetPath(flat_path)
364
        dark_flat_path = prev_settings.get('dark_flat_path', '')
365
        dark_flat_path = "" if dark_flat_path is None else dark_flat_path
366
        self.dark_flat_path_picker.SetPath(dark_flat_path)
367
        self.chk_to_align.Enable(False)
368
        self.to_debayer.Enable(False)
369
        self.chk_non_linear.Enable(False)
370
        self.btn_load_files.Enable(False)
371
        self.results_path_picker.Enable(False)
372
        self.results_label.Enable(False)
373
        self.magnitude_label.Enable(False)
374
        self.magnitude_input.Enable(False)
375
        self.btn_process.Enable(False)
376
        self.checkbox_list.SetObjects([])
377
        self.checkbox_list.Refresh(eraseBackground=False)
378
        self.results_path_picker.SetPath('')
379

380
    def on_exit(self, event: Event) -> None:
381
        """
382
        Event handler for the exit button.
383

384
        Args:
385
            event: The event object that triggered this function.
386

387
        Returns:
388
            None
389
        """
390
        logger.log.info("App exiting")
391
        self.Close()
392

393
    @staticmethod
394
    def _gen_short_file_names(file_list: list[str]) -> list[str]:
395
        """
396
        Generates short file names from a list of full file paths.
397
        Cuts off the common part of the file paths to make representation shorter.
398

399
        Args:
400
            file_list (list[str]): List of full file paths.
401

402
        Returns:
403
            list[str]: List of short file names.
404
        """
405
        folders = {os.path.split(fp)[0] for fp in file_list}
406
        split = os.path.split
407
        if len(folders) == 1:
408
            short_file_names = [split(fp)[1] for fp in file_list]
409
        else:
410
            short_file_names = [os.path.join(split(split(fp)[0])[1], split(fp)[1]) for fp in file_list]
411
        return short_file_names
412

413
    def on_header_loaded(self) -> None:
414
        """
415
        Event handler for the header loaded event.
416
        Displays the header information in the file list.
417
        Enables elements responsible for the next step: image loading, alignment and processing.
418

419
        Returns:
420
            None
421
        """
422
        self.source_data.stop_event.clear()
423
        self.progress_frame.Close()
424
        self.checkbox_list.SetObjects([])
425
        short_file_paths = self._gen_short_file_names([item.file_name for item in self.source_data.headers])
426
        self.checkbox_list.SetObjects([
427
            MyDataObject(
428
                fp, header.timestamp, header.exposure, checked=True
429
            ) for fp, header in zip(short_file_paths, self.source_data.headers)
430
        ])
431
        objects = self.checkbox_list.GetObjects()
432
        for obj in objects:
433
            self.checkbox_list.SetCheckState(obj, True)
434
        self.checkbox_list.RefreshObjects(objects)
435
        paths = [item.file_name for item in self.source_data.headers]
436
        logger.log.debug(f"Added the following files: {paths}")
437
        if paths:
438
            self.btn_load_files.Enable(True)
439
            self.chk_to_align.Enable(True)
440
            self.to_debayer.Enable(True)
441
            self.chk_non_linear.Enable(True)
442

443
    def handle_load_error(self, func: Callable, message: str, *args: Any, **kwargs: Any) -> Any:
444
        try:
445
            return func(*args, **kwargs)
446
        except Exception:
447
            self.progress_frame.label.SetLabel(message)
448
            self.progress_frame.label.Update()
449
            self.progress_frame.progress.SetValue(0)
450
            self.progress_frame.progress.SetRange(0)
451
            self.progress_frame.set_failed(True)
452
            raise
453

454
    def on_add_files(self, event: Event) -> None:
455
        """
456
        Event handler for the add files button.
457
        Opens dialog to select FIT(s) or XISF files.
458

459
        Args:
460
            event: The event object that triggered this function.
461

462
        Returns:
463
            None
464
        """
465
        wildcard = "Fits files (*.fits)|*.fits;*.FITS;*.fit;*.FIT|XISF files (*.xisf)|*.xisf;*.XISF|All files (*.*)|*.*"
466

467
        dialog = wx.FileDialog(self, message="Choose files to load",
468
                               wildcard=wildcard,
469
                               style=wx.FD_OPEN | wx.FD_MULTIPLE | wx.FD_FILE_MUST_EXIST)
470

471
        paths = []
472
        if dialog.ShowModal() == wx.ID_OK:
473
            paths = dialog.GetPaths()
474
            if self.source_data is None:
475
                self.source_data: SourceDataV2 = SourceDataV2()
476
            self.progress_frame = ProgressFrame(self, "Loading headers...", stop_event=self.source_data.stop_event)
477

478
            def load_headers():
479
                self.handle_load_error(
480
                    self.source_data.extend_headers, "Failed to load headers", file_list=paths,
481
                    progress_bar=ProgressBarFactory.create_progress_bar(self.progress_frame.progress))
482
                wx.CallAfter(self.on_header_loaded)
483
            self.process_thread = threading.Thread(target=load_headers)
484
            self.progress_frame.Show()
485
            self.process_thread.start()
486
        dialog.Destroy()
487

488
    def on_load_files(self, event: Event) -> None:
489
        """
490
        Event handler for loading files.
491
        Performs loading images, calibration, alignment, cropping and further processing.
492
        """
493

494
        self.progress_frame = ProgressFrame(self, "Loading progress", stop_event=self.source_data.stop_event)
495
        self.progress_frame.progress.SetValue(0)
496
        selected_headers = []
497
        objects = self.checkbox_list.GetObjects()
498
        for header, obj in zip(self.source_data.headers, objects):
499
            checked = self.checkbox_list.GetCheckState(obj)
500
            if checked:
501
                selected_headers.append(header)
502
            else:
503
                logger.log.debug(f"Excluding {header.file_name}")
504
        self.source_data.set_headers(selected_headers)
505
        short_file_paths = self._gen_short_file_names([item.file_name for item in self.source_data.headers])
506
        self.checkbox_list.SetObjects([
507
            MyDataObject(
508
                fp, header.timestamp, header.exposure, checked=True
509
            ) for fp, header in zip(short_file_paths, self.source_data.headers)
510
        ])
511
        objects = self.checkbox_list.GetObjects()
512
        for obj in objects:
513
            self.checkbox_list.SetCheckState(obj, True)
514
        self.checkbox_list.RefreshObjects(objects)
515

516
        to_align = self.chk_to_align.GetValue()
517

518
        logger.log.info(f"Loading {len(self.source_data.headers)} images...")
519
        dark_path = self.dark_path_picker.GetPath()
520
        dark_path = dark_path if dark_path else None
521
        flat_path = self.flat_path_picker.GetPath()
522
        flat_path = flat_path if flat_path else None
523
        dark_flat_path = self.dark_flat_path_picker.GetPath()
524
        dark_flat_path = dark_flat_path if dark_flat_path else None
525
        self.__write_previous_settings({
526
            'dark_path': dark_path,
527
            'flat_path': flat_path,
528
            'dark_flat_path': dark_flat_path
529
        })
530
        self.source_data.to_debayer = self.to_debayer.GetValue()
531
        self.source_data.linear = not self.chk_non_linear.GetValue()
532
        self.stop_event = Event()
533

534
        def load_images_calibrate_and_align() -> None:
535
            """
536
            This function loads, calibrates, and aligns images.
537
            It's run within the thread.
538
            """
539
            self.progress_frame.label.SetLabel("Loading images...")
540
            self.handle_load_error(
541
                self.source_data.load_images, "Failed to load images",
542
                progress_bar=ProgressBarFactory.create_progress_bar(self.progress_frame.progress))
543
            master_dark = None
544
            master_dark_flat = None
545
            master_flat = None
546
            if not self.source_data.stop_event.is_set():
547
                if dark_path:
548
                    if os.path.exists(dark_path):
549
                        self.progress_frame.label.SetLabel("Loading dark...")
550
                        master_dark = self.handle_load_error(
551
                            self.source_data.make_master_dark, "Failed to make master dark",
552
                            self.source_data.make_file_paths(dark_path),
553
                            progress_bar=ProgressBarFactory.create_progress_bar(self.progress_frame.progress))
554
                    else:
555
                        logger.log.warning(f"Dark path {dark_path} does not exist. Skipping make master dark...")
556
            if not self.source_data.stop_event.is_set():
557
                if dark_flat_path:
558
                    if os.path.exists(dark_flat_path):
559
                        self.progress_frame.label.SetLabel("Loading dark flat...")
560
                        master_dark_flat = self.handle_load_error(
561
                            self.source_data.make_master_dark, "Failed to make master dark flat",
562
                            self.source_data.make_file_paths(dark_flat_path),
563
                            progress_bar=ProgressBarFactory.create_progress_bar(self.progress_frame.progress))
564
                    else:
565
                        logger.log.warning(f"Dark flat path {dark_flat_path} does not exist. "
566
                                           f"Skipping make master dark flat...")
567

568
            if not self.source_data.stop_event.is_set():
569
                if flat_path:
570
                    if os.path.exists(flat_path):
571
                        self.progress_frame.label.SetLabel("Loading flats...")
572
                        flats = self.handle_load_error(
573
                            self.source_data.load_flats, "Failed to load flats",
574
                            self.source_data.make_file_paths(flat_path),
575
                            progress_bar=ProgressBarFactory.create_progress_bar(self.progress_frame.progress))
576
                        if master_dark_flat is not None:
577
                            for flat in flats:
578
                                flat -= master_dark_flat
579
                        master_flat = np.average(flats, axis=0)
580
            if master_dark is not None:
581
                self.source_data.original_frames -= master_dark
582
            if master_flat is not None:
583
                self.source_data.original_frames /= master_flat
584
            if to_align and not self.source_data.stop_event.is_set():
585
                self.progress_frame.label.SetLabel("Plate solving...")
586
                self.handle_load_error(
587
                    self.source_data.plate_solve_all, "Failed to plate solve images",
588
                    progress_bar=ProgressBarFactory.create_progress_bar(self.progress_frame.progress))
589
            if to_align and not self.source_data.stop_event.is_set():
590
                self.progress_frame.label.SetLabel("Aligning images...")
591
                self.handle_load_error(
592
                    self.source_data.align_images_wcs, "Failed to align images",
593
                    progress_bar=ProgressBarFactory.create_progress_bar(self.progress_frame.progress))
594
            if not self.source_data.stop_event.is_set():
595
                self.progress_frame.label.SetLabel("Cropping images...")
596
                self.handle_load_error(self.source_data.crop_images, "Failed to crop images")
597
            if not self.source_data.stop_event.is_set() and not self.chk_non_linear.GetValue():
598
                self.progress_frame.label.SetLabel("Stretching images... ")
599
                self.handle_load_error(
600
                    self.source_data.stretch_images, "Failed to stretch images",
601
                    progress_bar=ProgressBarFactory.create_progress_bar(self.progress_frame.progress))
602
            if not self.source_data.stop_event.is_set():
603
                self.source_data.images_from_buffer()
604
            else:
605
                self.source_data.original_frames = None
606
            self.source_data.stop_event.clear()
607
            self.progress_frame.Close()
608
            wx.CallAfter(self.on_load_finished)
609
        self.process_thread = threading.Thread(target=load_images_calibrate_and_align)
610
        self.progress_frame.Show()
611
        self.process_thread.start()
612

613
    def on_load_finished(self) -> None:
614
        """
615
        Event handler for the images loaded event.
616
        Displays the selected (first) image on image panel.
617
        Enables elements responsible for the next step: searching for moving objects and annotation.
618

619
        Returns:
620
            None
621
        """
622
        self.source_data.stop_event.clear()
623
        if self.source_data.images is not None and len(self.source_data.images) > 0:
624
            short_file_paths = self._gen_short_file_names([item.file_name for item in self.source_data.headers])
625
            self.checkbox_list.SetObjects([
626
                MyDataObject(
627
                    fp, header.timestamp, header.exposure, checked=True
628
                ) for fp, header in zip(short_file_paths, self.source_data.headers)
629
            ])
630
            objects = self.checkbox_list.GetObjects()
631
            for obj in objects:
632
                self.checkbox_list.SetCheckState(obj, True)
633
            self.checkbox_list.SelectObject(objects[0])
634
            self.checkbox_list.RefreshObjects(objects)
635
            img_to_draw = self.source_data.images[0]
636
            img_to_draw = (img_to_draw * 255).astype('uint8')
637
            self.draw_panel.image_array = img_to_draw
638
            self.draw_panel.Refresh(eraseBackground=False)
639
            self.results_label.Enable(True)
640
            self.magnitude_label.Enable(True)
641
            self.magnitude_input.Enable(True)
642
            self.results_path_picker.Enable(True)
643
            self.btn_process.Enable(True)
644

645
    def on_item_selected(self, event: Event) -> None:
646
        """
647
        Event handler for when an item is selected.
648
        Renders the selected item's image on the draw panel.
649
        """
650
        if self.source_data.images is not None and len(self.source_data.images) > 0:
651
            obj = self.checkbox_list.GetSelectedObject()
652
            file_paths = [item.file_name for item in self.source_data.headers]
653
            for num, item in enumerate(file_paths):
654
                if item.endswith(obj.file_path):
655
                    image_idx = num
656
                    break
657
            else:
658
                raise ValueError(f"{'obj.file_path'} is not in file list")
659
            img_to_draw = self.source_data.images[image_idx]
660
            img_to_draw = (img_to_draw * 255).astype('uint8')
661
            self.draw_panel.image_array = img_to_draw
662
            self.draw_panel.Refresh(eraseBackground=False)
663

664
    def on_process(self, event: Event) -> None:
665
        """
666
        Event handler for the 'Process' button click.
667
        Initiates the process of finding moving objects by AI model.
668

669
        Args:
670
            event (Event): The event object that triggered this function.
671

672
        Returns:
673
            None
674
        """
675
        self.progress_frame = ProgressFrame(self, "Finding moving objects", stop_event=self.source_data.stop_event)
676
        # self.progress_frame.label.SetLabel("Finding moving objects...")
677
        self.progress_frame.progress.SetValue(0)
678
        output_folder = self.results_path_picker.GetPath()
679
        objects = self.checkbox_list.GetObjects()
680
        use_img_mask = []
681
        for obj in objects:
682
            checked = self.checkbox_list.GetCheckState(obj)
683
            if checked:
684
                use_img_mask.append(True)
685
            else:
686
                use_img_mask.append(False)
687
        self.source_data.usage_map = np.array(use_img_mask, dtype=bool)
688

689
        def find_asteroids():
690
            """
691
            Predict asteroids in the given source data using AI model.
692
            The function to be run in processing thread.
693
            """
694

695
            self.progress_frame.label.SetLabel("Finding moving objects...")
696
            results = predict_asteroids(self.source_data, progress_bar=ProgressBarFactory.create_progress_bar(
697
                self.progress_frame.progress))
698
            if not self.source_data.stop_event.is_set():
699
                self.progress_frame.label.SetLabel("Saving results...")
700
                image_to_annotate = save_results(
701
                    source_data=self.source_data, results=results, output_folder=output_folder)
702
            if not self.source_data.stop_event.is_set():
703
                self.progress_frame.label.SetLabel("Annotating results...")
704
                magnitude_limit = float(self.magnitude_input.GetValue())
705
                annotate_results(self.source_data, image_to_annotate, output_folder, magnitude_limit=magnitude_limit)
706
            self.source_data.stop_event.clear()
707
            self.progress_frame.Close()
708

709
        self.progress_frame.Show()
710
        self.process_thread = threading.Thread(target=find_asteroids)
711
        self.process_thread.start()
712

713
    def on_process_finished(self) -> None:
714
        """
715
        Event handler for the process finished event.
716
        Closes the progress frame and resets the usage map in source data.
717
        """
718
        self.progress_frame.Close()
719
        self.source_data.usage_map = None
720

721
    def on_start_again(self, event: Event) -> None:
722
        """
723
        Event handler for the 'Start again' button click.
724
        Resets the source data and sets the startup states.
725

726
        Args:
727
            event (Event): The event object that triggered this function.
728

729
        Returns:
730
            None
731
        """
732
        self.source_data = None
733
        self.set_startup_states()
734

735

736
def start_ui():
737
    """
738
    Initializes the user interface and runs the main application loop.
739
    """
740
    app = wx.App(False)
741
    frame = MyFrame(None, wx.ID_ANY, "CelestialSurveyor", style=wx.DEFAULT_FRAME_STYLE)
742
    frame.Show()
743
    app.MainLoop()
744

745

746
if __name__ == '__main__':
747
    start_ui()
748