YOLOv8_Interface / interface /train_interface.py
Cippppy's picture
Initial Commit
6a7b87d verified
import gradio as gr
from interface.defaults import shared_theme
from interface.train_interface_methods import interface_train, interface_login
class TrainInterface():
def __init__(self):
self.demo = None
self.api_key = None
with gr.Blocks(theme=shared_theme) as demo:
gr.Markdown(
"""
# Training Interface for YOLOv8
Train your own YOLOv8 model!
""")
with gr.Row() as finetune_row:
is_finetune = gr.Checkbox(label="Finetune",info="Check this box if you want to finetune a model")
is_offical_pretrained = gr.Checkbox(label="Official",info="Check this box if you want to train an official model",visible=True,interactive=True,value=True)
custom_pretrained = gr.File(label="Pretrained Model Weights",file_count='single',type='binary',
file_types=['.pt'],visible=True,show_label=True,interactive=True)
offical_pretrained = gr.Dropdown(label="Pretrained Model",choices=["yolov8n.pt"],visible=True,interactive=True)
with gr.Row() as dataset_row:
is_official_dataset = gr.Checkbox(label="Official",info="Check this box if you want to use an official dataset",visible=True,interactive=True,value=True)
custom_dataset = gr.File(label="Custom Dataset",file_count='single',type='binary',
file_types=['.zip'],visible=True,show_label=True,interactive=True)
official_dataset = gr.Dropdown(label="Dataset",choices=["coco128"],visible=True,interactive=True)
# Row for start & clear buttons
with gr.Row() as buttons:
start_but = gr.Button(value="Start")
with gr.Accordion("Logger Options") as login_accordion:
use_logger = gr.Checkbox(label="Use Logger",info="Check this box if you want to use a logger",visible=True,interactive=True,value=False)
logger = gr.Radio(choices=['WANDB', 'ClearML', 'Tensorboard'],value='WANDB',show_label=True,interactive=True,visible=True,
label="Logger",info="Choose which logger to use")
wandb_key = gr.Textbox(label="WANDB Key",placeholder="Enter WANDB Key",visible=True,interactive=True)
login_but = gr.Button(value="Login")
def string_from_textbox(textbox):
self.api_key = textbox
wandb_key.change(fn=string_from_textbox,inputs=[wandb_key],outputs=[])
def logger_login(use_logger, logger):
if use_logger:
interface_login(logger, [self.api_key])
else:
gr.Warning("Not using logger, so no need to login")
start_but.click(fn=interface_train,inputs=[is_finetune, official_dataset],outputs=[])
login_but.click(fn=logger_login,inputs=[use_logger,logger],outputs=[])
self.demo = demo
def get_interface(self):
return self.demo
if __name__== "__main__" :
demo = TrainInterface().get_interface()
demo.queue().launch()