hoist-dataset-pipeline/upload_cvat.py

#!/usr/bin/env python3
"""
CVAT自动上传脚本
================
将pipeline生成的ZIP数据集自动上传到CVAT标注平台。

用法:
  python upload_cvat.py                  # 上传所有机位
  python upload_cvat.py --cameras cam01 cam02  # 只上传指定机位
"""
import os
import sys
import time
import zipfile
import argparse
import requests
from pathlib import Path

from config import *


class CVATClient:
    def __init__(self, url, username, password):
        self.url = url.rstrip('/')
        self.session = requests.Session()
        self.login(username, password)

    def login(self, username, password):
        self.session.get(f'{self.url}/api/v1/auth/login')
        csrf = self.session.cookies.get('csrftoken', '')
        resp = self.session.post(f'{self.url}/api/v1/auth/login', json={
            'username': username, 'password': password
        }, headers={'X-CSRFToken': csrf})
        resp.raise_for_status()
        self.csrf = self.session.cookies.get('csrftoken', '')
        print(f"Login OK")

    def _headers(self):
        return {'X-CSRFToken': self.csrf}

    def get_project_id(self, name):
        resp = self.session.get(f'{self.url}/api/v1/projects', params={'search': name})
        resp.raise_for_status()
        for p in resp.json()['results']:
            if p['name'] == name:
                return p['id']
        raise ValueError(f"Project '{name}' not found")

    def get_label_map(self, project_id):
        resp = self.session.get(f'{self.url}/api/v1/projects/{project_id}')
        resp.raise_for_status()
        labels = resp.json().get('labels', [])
        return {l['name']: l['id'] for l in labels}

    def list_tasks(self, project_id):
        resp = self.session.get(f'{self.url}/api/v1/tasks', params={'project': project_id})
        resp.raise_for_status()
        return resp.json()['results']

    def create_task(self, project_id, task_name):
        resp = self.session.post(f'{self.url}/api/v1/tasks', json={
            'name': task_name, 'project_id': project_id
        }, headers=self._headers())
        resp.raise_for_status()
        task = resp.json()
        print(f"  Created task: {task['id']} - {task_name}")
        return task['id']

    def delete_task(self, task_id):
        resp = self.session.delete(f'{self.url}/api/v1/tasks/{task_id}', headers=self._headers())
        return resp.status_code

    def upload_data(self, task_id, zip_path):
        with open(zip_path, 'rb') as f:
            resp = self.session.post(
                f'{self.url}/api/v1/tasks/{task_id}/data',
                files={'client_files[0]': (zip_path.name, f, 'application/zip')},
                data={'image_quality': '100'},
                headers=self._headers()
            )
        resp.raise_for_status()
        print(f"  Uploaded data")
        return resp.json()

    def upload_annotations(self, task_id, zip_path, class_names, label_map):
        shapes = []
        with zipfile.ZipFile(zip_path) as zf:
            ann_files = sorted([n for n in zf.namelist() if n.endswith('.txt') and 'obj_train_data' in n])
            for i, ann_file in enumerate(ann_files):
                with zf.open(ann_file) as f:
                    for line in f:
                        parts = line.decode().strip().split()
                        if len(parts) >= 5:
                            cls_id = int(parts[0])
                            if cls_id not in class_names:
                                continue
                            label_name = class_names[cls_id]
                            if label_name not in label_map:
                                continue
                            xc, yc, w, h = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4])
                            shapes.append({
                                'frame': i,
                                'label_id': label_map[label_name],
                                'type': 'rectangle',
                                'occluded': False,
                                'points': [
                                    round((xc - w/2) * IMG_WIDTH, 2),
                                    round((yc - h/2) * IMG_HEIGHT, 2),
                                    round((xc + w/2) * IMG_WIDTH, 2),
                                    round((yc + h/2) * IMG_HEIGHT, 2)
                                ],
                                'group': 0,
                                'attributes': []
                            })

        if not shapes:
            print(f"  No annotations to upload")
            return

        body = {'version': 0, 'tags': [], 'shapes': shapes, 'tracks': []}
        resp = self.session.put(
            f'{self.url}/api/v1/tasks/{task_id}/annotations',
            json=body, headers=self._headers()
        )
        if resp.status_code == 200:
            print(f"  Uploaded {len(shapes)} annotations")
        else:
            print(f"  Warning: {resp.status_code} {resp.text[:200]}")


def main():
    parser = argparse.ArgumentParser(description='Upload dataset to CVAT')
    parser.add_argument('--cameras', nargs='+', help='Specific cameras to upload (e.g. cam01 cam02)')
    parser.add_argument('--delete-test', action='store_true', help='Delete test tasks (czzt_*_reinfer)')
    args = parser.parse_args()

    print("Connecting to CVAT...")
    client = CVATClient(CVAT_URL, CVAT_USERNAME, CVAT_PASSWORD)

    project_id = client.get_project_id(CVAT_PROJECT_NAME)
    label_map = client.get_label_map(project_id)
    class_names = {0: 'hoist', 1: 'ear', 2: 'ready'}
    print(f"Project: {CVAT_PROJECT_NAME} (id={project_id}), labels: {label_map}")

    # 删除测试任务
    if args.delete_test:
        tasks = client.list_tasks(project_id)
        for t in tasks:
            if '_reinfer' in t['name']:
                print(f"  Deleting test task: {t['id']} - {t['name']}")
                client.delete_task(t['id'])
                time.sleep(0.5)

    # 确定要上传的机位
    if args.cameras:
        zips = [CVAT_ZIP_DIR / f'{cam}.zip' for cam in args.cameras]
        zips = [z for z in zips if z.exists()]
    else:
        zips = sorted(CVAT_ZIP_DIR.glob('*.zip'))

    print(f"\nFound {len(zips)} zips to upload\n")

    for zip_path in zips:
        cam_name = zip_path.stem
        print(f"\n{'='*40}")
        print(f"Processing {cam_name}...")
        print(f"{'='*40}")

        task_id = client.create_task(project_id, f'czzt_{cam_name}')
        client.upload_data(task_id, zip_path)
        time.sleep(1)
        client.upload_annotations(task_id, zip_path, class_names, label_map)
        time.sleep(0.5)

    print(f"\n{'='*40}")
    print("Done!")


if __name__ == '__main__':
    main()