Skip to content
This repository has been archived by the owner on Aug 21, 2024. It is now read-only.

feat: Import yolo semantic segmentation labels #247

Closed
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -263,13 +263,19 @@ Use cases: image object detection

### YOLO to Label Studio converter

Usage:
Object detection usage:

```
label-studio-converter import yolo -i /yolo/root/directory -o ls-tasks.json
```

Help:
Semantic segmentation usage:

```
label-studio-converter import yolo-seg -i /yolo/root/directory -o ls-tasks.json
```

Help (same for semantic segmentation):

```
label-studio-converter import yolo -h
Expand Down
214 changes: 214 additions & 0 deletions label_studio_converter/imports/yolo_seg.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,214 @@
import os
import json # better to use "imports ujson as json" for the best performance

import uuid
import logging

from PIL import Image
from typing import Optional, Tuple
from urllib.request import (
pathname2url,
) # for converting "+","*", etc. in file paths to appropriate urls

from label_studio_converter.utils import ExpandFullPath
from label_studio_converter.imports.label_config import generate_label_config

logger = logging.getLogger('root')


def convert_yolo_seg_to_ls(
input_dir,
out_file,
to_name='image',
from_name='label',
out_type="annotations",
image_root_url='/data/local-files/?d=',
image_ext='.jpg,.jpeg,.png',
image_dims: Optional[Tuple[int, int]] = None,
):
"""Convert YOLO labeling to Label Studio JSON
:param input_dir: directory with YOLO where images, labels, notes.json are located
:param out_file: output file with Label Studio JSON tasks
:param to_name: object name from Label Studio labeling config
:param from_name: control tag name from Label Studio labeling config
:param out_type: annotation type - "annotations" or "predictions"
:param image_root_url: root URL path where images will be hosted, e.g.: http://example.com/images
:param image_ext: image extension/s - single string or comma separated list to search, eg. .jpeg or .jpg, .png and so on.
:param image_dims: image dimensions - optional tuple of integers specifying the image width and height of *all* images in the dataset. Defaults to opening the image to determine it's width and height, which is slower. This should only be used in the special case where you dataset has uniform image dimesions.
"""

tasks = []
logger.info('Reading YOLO notes and categories from %s', input_dir)

# build categories=>labels dict
notes_file = os.path.join(input_dir, 'classes.txt')
with open(notes_file) as f:
lines = [line.strip() for line in f.readlines()]
categories = {i: line for i, line in enumerate(lines)}
logger.info(f'Found {len(categories)} categories')

# generate and save labeling config
label_config_file = out_file.replace('.json', '') + '.label_config.xml'
generate_label_config(
categories,
{from_name: 'PolygonLabels'},
to_name,
from_name,
label_config_file,
)

# define directories
labels_dir = os.path.join(input_dir, 'labels')
images_dir = os.path.join(input_dir, 'images')
logger.info('Converting labels from %s', labels_dir)

# build array out of provided comma separated image_extns (str -> array)
image_ext = [x.strip() for x in image_ext.split(",")]
logger.info(f'image extensions->, {image_ext}')

# loop through images
for f in os.listdir(images_dir):
image_file_found_flag = False
for ext in image_ext:
if f.endswith(ext):
image_file = f
image_file_base = f[0 : -len(ext)]
image_file_found_flag = True
break
if not image_file_found_flag:
continue

image_root_url += '' if image_root_url.endswith('/') else '/'
task = {
"data": {
# eg. '../../foo+you.py' -> '../../foo%2Byou.py'
"image": image_root_url
+ str(pathname2url(image_file))
}
}

# define coresponding label file and check existence
label_file = os.path.join(labels_dir, image_file_base + '.txt')

if os.path.exists(label_file):
task[out_type] = [
{
"result": [],
"ground_truth": False,
}
]

# read image sizes
if image_dims is None:
# default to opening file if we aren't given image dims. slow!
with Image.open(os.path.join(images_dir, image_file)) as im:
image_width, image_height = im.size
else:
image_width, image_height = image_dims

with open(label_file) as file:
# convert all polygons to Label Studio Results
lines = file.readlines()
for line in lines:
label_id, *label_coords = line.split()
label_coords = [float(coord) * 100 for coord in label_coords]
vertices_xy = [label_coords[i: i + 2] for i in range(0, len(label_coords), 2)]

item = {
"id": uuid.uuid4().hex[0:10],
"type": "polygonlabels",
"value": {
"points": vertices_xy,
"closed": True,
"polygonlabels": [categories[int(label_id)]],
},
"to_name": to_name,
"from_name": from_name,
"image_rotation": 0,
"original_width": image_width,
"original_height": image_height
}
task[out_type][0]["result"].append(item)

tasks.append(task)

if len(tasks) > 0:
logger.info('Saving Label Studio JSON to %s', out_file)
with open(out_file, 'w') as out:
json.dump(tasks, out)

print(
'\n'
f' 1. Create a new project in Label Studio\n'
f' 2. Use Labeling Config from "{label_config_file}"\n'
f' 3. Setup serving for images [e.g. you can use Local Storage (or others):\n'
f' https://labelstud.io/guide/storage.html#Local-storage]\n'
f' 4. Import "{out_file}" to the project\n'
)
else:
logger.error('No labels converted')


def add_parser(subparsers):
yolo_seg = subparsers.add_parser('yolo-seg')

yolo_seg.add_argument(
'-i',
'--input',
dest='input',
required=True,
help='directory with YOLO where images, labels, notes.json are located',
action=ExpandFullPath,
)
yolo_seg.add_argument(
'-o',
'--output',
dest='output',
help='output file with Label Studio JSON tasks',
default='output.json',
action=ExpandFullPath,
)
yolo_seg.add_argument(
'--to-name',
dest='to_name',
help='object name from Label Studio labeling config',
default='image',
)
yolo_seg.add_argument(
'--from-name',
dest='from_name',
help='control tag name from Label Studio labeling config',
default='label',
)
yolo_seg.add_argument(
'--out-type',
dest='out_type',
help='annotation type - "annotations" or "predictions"',
default='annotations',
)
yolo_seg.add_argument(
'--image-root-url',
dest='image_root_url',
help='root URL path where images will be hosted, e.g.: http://example.com/images',
default='/data/local-files/?d=',
)
yolo_seg.add_argument(
'--image-ext',
dest='image_ext',
help='image extension to search: .jpeg or .jpg, .png',
default='.jpg',
)
yolo_seg.add_argument(
'--image-dims',
dest='image_dims',
type=int,
nargs=2,
help=(
"optional tuple of integers specifying the image width and height of *all* "
"images in the dataset. Defaults to opening the image to determine it's width "
"and height, which is slower. This should only be used in the special "
"case where you dataset has uniform image dimesions. e.g. `--image-dims 600 800` "
"if all your images are of dimensions width=600, height=800"
),
default=None,
)
14 changes: 13 additions & 1 deletion label_studio_converter/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from label_studio_converter.converter import Converter, Format, FormatNotSupportedError
from label_studio_converter.exports.csv import ExportToCSV
from label_studio_converter.utils import ExpandFullPath
from label_studio_converter.imports import yolo as import_yolo, coco as import_coco
from label_studio_converter.imports import yolo as import_yolo, yolo_seg as import_yolo_seg, coco as import_coco

logging.basicConfig(level=logging.INFO)

Expand Down Expand Up @@ -97,6 +97,7 @@ def get_all_args():
)
import_format = parser_import.add_subparsers(dest='import_format')
import_yolo.add_parser(import_format)
import_yolo_seg.add_parser(import_format)
import_coco.add_parser(import_format)

return parser.parse_args()
Expand Down Expand Up @@ -166,6 +167,17 @@ def imports(args):
image_root_url=args.image_root_url,
image_ext=args.image_ext,
)

elif args.import_format == 'yolo-seg':
import_yolo_seg.convert_yolo_seg_to_ls(
input_dir=args.input,
out_file=args.output,
to_name=args.to_name,
from_name=args.from_name,
out_type=args.out_type,
image_root_url=args.image_root_url,
image_ext=args.image_ext,
)

elif args.import_format == 'coco':
import_coco.convert_coco_to_ls(
Expand Down
4 changes: 4 additions & 0 deletions tests/data/test_import_yolo_seg_data/classes.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
closed_door
opened_door
bus
number
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
4 changes: 4 additions & 0 deletions tests/data/test_import_yolo_seg_data/labels/img_1.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
2 0.0633721205070745 0.08960138500340462 0.05533468083300653 0.1070405807423223 0.046678976568625616 0.18461493420164576 0.037714140009088244 0.2970676791388046 0.03153149410595902 0.3938852830686579 0.028749303449550866 0.499121809079368 0.027203641973768566 0.5911286003915888 0.02844017115439441 0.6819326885494016 0.03400455246721071 0.7709327219756021 0.035241081647836556 0.8388377773093125 0.03678674312361886 0.8586823793570464 0.043587653617061 0.8683040045923114 0.10665064182897904 0.874317520364352 0.11190589084663888 0.8791283329819843 0.2559615403895497 0.8899526613716574 0.35921172697180764 0.8989729350297183 0.38270578140369876 0.9037837476473507 0.423202112069195 0.8977702318753102 0.4541153415848412 0.8911553645260656 0.4902838201181472 0.8779256298275763 0.5041947734001879 0.8779256298275763 0.5100682870081605 0.9031823960701466 0.5233609756998884 0.9200202402318602 0.5397449873431809 0.9152094276142279 0.5514920145591264 0.8941621224120857 0.5607659834138202 0.8610877856658626 0.5774591273522691 0.8496621056989855 0.6195011194935479 0.8334256131144762 0.6553604657316973 0.8213985815703948 0.6862736952473435 0.8099729016035176 0.718114321648459 0.8003512763682528 0.7450088313270711 0.7901282995557838 0.7802499129749075 0.781709377474927 0.812090539376023 0.7696823459308458 0.84640422413839 0.761263423849989 0.8525868700415196 0.7889255964013757 0.8658795587332472 0.8021553310998649 0.8748443952927846 0.7955404637506203 0.879790512015288 0.77028369750805 0.8835000995571655 0.7462296344198875 0.9147224613679683 0.7366080091846228 0.9379073835047026 0.728189087103766 0.9715598288456058 0.6957161019347466 0.9727963580262313 0.6229525610930557 0.9715598288456054 0.5435741529021202 0.9703232996649797 0.47020926048322503 0.9693959027795104 0.37519571128498397 0.968468505894041 0.282587568395559 0.9666137121231021 0.17795239396205292 0.9595036693345037 0.16472265926356364 0.9131338250610345 0.1569050887599109 0.8593448057038102 0.14608076037023787 0.8040101248708037 0.13645913513497296 0.7517667669893617 0.12743886147691208 0.6967412184515118 0.11726355755479129 0.6370786854863147 0.10643922916511822 0.5789618139969 0.09501354919824113 0.5242453977542065 0.08719597869458838 0.48776778692574413 0.08238516607695592 0.4383066197007104 0.07637165030491536 0.40028334739646565 0.07035813453287476 0.3585504875503434 0.06254056402922202 0.3189815537703164 0.0571283998343855 0.27415737097262954 0.0541216419483652 0.21855653267561792 0.0541216419483652 0.1765145405343392 0.05772975141158956 0.13509081298337347 0.06494597033803824 0.08995749789053015 0.0781757050365275
0 0.3913081473705955 0.2077395452853236 0.42526224388285544 0.21093553828971323 0.46140692726687405 0.21306620029263962 0.49207514347149584 0.21519686229556606 0.4945910723225618 0.3550514436146151 0.4940932116698893 0.48773357743321344 0.49508893297523426 0.5874872984793128 0.49508893297523426 0.6882095022540152 0.496084654280579 0.7618141896278362 0.497080375585924 0.848009152473495 0.44928575292937056 0.8605994279453328 0.39501894178807545 0.8848114961604054 0.3935253598300581 0.7569717759848218 0.3945210811354031 0.6194472285232089 0.3940232204827305 0.5138826111054919 0.39153391721936837 0.3821689600154965 0.391036056566696 0.2814467562407942
0 0.7565461408431807 0.23786503345366333 0.7874135013088714 0.23786503345366333 0.8217658863432692 0.24173896436807496 0.8222637469959416 0.34149268541417444 0.823757328953959 0.42381371734542156 0.824753050259304 0.5332522656775501 0.824753050259304 0.6572180549387222 0.8267444928699936 0.7482554314273956 0.7963749930569751 0.7608457068992333 0.7615247473699049 0.7695620514566595 0.7615247473699049 0.6562495722101194 0.7620226080225774 0.5690861266358577 0.7590354441065428 0.47223785377556693 0.7575418621485255 0.37054716727226167
3 0.26465981600281774 0.1129307614638883 0.26117479143411076 0.22140082706741387 0.31593946322807825 0.22140082706741387 0.32191379106014745 0.10808834782087376
4 changes: 4 additions & 0 deletions tests/data/test_import_yolo_seg_data/labels/img_2.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
2 0.8570140867160779 0.29644290501897036 0.7891087120622501 0.287059067307079 0.7298721086408257 0.27839706326533303 0.6706355052194016 0.2791188969354785 0.6118805001997775 0.27984073060562403 0.54493832234955 0.2863372336369334 0.4582506100255145 0.3180979151233352 0.3763788817194809 0.3476930955993005 0.30076793263685 0.3736791077245383 0.2434577228226265 0.3953341178289031 0.21648821232181542 0.4025524545303581 0.21070903150021308 0.45957731447185224 0.2092642362948125 0.5418663528684386 0.20733784268761168 0.6494195697201174 0.20733784268761168 0.715828267373503 0.23671534519742374 0.7461452815196137 0.25164489565322984 0.7504762835404866 0.25068169884962943 0.7721312936448517 0.2564608796712318 0.7721312936448517 0.259350470082033 0.785846133377616 0.26705604451083614 0.7930644700790708 0.2771696109486403 0.77718412933587 0.2805407997612417 0.7569727865717961 0.3368878127718647 0.7742967946552881 0.3797500705320823 0.7872898007179068 0.42550191870310106 0.7973954720999439 0.45247142920391215 0.8046138088013987 0.4563242164183137 0.837096323957946 0.4664377828561178 0.8616386687428929 0.488591309338927 0.8616386687428929 0.4982232773749309 0.8392618249683825 0.49918647417853124 0.8132758128431448 0.5304903702955441 0.8255469852356182 0.586355784904367 0.837096323957946 0.6301812394681849 0.8385399912982371 0.6985682125238131 0.8392618249683825 0.7236113294174232 0.8378181576280916 0.7332432974534272 0.8486456626802739 0.7650287919722402 0.854420332041438 0.7761055552136447 0.8378181576280916 0.7982590816964538 0.8349308229475095 0.8016302705090552 0.8125539791729992 0.842566134662072 0.8190504822043088 0.8748332275826853 0.8132758128431448 0.8859099908240898 0.8031701414611078 0.8878363844312903 0.7569727865717961 0.8690540467610828 0.7461452815196137 0.8700172435646834 0.6790147501960827 0.8690540467610828 0.6097187178621151 0.8685724483592826 0.5252641784550922 0.8661644563502816 0.4335913023466144 0.8637564643412806 0.37440094139468383
0 0.4857017189281258 0.3902812821378847 0.5280823782865431 0.37656644240512027 0.5300087718937438 0.4104926249019585 0.5256743862775421 0.5332043488266928 0.5261559846793423 0.6025003811606602 0.5271191814829427 0.7194374357242305 0.5309719686973443 0.8082229771521263 0.4857017189281258 0.7973954720999439 0.4881097109371267 0.6833457522169557 0.488591309338927 0.559912194622076 0.48666491573172604 0.4718484868643257
0 0.2949887518152476 0.43864413803763286 0.32051346711065803 0.4328694686764689 0.31810547510165704 0.5389790181878567 0.3161790814944563 0.6515850707305539 0.31762387669985687 0.7490326162001957 0.2940255550116472 0.7439797805091772 0.2925807598062466 0.6530287380708448 0.2940255550116472 0.575792535365277 0.293543956609847 0.49999999999999994
3 0.5849109896989664 0.3029394080502798 0.5844293912971663 0.35563326597090095 0.6070645161817755 0.35707693331119195 0.608509311387176 0.30149574070998886
4 changes: 4 additions & 0 deletions tests/data/test_import_yolo_seg_data/labels/img_3.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
3 0.2644538903363631 0.20499629894854596 0.26521803018928647 0.25593895581010573 0.29731190401206903 0.25729742665974725 0.29616569423268396 0.2043170635237252
0 0.1535921104376025 0.2518635432611809 0.179572865436998 0.24235424731368976 0.18110114514284475 0.3272586754162893 0.17919079551053627 0.4338986371131542 0.18186528499576815 0.5452932467837648 0.1574128097022195 0.5364631862610945 0.15168176080529402 0.4325401662635127 0.15550246006991103 0.33608873593895955
0 0.07717812514526295 0.2769952539795504 0.09398920190957767 0.27020289973134237 0.09246092220373088 0.374805155153745 0.09475334176250104 0.4705773500534773 0.08214503418926505 0.46989811462865644 0.07947054470403314 0.37548439057856586
2 0.4233334785195612 0.1866569424783845 0.39544237388785725 0.1689968214330438 0.35799952109461086 0.1608459963351942 0.3136794096250539 0.15473287751180706 0.27012343800842037 0.15677058378626946 0.23764749425917608 0.16763835058340215 0.19485566249546588 0.18597770705356367 0.14671485176129195 0.20975094692229157 0.1043050899240435 0.23012800966691543 0.08061675448341825 0.24303348273851058 0.07221121610126088 0.26001436835903047 0.07182914617479919 0.35239038613465873 0.07297535595418428 0.4325401662635127 0.07373949580710767 0.4739735271775812 0.07443104237400336 0.5188132537471261 0.09754627292493608 0.5278538772514909 0.10216931903512264 0.5664819958610495 0.12019919886485014 0.5845632428697792 0.1317568141403165 0.5673038707250827 0.14100290636068957 0.5541538729005522 0.2117355118465437 0.5853851177338124 0.21635855795673029 0.6182601122951389 0.23300152395340176 0.6281226106635368 0.24964448995007338 0.624013236343371 0.2547298406712786 0.5919601166460776 0.25889058217044647 0.573878869637348 0.3698436888149235 0.5516882483084526 0.37307982109205406 0.5722351199092817 0.38787356864465106 0.5845632428697792 0.4045165346413226 0.5796319936855803 0.41283801763965833 0.5541538729005522 0.414687236083733 0.5410038750760214 0.44843547268809475 0.5286757521155241 0.4488977772991134 0.43498201761574345 0.44288781735587096 0.37169765308518987 0.43364172513549787 0.28868829181784034
Loading