hbfreed commited on
Commit
6f1d1cd
•
1 Parent(s): df4c6a1

loading should be nicer

Browse files
Files changed (1) hide show
  1. app.py +17 -17
app.py CHANGED
@@ -28,7 +28,7 @@ def get_demo_call(video_name):
28
  ump_out = None
29
  return ump_out
30
 
31
- def call_pitch(pitch):
32
  start = time.time()
33
  std = (0.2104, 0.1986, 0.1829)
34
  mean = (0.3939, 0.3817, 0.3314)
@@ -42,14 +42,14 @@ def call_pitch(pitch):
42
  model = MobileNetSmall3D()
43
  model.load_state_dict(torch.load('weights/MobileNetSmall.pth',map_location=torch.device('cpu')))
44
  model.eval()
45
-
46
- #run the model
47
  with torch.no_grad():
48
  output = model(pitch_tensor)
49
  output = F.softmax(output,dim=1)
50
  final_call = classify_pitch(output)
51
- # time.sleep(video_length-1.5) #wait until the video is done to return the call
52
- while time.time() - start < (video_length-5)/15: #wait until the video is done to return the call and go 2 frames early to seem speedy
 
 
53
  time.sleep(0.1)
54
 
55
  if ump_out is not None:
@@ -74,26 +74,25 @@ def generate_random_pitch():
74
 
75
  random_pitch = download_pitch(video_link,ump_out)
76
 
77
- return random_pitch
78
 
79
 
80
  def download_pitch(video_link,ump_out=None):
81
- # progress = gr.Progress()
82
- # progress(1,"Downloading Video...")
83
  #download and process the video link using yt-dlp and ffmpeg
 
84
  if ump_out is not None:
85
  video_out = f"demo_files/downloaded_videos/{ump_out}_demo_video.mp4"
86
- os.system(f"yt-dlp -q --no-warnings --force-overwrites --progress -f mp4 {video_link} -o '{video_out}'")
87
- # progress(None,"Processing Video...")
88
  os.system(f'ffmpeg -y -nostats -loglevel 0 -i {video_out} -vf "crop=700:700:in_w/2-350:in_h/2-350,scale=224:224" -r 15 -c:v libx264 -crf 23 -c:a aac -an -strict experimental {video_out}_cropped_video.mp4')
89
  else:
90
  video_out = f"demo_files/downloaded_videos/own_demo_video.mp4"
91
  os.system(f"yt-dlp -q --no-warnings --force-overwrites -f mp4 {video_link} -o '{video_out}'")
92
- # progress(None,"Processing Video...")
93
  os.system(f'ffmpeg -y -nostats -loglevel 0 -i {video_out} -vf "crop=700:700:in_w/2-350:in_h/2-350,scale=224:224" -r 15 -c:v libx264 -crf 23 -c:a aac -an -strict experimental {video_out}_cropped_video.mp4')
 
 
 
94
  return f"{video_out}_cropped_video.mp4"
95
 
96
-
97
  demo_files = os.listdir("demo_files/")
98
  demo_files = [os.path.join("demo_files/", file) for file in demo_files if file.endswith(".mp4")]
99
  with gr.Blocks(title="Picklebot") as demo:
@@ -110,9 +109,10 @@ with gr.Blocks(title="Picklebot") as demo:
110
  with gr.Column(scale=2):
111
  ump_out = gr.Label(label="Umpire's Original Call")
112
  pb_out = gr.Label(label="Picklebot's Call")
 
113
  random_button = gr.Button("🔀 Load a Random Pitch")
114
- random_button.click(fn=generate_random_pitch,outputs=[inp],show_progress="minimal")
115
- inp.play(fn=call_pitch,inputs=inp,outputs=[pb_out,ump_out],show_progress="minimal")
116
  gr.ClearButton([inp,pb_out,ump_out])
117
  gr.Examples(demo_files,inputs=inp,outputs=[inp])
118
  with gr.Tab("Use Your Own Video!"):
@@ -133,11 +133,11 @@ with gr.Blocks(title="Picklebot") as demo:
133
  with gr.Column(scale=2):
134
  #make a textbox to take in the user's video link
135
  input_txt = gr.Textbox(placeholder="Paste your video link here",label="Video Link")
136
- input_txt.submit(fn=download_pitch,inputs=input_txt,outputs=vid_inp,show_progress="minimal")
137
  submit_button = gr.Button("Go!")
138
- submit_button.click(fn=download_pitch,inputs=input_txt,outputs=vid_inp,show_progress="minimal")
139
  pb_out = gr.Label(label="Picklebot's Call")
