Adding the separate class

app.py

@@ -2,101 +2,19 @@ import gradio as gr
 import spaces
 import sys
 import os
-from datasets import load_dataset
-from huggingface_hub import snapshot_download
-import tempfile
+import training
 
-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 + "!!"
+model = training.model_training()
 
 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 = 'Synaptics/COCO8'
+    return model.get_dataset(dataset)
 
- 
-
-    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
+    return model.run_training()
 
 with gr.Blocks() as demo:
     text1 = gr.Markdown("Starting to test SyNet")
@@ -110,9 +28,7 @@ with gr.Blocks() as demo:
     
     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()

training.py

@@ -21,7 +21,7 @@ except ImportError:
 backend = get_backend('ultralytics')
 backend.patch()
 
-class training:
+class model_training:
 
     def __init__(self):
 
@@ -31,7 +31,14 @@ class training:
 
         self.dataset_cfg = None
 
-    def get_dataset(self, dataset, token):
+        self.model_select = 'detect'
+        self.model_options = {
+            'detect' : { cfg :  'synet_package/synet/zoo/ultralytics/sabre-detect-vga.yaml',
+                         image_size: (480,640) }
+        }
+
+
+    def get_dataset(self, dataset, token=None):
     
         # Get the dataset name
         dataset_name = dataset.split('/')[1]
@@ -57,25 +64,26 @@ class training:
         return f"Loading the dataset in {dataset_path}"    
 
     @spaces.GPU
-    def run_training(self):
+    def run_training(self, model_type='detect'):
+
+        model_cfg = self.model_options[model_type]['cfg']
+        image_size = self.model_options[model_type]['image_size']
 
-        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')
+        #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)
+        #files = os.listdir(project_path)
         local_files = 'local_files: '
-        for file in files:
-            local_files += f'{file} '
+        #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, {})
+        #get_tflite(backend, image_size, 'runs/example_train/weights/best.pt',
+        #       'coco.yaml', 500, 3, {})