import os import shutil from models.training import Training from models.TrainingProject import TrainingProject def generate_yolox_exp(training_id): """Generate YOLOX exp.py file""" # Fetch training row from DB training = Training.query.get(training_id) if not training: training = Training.query.filter_by(project_details_id=training_id).first() if not training: raise Exception(f'Training not found for trainingId or project_details_id: {training_id}') # If transfer_learning is 'coco', copy default exp.py if training.transfer_learning == 'coco': selected_model = training.selected_model.lower().replace('-', '_') exp_source_path = f'/home/kitraining/Yolox/YOLOX-main/exps/default/{selected_model}.py' if not os.path.exists(exp_source_path): raise Exception(f'Default exp.py not found for model: {selected_model} at {exp_source_path}') # Copy to project folder project_details_id = training.project_details_id project_folder = os.path.join(os.path.dirname(__file__), '..', f'project_23/{project_details_id}') os.makedirs(project_folder, exist_ok=True) exp_dest_path = os.path.join(project_folder, 'exp.py') shutil.copyfile(exp_source_path, exp_dest_path) return {'type': 'default', 'expPath': exp_dest_path} # If transfer_learning is 'sketch', generate custom exp.py if training.transfer_learning == 'sketch': exp_content = generate_yolox_inference_exp(training_id) return {'type': 'custom', 'expContent': exp_content} raise Exception(f'Unknown transfer_learning type: {training.transfer_learning}') def save_yolox_exp(training_id, out_path): """Save YOLOX exp.py to specified path""" exp_result = generate_yolox_exp(training_id) if exp_result['type'] == 'custom' and 'expContent' in exp_result: with open(out_path, 'w') as f: f.write(exp_result['expContent']) return out_path elif exp_result['type'] == 'default' and 'expPath' in exp_result: # Optionally copy the file if outPath is different if exp_result['expPath'] != out_path: shutil.copyfile(exp_result['expPath'], out_path) return out_path else: raise Exception('Unknown expResult type or missing content') def generate_yolox_inference_exp(training_id, options=None): """Generate inference exp.py using DB values""" if options is None: options = {} training = Training.query.get(training_id) if not training: training = Training.query.filter_by(project_details_id=training_id).first() if not training: raise Exception(f'Training not found for trainingId or project_details_id: {training_id}') # Always use the training_id (project_details_id) for annotation file names project_details_id = training.project_details_id data_dir = options.get('data_dir', '/home/kitraining/To_Annotate/') train_ann = options.get('train_ann', f'coco_project_{training_id}_train.json') val_ann = options.get('val_ann', f'coco_project_{training_id}_valid.json') test_ann = options.get('test_ann', f'coco_project_{training_id}_test.json') # Get num_classes from TrainingProject.classes JSON num_classes = 80 try: training_project = TrainingProject.query.get(project_details_id) if training_project and training_project.classes: classes_arr = training_project.classes if isinstance(classes_arr, str): import json classes_arr = json.loads(classes_arr) if isinstance(classes_arr, list): num_classes = len([c for c in classes_arr if c not in [None, '']]) elif isinstance(classes_arr, dict): num_classes = len([k for k, v in classes_arr.items() if v not in [None, '']]) except Exception as e: print(f'Could not determine num_classes from TrainingProject.classes: {e}') depth = options.get('depth', training.depth or 1.00) width = options.get('width', training.width or 1.00) input_size = options.get('input_size', training.input_size or [640, 640]) mosaic_scale = options.get('mosaic_scale', training.mosaic_scale or [0.1, 2]) random_size = options.get('random_size', [10, 20]) test_size = options.get('test_size', training.test_size or [640, 640]) exp_name = options.get('exp_name', 'inference_exp') enable_mixup = options.get('enable_mixup', False) # Build exp content exp_content = f'''#!/usr/bin/env python3 # -*- coding:utf-8 -*- # Copyright (c) Megvii, Inc. and its affiliates. import os from yolox.exp import Exp as MyExp class Exp(MyExp): def __init__(self): super(Exp, self).__init__() self.data_dir = "{data_dir}" self.train_ann = "{train_ann}" self.val_ann = "{val_ann}" self.test_ann = "coco_project_{training_id}_test.json" self.num_classes = {num_classes} ''' # Set pretrained_ckpt if transfer_learning is 'coco' if training.transfer_learning and isinstance(training.transfer_learning, str) and training.transfer_learning.lower() == 'coco': yolox_base_dir = '/home/kitraining/Yolox/YOLOX-main' selected_model = training.selected_model.replace('.pth', '') if training.selected_model else '' if selected_model: exp_content += f" self.pretrained_ckpt = r'{yolox_base_dir}/pretrained/{selected_model}.pth'\n" # Format arrays input_size_str = ', '.join(map(str, input_size)) if isinstance(input_size, list) else str(input_size) mosaic_scale_str = ', '.join(map(str, mosaic_scale)) if isinstance(mosaic_scale, list) else str(mosaic_scale) random_size_str = ', '.join(map(str, random_size)) if isinstance(random_size, list) else str(random_size) test_size_str = ', '.join(map(str, test_size)) if isinstance(test_size, list) else str(test_size) exp_content += f''' self.depth = {depth} self.width = {width} self.input_size = ({input_size_str}) self.mosaic_scale = ({mosaic_scale_str}) self.random_size = ({random_size_str}) self.test_size = ({test_size_str}) self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] self.enable_mixup = {str(enable_mixup)} ''' return exp_content def save_yolox_inference_exp(training_id, out_path, options=None): """Save inference exp.py to custom path""" exp_content = generate_yolox_inference_exp(training_id, options) with open(out_path, 'w') as f: f.write(exp_content) return out_path