Optimizing the speed of a python program that uses PyQt5 and images

I’m trying to create a python program that let me visualize a stylized fictional political map where each country is represented by pixel of a specific color.
I’m using PyQt5 to create some kind of simple GUI that let me open the map, zoom and drag around and that shows on the side of the window some informations when you click on the map, like the country name and capital.
One of the functionalities is that once a country it’s clicked it’s color changes to a shade of bright green, so that it’s clear which country is selected. The way I implemented this works fine for small images, but it struggles with larger images, with up to 6-7 seconds passing between the moment I click on a country and the moment when the country color changes to green.

This is the code I used to implement the color changing upon clicking on the country:
(note that there are three timers included as I was trying to figure out what parts of the function were slower)

    def highlight_pixels(self, RGB):
        #BLOCK ONE
        #Timer 1 starts
        start_time = time.perf_counter()
        # Copy a numpy array that represent the base image. This is so that every time the highlight_pixels function is called it reverts previous changes
        temp_image = self.base_image_array.copy()
        # Print time 1
        end_time = time.perf_counter()
        elapsed_time = (end_time - start_time) * 1000
        print(f"Time taken to initialise image: {elapsed_time:.6f} milliseconds")
        
        #BLOCK TWO
        #Timer 2 starts
        start_time = time.perf_counter()
        # Select pixels that match the target color
        mask = (temp_image[:,:,0] == RGB[0]) & (temp_image[:,:,1] == RGB[1]) & (temp_image[:,:,2] == RGB[2])
        # Set color of the selected pixels to a bright green.
        temp_image[self.ignore_color_mask & mask, :] = (83, 255, 26)
        # Print time 2
        end_time = time.perf_counter()
        elapsed_time = (end_time - start_time) * 1000
        print(f"Time taken to change color: {elapsed_time:.6f} milliseconds")
        
        #BLOCK THREE
        #Timer 3 starts
        start_time = time.perf_counter()
        # Convert array back to image (qimage)
        temp_image = qimage2ndarray.array2qimage(temp_image)
        # convert the image back to pixmap
        self.map_image = QPixmap.fromImage(temp_image)
        # update the map scene
        self.view.scene().clear()
        self.view.scene().addPixmap(self.map_image)
        # Print time 3
        end_time = time.perf_counter()
        elapsed_time = (end_time - start_time) * 1000
        print(f"Time taken to show image: {elapsed_time:.6f} milliseconds")

In a test run with two images, one large and one small, I got this times for the three blocks. Obviously this change a bit everytime I run the program, but the results are always quite similar:

I tried a couple of different things. I first tought of somehow trying to create a new image with the green pixels and overlay that on the base image, instead of making changes to the entire image, but I don’t think that was implemented properly and decided to go back.

Then I thought I’d be able to speed things up by creating a dictionary that contained the coordinates of all pixel of a certain color. This worked somewhat, meaning that when tested on the small image (200×100 px) it showed some signs of improvement:

However, when trying to use this approach with the large image (16300×8150 px) it simply runs out of memory when trying to create the dictionary.

This is the updated function:

    def highlight_pixels(self, RGB):
        #BLOCK ONE
        #Timer 1 starts
        start_time = time.perf_counter()
        # Copy a numpy array that represent the base image. This is so that every time the highlight_pixels function is called it reverts previous changes
        temp_image = self.base_image_array.copy()
        # Print time 1
        end_time = time.perf_counter()
        elapsed_time = (end_time - start_time) * 1000
        print(f"Time taken to initialise image: {elapsed_time:.6f} milliseconds")
        
        #BLOCK TWO
        #Timer 2 starts
        start_time = time.perf_counter()
        # Select pixels that match the target color
        coordinates = self.color_dict.get(RGB)
        # Set their color to green
        temp_image[coordinates[:, 0], coordinates[:, 1], :] = (83, 255, 26)    
        # Print time 2
        end_time = time.perf_counter()
        elapsed_time = (end_time - start_time) * 1000
        print(f"Time taken to change color: {elapsed_time:.6f} milliseconds")
        
        #BLOCK THREE
        #Timer 3 starts
        start_time = time.perf_counter()
        #convert array back to image (qimage)
        temp_image = qimage2ndarray.array2qimage(temp_image)
        # convert the image back to pixmap
        self.map_image = QPixmap.fromImage(temp_image)
        # update the map scene
        self.view.scene().clear()
        self.view.scene().addPixmap(self.map_image)
        # Print time 3
        end_time = time.perf_counter()
        elapsed_time = (end_time - start_time) * 1000
        print(f"Time taken to show image: {elapsed_time:.6f} milliseconds")

