putyolo.py
Python script, ASCII text executable
1import argparse 2import sys 3import json 4from os import path, environ 5from PIL import Image 6 7parser = argparse.ArgumentParser(description="Put a YOLO annotation into a VIA-formatted annotation file") 8parser.add_argument("image", type=str, help="Path to the image") 9parser.add_argument("annotation", type=str, help="Path to the YOLO annotation file, automatically determined from the image path if not provided") 10parser.add_argument("via-location", type=str, help="Path to the VIA installation") 11parser.add_argument("path", type=str, help="Path to put the image in, relative to the VIA search path") 12 13args = parser.parse_args() 14 15via_json = json.load(sys.stdin) # we can pipe the JSON from the VIA tool into this script 16image_path = args.image 17annotation_path = args.annotation or path.splitext(image_path)[0] + ".txt" 18image_data = via_json["_via_img_metadata"] 19via_location = args.via_location or "../via" 20 21if not path.exists(image_path): 22raise FileNotFoundError(f"Image {image_path} not found") 23if not path.exists(annotation_path): 24raise FileNotFoundError(f"Annotation {annotation_path} not found. Please specify it manually if it's in a different location") 25 26with open(annotation_path, "r") as f: 27regions = [] 28for line in f: 29klass, 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 31image = Image.open(image_path) 32width, height = image.size 33region = { 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 44image_data.append({ 45"filename": path.join(args.path, path.basename(image_path)), 46"size": -1, 47"regions": regions, 48}) 49 50via_json["_via_img_metadata"] = image_data 51print(json.dumps(via_json)) 52