-
Notifications
You must be signed in to change notification settings - Fork 0
/
create_comet_artifact.py
71 lines (54 loc) · 2.78 KB
/
create_comet_artifact.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
"""
Download an exported dataset from Roboflow, and create a Comet data artifact with the files.
"""
from roboflow import Roboflow
from glob import glob
from comet_ml import Artifact, Experiment
import yaml
import os
import shutil
# These are all the "configuration parameters" to set before running this script
roboflow_workspace = "kristen-kehrer"
roboflow_project = "schoolbus-images"
roboflow_version = 4
staging_directory = f"Schoolbus-Images-{roboflow_version}"
roboflow_api_key_file = '.roboflow_api_key'
def create_artifact():
# Delete the staging directory if it already exists
if os.path.exists(staging_directory):
shutil.rmtree(staging_directory)
if not os.path.exists(roboflow_api_key_file):
raise f'Missing configuration file: {roboflow_api_key_file}'
# Load api key from saved credential file and strip whitespace characters from the ends
roboflow_api_key = open(roboflow_api_key_file).read().strip()
# With the Roboflow python SDK, download the desired dataset.
# After `.download("yolov5")` completes, the full data artifact will
# be on the local filesystem at `staging_directory`
rf = Roboflow(api_key=roboflow_api_key)
project = rf.workspace(roboflow_workspace).project(roboflow_project)
dataset = project.version(roboflow_version).download("yolov5")
# Start a comet experiment (project name / api key are already set in `.comet.config`)
experiment = Experiment()
# Load the `data.yaml` file (this is how data needs to be formatted for YOLOv5) from the Roboflow artifact
metadata = yaml.load(open(os.path.join(staging_directory, 'data.yaml')), Loader=yaml.Loader)
# Convert the `names` field in the metadata from a list of strings to an integer-keyed dictionary
# of strings (`['bus', 'car']` needs to be `{ '0': 'bus', '1': 'car'}` for the comet metadata
# I created a PR to add this to the yolov5 library, I'll remove this once it's merged.
names_dict = {}
for i in range(len(metadata['names'])):
names_dict[str(i)] = metadata['names'][i]
metadata['names'] = names_dict
# Create the comet Artifact object
artifact = Artifact(name="schoolbus-yolov5-take2", artifact_type="dataset", metadata=metadata)
# Iterate through all files in `staging_directory`, and add the file to the artifact
for file in glob(f'{staging_directory}/**/*.*', recursive=True):
artifact.add(file, logical_path=file.replace('\\', '/'))
# Log the artifact and end the experiment
experiment.log_artifact(artifact)
experiment.end()
def write_yolov5_data_file():
with open('comet_artifact.yaml', 'w') as file:
file.write(f'path: "comet://kristenkehrer/schoolbus-yolov5-take2:latest"\n')
# When this script is run, run these two functions:
create_artifact()
write_yolov5_data_file()