Create training.py

training.py

@@ -0,0 +1,118 @@
+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()