app.py

import gradio as gr
import spaces
import sys
import os
from datasets import load_dataset
from huggingface_hub import snapshot_download
import tempfile

temp_dir = tempfile.TemporaryDirectory()
print(f'Creating {temp_dir.name}')

try:
    from synet.backends import get_backend
    get_backend('ultralytics').patch()
    from synet.backends import get_backend
    from ultralytics import YOLO
except ImportError:
    import subprocess
    print('Installing synet package')
    subprocess.check_call([sys.executable, "-m", "pip", "install", "./synet_package[ultra]"])
    # Import the file
    from synet.backends import get_backend
    from synet.backends import get_backend
    from ultralytics import YOLO


# Setup the backend processing
backend = get_backend('ultralytics')
backend.patch()

dataset_cfg = None


@spaces.GPU
def greet(name):
    return "Hello " + name + "!!"

def get_dataset():
    global dataset_cfg
    
    dataset_name = 'Ultralytics/COCO8'

    dataset_path = f'{temp_dir.name}/COCO8'
    print(f'Writing to dataset {dataset_path}')
    snapshot_download(repo_id=dataset_name, repo_type='dataset', local_dir=dataset_path)
    
    #dataset = load_dataset(dataset_name)
    #print(dataset)
    #dataset.save_to_disk(dataset_path)

    files = os.listdir(dataset_path)
    local_files = 'local_files in {dataset_path}: '
    for file in files:
        local_files += f'{file} '
    print(local_files)

 

    dataset_cfg = f'{dataset_path}/dataset.yaml'
    print(f'Writing to dataset {dataset_cfg}')
    with open(dataset_cfg, 'r') as fp:
        contents = fp.read()
        print(contents)
        
    return f"Loading the dataset in {temp_dir.name}"    

@spaces.GPU
def run_training():

    model_cfg = 'synet_package/synet/zoo/ultralytics/sabre-detect-vga.yaml'
    model = YOLO(model=model_cfg)
    print(f'Loading model_cfg {model_cfg}')
    print(model)
    image_size = (480,640)

    # Run the initial training
    project_path = f'{temp_dir.name}/runs'
    print(f'Run the training in {project_path}')
    model.train(data=dataset_cfg, project=project_path, name='example_train')

    # PRint teh files
    files = os.listdir(project_path)
    local_files = 'local_files: '
    for file in files:
        local_files += f'{file} '

    return f"Done with training: {local_files}"
    #get_tflite(backend, image_size, 'runs/example_train/weights/best.pt',
    #       'coco.yaml', 500, 3, {})

# Run the ultralytics
#def run_ultralytics():
#
#    files = os.listdir()
#    local_files = 'local_files: '
#    for file in files:
#        local_files += f'{file} '
#
#    return local_files

with gr.Blocks() as demo:
    text1 = gr.Markdown("Starting to test SyNet")
    text2 = gr.Markdown("")

    load_btn = gr.Button("Load dataset")
    load_text = gr.Markdown("")

    train_btn = gr.Button("Train")
    train_text = gr.Markdown("")
    
    load_btn.click(get_dataset, inputs=None, outputs=[load_text])
    train_btn.click(run_training, inputs=None, outputs=[train_text])
    
    #demo.load(run_ultralytics, inputs=None, outputs=[text2])

if __name__ == "__main__":
    #demo.launch(share=True)
    demo.launch()