Python • Pygame display freezes during button press

Code:

import pygameimport randomimport math import numpy as npimport rembgfrom PIL import Imageimport RPi.GPIO as GPIO from time import sleep , localtime, strftimeimport subprocessimport osGPIO.setwarnings(False) GPIO.setmode(GPIO.BCM) GPIO.setup(2, GPIO.IN)pygame.init()screen_width = 500screen_height = 500screen = pygame.display.set_mode((screen_width, screen_height))background = pygame.image.load('back_down.png').convert()background = pygame.transform.scale(background, (500, 500))# Load multiple fish imagesfish_images = [    pygame.image.load('1437.png').convert_alpha(),]# Scale fish imagesfish_images = [pygame.transform.scale(fish, (25, 25)) for fish in fish_images]fish_list = []  # List to hold fish information# # Initialize fish positions, speeds, and timers# for fish_image in fish_images:    # fish_x = random.randint(0, screen_width - fish_image.get_width())    # fish_y = random.randint(0, screen_height - fish_image.get_height())    # speed_x_timer = 0    # speed_y_timer = 0    # speed_x = random.uniform(-1, 1) * 0.2    # speed_y = random.uniform(-1, 1) * 0.2    # fish_list.append({        # 'image': fish_image,        # 'x': fish_x,        # 'y': fish_y,        # 'speed_x': speed_x,        # 'speed_y': speed_y,        # 'speed_x_timer': speed_x_timer,        # 'speed_y_timer': speed_y_timer    # })clock = pygame.time.Clock()run = Truecount = 0while run:    dt = clock.tick(60)  # Target 60 FPS    for event in pygame.event.get():        if event.type == pygame.QUIT:            run = False    #if the button is pressed, capture an image, extractor outline and     #add to array of images      if GPIO.input(2) == 0:        name = strftime("%H%M",localtime())        os.system("fswebcam -d /dev/video0 -r 1280x720 /home/raspberry/python/rpitips/" + name + '.png')        input_image = Image.open(name + '.png')        input_array = np.array(input_image)        output_array = rembg.remove(input_array)        output_image = Image.fromarray(output_array)        output_image.save(name + '.png')        fish_images.append(pygame.image.load(name+'.png').convert_alpha())        fish_images = [pygame.transform.scale(fish, (25, 25)) for fish in fish_images]        count = 2        sleep(0.25)            if (count==0):        # Initialize fish positions, speeds, and timers        for fish_image in fish_images:            fish_x = random.randint(0, screen_width - fish_image.get_width())            fish_y = random.randint(0, screen_height - fish_image.get_height())                    speed_x_timer = 0            speed_y_timer = 0                    speed_x = random.uniform(-1, 1) * 0.2            speed_y = random.uniform(-1, 1) * 0.2                    fish_list.append({                'image': fish_image,                'x': fish_x,                'y': fish_y,                'speed_x': speed_x,                'speed_y': speed_y,                'speed_x_timer': speed_x_timer,                'speed_y_timer': speed_y_timer            })            count = count+1        for fish_info in fish_list:        fish_x = fish_info['x']        fish_y = fish_info['y']        speed_x = fish_info['speed_x']        speed_y = fish_info['speed_y']        speed_x_timer = fish_info['speed_x_timer']        speed_y_timer = fish_info['speed_y_timer']        fish_image = fish_info['image']        fish_x += speed_x * 18        fish_y += speed_y * 14        speed_x_timer += 1        speed_y_timer += 1        if speed_x_timer >= random.randint(500, 1200):            speed_x *= -1            speed_x_timer = 0        if speed_y_timer >= random.randint(500, 1200):            speed_y *= -1            speed_y_timer = 0                    if speed_x < 0:            flipped_fish = pygame.transform.flip(fish_image, True, False)        else:            flipped_fish = fish_image                    if fish_x < 0:            speed_x = random.uniform(-1, 1) * 0.2            fish_x = 0        elif fish_x > screen_width - fish_image.get_width():            speed_x = random.uniform(-1, 1) * 0.2            fish_x = screen_width - fish_image.get_width()        if fish_y < 0:            speed_y = random.uniform(-1, 1) * 0.2            fish_y = 0        elif fish_y > screen_height - fish_image.get_height():            speed_y = random.uniform(-1, 1) * 0.2            fish_y = screen_height - fish_image.get_height()                fish_info['x'] = fish_x        fish_info['y'] = fish_y        fish_info['speed_x'] = speed_x        fish_info['speed_y'] = speed_y        fish_info['speed_x_timer'] = speed_x_timer        fish_info['speed_y_timer'] = speed_y_timer    screen.blit(background, (0, 0))        for fish_info in fish_list:        fish_x = fish_info['x']        fish_y = fish_info['y']        fish_image = fish_info['image']        screen.blit(fish_image, (fish_x, fish_y))        #if the button is pressed, need to add the new fish to the list and give info    if (count==2):        fish_x = random.randint(0, screen_width - fish_images[-1].get_width())        fish_y = random.randint(0, screen_height - fish_images[-1].get_height())            speed_x_timer = 0        speed_y_timer = 0            speed_x = random.uniform(-1, 1) * 0.2        speed_y = random.uniform(-1, 1) * 0.2            fish_list.append({            'image': fish_images[-1],            'x': fish_x,            'y': fish_y,            'speed_x': speed_x,            'speed_y': speed_y,            'speed_x_timer': speed_x_timer,            'speed_y_timer': speed_y_timer        })        count=1       pygame.display.update()pygame.quit()