Spaces:
Sleeping
Sleeping
File size: 1,826 Bytes
6a7b87d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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() |