Update app.py

app.py

@@ -17,15 +17,24 @@ def video_change(v, vc):
     fms = vc.get(cv.CAP_PROP_FRAME_COUNT)
     w = vc.get(cv.CAP_PROP_FRAME_WIDTH)
     h = vc.get(cv.CAP_PROP_FRAME_HEIGHT)
-    return fps, fms, w, h, gr.update(maximum=int(fms)-1), vc
-
-def get_nth_frame(vc, n):
+    return [fps, fms, w, h, 
+            gr.update(maximum=int(fms)-1), 
+            gr.update(maximum=int(w)), 
+            gr.update(maximum=int(w), value=int(w)), 
+            gr.update(maximum=int(h)), 
+            gr.update(maximum=int(h), value=int(h)), 
+            vc]
+
+def _cut_frame(img, input_w1, input_w2, input_h1, input_h2):
+    return img[input_h1:input_h2, input_w1:input_w2]
+
+def get_nth_frame(vc, n, input_w1, input_w2, input_h1, input_h2):
     vc.set(cv.CAP_PROP_POS_FRAMES, n)
     ok, img = vc.read()
     img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
-    return img
+    return _cut_frame(img, input_w1, input_w2, input_h1, input_h2)
 
-def make_dataset(vc, df):
+def make_dataset(vc, df, input_w1, input_w2, input_h1, input_h2):
     path = os.path.dirname(__file__)+"/"+str(uuid.uuid1())
     os.mkdir(path)
     os.mkdir(path+"/train")
@@ -42,7 +51,7 @@ def make_dataset(vc, df):
                 if not ok:
                     break
                 name = "image_%08d.jpg"%(idx)
-                cv.imwrite(path+"/train/"+name, img)
+                cv.imwrite(path+"/train/"+name, _cut_frame(img, input_w1, input_w2, input_h1, input_h2))
                 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)])
@@ -53,31 +62,45 @@ def make_dataset(vc, df):
 
 with gr.Blocks() as demo:
     with gr.Tab("Create Image Classify Dataset"):
-        with gr.Row():
-            with gr.Column():
-                state_vc = gr.State(value=None)
-                input_video = gr.Video()
-                with gr.Accordion('Video Info'):
-                    output_fps = gr.Textbox(lines=1, label="fps")
-                    output_fms = gr.Textbox(lines=1, label="frame count")
-                    output_w = gr.Textbox(lines=1, label="width")
-                    output_h = gr.Textbox(lines=1, label="height")
-            with gr.Column():
-                output_img = gr.Image()
-                with gr.Accordion('Show Frame'):
-                    input_n = gr.Slider(0, 9999, value=0, step=1, label="n")
+
+        with gr.Accordion('Step 1: Video Info'):
+            with gr.Row():
+                with gr.Column():
+                    state_vc = gr.State(value=None)
+                    input_video = gr.Video()
+                with gr.Column():
+                    output_fps = gr.Number(label="fps")
+                    output_fms = gr.Number(label="frame count")
+                    output_w = gr.Number(label="width")
+                    output_h = gr.Number(label="height")
+        
+        gr.Markdown("*****")
+
+        with gr.Accordion('Step 2: Frame Info'):
+            with gr.Row():
+                with gr.Column():
+                    input_n = gr.Slider(0, 9999, value=0, step=1, label="nth frame")
+                    input_w1 = gr.Slider(0, 9999, value=0, step=1, label="w1")
+                    input_w2 = gr.Slider(0, 9999, value=0, step=1, label="w2")
+                    input_h1 = gr.Slider(0, 9999, value=0, step=1, label="h1")
+                    input_h2 = gr.Slider(0, 9999, value=0, step=1, label="h2")
                     btn = gr.Button(value="Submit")
+                with gr.Column():
+                    output_img = gr.Image()
 
-            input_video.change(
+        input_video.change(
                 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)
+                outputs=[output_fps, output_fms, output_w, output_h, 
+                         input_n, input_w1, input_w2, input_h1, input_h2, 
+                         state_vc]
+        )
+        btn.click(get_nth_frame, 
+                  inputs=[state_vc, input_n, input_w1, input_w2, input_h1, input_h2], 
+                  outputs=output_img)
 
         gr.Markdown("*****")
 
-
         with gr.Row():
             df = gr.Dataframe(
                 headers=["start_index", "end_index", "label"],
@@ -86,10 +109,14 @@ with gr.Blocks() as demo:
                 interactive=True,
             )
             file = gr.File(file_count='multiple')
+
         btn_make_dataset = gr.Button(value="make dataset")
-        btn_make_dataset.click(make_dataset, inputs=[state_vc, df], outputs=file)
+        btn_make_dataset.click(make_dataset, 
+                               inputs=[state_vc, df, input_w1, input_w2, input_h1, input_h2 ], 
+                               outputs=file)
 
         gr.Markdown("*****")
+
         gr.Examples([os.path.join(os.path.dirname(__file__), "test.mp4")], inputs=input_video)