|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import argparse |
|
|
from rf100vl import get_rf100vl_projects |
|
|
import roboflow |
|
|
from rfdetr import RFDETRBase |
|
|
import torch |
|
|
import os |
|
|
|
|
|
def download_dataset(rf_project: roboflow.Project, dataset_version: int): |
|
|
versions = rf_project.versions() |
|
|
if dataset_version is not None: |
|
|
versions = [v for v in versions if v.version == str(dataset_version)] |
|
|
if len(versions) == 0: |
|
|
raise ValueError(f"Dataset version {dataset_version} not found") |
|
|
version = versions[0] |
|
|
else: |
|
|
version = max(versions, key=lambda v: v.id) |
|
|
location = os.path.join("datasets/", rf_project.name + "_v" + version.version) |
|
|
if not os.path.exists(location): |
|
|
location = version.download( |
|
|
model_format="coco", location=location, overwrite=False |
|
|
).location |
|
|
|
|
|
return location |
|
|
|
|
|
|
|
|
def train_from_rf_project(rf_project: roboflow.Project, dataset_version: int): |
|
|
location = download_dataset(rf_project, dataset_version) |
|
|
print(location) |
|
|
rf_detr = RFDETRBase() |
|
|
device_supports_cuda = torch.cuda.is_available() |
|
|
rf_detr.train( |
|
|
dataset_dir=location, |
|
|
epochs=1, |
|
|
device="cuda" if device_supports_cuda else "cpu", |
|
|
) |
|
|
|
|
|
|
|
|
def train_from_coco_dir(coco_dir: str): |
|
|
rf_detr = RFDETRBase() |
|
|
rf_detr.train( |
|
|
dataset_dir=coco_dir, |
|
|
epochs=1, |
|
|
device="cuda" if device_supports_cuda else "cpu", |
|
|
) |
|
|
|
|
|
|
|
|
def trainer(): |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument("--coco_dir", type=str, required=False) |
|
|
parser.add_argument("--api_key", type=str, required=False) |
|
|
parser.add_argument("--workspace", type=str, required=False, default=None) |
|
|
parser.add_argument("--project_name", type=str, required=False, default=None) |
|
|
parser.add_argument("--dataset_version", type=int, required=False, default=None) |
|
|
args = parser.parse_args() |
|
|
|
|
|
if args.coco_dir is not None: |
|
|
train_from_coco_dir(args.coco_dir) |
|
|
return |
|
|
|
|
|
if (args.workspace is None and args.project_name is not None) or ( |
|
|
args.workspace is not None and args.project_name is None |
|
|
): |
|
|
raise ValueError( |
|
|
"Either both workspace and project_name must be provided or none of them" |
|
|
) |
|
|
|
|
|
if args.workspace is not None: |
|
|
rf = roboflow.Roboflow(api_key=args.api_key) |
|
|
project = rf.workspace(args.workspace).project(args.project_name) |
|
|
else: |
|
|
projects = get_rf100vl_projects(api_key=args.api_key) |
|
|
project = projects[0].rf_project |
|
|
|
|
|
train_from_rf_project(project, args.dataset_version) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
trainer() |
|
|
|
