Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| from wandb.integration.ultralytics import add_wandb_callback | |
| import wandb | |
| def interface_login(logger, args): | |
| if logger == 'WANDB': | |
| result = False | |
| wandb_key = args[0] | |
| if (wandb_key is not None) & isinstance(wandb_key, str): | |
| try: | |
| result = wandb.login(key=wandb_key,relogin=True,timeout=15) | |
| except: | |
| gr.Warning("Issue with the WANDB key") | |
| else: | |
| gr.Warning("Issue with the WANDB key") | |
| if result: | |
| gr.Info("Logged in to WANDB") | |
| else: | |
| gr.Warning("Failed to log in to WANDB") | |
| elif logger == 'ClearML': | |
| pass | |
| elif logger == 'Tensorboard': | |
| pass | |
| def interface_finetune(): | |
| # Load a pretrained YOLOv8n model | |
| model = YOLO('yolov8n.pt') # Load an official Detect model | |
| return model | |
| def interface_train(is_fintune=False, dataset=None, epochs=2, imgsz=640): | |
| model = YOLO('yolov8n.yaml') | |
| if is_fintune: | |
| model = interface_finetune() | |
| results = model.train(data=dataset, epochs=epochs, imgsz=imgsz) | |
| def interface_train_wandb(project_name, model_name, dataset_name, epochs=2, imgsz=640): | |
| # Step 1: Initialize a Weights & Biases run | |
| wandb.init(project=project_name, job_type="training") | |
| model = YOLO(f"{model_name}.pt") | |
| # Step 3: Add W&B Callback for Ultralytics | |
| add_wandb_callback(model, enable_model_checkpointing=True) | |
| # Step 4: Train and Fine-Tune the Model | |
| model.train(project=project_name, data=dataset_name, epochs=epochs, imgsz=imgsz) | |
| # Step 5: Validate the Model | |
| model.val() | |
| # # Step 6: Perform Inference and Log Results | |
| # model(["Images\Craig.jpg", "Images\WalterWhite.jpg"]) | |
| # Step 7: Finalize the W&B Run | |
| wandb.finish() |