Update app.py

app.py

@@ -383,6 +383,29 @@ def switch_rows(v):
         print(frames[0])
         return frames
 
+def optimize(v, d):
+    global frames
+    global depths
+    
+    if d == False:
+      if v == True:
+        scale = 1
+        delta = 0
+        ddepth = cv2.CV_16S
+    
+        indices = np.arange(0,256)   # List of all colors 
+        divider = np.linspace(0,255,48+1)[1] # we get a divider
+        quantiz = np.int0(np.linspace(0,255,48)) # we get quantization colors
+        color_levels = np.clip(np.int0(indices/divider),0,48-1) # color levels 0,1,2..
+        palette = quantiz[color_levels]
+        for k, f in enumerate(frames):
+            frames[k] = cv2.convertScaleAbs(palette[frames[k]]).astype(np.uint8) # Converting image back to uint
+            frames[k] = cv2.convertScaleAbs(cv2.Sobel(frames[k], ddepth, 1, 1, ksize=3, scale=scale, delta=delta, borderType=cv2.BORDER_DEFAULT))
+
+      return frames
+    else:
+      return depths
+
 def reset_mask():
     global frame_selected
     global masks
@@ -590,7 +613,9 @@ with gr.Blocks(css=css, js=js) as demo:
             submit = gr.Button("Submit")
             output_frame = gr.Gallery(label="Frames", preview=True, columns=8192, type="numpy")
             output_switch = gr.Checkbox(label="Show depths")
+            optimize_switch = gr.Checkbox(label="Optimize")
             output_switch.input(fn=switch_rows, inputs=[output_switch], outputs=[output_frame])
+            optimize_switch.input(fn=optimize, inputs=[optimize_switch, output_switch], outputs=[output_frame])
             output_mask = gr.ImageEditor(layers=False, sources=('upload', 'clipboard'), show_download_button=True, type="numpy", interactive=True, transforms=(None,), eraser=gr.Eraser(), brush=gr.Brush(default_size=0, colors=['black', '#505050', '#a0a0a0', 'white']), elem_id="image_edit")
             with gr.Row():
               selector = gr.HTML(value="""