def post(self):
""" Creates a category """
args = create_category.parse_args()
name = args.get('name')
supercategory = args.get('supercategory')
metadata = args.get('metadata', {})
color = args.get('color')
keypoint_edges = args.get('keypoint_edges')
keypoint_labels = args.get('keypoint_labels')
keypoint_colors = args.get('keypoint_colors')
try:
category = CategoryModel(
name=name,
supercategory=supercategory,
color=color,
metadata=metadata,
keypoint_edges=keypoint_edges,
keypoint_labels=keypoint_labels,
keypoint_colors=keypoint_colors,
)
category.save()
except NotUniqueError as e:
return {
'message':
'Category already exists. Check the undo tab to fully delete the category.'
}, 400
return query_util.fix_ids(category)
def create_category():
category = CategoryModel(
name="category",
supercategory="supercategory",
color="red",
metadata={},
keypoint_edges=[],
keypoint_labels=[],
keypoint_colors=[],
)
category.save()
return category.id
def post(self, dataset_id):
""" Updates dataset by ID """
dataset = current_user.datasets.filter(id=dataset_id,
deleted=False).first()
if dataset is None:
return {"message": "Invalid dataset id"}, 400
args = update_dataset.parse_args()
categories = args.get('categories')
default_annotation_metadata = args.get('default_annotation_metadata')
set_default_annotation_metadata = args.get(
'set_default_annotation_metadata')
if categories is not None:
dataset.categories = CategoryModel.bulk_create(categories)
if default_annotation_metadata is not None:
update = {}
for key, value in default_annotation_metadata.items():
if key not in dataset.default_annotation_metadata:
update[f'set__metadata__{key}'] = value
dataset.default_annotation_metadata = default_annotation_metadata
if len(update.keys()) > 0:
AnnotationModel.objects(dataset_id=dataset.id, deleted=False)\
.update(**update)
dataset.update(
categories=dataset.categories,
default_annotation_metadata=dataset.default_annotation_metadata)
return {"success": True}
def get(self, image_id):
""" Called when loading from the annotator client """
image = ImageModel.objects(id=image_id)\
.exclude('events').first()
if image is None:
return {'success': False, 'message': 'Could not load image'}, 400
dataset = current_user.datasets.filter(id=image.dataset_id).first()
if dataset is None:
return {
'success': False,
'message': 'Could not find associated dataset'
}, 400
categories = CategoryModel.objects(deleted=False)\
.in_bulk(dataset.categories).items()
# Get next and previous image
images = ImageModel.objects(dataset_id=dataset.id, deleted=False)
pre = images.filter(
file_name__lt=image.file_name).order_by('-file_name').first()
nex = images.filter(
file_name__gt=image.file_name).order_by('file_name').first()
preferences = {}
if not Config.LOGIN_DISABLED:
preferences = current_user.preferences
# Generate data about the image to return to client
data = {
'image': query_util.fix_ids(image),
'categories': [],
'dataset': query_util.fix_ids(dataset),
'preferences': preferences,
'permissions': {
'dataset': dataset.permissions(current_user),
'image': image.permissions(current_user)
}
}
data['image']['previous'] = pre.id if pre else None
data['image']['next'] = nex.id if nex else None
for category in categories:
category = query_util.fix_ids(category[1])
category_id = category.get('id')
annotations = AnnotationModel.objects(image_id=image_id, category_id=category_id, deleted=False)\
.exclude('events').all()
category['show'] = True
category['visualize'] = False
category[
'annotations'] = [] if annotations is None else query_util.fix_ids(
annotations)
data.get('categories').append(category)
return data
def create_category(self, category_data):
exist_category = CategoryModel.find_by_name(category_data['name'])
if exist_category is not None:
return exist_category.id
category = CategoryModel(
name=category_data['name'],
supercategory="supercategory",
color="color",
metadata={},
keypoint_edges=[],
keypoint_labels=[],
keypoint_colors=[],
)
category.save()
return category.id
def get(self, dataset_id):
""" All users in the dataset """
args = dataset_generate.parse_args()
dataset = current_user.datasets.filter(id=dataset_id,
deleted=False).first()
if dataset is None:
return {"message": "Invalid dataset id"}, 400
images = ImageModel.objects(dataset_id=dataset.id, deleted=False)
annotated_images = images.filter(annotated=True)
annotations = AnnotationModel.objects(dataset_id=dataset_id,
deleted=False)
# Calculate annotation counts by category in this dataset
category_count = dict()
image_category_count = dict()
for category in dataset.categories:
# Calculate the annotation count in the current category in this dataset
cat_name = CategoryModel.objects(id=category).first()['name']
cat_count = AnnotationModel.objects(dataset_id=dataset_id,
category_id=category,
deleted=False).count()
category_count.update({str(cat_name): cat_count})
# Calculate the annotated images count in the current category in this dataset
image_count = len(
AnnotationModel.objects(dataset_id=dataset_id,
category_id=category,
deleted=False).distinct('image_id'))
image_category_count.update({str(cat_name): image_count})
stats = {
'total': {
'Users': dataset.get_users().count(),
'Images': images.count(),
'Annotated Images': annotated_images.count(),
'Annotations': annotations.count(),
'Categories': len(dataset.categories),
'Time Annotating (s)': (images.sum('milliseconds') or 0) / 1000
},
'average': {
'Image Size (px)':
images.average('width'),
'Image Height (px)':
images.average('height'),
'Annotation Area (px)':
annotations.average('area'),
'Time (ms) per Image':
images.average('milliseconds') or 0,
'Time (ms) per Annotation':
annotations.average('milliseconds') or 0
},
'categories': category_count,
'images_per_category': image_category_count
}
return stats
def test_create_category(self):
query = {"name": category1.get("name")}
create_category1 = upsert(CategoryModel, query=query, update=category1)
assert create_category1.name == category1.get("name")
assert create_category1.color == category1.get("color")
found = CategoryModel.objects(**query).first()
assert found.name == category1.get("name")
assert found.color == category1.get("color")
def post(self):
""" Creates a dataset """
args = dataset_create.parse_args()
name = args['name']
categories = args.get('categories', [])
category_ids = CategoryModel.bulk_create(categories)
try:
dataset = DatasetModel(name=name, categories=category_ids)
dataset.save()
except NotUniqueError:
return {'message': 'Dataset already exists. Check the undo tab to fully delete the dataset.'}, 400
return query_util.fix_ids(dataset)
def get(self, refset_id):
""" All users in the refset """
args = refset_generate.parse_args()
refset = current_user.refsets.filter(id=refset_id, deleted=False).first()
if refset is None:
return {"message": "Invalid refset id"}, 400
pixels = PixelModel.objects(refset_id=refset.id, deleted=False)
annotated_pixels = pixels.filter(annotated=True)
annotations = AnnotationModel.objects(refset_id=refset_id, deleted=False)
# Calculate annotation counts by category in this refset
category_count = dict()
image_category_count = dict()
for category in refset.categories:
# Calculate the annotation count in the current category in this refset
cat_name = CategoryModel.objects(id=category).first()['name']
cat_count = AnnotationModel.objects(refset_id=refset_id, category_id=category, deleted=False).count()
category_count.update({str(cat_name): cat_count})
# Calculate the annotated pixels count in the current category in this refset
image_count = len(AnnotationModel.objects(refset_id=refset_id, category_id=category, deleted=False).distinct('image_id'))
image_category_count.update({str(cat_name): image_count})
stats = {
'total': {
'Users': refset.get_users().count(),
'Pixels': pixels.count(),
'Annotated Pixels': annotated_pixels.count(),
'Annotations': annotations.count(),
'Categories': len(refset.categories),
'Time Annotating (s)': (pixels.sum('milliseconds') or 0) / 1000
},
'average': {
'Pixel Size (px)': pixels.average('width'),
'Pixel Height (px)': pixels.average('height'),
'Annotation Area (px)': annotations.average('area'),
'Time (ms) per Pixel': pixels.average('milliseconds') or 0,
'Time (ms) per Annotation': annotations.average('milliseconds') or 0
},
'categories': category_count,
'pixels_per_category': image_category_count
}
return stats
def get(self, dataset_id):
""" All users in the dataset """
dataset = current_user.datasets.filter(id=dataset_id,
deleted=False).first()
if dataset is None:
return {"message": "Invalid dataset id"}, 400
category_dict = dict()
for category in dataset.categories:
cat_name = CategoryModel.objects(id=category).first()['name']
category_dict.update({str(cat_name): category})
res = {
'categories': category_dict,
}
return res
def post(self, image_id):
""" COCO data test """
print("Converter Model")
image_model = ImageModel.objects(id=image_id).first()
if not image_model:
return {"message": "Invalid image ID"}, 400
#image = Image.open(image_model.path)
annotations = AnnotationModel.objects(image_id=image_id)
for annotation in annotations:
category_id = annotation.category_id
category = CategoryModel.objects(id=category_id).first()
for segmentation in annotation.segmentation:
print(segmentation)
convert(image_model.path, segmentation, category.color)
AnnotationModel.objects(image_id=image_id).delete()
return {"success": True}
def get(self, image_id):
""" Called when loading from the annotator client """
image = ImageModel.objects(id=image_id)\
.exclude('events').first()
if image is None:
return {'success': False, 'message': 'Could not load image'}, 400
dataset = current_user.datasets.filter(id=image.dataset_id).first()
if dataset is None:
return {
'success': False,
'message': 'Could not find associated dataset'
}, 400
categories = CategoryModel.objects(deleted=False)\
.in_bulk(dataset.categories).items()
# Generate data about the image to return to client
data = {'categories': []}
for category in categories:
category = query_util.fix_ids(category[1])
category_id = category.get('id')
annotations = AnnotationModel.objects(image_id=image_id, category_id=category_id, deleted=False)\
.exclude('events').all()
category['show'] = True
category['visualize'] = False
category[
'annotations'] = [] if annotations is None else query_util.fix_ids(
annotations)
data.get('categories').append(category)
return data
def export_annotations(task_id,
dataset_id,
categories,
with_empty_images=False):
task = TaskModel.objects.get(id=task_id)
dataset = DatasetModel.objects.get(id=dataset_id)
task.update(status="PROGRESS")
socket = create_socket()
task.info("Beginning Export (COCO Format)")
db_categories = CategoryModel.objects(id__in=categories, deleted=False) \
.only(*CategoryModel.COCO_PROPERTIES)
db_images = ImageModel.objects(
deleted=False, dataset_id=dataset.id).only(*ImageModel.COCO_PROPERTIES)
db_annotations = AnnotationModel.objects(deleted=False,
category_id__in=categories)
total_items = db_categories.count()
coco = {'images': [], 'categories': [], 'annotations': []}
total_items += db_images.count()
progress = 0
# iterate though all categoires and upsert
category_names = []
for category in fix_ids(db_categories):
if len(category.get('keypoint_labels', [])) > 0:
category['keypoints'] = category.pop('keypoint_labels', [])
category['skeleton'] = category.pop('keypoint_edges', [])
else:
if 'keypoint_edges' in category:
del category['keypoint_edges']
if 'keypoint_labels' in category:
del category['keypoint_labels']
task.info(f"Adding category: {category.get('name')}")
coco.get('categories').append(category)
category_names.append(category.get('name'))
progress += 1
task.set_progress((progress / total_items) * 100, socket=socket)
total_annotations = db_annotations.count()
total_images = db_images.count()
for image in db_images:
image = fix_ids(image)
progress += 1
task.set_progress((progress / total_items) * 100, socket=socket)
annotations = db_annotations.filter(image_id=image.get('id'))\
.only(*AnnotationModel.COCO_PROPERTIES)
annotations = fix_ids(annotations)
if len(annotations) == 0:
if with_empty_images:
coco.get('images').append(image)
continue
num_annotations = 0
for annotation in annotations:
has_keypoints = len(annotation.get('keypoints', [])) > 0
has_segmentation = len(annotation.get('segmentation', [])) > 0
if has_keypoints or has_segmentation:
if not has_keypoints:
if 'keypoints' in annotation:
del annotation['keypoints']
else:
arr = np.array(annotation.get('keypoints', []))
arr = arr[2::3]
annotation['num_keypoints'] = len(arr[arr > 0])
num_annotations += 1
coco.get('annotations').append(annotation)
task.info(
f"Exporting {num_annotations} annotations for image {image.get('id')}"
)
coco.get('images').append(image)
task.info(
f"Done export {total_annotations} annotations and {total_images} images from {dataset.name}"
)
timestamp = time.time()
directory = f"{dataset.directory}.exports/"
file_path = f"{directory}coco-{timestamp}.json"
if not os.path.exists(directory):
os.makedirs(directory)
task.info(f"Writing export to file {file_path}")
with open(file_path, 'w') as fp:
json.dump(coco, fp)
task.info("Creating export object")
export = ExportModel(dataset_id=dataset.id,
path=file_path,
tags=["COCO", *category_names])
export.save()
task.set_progress(100, socket=socket)
def get(self, dataset_id):
""" Endpoint called by image viewer client """
parsed_args = page_data.parse_args()
per_page = parsed_args.get('limit')
page = parsed_args.get('page') - 1
folder = parsed_args.get('folder')
order = parsed_args.get('order')
args = dict(request.args)
# Check if dataset exists
dataset = current_user.datasets.filter(id=dataset_id,
deleted=False).first()
if dataset is None:
return {'message', 'Invalid dataset id'}, 400
# Make sure folder starts with is in proper format
if len(folder) > 0:
folder = folder[0].strip('/') + folder[1:]
if folder[-1] != '/':
folder = folder + '/'
# Get directory
directory = os.path.join(dataset.directory, folder)
if not gfile.exists(directory):
return {'message': 'Directory does not exist.'}, 400
# Remove parsed arguments
for key in parsed_args:
args.pop(key, None)
# Generate query from remaining arugments
query = {}
for key, value in args.items():
lower = value.lower()
if lower in ["true", "false"]:
value = json.loads(lower)
if len(lower) != 0:
query[key] = value
# Change category_ids__in to list
if 'category_ids__in' in query.keys():
query['category_ids__in'] = [
int(x) for x in query['category_ids__in'].split(',')
]
# Initialize mongo query with required elements:
query_build = Q(dataset_id=dataset_id)
query_build &= Q(path__startswith=directory)
query_build &= Q(deleted=False)
# Define query names that should use complex logic:
complex_query = ['annotated', 'category_ids__in']
# Add additional 'and' arguments to mongo query that do not require complex_query logic
for key in query.keys():
if key not in complex_query:
query_dict = {}
query_dict[key] = query[key]
query_build &= Q(**query_dict)
# Add additional arguments to mongo query that require more complex logic to construct
if 'annotated' in query.keys():
if 'category_ids__in' in query.keys() and query['annotated']:
# Only show annotated images with selected category_ids
query_dict = {}
query_dict['category_ids__in'] = query['category_ids__in']
query_build &= Q(**query_dict)
else:
# Only show non-annotated images
query_dict = {}
query_dict['annotated'] = query['annotated']
query_build &= Q(**query_dict)
elif 'category_ids__in' in query.keys():
# Ahow annotated images with selected category_ids or non-annotated images
query_dict_1 = {}
query_dict_1['category_ids__in'] = query['category_ids__in']
query_dict_2 = {}
query_dict_2['annotated'] = False
query_build &= (Q(**query_dict_1) | Q(**query_dict_2))
# Perform mongodb query
images = current_user.images \
.filter(query_build) \
.order_by(order).only('id', 'file_name', 'annotating', 'annotated', 'num_annotations')
total = images.count()
pages = int(total / per_page) + 1
images = images.skip(page * per_page).limit(per_page)
images_json = query_util.fix_ids(images)
# for image in images:
# image_json = query_util.fix_ids(image)
# query = AnnotationModel.objects(image_id=image.id, deleted=False)
# category_ids = query.distinct('category_id')
# categories = CategoryModel.objects(id__in=category_ids).only('name', 'color')
# image_json['annotations'] = query.count()
# image_json['categories'] = query_util.fix_ids(categories)
# images_json.append(image_json)
subdirectories = [
f for f in sorted(gfile.listdir(directory))
if gfile.isdir(directory + f) and not f.startswith('.')
]
categories = CategoryModel.objects(id__in=dataset.categories).only(
'id', 'name')
return {
"total": total,
"per_page": per_page,
"pages": pages,
"page": page,
"images": images_json,
"folder": folder,
"directory": directory,
"dataset": query_util.fix_ids(dataset),
"categories": query_util.fix_ids(categories),
"subdirectories": subdirectories
}
def execute(self, dataset_id, categories):
dataset = DatasetModel.find_by(dataset_id)
category_id_list = CategoryModel.bulk_create(categories)
# dataset.update(set__categories=category_id_list)
dataset.update_categories(category_id_list)
def import_annotations(task_id, dataset_id, coco_json):
task = TaskModel.objects.get(id=task_id)
dataset = DatasetModel.objects.get(id=dataset_id)
task.update(status="PROGRESS")
socket = create_socket()
task.info("Beginning Import")
images = ImageModel.objects(dataset_id=dataset.id)
categories = CategoryModel.objects
coco_images = coco_json.get('images', [])
coco_annotations = coco_json.get('annotations', [])
coco_categories = coco_json.get('categories', [])
task.info(f"Importing {len(coco_categories)} categories, "
f"{len(coco_images)} images, and "
f"{len(coco_annotations)} annotations")
total_items = sum(
[len(coco_categories),
len(coco_annotations),
len(coco_images)])
progress = 0
task.info("===== Importing Categories =====")
# category id mapping ( file : database )
categories_id = {}
# Create any missing categories
for category in coco_categories:
category_name = category.get('name')
category_id = category.get('id')
category_model = categories.filter(name__iexact=category_name).first()
if category_model is None:
task.warning(
f"{category_name} category not found (creating a new one)")
new_category = CategoryModel(
name=category_name,
keypoint_edges=category.get('skeleton', []),
keypoint_labels=category.get('keypoints', []))
new_category.save()
category_model = new_category
dataset.categories.append(new_category.id)
task.info(f"{category_name} category found")
# map category ids
categories_id[category_id] = category_model.id
# update progress
progress += 1
task.set_progress((progress / total_items) * 100, socket=socket)
dataset.update(set__categories=dataset.categories)
task.info("===== Loading Images =====")
# image id mapping ( file: database )
images_id = {}
categories_by_image = {}
# Find all images
for image in coco_images:
image_id = image.get('id')
image_filename = image.get('file_name')
# update progress
progress += 1
task.set_progress((progress / total_items) * 100, socket=socket)
image_model = images.filter(file_name__exact=image_filename).all()
if len(image_model) == 0:
task.warning(f"Could not find image {image_filename}")
continue
if len(image_model) > 1:
task.error(
f"Too many images found with the same file name: {image_filename}"
)
continue
task.info(f"Image {image_filename} found")
image_model = image_model[0]
images_id[image_id] = image_model
categories_by_image[image_id] = list()
task.info("===== Import Annotations =====")
for annotation in coco_annotations:
image_id = annotation.get('image_id')
category_id = annotation.get('category_id')
segmentation = annotation.get('segmentation', [])
keypoints = annotation.get('keypoints', [])
# is_crowd = annotation.get('iscrowed', False)
area = annotation.get('area', 0)
bbox = annotation.get('bbox', [0, 0, 0, 0])
isbbox = annotation.get('isbbox', False)
progress += 1
task.set_progress((progress / total_items) * 100, socket=socket)
has_segmentation = len(segmentation) > 0
has_keypoints = len(keypoints) > 0
if not has_segmentation and not has_keypoints:
task.warning(
f"Annotation {annotation.get('id')} has no segmentation or keypoints"
)
continue
try:
image_model = images_id[image_id]
category_model_id = categories_id[category_id]
image_categories = categories_by_image[image_id]
except KeyError:
task.warning(
f"Could not find image assoicated with annotation {annotation.get('id')}"
)
continue
annotation_model = AnnotationModel.objects(
image_id=image_model.id,
category_id=category_model_id,
segmentation=segmentation,
keypoints=keypoints).first()
if annotation_model is None:
task.info(f"Creating annotation data ({image_id}, {category_id})")
annotation_model = AnnotationModel(image_id=image_model.id)
annotation_model.category_id = category_model_id
annotation_model.color = annotation.get('color')
annotation_model.metadata = annotation.get('metadata', {})
if has_segmentation:
annotation_model.segmentation = segmentation
annotation_model.area = area
annotation_model.bbox = bbox
if has_keypoints:
annotation_model.keypoints = keypoints
annotation_model.isbbox = isbbox
annotation_model.save()
image_categories.append(category_id)
else:
annotation_model.update(deleted=False, isbbox=isbbox)
task.info(
f"Annotation already exists (i:{image_id}, c:{category_id})")
for image_id in images_id:
image_model = images_id[image_id]
category_ids = categories_by_image[image_id]
all_category_ids = list(image_model.category_ids)
all_category_ids += category_ids
num_annotations = AnnotationModel.objects(
Q(image_id=image_id) & Q(deleted=False)
& (Q(area__gt=0) | Q(keypoints__size__gt=0))).count()
image_model.update(set__annotated=True,
set__category_ids=list(set(all_category_ids)),
set__num_annotations=num_annotations)
task.set_progress(100, socket=socket)
def get_dataset_coco(dataset):
"""
Generates coco for all images in dataset
:param dataset: DatasetModel
:return: Coco in dictionary format
"""
categories = CategoryModel.objects(deleted=False) \
.exclude('deleted_date').in_bulk(dataset.categories).items()
dataset = fix_ids(dataset)
images = ImageModel.objects(deleted=False, dataset_id=dataset.get('id')).exclude('deleted_date')
all_annotations = AnnotationModel.objects(deleted=False).exclude('deleted_date', 'paper_object')
coco = {
'images': [],
'categories': [],
'annotations': []
}
for category in categories:
category = fix_ids(category[1])
del category['deleted']
if len(category.get('keypoint_labels', [])) > 0:
category['keypoints'] = category.pop('keypoint_labels')
category['skeleton'] = category.pop('keypoint_edges')
else:
del category['keypoint_edges']
del category['keypoint_labels']
coco.get('categories').append(category)
for image in images:
annotations = all_annotations.filter(image_id=image.id)
if annotations.count() == 0:
continue
annotations = fix_ids(annotations.all())
for annotation in annotations:
has_keypoints = len(annotation.get('keypoints', [])) > 0
has_segmentation = len(annotation.get('segmentation', [])) > 0
if has_keypoints or has_segmentation:
del annotation['deleted']
if not has_keypoints:
del annotation['keypoints']
else:
arr = np.array(annotation.get('keypoints', []))
arr = arr[2::3]
annotation['num_keypoints'] = len(arr[arr > 0])
coco.get('annotations').append(annotation)
image = fix_ids(image)
del image['deleted']
coco.get('images').append(image)
return coco
def get_image_coco(image_id):
"""
Generates coco for an image
:param image: ImageModel
:return: Coco in dictionary format
"""
image = ImageModel.objects(id=image_id)\
.only(*ImageModel.COCO_PROPERTIES)
image = fix_ids(image)[0]
dataset = DatasetModel.objects(id=image.get('dataset_id')).first()
bulk_categories = CategoryModel.objects(id__in=dataset.categories, deleted=False) \
.only(*CategoryModel.COCO_PROPERTIES)
print(bulk_categories)
db_annotations = AnnotationModel.objects(deleted=False, image_id=image_id)
categories = []
annotations = []
for category in fix_ids(bulk_categories):
category_annotations = db_annotations\
.filter(category_id=category.get('id'))\
.only(*AnnotationModel.COCO_PROPERTIES)
if category_annotations.count() == 0:
continue
category_annotations = fix_ids(category_annotations)
for annotation in category_annotations:
has_segmentation = len(annotation.get('segmentation', [])) > 0
has_keypoints = len(annotation.get('keypoints', [])) > 0
if has_segmentation or has_keypoints:
if has_keypoints:
arr = np.array(annotation.get('keypoints', []))
arr = arr[2::3]
annotation['num_keypoints'] = len(arr[arr > 0])
annotations.append(annotation)
if len(category.get('keypoint_labels')) > 0:
category['keypoints'] = category.pop('keypoint_labels')
category['skeleton'] = category.pop('keypoint_edges')
else:
del category['keypoint_edges']
del category['keypoint_labels']
categories.append(category)
coco = {
"images": [image],
"categories": categories,
"annotations": annotations
}
return coco
def get(self, image_id):
""" Called when loading from the annotator client """
data = folder_data.parse_args()
folder = data.get('folder')
order = data.get('order')
image = ImageModel.objects(id=image_id)\
.exclude('events').first()
if image is None:
return {'success': False, 'message': 'Could not load image'}, 400
dataset = current_user.datasets.filter(id=image.dataset_id).first()
if dataset is None:
return {
'success': False,
'message': 'Could not find associated dataset'
}, 400
categories = CategoryModel.objects(deleted=False)\
.in_bulk(dataset.categories).items()
if len(folder) > 0:
folder = folder[0].strip('/') + folder[1:]
if folder[-1] != '/':
folder = folder + '/'
else:
folder = ''
# Get directory
directory = os.path.join(dataset.directory, folder)
if not os.path.exists(directory):
return {'message': 'Directory does not exist.'}, 400
# Get next and previous image
# images = ImageModel.objects(dataset_id=dataset.id, deleted=False)
# pre = images.filter(file_name__lt=image.file_name).order_by('-file_name').first()
# nex = images.filter(file_name__gt=image.file_name).order_by('file_name').first()
images = ImageModel.objects(dataset_id=dataset.id,
deleted=False,
path__startswith=directory)
pre = images.filter(
file_name__lt=image.file_name).order_by('-' + order).first()
nex = images.filter(
file_name__gt=image.file_name).order_by(order).first()
preferences = {}
if not Config.LOGIN_DISABLED:
preferences = current_user.preferences
# Generate data about the image to return to client
data = {
'image': query_util.fix_ids(image),
'categories': [],
'dataset': query_util.fix_ids(dataset),
'preferences': preferences,
'permissions': {
'dataset': dataset.permissions(current_user),
'image': image.permissions(current_user)
}
}
data['image']['previous'] = pre.id if pre else None
data['image']['next'] = nex.id if nex else None
for category in categories:
category = query_util.fix_ids(category[1])
category_id = category.get('id')
annotations = AnnotationModel.objects(image_id=image_id, category_id=category_id, deleted=False)\
.exclude('events').all()
category['show'] = True
category['visualize'] = False
category[
'annotations'] = [] if annotations is None else query_util.fix_ids(
annotations)
data.get('categories').append(category)
return data
def get(self, dataset_id):
""" Endpoint called by image viewer client """
parsed_args = page_data.parse_args()
per_page = parsed_args.get('limit')
page = parsed_args.get('page') - 1
folder = parsed_args.get('folder')
order = parsed_args.get('order')
args = dict(request.args)
# Check if dataset exists
dataset = current_user.datasets.filter(id=dataset_id,
deleted=False).first()
if dataset is None:
return {'message', 'Invalid dataset id'}, 400
# Make sure folder starts with is in proper format
if len(folder) > 0:
folder = folder[0].strip('/') + folder[1:]
if folder[-1] != '/':
folder = folder + '/'
# Get directory
directory = os.path.join(dataset.directory, folder)
if not os.path.exists(directory):
return {'message': 'Directory does not exist.'}, 400
# Remove parsed arguments
for key in parsed_args:
args.pop(key, None)
# Generate query from remaining arugments
query = {}
for key, value in args.items():
lower = value.lower()
if lower in ["true", "false"]:
value = json.loads(lower)
if len(lower) != 0:
query[key] = value
images = current_user.images \
.filter(dataset_id=dataset_id, path__startswith=directory, deleted=False, **query) \
.order_by(order).only('id', 'file_name', 'annotating', 'annotated', 'num_annotations')
total = images.count()
pages = int(total / per_page) + 1
images = images.skip(page * per_page).limit(per_page)
images_json = query_util.fix_ids(images)
# for image in images:
# image_json = query_util.fix_ids(image)
# query = AnnotationModel.objects(image_id=image.id, deleted=False)
# category_ids = query.distinct('category_id')
# categories = CategoryModel.objects(id__in=category_ids).only('name', 'color')
# image_json['annotations'] = query.count()
# image_json['categories'] = query_util.fix_ids(categories)
# images_json.append(image_json)
subdirectories = [
f for f in sorted(os.listdir(directory))
if os.path.isdir(directory + f) and not f.startswith('.')
]
categories = CategoryModel.objects(id__in=dataset.categories).only(
'id', 'name')
return {
"total": total,
"per_page": per_page,
"pages": pages,
"page": page,
"images": images_json,
"folder": folder,
"directory": directory,
"dataset": query_util.fix_ids(dataset),
"categories": query_util.fix_ids(categories),
"subdirectories": subdirectories
}
def update_all_category(self, data, current_user):
for category_data in data.get('categories', []):
db_category = CategoryModel.find_by(category_data.get('id'))
if db_category is None:
continue
self.update_category(db_category, category_data, current_user)
def collect_coco_annotations(task, categories, dataset, socket):
"""
Getting all coco labels from current dataset and creating a dict from it
:return: COCO labels from current dataset as dict with fields: "images", "categories",
"annotations".
"""
# Getting coco annotations
task.info("===== Getting COCO annotations =====")
db_categories = CategoryModel.objects(id__in=categories, deleted=False) \
.only(*CategoryModel.COCO_PROPERTIES)
db_images = ImageModel.objects(deleted=False, dataset_id=dataset.id) \
.only(*ImageModel.COCO_PROPERTIES)
db_annotations = AnnotationModel.objects(deleted=False, category_id__in=categories)
total_items = db_categories.count()
coco = {
'images': [],
'categories': [],
'annotations': []
}
total_items += db_images.count()
progress = 0
# iterate though all categoires and upsert
category_names = []
for category in fix_ids(db_categories):
if len(category.get('keypoint_labels', [])) > 0:
category['keypoints'] = category.pop('keypoint_labels', [])
category['skeleton'] = category.pop('keypoint_edges', [])
else:
if 'keypoint_edges' in category:
del category['keypoint_edges']
if 'keypoint_labels' in category:
del category['keypoint_labels']
task.info(f"Adding category: {category.get('name')}")
coco.get('categories').append(category)
category_names.append(category.get('name'))
progress += 1
task.set_progress((progress / total_items) * 50, socket=socket)
total_annotations = db_annotations.count()
total_images = db_images.count()
for img_counter, image in enumerate(fix_ids(db_images)):
progress += 1
task.set_progress((progress / total_items) * 50, socket=socket)
annotations = db_annotations.filter(image_id=image.get('id')) \
.only(*AnnotationModel.COCO_PROPERTIES)
annotations = fix_ids(annotations)
num_annotations = 0
for annotation in annotations:
has_keypoints = len(annotation.get('keypoints', [])) > 0
has_segmentation = len(annotation.get('segmentation', [])) > 0
if has_keypoints or has_segmentation:
if not has_keypoints:
if 'keypoints' in annotation:
del annotation['keypoints']
else:
arr = np.array(annotation.get('keypoints', []))
arr = arr[2::3]
annotation['num_keypoints'] = len(arr[arr > 0])
num_annotations += 1
coco.get('annotations').append(annotation)
task.info(f"Exporting {num_annotations} annotations for image {image.get('id')} ({img_counter+1}/{total_images})")
coco.get('images').append(image)
task.info(f"Done export {total_annotations} annotations and {total_images} images from {dataset.name}")
return coco, category_names
def get(self, dataset_id):
""" Endpoint called by image viewer client """
args = page_data.parse_args()
limit = args['limit']
page = args['page']
folder = args['folder']
args = dict(request.args)
if args.get('limit') != None:
del args['limit']
if args.get('page') != None:
del args['page']
if args.get('folder') != None:
del args['folder']
query = {}
for key, value in args.items():
lower = value.lower()
if lower in ["true", "false"]:
value = json.loads(lower)
if len(lower) != 0:
query[key] = value
# print(query, flush=True)
# Check if dataset exists
dataset = current_user.datasets.filter(id=dataset_id,
deleted=False).first()
if dataset is None:
return {'message', 'Invalid dataset id'}, 400
# Make sure folder starts with is in proper format
if len(folder) > 0:
folder = folder[0].strip('/') + folder[1:]
if folder[-1] != '/':
folder = folder + '/'
# Get directory
directory = os.path.join(dataset.directory, folder)
if not os.path.exists(directory):
return {'message': 'Directory does not exist.'}, 400
images = ImageModel.objects(dataset_id=dataset_id, path__startswith=directory, deleted=False, **query) \
.order_by('file_name').only('id', 'file_name', 'annotating')
pagination = Pagination(images.count(), limit, page)
images = images[pagination.start:pagination.end]
images_json = []
for image in images:
image_json = query_util.fix_ids(image)
query = AnnotationModel.objects(image_id=image.id, deleted=False)
category_ids = query.distinct('category_id')
image_json['annotations'] = query.count()
image_json['categories'] = query_util.fix_ids(
CategoryModel.objects(id__in=category_ids).only(
'name', 'color'))
image_json['permissions'] = image.permissions(current_user)
images_json.append(image_json)
subdirectories = [
f for f in sorted(os.listdir(directory))
if os.path.isdir(directory + f) and not f.startswith('.')
]
categories = CategoryModel.objects(id__in=dataset.categories).only(
'id', 'name')
return {
"pagination": pagination.export(),
"images": images_json,
"folder": folder,
"directory": directory,
"dataset": query_util.fix_ids(dataset),
"categories": query_util.fix_ids(categories),
"subdirectories": subdirectories
}