BigShoe/QuicksortVisualization

1import gradio as gr2import random3import matplotlib4matplotlib.use('agg')5import matplotlib.pyplot as plt6from PIL import Image7from io import BytesIO8import cv29import numpy as np10from multiprocessing import Pool, cpu_count11import warnings12warnings.filterwarnings("ignore")13 14def new_array(size, numRange):15    if numRange<=0:16        # generate random non repeating list starting from 117        arr = [i+1 for i in range(size)]18        random.shuffle(arr)19        return arr20    else:21        # generate random repeating list from 1 to numRange22        return [random.randint(1, numRange) for i in range(size)]23    24def swap(arr, index1, index2):25    # swap 2 elements in a list26    arr[index1], arr[index2] = arr[index2], arr[index1]27 28def quicksort(arr, ascending, leftBound=None, rightBound=None):29    if leftBound is None or rightBound is None:30        leftBound = 031        rightBound = len(arr)-132        steps.append([arr[:], leftBound, rightBound, -1, -1, "start", [-1, -1, -1], 1]) # store step -----33    else:34        steps.append([arr[:], leftBound, rightBound, -1, -1, "new sublist", [-1, -1, -1], steps[-1][-1]+1]) # store step -----35 36    # base case37    if rightBound-leftBound<1:38        return39 40    # get middle element41    mid = (rightBound+leftBound)//242 43    # get median of three44    l = arr[leftBound]45    m = arr[mid]46    r = arr[rightBound]47    if (l < m < r) or (r < m < l):48        pivotIndex = mid49    elif (m < l < r) or (r < l < m):50        pivotIndex = leftBound51    else:52        pivotIndex = rightBound53 54    steps.append([arr[:], leftBound, rightBound, -1, -1, "get left, middle, and right elements", [leftBound, mid, rightBound], steps[-1][-1]]) # store step -----55    steps.append([arr[:], leftBound, rightBound, -1, -1, "get median of three for pivot", [pivotIndex, -1, -1], steps[-1][-1]]) # store step -----56    57    swap(arr, rightBound, pivotIndex)58 59    steps.append([arr[:], leftBound, rightBound, -1, -1, "move pivot to right", [rightBound, -1, -1], steps[-1][-1]]) # store step -----60 61    # left and right pointers62    left = leftBound63    right = rightBound-164 65    steps.append([arr[:], leftBound, rightBound, left, right, "initialize left and right pointers", [rightBound, -1, -1], steps[-1][-1]]) # store step -----66 67    while left<right:68        # move left pointer69        while ((ascending and arr[left]<=arr[rightBound]) or (not ascending and arr[left]>=arr[rightBound])) and left<right:70            left+=171            steps.append([arr[:], leftBound, rightBound, left, right, "move left pointer", [rightBound, -1, -1], steps[-1][-1]]) # store step -----72 73        # more right pointer74        while ((ascending and arr[right]>=arr[rightBound]) or (not ascending and arr[right]<=arr[rightBound])) and left<right:75            right-=176            steps.append([arr[:], leftBound, rightBound, left, right, "move right pointer", [rightBound, -1, -1], steps[-1][-1]]) # store step -----77 78        swap(arr, left, right)79        steps.append([arr[:], leftBound, rightBound, left, right, "swap left and right pointer", [rightBound, -1, -1], steps[-1][-1]]) # store step -----80 81    # move pivot back82    if (ascending and arr[right]>arr[rightBound]) or (not ascending and arr[right]<arr[rightBound]):83        swap(arr, rightBound, right)84        steps.append([arr[:], leftBound, rightBound, left, right, "swap pivot and right pointer", [rightBound, -1, -1], steps[-1][-1]]) # store step -----85 86    # quicksort on the 2 sub lists87    quicksort(arr, ascending, leftBound, right)88    quicksort(arr, ascending, right+1, rightBound)89 90# GUI ------------------------------------------------------------91 92outputFilepath = "output.mp4"93 94def draw_step(stepIndex, stepsCopy):95    arr, leftBound, rightBound, left, right, stepText, pivotIndex, calls = stepsCopy[stepIndex]96 97    # set colors of bars98    colors = []99    for i in range(len(arr)):100        if stepIndex >= len(stepsCopy)-1:101            colors.append("green")102        elif i==pivotIndex[0] or i==pivotIndex[1] or i==pivotIndex[2]:103            colors.append("red")104        elif i==left and i==right:105            colors.append("brown")106        elif i==left:107            colors.append("orange")108        elif i==right:109            colors.append("blue")110        else:111            colors.append("white")112    113    # make bar graph114    fig, ax = plt.subplots()115    fig.patch.set_facecolor('black')116    ax.set_xticks([])117    ax.set_yticks([])118    ax.bar(range(len(arr)), arr, color=colors)119    ax.set_facecolor("black")120    ax.set_title("Calls " + str(calls) + "     Steps " + str(stepIndex) + "\n" + stepText, color="white")121 122    # make leftBound and rightBound lines123    if leftBound != -1:124        ax.axvline(x=leftBound-0.5, color="red", linestyle="solid", linewidth=2)125        ax.axvline(x=rightBound+0.5, color="red", linestyle="solid", linewidth=2)126 127    # save plot as image (https://stackoverflow.com/questions/57316491/how-to-convert-matplotlib-figure-to-pil-image-object-without-saving-image)128    buf = BytesIO()129    fig.savefig(buf)130    buf.seek(0)131    plt.close()132    return Image.open(buf)133 134def render_sort(size, numRange, ascending, fps):135    # reset program with new list and sort136    global steps137    steps = []138    arr = new_array(size, numRange)139    quicksort(arr, ascending)140    steps.append([arr[:], -1, -1, -1, -1, "sorted", [-1, -1, -1], steps[-1][-1]]) # store step -----141 142    # multiprocessing to draw steps (https://stackoverflow.com/questions/44660676/python-using-multiprocessing)143    results = []144    pool = Pool(processes=(cpu_count() - 1))145    for i in range(len(steps)):146        result = pool.apply_async(draw_step, args=(i,steps))147        results.append(result)148    pool.close()149    pool.join()150 151    # get images from multiprocessing results152    stepImages = []153    for i in range(len(steps)):154        stepImages.append(results[i].get())155 156    # generate video (https://stackoverflow.com/questions/52414148/turn-pil-images-into-video-on-linux)157    videodims = (640,480)158    fourcc = cv2.VideoWriter_fourcc(*'mp4v')    159    video = cv2.VideoWriter(outputFilepath,fourcc, fps,videodims)160    for img in stepImages:161        video.write(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))162 163    video.release()164 165    return outputFilepath166 167def returnInput(input):168    return input169 170with gr.Blocks() as demo:171    video = gr.Video(width=640, height=480)172 173    # variables174    fps = gr.State(10)175    size = gr.State(10)176    numRange = gr.State(0)177    ascending = gr.State(True)178 179    # ui180    fpsSlider = gr.Slider(label="FPS", minimum=1, maximum=120, step=1, value=10)181    sizeSlider = gr.Slider(label="Array Size", minimum=0, maximum=100, step=1, value=10)182    numRangeSlider = gr.Slider(label="Range of Numbers (0 for no repeat)", minimum=0, maximum=100, step=1, value=0)183    ascendingCheckBox = gr.Checkbox(label="Ascending", value=True)184    renderButton = gr.Button("Render Sort")185 186    # fps slider187    fpsSlider.release(188        fn=returnInput,189        inputs=fpsSlider,190        outputs=fps,191        show_progress=False192    )193 194    # sort settings195    sizeSlider.release(196        fn=returnInput,197        inputs=sizeSlider,198        outputs=size,199        show_progress=False200    )201 202    numRangeSlider.release(203        fn=returnInput,204        inputs=numRangeSlider,205        outputs=numRange,206        show_progress=False207    )208 209    ascendingCheckBox.change(210        fn=returnInput,211        inputs=ascendingCheckBox,212        outputs=ascending,213        show_progress=False214    )215 216    # render sort217    renderButton.click(218        fn=render_sort,219        inputs=[size, numRange, ascending, fps],220        outputs=video,221        show_progress=True222    )223 224if __name__ == "__main__":225    demo.launch()226