Spaces:
Sleeping
Sleeping
| import json | |
| import os | |
| import shutil | |
| import subprocess | |
| import sys | |
| import threading | |
| import uuid | |
| from pathlib import Path | |
| from typing import Optional | |
| ROOT = Path(__file__).resolve().parent | |
| PD_DIR = ROOT / '.paddledetection' | |
| PD_REPO = 'https://github.com/PaddlePaddle/PaddleDetection.git' | |
| PD_BRANCH = os.getenv('PADDLEDETECTION_BRANCH', 'release/2.9') | |
| _backend_lock = threading.RLock() | |
| _training = {"running": False, "message": "not started", "error": None, "epoch": None, "log": ""} | |
| def _run(cmd, cwd=None, timeout=None): | |
| p = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=timeout) | |
| if p.returncode != 0: | |
| raise RuntimeError(p.stdout[-16000:] or f'command failed: {p.returncode}') | |
| return p.stdout | |
| def ensure_paddledetection(): | |
| """Clone official PaddleDetection only when first needed; no Hugging Face dependency.""" | |
| marker = PD_DIR / 'tools' / 'train.py' | |
| if marker.exists(): | |
| return PD_DIR | |
| PD_DIR.parent.mkdir(parents=True, exist_ok=True) | |
| tmp = ROOT / '.paddledetection_tmp' | |
| if tmp.exists(): | |
| shutil.rmtree(tmp, ignore_errors=True) | |
| _run(['git', 'clone', '--depth', '1', '--branch', PD_BRANCH, PD_REPO, str(tmp)]) | |
| if marker.exists(): | |
| return PD_DIR | |
| if PD_DIR.exists(): | |
| shutil.rmtree(PD_DIR, ignore_errors=True) | |
| tmp.rename(PD_DIR) | |
| return PD_DIR | |
| def paddle_status(): | |
| try: | |
| import paddle | |
| return {"installed": True, "version": paddle.__version__, "device": paddle.device.get_device(), "cuda": bool(paddle.is_compiled_with_cuda())} | |
| except Exception as e: | |
| return {"installed": False, "error": str(e)} | |
| def prepare_config(dataset_root: Path, classes_file: Path, output_dir: Path, | |
| epochs: int, batch_size: int, learning_rate: float): | |
| """Create a native PaddleDetection RT-DETR R50vd config.""" | |
| pd = ensure_paddledetection() | |
| cfg_dir = pd / 'configs' / 'rtdetr' | |
| ds_cfg = pd / 'configs' / 'datasets' / 'icecream_coco.yml' | |
| model_cfg = cfg_dir / 'icecream_rtdetr_r50vd.yml' | |
| cfg_dir.mkdir(parents=True, exist_ok=True) | |
| classes = [x.strip() for x in classes_file.read_text(encoding='utf-8').splitlines() if x.strip()] | |
| if not classes: | |
| raise RuntimeError('No classes configured.') | |
| train_dir = (dataset_root / 'train').resolve().as_posix() | |
| val_dir = (dataset_root / 'val').resolve().as_posix() | |
| save_dir = output_dir.resolve().as_posix() | |
| ds_cfg.write_text(f"""metric: COCO | |
| num_classes: {len(classes)} | |
| TrainDataset: | |
| !COCODataSet | |
| image_dir: images | |
| anno_path: annotations.json | |
| dataset_dir: {train_dir} | |
| data_fields: ['image', 'gt_bbox', 'gt_class', 'is_crowd'] | |
| EvalDataset: | |
| !COCODataSet | |
| image_dir: images | |
| anno_path: annotations.json | |
| dataset_dir: {val_dir} | |
| TestDataset: | |
| !ImageFolder | |
| anno_path: annotations.json | |
| dataset_dir: {val_dir} | |
| """, encoding='utf-8') | |
| model_cfg.write_text(f"""_BASE_: [ | |
| '../datasets/icecream_coco.yml', | |
| '../runtime.yml', | |
| '_base_/optimizer_6x.yml', | |
| '_base_/rtdetr_r50vd.yml', | |
| '_base_/rtdetr_reader.yml', | |
| ] | |
| architecture: DETR | |
| num_classes: {len(classes)} | |
| use_gpu: true | |
| find_unused_parameters: True | |
| log_iter: 1 | |
| save_dir: {save_dir} | |
| weights: {save_dir}/model_final | |
| pretrain_weights: https://paddledet.bj.bcebos.com/models/pretrained/ResNet50_vd_ssld_v2_pretrained.pdparams | |
| epoch: {int(epochs)} | |
| LearningRate: | |
| base_lr: {float(learning_rate)} | |
| TrainReader: | |
| batch_size: {int(batch_size)} | |
| worker_num: 0 | |
| use_shared_memory: false | |
| EvalReader: | |
| batch_size: 1 | |
| worker_num: 0 | |
| TestReader: | |
| batch_size: 1 | |
| worker_num: 0 | |
| """, encoding='utf-8') | |
| return pd, model_cfg, classes | |
| def train(dataset_root: Path, classes_file: Path, output_dir: Path, epochs: int, batch_size: int, learning_rate: float): | |
| if 'transformers' in sys.modules or 'torch' in sys.modules: | |
| raise RuntimeError( | |
| 'Wrong training engine loaded: this package uses PaddleDetection RT-DETR, ' | |
| 'not PyTorch or Hugging Face Transformers. Restart with this package.' | |
| ) | |
| global _training | |
| with _backend_lock: | |
| if _training['running']: | |
| return _training | |
| _training = {"running": True, "message": "preparing PaddleDetection", "error": None, "epoch": None, "log": ""} | |
| try: | |
| pd, cfg, classes = prepare_config(dataset_root, classes_file, output_dir, epochs, batch_size, learning_rate) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| status = paddle_status() | |
| if not status.get('installed'): | |
| raise RuntimeError(f'PaddlePaddle is not installed: {status.get("error")}') | |
| if not status.get('cuda'): | |
| raise RuntimeError( | |
| 'GPU PaddlePaddle is not available. ' | |
| f'Paddle device={status.get("device")}. Install the GPU PaddlePaddle build ' | |
| 'matching the cloud CUDA runtime.' | |
| ) | |
| env = os.environ.copy() | |
| env['CUDA_VISIBLE_DEVICES'] = env.get('CUDA_VISIBLE_DEVICES', '0') | |
| env['FLAGS_selected_gpus'] = env.get('FLAGS_selected_gpus', '0') | |
| cmd = [os.sys.executable, '-u', 'tools/train.py', '-c', str(cfg), '--eval', '-o', 'use_gpu=true'] | |
| _training['message'] = 'training RT-DETR with PaddleDetection' | |
| proc = subprocess.Popen(cmd, cwd=pd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1) | |
| lines = [] | |
| for line in proc.stdout: | |
| lines.append(line.rstrip()) | |
| if len(lines) > 250: | |
| lines = lines[-250:] | |
| _training['log'] = '\n'.join(lines[-80:]) | |
| low = line.lower() | |
| if 'epoch' in low: | |
| import re | |
| m = re.search(r'epoch[^0-9]*(\d+)', low) | |
| if m: | |
| _training['epoch'] = int(m.group(1)) | |
| code = proc.wait() | |
| if code != 0: | |
| raise RuntimeError('\n'.join(lines[-120:]) or f'PaddleDetection training exited with {code}') | |
| final = output_dir / 'model_final.pdparams' | |
| if not final.exists(): | |
| # Some PaddleDetection versions save model_final as a prefix with .pdparams. | |
| candidates = list(output_dir.glob('model_final*')) | |
| if not candidates: | |
| raise RuntimeError('Training finished but model_final was not produced.') | |
| (output_dir / 'classes.json').write_text(json.dumps({str(i): c for i, c in enumerate(classes)}, indent=2), encoding='utf-8') | |
| _training = {"running": False, "message": "training complete", "error": None, "epoch": int(epochs), "log": '\n'.join(lines[-80:])} | |
| except Exception as e: | |
| _training = {"running": False, "message": "training failed", "error": str(e), "epoch": _training.get('epoch'), "log": _training.get('log', '')} | |
| return _training | |
| def training_status(): | |
| return dict(_training) | |
| def model_ready(model_dir: Path): | |
| return any((model_dir / name).exists() for name in ('model_final.pdparams', 'model_final.pdiparams')) | |
| def predict(image_path: Path, model_dir: Path, threshold: float, classes_file: Path, output_dir: Path): | |
| pd = ensure_paddledetection() | |
| if not model_ready(model_dir): | |
| raise RuntimeError('No trained model yet. Train the model first.') | |
| cfg = pd / 'configs' / 'rtdetr' / 'icecream_rtdetr_r50vd.yml' | |
| if not cfg.exists(): | |
| raise RuntimeError('PaddleDetection model config is missing; train once to generate it.') | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| for p in output_dir.iterdir(): | |
| if p.is_file(): | |
| p.unlink() | |
| cmd = [os.sys.executable, '-u', 'tools/infer.py', '-c', str(cfg), '--infer_img', str(image_path), '--output_dir', str(output_dir), '--draw_threshold', str(threshold), '--save_threshold', str(threshold), '--save_results=True', '--visualize=True', '-o', f'use_gpu=true', f'weights={str(model_dir / "model_final")}'] | |
| log = _run(cmd, cwd=pd, timeout=int(os.getenv('INFER_TIMEOUT', '300'))) | |
| bbox_file = output_dir / 'bbox.json' | |
| rows = json.loads(bbox_file.read_text(encoding='utf-8')) if bbox_file.exists() else [] | |
| classes = [x.strip() for x in classes_file.read_text(encoding='utf-8').splitlines() if x.strip()] | |
| detections = [] | |
| for r in rows: | |
| cid = int(r.get('category_id', 0)) | |
| detections.append({ | |
| 'class': classes[cid] if 0 <= cid < len(classes) else str(cid), | |
| 'class_id': cid, | |
| 'confidence': round(float(r.get('score', 0.0)), 4), | |
| 'box': [round(float(r['bbox'][0]), 2), round(float(r['bbox'][1]), 2), round(float(r['bbox'][0] + r['bbox'][2]), 2), round(float(r['bbox'][1] + r['bbox'][3]), 2)] | |
| }) | |
| counts = {} | |
| for d in detections: | |
| counts[d['class']] = counts.get(d['class'], 0) + 1 | |
| visual = next((p for p in output_dir.iterdir() if p.suffix.lower() in {'.jpg','.jpeg','.png'}), None) | |
| return {'total': len(detections), 'counts': dict(sorted(counts.items())), 'detections': detections, 'visual_path': str(visual) if visual else None, 'log': log[-4000:]} | |