# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/25f_converter.yolo.ipynb.

# %% auto 0
__all__ = ['YOLO']

# %% ../../nbs/25f_converter.yolo.ipynb 3
from fastcore.all import *
import shutil
import polvo as pv

# %% ../../nbs/25f_converter.yolo.ipynb 4
class YOLO(pv.Visitor):
    def convert_dataset(self, save_dir, classes, train=None, valid=None, test=None, n_workers=defaults.cpus):
        save_dir = Path(save_dir)
        dirs = {}
        if train: 
            train_dir = dirs['train_dir'] = pv.mkdir(save_dir/'train')
            self.convert_records(train_dir, train, n_workers=n_workers)
        if valid:
            valid_dir = dirs['valid_dir'] = pv.mkdir(save_dir/'valid')
            self.convert_records(valid_dir, valid, n_workers=n_workers)
        if test: 
            test_dir = dirs['test_dir'] = pv.mkdir(save_dir/'test')
            self.convert_records(test_dir, test, n_workers=n_workers)
            
        yolo_yaml = self.yolo_yaml(classes, save_dir, **dirs)
        pv.save_txt(yolo_yaml, save_dir/'dataset.yaml')
        
    def convert_records(self, save_dir, records, n_workers=defaults.cpus):
        image_dir = pv.mkdir(Path(save_dir)/'images')
        ann_dir = pv.mkdir(Path(save_dir)/'labels')
        parallel(pv.partial(self.convert_record, image_dir=image_dir, ann_dir=ann_dir),
                 records, progress=pv.pbar, n_workers=n_workers)
        
    def convert_record(self, record, image_dir, ann_dir):
        self._labels, self._bboxes = [], []
        self.visit_all(record)
        lines = [' '.join(o) for o in pv.safe_zip(self._labels, self._bboxes)]
        pv.save_txt('\n'.join(lines), ann_dir/self._image_file.with_suffix('.txt').name)
        shutil.copy(self._image_file.absolute(), image_dir/self._image_file.name)
        
    def yolo_yaml(self, class_map, save_dir, train_dir=None, valid_dir=None, test_dir=None):
        classes = '\n'.join([f'  {k}: {v}' for k,v in class_map.items()])
        train_str = f"train: {train_dir.relative_to(save_dir)}\n" if train_dir is not None else ""
        valid_str = f"val: {valid_dir.relative_to(save_dir)}\n" if valid_dir is not None else ""
        test_str = f"test: {test_dir.relative_to(save_dir)}\n" if test_dir is not None else ""
        return (
            f"path: {save_dir.absolute()}\n"
            f"{train_str}"
            f"{valid_str}"
            f"{test_str}"
            f"names:\n"
            f"{classes}"
        )
    
    def _visit_image_file(self, image_file): 
        self._w,self._h = pv.image_size(image_file)
        self._image_file = image_file
    
    def _visit_bbox_labelled(self, bbox):
        self._labels.append(str(bbox.label.id))
        points = bbox.normalized_points(self._w,self._h).flatten()
        self._bboxes.append(' '.join(map(str, points)))
