Spaces:
Sleeping
Sleeping
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() |