140
- vid_inp.play(fn=call_pitch,inputs=vid_inp,outputs=pb_out,show_progress="minimal")
141
  gr.ClearButton([vid_inp,pb_out,input_txt])
142
 
143
  with gr.Tab("About"):
 
28
  ump_out = None
29
  return ump_out
30
 
31
+ def call_pitch(pitch,progress=gr.Progress()):
32
  start = time.time()
33
  std = (0.2104, 0.1986, 0.1829)
34
  mean = (0.3939, 0.3817, 0.3314)
 
42
  model = MobileNetSmall3D()
43
  model.load_state_dict(torch.load('weights/MobileNetSmall.pth',map_location=torch.device('cpu')))
44
  model.eval()
 
 
45
  with torch.no_grad():
46
  output = model(pitch_tensor)
47
  output = F.softmax(output,dim=1)
48
  final_call = classify_pitch(output)
49
+
50
+ differential = time.time() - start
51
+ wait_time = (int(video_length-(5+differential))//15)*10
52
+ for _ in progress.tqdm(range(wait_time)): #wait until the video is done to return the call and go 5 frames early to seem speedy
53
  time.sleep(0.1)
54
 
55
  if ump_out is not None:
 
74
 
75
  random_pitch = download_pitch(video_link,ump_out)
76
 
77
+ return random_pitch,None,None #none and none to write none to ump out and pb out
78
 
79
 
80
  def download_pitch(video_link,ump_out=None):
 
 
81
  #download and process the video link using yt-dlp and ffmpeg
82
+ start = time.time()
83
  if ump_out is not None:
84
  video_out = f"demo_files/downloaded_videos/{ump_out}_demo_video.mp4"
85
+ os.system(f"yt-dlp -q --no-warnings --force-overwrites -f mp4 {video_link} -o '{video_out}'")
 
86
  os.system(f'ffmpeg -y -nostats -loglevel 0 -i {video_out} -vf "crop=700:700:in_w/2-350:in_h/2-350,scale=224:224" -r 15 -c:v libx264 -crf 23 -c:a aac -an -strict experimental {video_out}_cropped_video.mp4')
87
  else:
88
  video_out = f"demo_files/downloaded_videos/own_demo_video.mp4"
89
  os.system(f"yt-dlp -q --no-warnings --force-overwrites -f mp4 {video_link} -o '{video_out}'")
 
90
  os.system(f'ffmpeg -y -nostats -loglevel 0 -i {video_out} -vf "crop=700:700:in_w/2-350:in_h/2-350,scale=224:224" -r 15 -c:v libx264 -crf 23 -c:a aac -an -strict experimental {video_out}_cropped_video.mp4')
91
+ #if the process is taking too long, throw an error
92
+ if time.time() - start > 30:
93
+ raise gr.Error("Video download is taking too long. Please try again.")
94
  return f"{video_out}_cropped_video.mp4"
95
 
 
96
  demo_files = os.listdir("demo_files/")
97
  demo_files = [os.path.join("demo_files/", file) for file in demo_files if file.endswith(".mp4")]
98
  with gr.Blocks(title="Picklebot") as demo:
 
109
  with gr.Column(scale=2):
110
  ump_out = gr.Label(label="Umpire's Original Call")
111
  pb_out = gr.Label(label="Picklebot's Call")
112
+ #when random is clicked, we want to somehow have a loading something to show that it's working
113
  random_button = gr.Button("🔀 Load a Random Pitch")
114
+ random_button.click(fn=generate_random_pitch,outputs=[inp,ump_out,pb_out])
115
+ inp.play(fn=call_pitch,inputs=inp,outputs=[pb_out,ump_out],trigger_mode="once")
116
  gr.ClearButton([inp,pb_out,ump_out])
117
  gr.Examples(demo_files,inputs=inp,outputs=[inp])
118
  with gr.Tab("Use Your Own Video!"):
 
133
  with gr.Column(scale=2):
134
  #make a textbox to take in the user's video link
135
  input_txt = gr.Textbox(placeholder="Paste your video link here",label="Video Link")
136
+ input_txt.submit(fn=download_pitch,inputs=input_txt,outputs=vid_inp)
137
  submit_button = gr.Button("Go!")
138
+ submit_button.click(fn=download_pitch,inputs=input_txt,outputs=vid_inp)
139
  pb_out = gr.Label(label="Picklebot's Call")
140
+ vid_inp.play(fn=call_pitch,inputs=vid_inp,outputs=pb_out,trigger_mode="once")
141
  gr.ClearButton([vid_inp,pb_out,input_txt])
142
 
143
  with gr.Tab("About"):