Is there any way to convert the screen of a pygame gui into an image somehow?

So I'm working on an interactive version of the MNIST handwritten digit image classification project with pygame, where the user draws on the gui and based off of that, the model which I have already made will look at the screen and output a prediction as to what number the image contains. My problem is that I'm not sure what approach to take to make whatever is displayed on the gui into an input of sorts for my model to predict (an image is required as an input)

Here's the code that I made for the gui:

import pygame as pg
#I'm importing a function that I made with my model
#Takes an image input and spits out a prediction as to what the number displayed in the image should be
from MNIST_Classification_GUI import makePrediction

pg.init()
screen = pg.display.set_mode([800, 600])
pg.display.set_caption("Draw a Number")

radius = 10
black = (0, 0, 0)
isGoing = True
screen.fill((255, 255, 255))
last_pos = (0, 0)


def roundline(srf, color, start, end, radius=1):
    dx = end[0]-start[0]
    dy = end[1]-start[1]
    distance = max(abs(dx), abs(dy))
    for i in range(distance):
        x = int( start[0]+float(i)/distance*dx)
        y = int( start[1]+float(i)/distance*dy)
        pg.draw.circle(srf, color, (x, y), radius)

#To be used for the popup text containing the prediction
pg.font.init()
myFont = pg.font.SysFont("Sans Serif", 10)

draw_on = False
while isGoing:
    for event in pg.event.get():
        if event.type == pg.QUIT:
            isGoing = False

        if event.type == pg.MOUSEBUTTONDOWN:
            spot = event.pos
            pg.draw.circle(screen, black, spot, radius)
            draw_on = True

        #This is the part where I want to somehow obtain an image from the gui
        #So when the user stops drawing, the popup text appears with the prediction
        if event.type == pg.MOUSEBUTTONUP:
            draw_on = False
            #The makePrediction takes an image input and returns the predicted value
            prediction = makePrediction(screen)
            textSurface = myFont.render(f"The number should be {prediction}", False, black)
            screen.blit(textSurface, (0, 0))

        if event.type == pg.MOUSEMOTION:
            if draw_on:
                pg.draw.circle(screen, black, event.pos, radius)
                roundline(screen, black, event.pos, last_pos, radius)
            last_pos = event.pos
        pg.display.flip()

pg.quit()