And here is the function that creates the dictionary I mentioned:

def create_color_dict(self, image_array):
        self.color_dict = {}
        # Create an array of indices for each pixel
        pixels = np.indices(image_array.shape[:2]).T.reshape(-1, 2)
        # Get the color of each pixel
        pixel_colors = image_array[pixels[:, 0], pixels[:, 1], :]
        # Convert pixel colors to tuples
        pixel_colors = tuple(map(tuple, pixel_colors))
        # Create dictionary of pixels by color
        for color, pixel in zip(pixel_colors, pixels):
            if color not in self.color_dict:
                self.color_dict[color] = np.array([pixel])
            else:
                self.color_dict[color] = np.append(self.color_dict[color], [pixel], axis=0)
        return self.color_dict

I also tried to make a separate program to create the dictionary (as I would only need to run that part once, really), but it obviously runs out of memory as well.

So my question really is how could this be optimized? Is the approach I attempted valid? If that is the case, how can I circumvent the running out of memory?

I would be more than happy to provide any additional information, the entire code or the test images. Bare with me for the jankiness or anything weird: I’m an absolute beginner and it’s really hard to understand how not to make something just work, but how to make it work better.

edit:
Here is the entire code in case anyone asks, since I won’t be able to have access to a pc until tomorrow. Note that it might be a little confusing as I’m a noob at this.

import sys
from PyQt5.QtCore import Qt, QEvent
from PyQt5.QtGui import QImage, QPixmap, QColor, QPainter, QPen, QBrush, qRed, qGreen, qBlue, QFont
from PyQt5.QtWidgets import QApplication, QGraphicsView, QGraphicsScene, QLabel, QVBoxLayout, QWidget, QGraphicsRectItem, QMainWindow, QDockWidget, QHBoxLayout, QSizePolicy, QGraphicsPixmapItem
import time
import qimage2ndarray
from qimage2ndarray import rgb_view
import numpy as np

