andrew.ng
/
YoLo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update

This commit is contained in:
nguentrungthat 2025-03-21 14:57:23 +07:00
parent c1214662a3
commit f770973b5e
2 changed files with 48 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/App.js
View File

@ -311,9 +311,9 @@ function App() {
</Tabs.List>
</Box>
<Box className="px-4 flex items-center gap-4 w-fit">
<Button disabled={isLoading} onClick={() => setOpenConfirm(true)} leftSection={<IconReportAnalytics size={14} />} color="#caa32c">
{/* <Button disabled={isLoading} onClick={() => setOpenConfirm(true)} leftSection={<IconReportAnalytics size={14} />} color="#caa32c">
Training
</Button>
</Button> */}
<Button disabled={!clickData || isLoading} onClick={handleClearSelect} leftSection={<IconRefreshDot size={14} />} color="orange">
Reset select
</Button>

85
src/server/train.py
View File

@ -15,6 +15,8 @@ IMAGE_EXTENSION = ".png"
LOG_FILE = "training_logs.log"
LATEST_DETECT_IMAGES = "latest_detect/images"
LATEST_DETECT_LABELS = "latest_detect/labels"
image_folder=""
label_folder=""
def get_label_names():
conn = sqlite3.connect("my_database.db")
@ -75,6 +77,7 @@ for folder in [train_img_folder, train_lbl_folder, val_img_folder, val_lbl_folde
# 📂 Get image list
image_files = [f for f in os.listdir(image_folder) if f.endswith(IMAGE_EXTENSION)]
label_files = [f for f in os.listdir(label_folder) if os.path.exists(label_folder)]
# 🌀 Shuffle images
random.shuffle(image_files)
@ -86,18 +89,18 @@ split_idx = int(len(image_files) * train_ratio)
train_files = image_files[:split_idx]
val_files = image_files[split_idx:]
latest_img_folder = [f for f in os.listdir(LATEST_DETECT_IMAGES) if f.endswith(IMAGE_EXTENSION)]
train_files += latest_img_folder
# Copy latest detect images to train folder
for img_file in latest_img_folder:
img_path = os.path.join(LATEST_DETECT_IMAGES, img_file)
label_file = os.path.splitext(img_file)[0] + ".txt"
label_path = os.path.join(LATEST_DETECT_LABELS, label_file)
if os.path.exists(img_path) and os.path.exists(label_path) and not os.path.exists(os.path.join(image_folder, img_file)):
shutil.copy(img_path, os.path.join(image_folder, img_file))
shutil.copy(label_path, os.path.join(label_folder, label_file))
def copy_file_latest_training():
latest_img_folder = [f for f in os.listdir(LATEST_DETECT_IMAGES) if f.endswith(IMAGE_EXTENSION)]
train_files += latest_img_folder
# Copy latest detect images to train folder
for img_file in latest_img_folder:
img_path = os.path.join(LATEST_DETECT_IMAGES, img_file)
label_file = os.path.splitext(img_file)[0] + ".txt"
label_path = os.path.join(LATEST_DETECT_LABELS, label_file)
if os.path.exists(img_path) and os.path.exists(label_path) and not os.path.exists(os.path.join(image_folder, img_file)):
shutil.copy(img_path, os.path.join(image_folder, img_file))
shutil.copy(label_path, os.path.join(label_folder, label_file))
def log_message(message: str):
"""Ghi log vào file"""
@ -266,35 +269,39 @@ def train_yolo_model(
call_reload_model_api()
clear_images_source()
# 🚀 Copy files
train_copied_imgs, train_copied_lbls = move_files(train_files, train_img_folder, train_lbl_folder)
val_copied_imgs, val_copied_lbls = move_files(val_files, val_img_folder, val_lbl_folder)
if(len(label_folder) >= 20):
copy_file_latest_training()
# Create yml file
create_dataset_yaml(dataset_folder, class_names)
# 🚀 Copy files
train_copied_imgs, train_copied_lbls = move_files(train_files, train_img_folder, train_lbl_folder)
val_copied_imgs, val_copied_lbls = move_files(val_files, val_img_folder, val_lbl_folder)
# Create yml file
create_dataset_yaml(dataset_folder, class_names)
# 🏁 Log summary
with open(LOG_FILE, "a") as log:
log.write(f"\n[{datetime.datetime.now()}] Dataset split completed\n")
log.write(f"Source folder: {image_folder}\n")
log.write(f"Dataset saved in: {dataset_folder}\n")
log.write(f"Planned Train: {len(train_files)} images, Val: {len(val_files)} images\n")
log.write(f"Actual Train: {train_copied_imgs} images, {train_copied_lbls} labels\n")
log.write(f"Actual Val: {val_copied_imgs} images, {val_copied_lbls} labels\n")
if(train_copied_imgs <=0 or train_copied_lbls <= 0 or val_copied_imgs <=0 or val_copied_lbls <=0):
# 🏁 Log summary
with open(LOG_FILE, "a") as log:
log.write(f"\n[{datetime.datetime.now()}] Data not qualified\n")
log.write("=" * 50 + "\n")
else:
train_yolo_model(pretrained_model = get_latest_model(
trained_model_folder=TRAINED_MODEL_FOLDER,
default_model=PRETRAINED_MODEL),
dataset_folder = dataset_folder,
epochs = 50,
name=today_str,
project_name=model_folder_name,
)
log.write(f"\n[{datetime.datetime.now()}] Dataset split completed\n")
log.write(f"Source folder: {image_folder}\n")
log.write(f"Dataset saved in: {dataset_folder}\n")
log.write(f"Planned Train: {len(train_files)} images, Val: {len(val_files)} images\n")
log.write(f"Actual Train: {train_copied_imgs} images, {train_copied_lbls} labels\n")
log.write(f"Actual Val: {val_copied_imgs} images, {val_copied_lbls} labels\n")
if(train_copied_imgs <=0 or train_copied_lbls <= 0 or val_copied_imgs <=0 or val_copied_lbls <=0):
with open(LOG_FILE, "a") as log:
log.write(f"\n[{datetime.datetime.now()}] Data not qualified\n")
log.write("=" * 50 + "\n")
else:
train_yolo_model(pretrained_model = get_latest_model(
trained_model_folder=TRAINED_MODEL_FOLDER,
default_model=PRETRAINED_MODEL),
dataset_folder = dataset_folder,
epochs = 50,
name=today_str,
project_name=model_folder_name,
)
else:
log_message(f"Date: {today_str}. Need more image for training!!")