Joshua Lochner commited on
Commit
e926596
1 Parent(s): 0b48a99

Hide previous output on run

Browse files
Files changed (1) hide show
  1. app.py +73 -65
app.py CHANGED
@@ -159,82 +159,90 @@ def load_predict(model_id):
159
 
160
 
161
  def main():
 
 
162
 
163
  # Display heading and subheading
164
- st.write('# SponsorBlock ML')
165
- st.write('##### Automatically detect in-video YouTube sponsorships, self/unpaid promotions, and interaction reminders.')
166
 
167
- model_id = st.selectbox('Select model', MODELS.keys(), index=0)
 
 
168
 
169
- # Load prediction function
170
- predict = load_predict(model_id)
171
-
172
- video_input = st.text_input('Video URL/ID:')
173
 
174
- categories = st.multiselect('Categories:',
175
- CATGEGORY_OPTIONS.keys(),
176
- CATGEGORY_OPTIONS.keys(),
177
- format_func=CATGEGORY_OPTIONS.get
178
- )
179
 
180
  # Hide segments with a confidence lower than
181
- confidence_threshold = st.slider(
182
- 'Confidence Threshold (%):', min_value=0, max_value=100)
183
 
184
- if len(video_input) == 0: # No input, do not continue
185
- return
186
 
187
- video_id = regex_search(video_input, YT_VIDEO_REGEX)
188
- if video_id is None:
189
- st.exception(ValueError('Invalid YouTube URL/ID'))
190
- return
191
-
192
- with st.spinner('Running model...'):
193
- predictions = predict(video_id)
194
-
195
- if len(predictions) == 0:
196
- st.success('No segments found!')
197
  return
198
 
199
- submit_segments = []
200
- for index, prediction in enumerate(predictions, start=1):
201
- if prediction['category'] not in categories:
202
- continue # Skip
203
-
204
- confidence = prediction['probability'] * 100
205
-
206
- if confidence < confidence_threshold:
207
- continue
208
-
209
- submit_segments.append({
210
- 'segment': [prediction['start'], prediction['end']],
211
- 'category': prediction['category'].lower(),
212
- 'actionType': 'skip'
213
- })
214
- start_time = seconds_to_time(prediction['start'])
215
- end_time = seconds_to_time(prediction['end'])
216
- with st.expander(
217
- f"[{prediction['category']}] Prediction #{index} ({start_time} \u2192 {end_time})"
218
- ):
219
-
220
- url = f"https://www.youtube-nocookie.com/embed/{video_id}?&start={floor(prediction['start'])}&end={ceil(prediction['end'])}"
221
- # autoplay=1controls=0&&modestbranding=1&fs=0
222
-
223
- # , width=None, height=None, scrolling=False
224
- components.iframe(url, width=670, height=376)
225
-
226
- text = ' '.join(w['text'] for w in prediction['words'])
227
- st.write(f"**Times:** {start_time} \u2192 {end_time}")
228
- st.write(
229
- f"**Category:** {CATGEGORY_OPTIONS[prediction['category']]}")
230
- st.write(f"**Confidence:** {confidence:.2f}%")
231
- st.write(f'**Text:** "{text}"')
232
-
233
- json_data = quote(json.dumps(submit_segments))
234
- link = f'[Submit Segments](https://www.youtube.com/watch?v={video_id}#segments={json_data})'
235
- st.markdown(link, unsafe_allow_html=True)
236
- wiki_link = '[Review generated segments before submitting!](https://wiki.sponsor.ajay.app/w/Automating_Submissions)'
237
- st.markdown(wiki_link, unsafe_allow_html=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
238
 
239
 
240
  if __name__ == '__main__':
 
159
 
160
 
161
  def main():
162
+ top = st.container()
163
+ output = st.empty()
164
 
165
  # Display heading and subheading
166
+ top.markdown('# SponsorBlock ML')
167
+ top.markdown('##### Automatically detect in-video YouTube sponsorships, self/unpaid promotions, and interaction reminders.')
168
 
169
+ # Add controls
170
+ model_id = top.selectbox(
171
+ 'Select model', MODELS.keys(), index=0, on_change=output.empty)
172
 
173
+ video_input = top.text_input(
174
+ 'Video URL/ID:', on_change=output.empty)
 
 
175
 
176
+ categories = top.multiselect('Categories:',
177
+ CATGEGORY_OPTIONS.keys(),
178
+ CATGEGORY_OPTIONS.keys(),
179
+ format_func=CATGEGORY_OPTIONS.get, on_change=output.empty
180
+ )
181
 
182
  # Hide segments with a confidence lower than
183
+ confidence_threshold = top.slider(
184
+ 'Confidence Threshold (%):', min_value=0, max_value=100, on_change=output.empty)
185
 
 
 
186
 
187
+ if len(video_input) == 0: # No input, do not continue
 
 
 
 
 
 
 
 
 
188
  return
189
 
190
+ # Load prediction function
191
+ with st.spinner('Loading model...'):
192
+ predict = load_predict(model_id)
193
+
194
+ with output.container(): # Place all content in output container
195
+ video_id = regex_search(video_input, YT_VIDEO_REGEX)
196
+ if video_id is None:
197
+ st.exception(ValueError('Invalid YouTube URL/ID'))
198
+ return
199
+
200
+ with st.spinner('Running model...'):
201
+ predictions = predict(video_id)
202
+
203
+ if len(predictions) == 0:
204
+ st.success('No segments found!')
205
+ return
206
+
207
+ submit_segments = []
208
+ for index, prediction in enumerate(predictions, start=1):
209
+ if prediction['category'] not in categories:
210
+ continue # Skip
211
+
212
+ confidence = prediction['probability'] * 100
213
+
214
+ if confidence < confidence_threshold:
215
+ continue
216
+
217
+ submit_segments.append({
218
+ 'segment': [prediction['start'], prediction['end']],
219
+ 'category': prediction['category'].lower(),
220
+ 'actionType': 'skip'
221
+ })
222
+ start_time = seconds_to_time(prediction['start'])
223
+ end_time = seconds_to_time(prediction['end'])
224
+ with st.expander(
225
+ f"[{prediction['category']}] Prediction #{index} ({start_time} \u2192 {end_time})"
226
+ ):
227
+
228
+ url = f"https://www.youtube-nocookie.com/embed/{video_id}?&start={floor(prediction['start'])}&end={ceil(prediction['end'])}"
229
+ # autoplay=1controls=0&&modestbranding=1&fs=0
230
+
231
+ # , width=None, height=None, scrolling=False
232
+ components.iframe(url, width=670, height=376)
233
+
234
+ text = ' '.join(w['text'] for w in prediction['words'])
235
+ st.write(f"**Times:** {start_time} \u2192 {end_time}")
236
+ st.write(
237
+ f"**Category:** {CATGEGORY_OPTIONS[prediction['category']]}")
238
+ st.write(f"**Confidence:** {confidence:.2f}%")
239
+ st.write(f'**Text:** "{text}"')
240
+
241
+ json_data = quote(json.dumps(submit_segments))
242
+ link = f'[Submit Segments](https://www.youtube.com/watch?v={video_id}#segments={json_data})'
243
+ st.markdown(link, unsafe_allow_html=True)
244
+ wiki_link = '[Review generated segments before submitting!](https://wiki.sponsor.ajay.app/w/Automating_Submissions)'
245
+ st.markdown(wiki_link, unsafe_allow_html=True)
246
 
247
 
248
  if __name__ == '__main__':