Spaces:
Sleeping
Sleeping
File size: 4,545 Bytes
38c70a2 de13994 6e20ef2 de13994 eb9c960 38c70a2 b819df9 2ccc721 b42cdae 2ccc721 b42cdae a9eb8ce b42cdae 6f64a48 de13994 269a7d4 de13994 b42cdae de13994 460bbce de13994 a9eb8ce 4e9dd30 a9eb8ce b42cdae 2ccc721 b42cdae de13994 b42cdae de13994 b42cdae de13994 ad6a224 8b2e8f9 315c7c6 c0f6513 8b2e8f9 4e9dd30 b42cdae 6a0d672 b42cdae 6d48d5d d508011 b42cdae a9eb8ce de13994 6e20ef2 96ad467 38c70a2 4e9dd30 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import os
import uuid
import cv2 as cv
import pandas as pd
from zipfile import ZipFile
import gradio as gr
def video_change(v, vc):
if vc:
vc.release()
vc = cv.VideoCapture(v)
fps = vc.get(cv.CAP_PROP_FPS)
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),
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 _cut_frame(img, input_w1, input_w2, input_h1, input_h2)
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")
headers = ["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, _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)])
label_df.to_csv(path+'/label.csv', index=False)
return [path+"/label.csv", path+"/train.zip"]
with gr.Blocks() as demo:
with gr.Tab("Create Image Classify Dataset"):
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(
video_change,
inputs=[input_video, state_vc],
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"],
datatype=["number", "number", "number"],
height=500,
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, 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)
if __name__ == "__main__":
demo.queue().launch(debug=True) |