mirror of https://github.com/vladmandic/automatic
138 lines
6.1 KiB
Python
138 lines
6.1 KiB
Python
from collections import defaultdict
|
|
from csv import DictReader, reader as TupleReader
|
|
from pathlib import Path
|
|
from typing import Dict, List, Any
|
|
import warnings
|
|
|
|
from taming.data.annotated_objects_dataset import AnnotatedObjectsDataset
|
|
from taming.data.helper_types import Annotation, Category
|
|
from tqdm import tqdm
|
|
|
|
OPEN_IMAGES_STRUCTURE = {
|
|
'train': {
|
|
'top_level': '',
|
|
'class_descriptions': 'class-descriptions-boxable.csv',
|
|
'annotations': 'oidv6-train-annotations-bbox.csv',
|
|
'file_list': 'train-images-boxable.csv',
|
|
'files': 'train'
|
|
},
|
|
'validation': {
|
|
'top_level': '',
|
|
'class_descriptions': 'class-descriptions-boxable.csv',
|
|
'annotations': 'validation-annotations-bbox.csv',
|
|
'file_list': 'validation-images.csv',
|
|
'files': 'validation'
|
|
},
|
|
'test': {
|
|
'top_level': '',
|
|
'class_descriptions': 'class-descriptions-boxable.csv',
|
|
'annotations': 'test-annotations-bbox.csv',
|
|
'file_list': 'test-images.csv',
|
|
'files': 'test'
|
|
}
|
|
}
|
|
|
|
|
|
def load_annotations(descriptor_path: Path, min_object_area: float, category_mapping: Dict[str, str],
|
|
category_no_for_id: Dict[str, int]) -> Dict[str, List[Annotation]]:
|
|
annotations: Dict[str, List[Annotation]] = defaultdict(list)
|
|
with open(descriptor_path) as file:
|
|
reader = DictReader(file)
|
|
for i, row in tqdm(enumerate(reader), total=14620000, desc='Loading OpenImages annotations'):
|
|
width = float(row['XMax']) - float(row['XMin'])
|
|
height = float(row['YMax']) - float(row['YMin'])
|
|
area = width * height
|
|
category_id = row['LabelName']
|
|
if category_id in category_mapping:
|
|
category_id = category_mapping[category_id]
|
|
if area >= min_object_area and category_id in category_no_for_id:
|
|
annotations[row['ImageID']].append(
|
|
Annotation(
|
|
id=i,
|
|
image_id=row['ImageID'],
|
|
source=row['Source'],
|
|
category_id=category_id,
|
|
category_no=category_no_for_id[category_id],
|
|
confidence=float(row['Confidence']),
|
|
bbox=(float(row['XMin']), float(row['YMin']), width, height),
|
|
area=area,
|
|
is_occluded=bool(int(row['IsOccluded'])),
|
|
is_truncated=bool(int(row['IsTruncated'])),
|
|
is_group_of=bool(int(row['IsGroupOf'])),
|
|
is_depiction=bool(int(row['IsDepiction'])),
|
|
is_inside=bool(int(row['IsInside']))
|
|
)
|
|
)
|
|
if 'train' in str(descriptor_path) and i < 14000000:
|
|
warnings.warn(f'Running with subset of Open Images. Train dataset has length [{len(annotations)}].')
|
|
return dict(annotations)
|
|
|
|
|
|
def load_image_ids(csv_path: Path) -> List[str]:
|
|
with open(csv_path) as file:
|
|
reader = DictReader(file)
|
|
return [row['image_name'] for row in reader]
|
|
|
|
|
|
def load_categories(csv_path: Path) -> Dict[str, Category]:
|
|
with open(csv_path) as file:
|
|
reader = TupleReader(file)
|
|
return {row[0]: Category(id=row[0], name=row[1], super_category=None) for row in reader}
|
|
|
|
|
|
class AnnotatedObjectsOpenImages(AnnotatedObjectsDataset):
|
|
def __init__(self, use_additional_parameters: bool, **kwargs):
|
|
"""
|
|
@param data_path: is the path to the following folder structure:
|
|
open_images/
|
|
│ oidv6-train-annotations-bbox.csv
|
|
├── class-descriptions-boxable.csv
|
|
├── oidv6-train-annotations-bbox.csv
|
|
├── test
|
|
│ ├── 000026e7ee790996.jpg
|
|
│ ├── 000062a39995e348.jpg
|
|
│ └── ...
|
|
├── test-annotations-bbox.csv
|
|
├── test-images.csv
|
|
├── train
|
|
│ ├── 000002b66c9c498e.jpg
|
|
│ ├── 000002b97e5471a0.jpg
|
|
│ └── ...
|
|
├── train-images-boxable.csv
|
|
├── validation
|
|
│ ├── 0001eeaf4aed83f9.jpg
|
|
│ ├── 0004886b7d043cfd.jpg
|
|
│ └── ...
|
|
├── validation-annotations-bbox.csv
|
|
└── validation-images.csv
|
|
@param: split: one of 'train', 'validation' or 'test'
|
|
@param: desired image size (returns square images)
|
|
"""
|
|
|
|
super().__init__(**kwargs)
|
|
self.use_additional_parameters = use_additional_parameters
|
|
|
|
self.categories = load_categories(self.paths['class_descriptions'])
|
|
self.filter_categories()
|
|
self.setup_category_id_and_number()
|
|
|
|
self.image_descriptions = {}
|
|
annotations = load_annotations(self.paths['annotations'], self.min_object_area, self.category_mapping,
|
|
self.category_number)
|
|
self.annotations = self.filter_object_number(annotations, self.min_object_area, self.min_objects_per_image,
|
|
self.max_objects_per_image)
|
|
self.image_ids = list(self.annotations.keys())
|
|
self.clean_up_annotations_and_image_descriptions()
|
|
|
|
def get_path_structure(self) -> Dict[str, str]:
|
|
if self.split not in OPEN_IMAGES_STRUCTURE:
|
|
raise ValueError(f'Split [{self.split} does not exist for Open Images data.]')
|
|
return OPEN_IMAGES_STRUCTURE[self.split]
|
|
|
|
def get_image_path(self, image_id: str) -> Path:
|
|
return self.paths['files'].joinpath(f'{image_id:0>16}.jpg')
|
|
|
|
def get_image_description(self, image_id: str) -> Dict[str, Any]:
|
|
image_path = self.get_image_path(image_id)
|
|
return {'file_path': str(image_path), 'file_name': image_path.name}
|