app.py

import gradio as gr
import pandas as pd
import random
import os
import torch
import torch.nn.functional as F
from mobilenet import MobileNetLarge3D, MobileNetSmall3D
from torchvision.io import read_video
import time

def classify_pitch(confidence_scores):
    if torch.argmax(confidence_scores) == 0:
        call = 'Ball'
    elif torch.argmax(confidence_scores) == 1:
        call = 'Strike'
    else:
        print("that's odd, something is wrong")
        pass
    return call

def get_demo_call(video_name):
    video_name = os.path.basename(video_name).lower()
    if video_name.startswith("strike"):
        ump_out = "Strike"
    elif video_name.startswith("ball"):
        ump_out = "Ball"
    else:
        ump_out = None
    return ump_out

def call_pitch(pitch,progress=gr.Progress()):
    start = time.time()
    std = (0.2104, 0.1986, 0.1829)
    mean = (0.3939, 0.3817, 0.3314)
    # Convert the mean and std to tensors
    mean = torch.tensor(mean).view(1, 3, 1, 1, 1)
    std = torch.tensor(std).view(1, 3, 1, 1, 1)
    ump_out = get_demo_call(pitch)
    pitch_tensor = (read_video(pitch,pts_unit='sec')[0].permute(-1,0,1,2)).unsqueeze(0)/255
    pitch_tensor = (pitch_tensor-mean)/std #normalize the pitch tensor
    video_length = pitch_tensor.shape[2] #get video length to wait for the video to finish
    model = MobileNetSmall3D()
    model.load_state_dict(torch.load('weights/MobileNetSmall.pth',map_location=torch.device('cpu')))
    model.eval()
    with torch.no_grad():
        output = model(pitch_tensor)
    output = F.softmax(output,dim=1)
    final_call = classify_pitch(output)

    differential = time.time() - start
    wait_time = (int(video_length-(10+differential))//15)*10
    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
        time.sleep(0.1)

    if ump_out is not None:
        return final_call,ump_out
    else:
        return final_call

def generate_random_pitch():
    random_number = random.randint(1,2645342) #random number in our range to select a random pitch

    df = pd.read_csv('picklebot_2m.csv',skiprows=random_number,nrows=1,header=None) #just load one row
    video_link = df.values[0][0]
    label = df.values[0][1]
    
    if label == 0:
        ump_out = "Ball"
    elif label == 1:
        ump_out = "Strike"
    else:
        #error if it's not a 0 or 1
        ump_out = "Error"
    
    random_pitch = download_pitch(video_link,ump_out)

    return random_pitch,None,None #none and none to write none to ump out and pb out


def download_pitch(video_link,ump_out=None):
    #download and process the video link using yt-dlp and ffmpeg
    start = time.time()
    if ump_out is not None:
        video_out = f"demo_files/downloaded_videos/{ump_out}_demo_video.mp4"
        os.system(f"yt-dlp -q --no-warnings --force-overwrites -f mp4 {video_link} -o '{video_out}'")
        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')
    else:
        video_out = f"demo_files/downloaded_videos/own_demo_video.mp4"
        os.system(f"yt-dlp -q --no-warnings --force-overwrites -f mp4 {video_link} -o '{video_out}'")
        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')
    #if the process is taking too long, throw an error
    if time.time() - start > 30:
        raise gr.Error("Video download is taking too long. Please try again.")
    return f"{video_out}_cropped_video.mp4"

demo_files = os.listdir("demo_files/")
demo_files = [os.path.join("demo_files/", file) for file in demo_files if file.endswith(".mp4")]
with gr.Blocks(title="Picklebot") as demo:
    with gr.Tab("Picklebot"):
        gr.Markdown(value="""
                    To load a video, click the random button or choose from the examples. 
                    Play the video by clicking anywhere on the video; Picklebot's call will appear after the video is done playing.

                    
                    To use your own video, click the "Use Your Own Video!" tab and follow the instructions.""")
        with gr.Row():
            with gr.Column(scale=3):    
                inp = gr.Video(interactive=False, label="Pitch Video")
            with gr.Column(scale=2):
                ump_out = gr.Label(label="Umpire's Original Call")
                pb_out = gr.Label(label="Picklebot's Call")
                #when random is clicked, we want to somehow have a loading something to show that it's working
                random_button = gr.Button("🔀 Load a Random Pitch")
                random_button.click(fn=generate_random_pitch,outputs=[inp,ump_out,pb_out])
                inp.play(fn=call_pitch,inputs=inp,outputs=[pb_out,ump_out],trigger_mode="once")
                gr.ClearButton([inp,pb_out,ump_out])
                gr.Examples(demo_files,inputs=inp,outputs=[inp])
    with gr.Tab("Use Your Own Video!"):
        with gr.Row():
            gr.Markdown(value=
                    """
                    # Here\'s how to use your own video: 
                    1. Navigate to [Baseball Savant's statcast search](https://baseballsavant.mlb.com/statcast_search)
                    2. Choose the situation and pitch you want (just keep in mind that the network was only trained on called balls and called strikes) 
                    3. Click on the pitcher or batter's name to get to list of pitches
                    4. Right click on the camera icon to the right, and copy the link
                    5. Paste the video url in below, and press go to download the pitch
                    6. Watch the video and see what Picklebot thinks!
                    """)
        with gr.Row():
            with gr.Column(scale=3):    
                vid_inp = gr.Video(interactive=False, label="Pitch Video")
            with gr.Column(scale=2):
                #make a textbox to take in the user's video link
                input_txt = gr.Textbox(placeholder="Paste your video link here",label="Video Link")
                input_txt.submit(fn=download_pitch,inputs=input_txt,outputs=vid_inp)
                submit_button = gr.Button("Go!")
                submit_button.click(fn=download_pitch,inputs=input_txt,outputs=vid_inp)
                pb_out = gr.Label(label="Picklebot's Call")
                vid_inp.play(fn=call_pitch,inputs=vid_inp,outputs=pb_out,trigger_mode="once")
                gr.ClearButton([vid_inp,pb_out,input_txt])

    with gr.Tab("About"):
        gr.Markdown(value=
                    """Picklebot is a 3D Convolutional Neural Network based on MobileNetV3 that classifies baseball pitches as balls or strikes.
                    The network was trained on the [Picklebot-50K Dataset](https://huggingface.co/datasets/hbfreed/Picklebot-50K),
                    comprised of over fifty thousand pitches to achieve ~80% accuracy.

                    Here's the [GitHub](https://github.com/hbfreed/Picklebot) for the project.
                    
                    Here's my [LinkedIn](https://www.linkedin.com/in/hbfreed/) if you want to connect.""")
if __name__ == "__main__":
  demo.launch()