class MapViewer(QMainWindow):
    def __init__(self, map_path):
        super().__init__()
        self.setWindowTitle("Map Viewer")
        self.setGeometry(100, 100, 800, 600)
        self.dragCheck = False #False if not dragging, True if dragging

        # Create the QGraphicsView widget and set it as the central widget
        self.view = QGraphicsView(self)
        self.setCentralWidget(self.view)
        self.view.setRenderHint(QPainter.Antialiasing)
        self.view.setRenderHint(QPainter.SmoothPixmapTransform)
        self.view.setDragMode(QGraphicsView.NoDrag)
        self.view.setTransformationAnchor(QGraphicsView.AnchorUnderMouse)
        self.view.setResizeAnchor(QGraphicsView.AnchorUnderMouse)
        self.view.setViewportUpdateMode(QGraphicsView.FullViewportUpdate)
        self.view.setOptimizationFlag(QGraphicsView.DontAdjustForAntialiasing, True)
        self.view.setOptimizationFlag(QGraphicsView.DontSavePainterState, True)

        # Load the map image and set it as the background image
        self.map_image = QPixmap(map_path)
        #self.original_map_image = self.map_image.copy() # Assign the original image to this variable
        #self.map_item = QGraphicsPixmapItem(self.map_image)
        self.view.setScene(QGraphicsScene())
        self.view.scene().addPixmap(self.map_image)
        #self.view.fitInView(self.view.sceneRect(), Qt.KeepAspectRatio)
        
        # convert the pixmap to image
        self.base_image = QImage(map_path)
        # Convert image into numpy array
        self.base_image_array = rgb_view(self.base_image)
        
        # Create a mask with the same shape as the image, filled with True
        self.ignore_color_mask = np.ones(self.base_image_array.shape[:2], dtype=bool)
        # Set False to pixels that have black color (0, 0, 0) or sea color (172, 208, 239)
        self.ignore_color_mask[np.where((self.base_image_array[:,:,0]==0) & (self.base_image_array[:,:,1]==0) & (self.base_image_array[:,:,2]==0)| (self.base_image_array[:,:,0] == 172) & (self.base_image_array[:,:,1] == 208) & (self.base_image_array[:,:,2] == 239))] = False
        
        # Make it so the wheel zooms in and out
        self.view.wheelEvent = self.handle_mouse_wheel
        
        # Install an eventFilter to handle the mouse clicks under certain conditions only (to allow panning when pressing control)
        self.view.installEventFilter(self)
        
        # Handle ctrl beign pressed or released
        self.view.keyPressEvent = self.CtrlPressEvent
        self.view.keyReleaseEvent = self.CtrlReleaseEvent
        
        # Display the default country information in the dock widget
        RGB = (255, 0, 0)
        self.countries = {
            (255, 0, 0): {"name": "Red", "rgb": (255, 0, 0), "capital": "Red Capital", "government": "Red Government", "size": "Red Size", "population": "Red Population"},
            (0, 255, 0): {"name": "Green", "rgb": (0, 255, 0), "capital": "Green Capital", "government": "Green Government", "size": "Green Size", "population": "Green Population"},
            (0, 0, 255): {"name": "Blue", "rgb": (0, 0, 255), "capital": "Blue Capital", "government": "Blue Government", "size": "Blue Size", "population": "Blue Population"}
        }
        self.display_country_info(RGB)
        self.create_color_dict(self.base_image_array)

    def create_color_dict(self, image_array):
        self.color_dict = {}
        # Create an array of indices for each pixel
        pixels = np.indices(image_array.shape[:2]).T.reshape(-1, 2)
        # Get the color of each pixel
        pixel_colors = image_array[pixels[:, 0], pixels[:, 1], :]
        # Convert pixel colors to tuples
        pixel_colors = tuple(map(tuple, pixel_colors))
        # Create dictionary of pixels by color
        for color, pixel in zip(pixel_colors, pixels):
            if color not in self.color_dict:
                self.color_dict[color] = np.array([pixel])
            else:
                self.color_dict[color] = np.append(self.color_dict[color], [pixel], axis=0)
        return self.color_dict

        
        
    def handle_mouse_wheel(self, event):
        if event.angleDelta().y() > 0:
            # Scroll up, zoom in
            self.view.scale(1.1, 1.1)
        else:
            # Scroll down, zoom out
            self.view.scale(1 / 1.1, 1 / 1.1)

    def display_country_info(self, RGB):
        # Look up the country information based on the RGB value
        country_info = self.countries.get(RGB)
        if country_info is None:
            # Handle the case where the RGB value is not found in the dictionary
            print("Sorry, no country found with that RGB value.")
            return
        else:
            #Remove any existing dock widgets
            for dock_widget in self.findChildren(QDockWidget):
                self.removeDockWidget(dock_widget)

            # Create a QVBoxLayout to hold the labels
            layout = QVBoxLayout()
            
            # Create a QLabel for each piece of information and add it to the layout
            name_label = QLabel("Name:")
            name_label.setStyleSheet("font-weight: bold;")
            name_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)
            value_label = QLabel(f"{country_info['name']}")
            value_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

            # create a new layout for the name and value labels
            layout_name_value = QHBoxLayout()

            # add both labels to the new layout
            layout_name_value.addWidget(name_label)
            layout_name_value.addWidget(value_label)

            # add this new layout to the main layout
            layout.addLayout(layout_name_value)
         
            # Create a QLabel for each piece of information and add it to the layout
            name_label = QLabel("RGB:")
            name_label.setStyleSheet("font-weight: bold;")
            name_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)
            value_label = QLabel(f"{country_info['rgb']}")
            value_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

            # create a new layout for the name and value labels
            layout_name_value = QHBoxLayout()

            # add both labels to the new layout
            layout_name_value.addWidget(name_label)
            layout_name_value.addWidget(value_label)

            # add this new layout to the main layout
            layout.addLayout(layout_name_value)
            
            # Create a QLabel for each piece of information and add it to the layout
            name_label = QLabel("Capital:")
            name_label.setStyleSheet("font-weight: bold;")
            name_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)
            value_label = QLabel(f"{country_info['capital']}")
            value_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

            # create a new layout for the name and value labels
            layout_name_value = QHBoxLayout()

            # add both labels to the new layout
            layout_name_value.addWidget(name_label)
            layout_name_value.addWidget(value_label)

            # add this new layout to the main layout
            layout.addLayout(layout_name_value)
            
            # Create a QLabel for each piece of information and add it to the layout
            name_label = QLabel("Government:")
            name_label.setStyleSheet("font-weight: bold;")
            name_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)
            value_label = QLabel(f"{country_info['government']}")
            value_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

            # create a new layout for the name and value labels
            layout_name_value = QHBoxLayout()

            # add both labels to the new layout
            layout_name_value.addWidget(name_label)
            layout_name_value.addWidget(value_label)

            # add this new layout to the main layout
            layout.addLayout(layout_name_value)
            
            # Create a QLabel for each piece of information and add it to the layout
            name_label = QLabel("Size:")
            name_label.setStyleSheet("font-weight: bold;")
            name_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)
            value_label = QLabel(f"{country_info['size']}")
            value_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

            # create a new layout for the name and value labels
            layout_name_value = QHBoxLayout()

            # add both labels to the new layout
            layout_name_value.addWidget(name_label)
            layout_name_value.addWidget(value_label)

            # add this new layout to the main layout
            layout.addLayout(layout_name_value) 

            # Create a QLabel for each piece of information and add it to the layout
            name_label = QLabel("Population:")
            name_label.setStyleSheet("font-weight: bold;")
            name_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)
            value_label = QLabel(f"{country_info['population']}")
            value_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

            # create a new layout for the name and value labels
            layout_name_value = QHBoxLayout()

            # add both labels to the new layout
            layout_name_value.addWidget(name_label)
            layout_name_value.addWidget(value_label)

            # add this new layout to the main layout
            layout.addLayout(layout_name_value) 

            # Create a QWidget to hold the layout and add it to the right dock
            
            widget = QWidget()
            widget.setLayout(layout)
            dock_widget = QDockWidget()
            dock_widget.setWindowTitle("Country Informations")
            dock_widget.setWidget(widget)
            self.addDockWidget(Qt.RightDockWidgetArea, dock_widget)
            dock_widget.setFeatures(QDockWidget.NoDockWidgetFeatures) # removes close button and full screen button
            #dock_widget.setFixedSize(int(self.width()*0.20), self.height())
            dock_widget.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)

    def highlight_pixels(self, RGB):
        #BLOCK ONE
        #Timer 1 starts
        start_time = time.perf_counter()
        # Copy a numpy array that represent the base image. This is so that every time the highlight_pixels function is called it reverts previous changes
        temp_image = self.base_image_array.copy()
        # Print time 1
        end_time = time.perf_counter()
        elapsed_time = (end_time - start_time) * 1000
        print(f"Time taken to initialise image: {elapsed_time:.6f} milliseconds")
        
        #BLOCK TWO
        #Timer 2 starts
        start_time = time.perf_counter()
        # Select pixels that match the target color
        coordinates = self.color_dict.get(RGB)
        # Set their color to green
        temp_image[coordinates[:, 0], coordinates[:, 1], :] = (83, 255, 26)    
        # Print time 2
        end_time = time.perf_counter()
        elapsed_time = (end_time - start_time) * 1000
        print(f"Time taken to change color: {elapsed_time:.6f} milliseconds")
        
        #BLOCK THREE
        #Timer 3 starts
        start_time = time.perf_counter()
        #convert array back to image (qimage)
        temp_image = qimage2ndarray.array2qimage(temp_image)
        # convert the image back to pixmap
        self.map_image = QPixmap.fromImage(temp_image)
        # update the map scene
        self.view.scene().clear()
        self.view.scene().addPixmap(self.map_image)
        # Print time 3
        end_time = time.perf_counter()
        elapsed_time = (end_time - start_time) * 1000
        print(f"Time taken to show image: {elapsed_time:.6f} milliseconds")
        
    # Set up an event filter to recognize whether it should pan around or print informations.
    def eventFilter(self, obj, event):
        if event.type() == QEvent.MouseButtonPress and event.button() == Qt.LeftButton:
            if self.dragCheck == False:
                # Get the position of the mouse click in viewport coordinates
                pos = event.pos()
                # Convert the viewport coordinates to scene coordinates
                scene_pos = self.view.mapToScene(pos)
                # Get the pixel at the scene coordinates
                pixel = self.map_image.toImage().pixel(int(scene_pos.x()), int(scene_pos.y()))
                # Get the red, green, and blue components of the pixel
                red = qRed(pixel)
                green = qGreen(pixel)
                blue = qBlue(pixel)
                # Assign the RGB values to the RGB variable
                RGB = (red, green, blue)
                print("RGB:", RGB) # You can replace this with the call to display_country_info with the RGB variable
                self.display_country_info(RGB)
                self.highlight_pixels(RGB)
                return True
        return QMainWindow.eventFilter(self, obj, event)
    
    # Check if Ctrl is beign pressed
    def CtrlPressEvent(self, event):
        if event.key() == Qt.Key_Control:
            self.dragCheck = True
            self.view.setDragMode(QGraphicsView.ScrollHandDrag)
            #print("drag")
            
            
    # Check if Ctrl is beign released        
    def CtrlReleaseEvent(self, event):
        if event.key() == Qt.Key_Control:
            self.dragCheck = False
            self.view.setDragMode(QGraphicsView.NoDrag)
            #print("nodrag", self.dragCheck)
            

          
if __name__ == "__main__":
    app = QApplication(sys.argv)
    viewer = MapViewer("map.png")
    #viewer = MapViewer("countries_ee.png")
    viewer.show()
    
    sys.exit(app.exec_())