Update app.py

app.py

@@ -1,5 +1,10 @@
 import os
+import uuid
 import cv2 as cv
+import pandas as pd
+
+from zipfile import ZipFile
+
 
 import gradio as gr
 
@@ -20,6 +25,32 @@ def get_nth_frame(vc, n):
     img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
     return img
 
+def make_dataset(vc, df):
+    path = __file__+"/"+str(uuid.uuid1())
+    os.mkdir(path)
+    os.mkdir(path+"/train")
+    headers = ["image_name", "label"]
+    label_df = pd.DataFrame(columns=headers)
+    with ZipFile(path+"/train.zip", "w") as zf:
+        for line in df.itertuples():
+            start_idx = int(line[1])
+            end_idx = int(line[2])
+            label = int(line[3])
+            vc.set(cv.CAP_PROP_POS_FRAMES, start_idx)
+            for idx in range(start_idx, end_idx):
+                ok, img = vc.read()
+                if not ok:
+                    break
+                name = "image_%08d.jpg"%(idx)
+                cv.imwrite(path+"/train/"+name, img)
+                zf.write(path+"/train/"+name, arcname="train/"+name)
+                os.remove(path+"/train/"+name)
+                label_df = pd.concat([label_df, pd.DataFrame([[name, label]], columns=headers)])
+    label_df.to_csv(path+'/label.csv')
+    return [path+"/label.csv", path+"/train.zip"]
+    #return [ path+"/test.mp4", path+"/label.csv"]
+
+
 
 
 with gr.Blocks() as demo:
@@ -38,21 +69,17 @@ with gr.Blocks() as demo:
                 with gr.Accordion('Show Frame'):
                     input_n = gr.Slider(0, 9999, value=0, step=1, label="n")
                     btn = gr.Button(value="Submit")
-                    
+
             input_video.change(
-                video_change, 
-                inputs=[input_video, state_vc], 
+                video_change,
+                inputs=[input_video, state_vc],
                 outputs=[output_fps, output_fms, output_w, output_h, input_n, state_vc]
             )
             btn.click(get_nth_frame, inputs=[state_vc, input_n], outputs=output_img)
-        
+
         gr.Markdown("*****")
 
-        def make_dataset(df):
-            path = os.path.dirname(__file__)
-            print("path:", path)
-            df.to_csv('label.csv')
-            return [ path+"/test.mp4", path+"/label.csv"]
+
         with gr.Row():
             df = gr.Dataframe(
                 headers=["start_index", "end_index", "label"],
@@ -62,11 +89,11 @@ with gr.Blocks() as demo:
             )
             file = gr.File(file_count='multiple')
         btn_make_dataset = gr.Button(value="make dataset")
-        btn_make_dataset.click(make_dataset, inputs=df, outputs=file)
+        btn_make_dataset.click(make_dataset, inputs=[state_vc, df], outputs=file)
 
         gr.Markdown("*****")
         gr.Examples([os.path.join(os.path.dirname(__file__), "test.mp4")], inputs=input_video)
-            
+
 
 
 if __name__ == "__main__":