By using this site, you agree to have cookies stored on your device, strictly for functional purposes, such as storing your session and preferences.

Dismiss

 putyolo.py

View raw Download
text/x-script.python • 2.0 kiB
Python script, ASCII text executable
        
            
1
import argparse
2
import sys
3
import json
4
from os import path, environ
5
from PIL import Image
6
7
parser = argparse.ArgumentParser(description="Put a YOLO annotation into a VIA-formatted annotation file")
8
parser.add_argument("image", type=str, help="Path to the image")
9
parser.add_argument("annotation", type=str, help="Path to the YOLO annotation file, automatically determined from the image path if not provided")
10
parser.add_argument("via-location", type=str, help="Path to the VIA installation")
11
parser.add_argument("path", type=str, help="Path to put the image in, relative to the VIA search path")
12
13
args = parser.parse_args()
14
15
via_json = json.load(sys.stdin) # we can pipe the JSON from the VIA tool into this script
16
image_path = args.image
17
annotation_path = args.annotation or path.splitext(image_path)[0] + ".txt"
18
image_data = via_json["_via_img_metadata"]
19
via_location = args.via_location or "../via"
20
21
if not path.exists(image_path):
22
raise FileNotFoundError(f"Image {image_path} not found")
23
if not path.exists(annotation_path):
24
raise FileNotFoundError(f"Annotation {annotation_path} not found. Please specify it manually if it's in a different location")
25
26
with open(annotation_path, "r") as f:
27
regions = []
28
for line in f:
29
klass, cx, cy, w, h = [float(i) for i in line.split()]
30
# Since our project considers the class to be the directory, and classes may change, we discard the class index
31
image = Image.open(image_path)
32
width, height = image.size
33
region = {
34
"shape_attributes": {
35
"name": "rect",
36
"x": int((cx - w / 2) * width),
37
"y": int((cy - h / 2) * height),
38
"width": int(w * width),
39
"height": int(h * height),
40
},
41
"region_attributes": {}
42
}
43
44
image_data.append({
45
"filename": path.join(args.path, path.basename(image_path)),
46
"size": -1,
47
"regions": regions,
48
})
49
50
via_json["_via_img_metadata"] = image_data
51
print(json.dumps(via_json))
52