Update app.py

app.py

@@ -11,39 +11,47 @@ def encode(file_path):
     """
     Encodes a file by splitting it into chunks, rendering frames, and creating a video from the frames.
     """
-    yield "Splitting file into chunks..."
+    steps = gr.Steps()
+    steps.append("Splitting file into chunks...")
     split_file(input_file=file_path)
     
-    yield "Rendering frames from chunks..."
+    steps.append("Rendering frames from chunks...")
     render_frame()
     
-    yield "Creating video from frames..."
+    steps.append("Creating video from frames...")
     create_video_from_images()
     
-    yield "Cleaning up temporary files..."
-
+    steps.append("Cleaning up temporary files...")
+    shutil.rmtree("./chunks", ignore_errors=True)
+    shutil.rmtree("./images", ignore_errors=True)
+    shutil.rmtree("./extracted_frames", ignore_errors=True)
+    shutil.rmtree("./recovered", ignore_errors=True)
     
-    yield "Encoding completed! Encoded video is ready."
-    return "encoded_video.mp4"
+    steps.append("Encoding completed! Encoded video is ready.")
+    return steps, "encoded_video.mp4"
 
 def decode(output_file):
     """
     Decodes a video by extracting frames, recovering binary files, and merging the chunks back into the original file.
     """
-    yield "Extracting frames from video..."
+    steps = gr.Steps()
+    steps.append("Extracting frames from video...")
     extract_frames_from_video()
     
-    yield "Recovering binary files from frames..."
+    steps.append("Recovering binary files from frames...")
     recover_binary_files()
     
-    yield "Merging chunks to reconstruct the original file..."
+    steps.append("Merging chunks to reconstruct the original file...")
     merge_cunks("./recovered", output_file)
     
-    yield "Cleaning up temporary files..."
-  
+    steps.append("Cleaning up temporary files...")
+    shutil.rmtree("./extracted_frames", ignore_errors=True)
+    shutil.rmtree("./chunks", ignore_errors=True)
+    shutil.rmtree("./recovered", ignore_errors=True)
+    shutil.rmtree("./images", ignore_errors=True)
     
-    yield "Decoding completed! File successfully reconstructed."
-    return f"Decoding completed! Output file: {output_file}"
+    steps.append("Decoding completed! File successfully reconstructed.")
+    return steps, f"Decoding completed! Output file: {output_file}"
 
 def encode_interface(file):
     return encode(file.name)
@@ -58,20 +66,20 @@ with gr.Blocks() as demo:
     with gr.Tab("Encode"):
         file_input = gr.File(label="Upload file to encode")
         encode_button = gr.Button("Encode")
-        status_output = gr.Textbox(label="Status", interactive=False)
+        steps_output = gr.Steps(label="Encoding Steps")
         video_output = gr.Video(label="Encoded Video")
         download_button = gr.Button("Download Encoded Video")
         
-        encode_button.click(encode_interface, inputs=file_input, outputs=[status_output, video_output])
+        encode_button.click(encode_interface, inputs=file_input, outputs=[steps_output, video_output])
         download_button.click(lambda: "encoded_video.mp4", inputs=None, outputs=gr.File(label="Download Encoded Video"))
     
     with gr.Tab("Decode"):
         video_input = gr.Video(label="Upload video to decode")
         output_file_input = gr.Textbox(label="Output file name")
         decode_button = gr.Button("Decode")
-        decode_status = gr.Textbox(label="Status", interactive=False)
+        decode_steps = gr.Steps(label="Decoding Steps")
         
-        decode_button.click(decode_interface, inputs=[video_input, output_file_input], outputs=decode_status)
+        decode_button.click(decode_interface, inputs=[video_input, output_file_input], outputs=decode_steps)
 
 # Launch the Gradio app
 demo.launch()