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()