first push

backend/services/generate_yolox_exp.py

@@ -0,0 +1,152 @@
+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