diff --git a/requirements.txt b/requirements.txt index f6d336cbb..5651e7e0b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ einops faster-coco-eval graphviz hydra-core +lazy-loader lightning loguru numpy diff --git a/train_kwcoco_demo.sh b/train_kwcoco_demo.sh new file mode 100644 index 000000000..e631203a9 --- /dev/null +++ b/train_kwcoco_demo.sh @@ -0,0 +1,143 @@ +#!/bin/bash +__doc__=" +YOLO Training Tutorial with KWCOCO DemoData +=========================================== + +This demonstrates an end-to-end YOLO pipeline on toydata generated with kwcoco. +" + +# Define where we will store results +BUNDLE_DPATH=$HOME/demo-yolo-kwcoco-train +mkdir -p "$BUNDLE_DPATH" + +echo " +Generate Toy Data +----------------- + +Now that we know where the data and our intermediate files will go, lets +generate the data we will use to train and evaluate with. + +The kwcoco package comes with a commandline utility called 'kwcoco toydata' to +accomplish this. +" + +# Define the names of the kwcoco files to generate +TRAIN_FPATH=$BUNDLE_DPATH/vidshapes_rgb_train/data.kwcoco.json +VALI_FPATH=$BUNDLE_DPATH/vidshapes_rgb_vali/data.kwcoco.json +TEST_FPATH=$BUNDLE_DPATH/vidshapes_rgb_test/data.kwcoco.json + +# Generate toy datasets using the "kwcoco toydata" tool +kwcoco toydata vidshapes32-frames10 --dst "$TRAIN_FPATH" +kwcoco toydata vidshapes4-frames10 --dst "$VALI_FPATH" +kwcoco toydata vidshapes2-frames6 --dst "$TEST_FPATH" + +# Ensure legacy COCO structure for now +kwcoco conform "$TRAIN_FPATH" --inplace --legacy=True +kwcoco conform "$VALI_FPATH" --inplace --legacy=True +kwcoco conform "$TEST_FPATH" --inplace --legacy=True + + +echo " +Create the YOLO Configuration +----------------------------- + +Constructing the YOLO configuration is not entirely kwcoco aware +so we need to set +" +# In the current version we need to write configs to the repo itself. +# Its a bit gross, but this should be somewhat robust. +# Find where the yolo repo is installed (we need to be careful that this is the +# our fork of the WongKinYiu variant +REPO_DPATH=$(python -c "import yolo, pathlib; print(pathlib.Path(yolo.__file__).parent.parent)") +MODULE_DPATH=$(python -c "import yolo, pathlib; print(pathlib.Path(yolo.__file__).parent)") +CONFIG_DPATH=$(python -c "import yolo.config, pathlib; print(pathlib.Path(yolo.config.__file__).parent / 'dataset')") +echo "REPO_DPATH = $REPO_DPATH" +echo "MODULE_DPATH = $MODULE_DPATH" +echo "CONFIG_DPATH = $CONFIG_DPATH" + +DATASET_CONFIG_FPATH=$CONFIG_DPATH/kwcoco-demo.yaml + +# Hack to construct the class part of the YAML +CLASS_YAML=$(python -c "if 1: + import kwcoco + train_fpath = kwcoco.CocoDataset('$TRAIN_FPATH') + categories = train_fpath.categories().objs + # It would be nice to have better class introspection, but in the meantime + # do the same sorting as yolo.tools.data_conversion.discretize_categories + categories = sorted(categories, key=lambda cat: cat['id']) + class_num = len(categories) + class_list = [c['name'] for c in categories] + print(f'class_num: {class_num}') + print(f'class_list: {class_list}') +") + + +CONFIG_YAML=" +path: $BUNDLE_DPATH +train: $TRAIN_FPATH +validation: $VALI_FPATH + +$CLASS_YAML +" + +echo "$CONFIG_YAML" > "$DATASET_CONFIG_FPATH" + + +# This might only work in development mode, otherwise we will get site packages +# That still might be fine, but we do want to fix this to run anywhere. +cd "$REPO_DPATH" +LOG_BATCH_VIZ_TO_DISK=1 python -m yolo.lazy \ + task=train \ + dataset=kwcoco-demo \ + use_wandb=False \ + out_path="$BUNDLE_DPATH"/training \ + name=kwcoco-demo \ + cpu_num=0 \ + device=0 \ + accelerator=auto \ + task.data.batch_size=2 \ + "image_size=[224,224]" \ + task.optimizer.args.lr=0.0003 + +LOG_BATCH_VIZ_TO_DISK=1 python -m yolo.lazy \ + task=train \ + dataset=kwcoco-demo \ + use_wandb=False \ + out_path="$BUNDLE_DPATH"/training \ + name=kwcoco-demo \ + cpu_num=0 \ + device=0 \ + accelerator=auto \ + task.data.batch_size=2 \ + "image_size=[224,224]" \ + task.optimizer.args.lr=0.0003 + + +### TODO: show how to validate + +# Grab a checkpoint +CKPT_FPATH=$(python -c "if 1: + import pathlib + ckpt_dpath = pathlib.Path('$BUNDLE_DPATH') / 'training/train/kwcoco-demo/checkpoints' + checkpoints = sorted(ckpt_dpath.glob('*')) + print(checkpoints[-1]) +") +echo "CKPT_FPATH = $CKPT_FPATH" + + +#DISABLE_RICH_HANDLER=1 +LOG_BATCH_VIZ_TO_DISK=1 python -m yolo.lazy \ + task=validation \ + dataset=kwcoco-demo \ + use_wandb=False \ + out_path="$BUNDLE_DPATH"/training \ + name=kwcoco-demo \ + cpu_num=0 \ + device=0 \ + weight="'$CKPT_FPATH'" \ + accelerator=auto \ + "task.data.batch_size=2" \ + "image_size=[224,224]" + + +### TODO: show how to run inference diff --git a/yolo/__init__.py b/yolo/__init__.py index b4b98d7f8..77da671a6 100644 --- a/yolo/__init__.py +++ b/yolo/__init__.py @@ -1,33 +1,78 @@ -from yolo.config.config import Config, NMSConfig -from yolo.model.yolo import create_model -from yolo.tools.data_loader import AugmentationComposer, create_dataloader -from yolo.tools.drawer import draw_bboxes -from yolo.tools.solver import TrainModel -from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box, bbox_nms, create_converter -from yolo.utils.deploy_utils import FastModelLoader -from yolo.utils.logging_utils import ( - ImageLogger, - YOLORichModelSummary, - YOLORichProgressBar, +""" +The MIT YOLO rewrite +""" + +__autogen__ = """ +mkinit ~/code/YOLO-v9/yolo/__init__.py --nomods --write --lazy-loader + +# Check to see how long it takes to run a simple help command +time python -m yolo.lazy --help +""" + +__submodules__ = { + 'config.config': ['Config', 'NMSConfig'], + 'model.yolo': ['create_model'], + 'tools.data_loader': ['AugmentationComposer', 'create_dataloader'], + 'tools.drawer': ['draw_bboxes'], + 'tools.solver': ['TrainModel'], + 'utils.bounding_box_utils': ['Anc2Box', 'Vec2Box', 'bbox_nms', 'create_converter'], + 'utils.deploy_utils': ['FastModelLoader'], + 'utils.logging_utils': [ + 'ImageLogger', 'YOLORichModelSummary', + 'YOLORichProgressBar', + 'validate_log_directory' + ], + 'utils.model_utils': ['PostProcess'], +} + + +import lazy_loader + + +__getattr__, __dir__, __all__ = lazy_loader.attach( + __name__, + submodules={}, + submod_attrs={ + 'config.config': [ + 'Config', + 'NMSConfig', + ], + 'model.yolo': [ + 'create_model', + ], + 'tools.data_loader': [ + 'AugmentationComposer', + 'create_dataloader', + ], + 'tools.drawer': [ + 'draw_bboxes', + ], + 'tools.solver': [ + 'TrainModel', + ], + 'utils.bounding_box_utils': [ + 'Anc2Box', + 'Vec2Box', + 'bbox_nms', + 'create_converter', + ], + 'utils.deploy_utils': [ + 'FastModelLoader', + ], + 'utils.logging_utils': [ + 'ImageLogger', + 'YOLORichModelSummary', + 'YOLORichProgressBar', + 'validate_log_directory', + ], + 'utils.model_utils': [ + 'PostProcess', + ], + }, ) -from yolo.utils.model_utils import PostProcess - -all = [ - "create_model", - "Config", - "YOLORichProgressBar", - "NMSConfig", - "YOLORichModelSummary", - "validate_log_directory", - "draw_bboxes", - "Vec2Box", - "Anc2Box", - "bbox_nms", - "create_converter", - "AugmentationComposer", - "ImageLogger", - "create_dataloader", - "FastModelLoader", - "TrainModel", - "PostProcess", -] + +__all__ = ['Anc2Box', 'AugmentationComposer', 'Config', 'FastModelLoader', + 'ImageLogger', 'NMSConfig', 'PostProcess', 'TrainModel', 'Vec2Box', + 'YOLORichModelSummary', 'YOLORichProgressBar', 'bbox_nms', + 'create_converter', 'create_dataloader', 'create_model', + 'draw_bboxes', 'validate_log_directory'] diff --git a/yolo/config/config.py b/yolo/config/config.py index 9dd85b07a..a0a3f2512 100644 --- a/yolo/config/config.py +++ b/yolo/config/config.py @@ -157,6 +157,7 @@ class Config: use_tensorboard: bool weight: Optional[str] + accelerator: str @dataclass diff --git a/yolo/config/general.yaml b/yolo/config/general.yaml index c3380a799..0a4af6ca0 100644 --- a/yolo/config/general.yaml +++ b/yolo/config/general.yaml @@ -11,3 +11,4 @@ use_wandb: True use_tensorboard: False weight: True # Path to weight or True for auto, False for no pretrained weight +accelerator: 'auto' diff --git a/yolo/lazy.py b/yolo/lazy.py index 0f1cc55b6..c0b0d6a1b 100644 --- a/yolo/lazy.py +++ b/yolo/lazy.py @@ -2,22 +2,22 @@ from pathlib import Path import hydra -from lightning import Trainer +from omegaconf.dictconfig import DictConfig +# FIXME: messing with sys.path is a bad idea. Factor this out. project_root = Path(__file__).resolve().parent.parent sys.path.append(str(project_root)) -from yolo.config.config import Config -from yolo.tools.solver import InferenceModel, TrainModel, ValidateModel -from yolo.utils.logging_utils import setup - @hydra.main(config_path="config", config_name="config", version_base=None) -def main(cfg: Config): +def main(cfg: DictConfig): + from yolo.utils.logging_utils import setup callbacks, loggers, save_path = setup(cfg) + from lightning import Trainer + from yolo.tools.solver import InferenceModel, TrainModel, ValidateModel trainer = Trainer( - accelerator="auto", + accelerator=cfg.accelerator, max_epochs=getattr(cfg.task, "epoch", None), precision="16-mixed", callbacks=callbacks, @@ -25,9 +25,10 @@ def main(cfg: Config): log_every_n_steps=1, gradient_clip_val=10, gradient_clip_algorithm="value", - deterministic=True, + # deterministic=True, enable_progress_bar=not getattr(cfg, "quite", False), default_root_dir=save_path, + num_sanity_val_steps=0, ) if cfg.task.task == "train": diff --git a/yolo/model/__init__.py b/yolo/model/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/yolo/model/yolo.py b/yolo/model/yolo.py index cc9ce20ba..bb43016f7 100644 --- a/yolo/model/yolo.py +++ b/yolo/model/yolo.py @@ -32,7 +32,7 @@ def __init__(self, model_cfg: ModelConfig, class_num: int = 80): def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]): self.layer_index = {} output_dim, layer_idx = [3], 1 - logger.info(f":tractor: Building YOLO") + logger.info(":tractor: Building YOLO") for arch_name in model_arch: if model_arch[arch_name]: logger.info(f" :building_construction: Building {arch_name}") @@ -129,27 +129,99 @@ def save_load_weights(self, weights: Union[Path, OrderedDict]): weights = torch.load(weights, map_location=torch.device("cpu"), weights_only=False) if "model_state_dict" in weights: weights = weights["model_state_dict"] + if "state_dict" in weights: + weights = weights["state_dict"] + + if 0: + # Debug the state of the model and the loaded weights + import networkx as nx + graph_src = nx.DiGraph() + for key in list(weights.keys()): + graph_src.add_node(key) + graph_src.add_node('__root__') + for key in list(weights.keys()): + parts = key.split('.') + graph_src.add_edge('__root__', parts[0]) + for i in range(1, len(parts)): + parent = '.'.join(parts[:i - 1]) + child = '.'.join(parts[:i]) + graph_src.add_edge(parent, child) + nx.write_network_text(graph_src, max_depth=4, sources=['__root__']) + + graph_dst = nx.DiGraph() + dst_weights = self.state_dict() + for key in list(dst_weights.keys()): + graph_dst.add_node(key) + graph_dst.add_node('__root__') + for key in list(dst_weights.keys()): + parts = key.split('.') + graph_dst.add_edge('__root__', parts[0]) + for i in range(1, len(parts)): + parent = '.'.join(parts[:i - 1]) + child = '.'.join(parts[:i]) + graph_dst.add_edge(parent, child) + nx.write_network_text(graph_dst, max_depth=3, sources=['__root__']) + + USE_TORCH_LIBERATOR = False + if USE_TORCH_LIBERATOR: + + # Torch liberator will figure out the mapping in most cases but it + # is slow. + HACK_DONT_LOAD_EMA_WEIGHTS = True + if HACK_DONT_LOAD_EMA_WEIGHTS: + for key in list(weights.keys()): + if key.startswith('ema.model'): + weights.pop(key) + from torch_liberator.initializer import load_partial_state + load_partial_state(self, weights, verbose=3) - model_state_dict = self.model.state_dict() - - # TODO1: autoload old version weight - # TODO2: weight transform if num_class difference - - error_dict = {"Mismatch": set(), "Not Found": set()} - for model_key, model_weight in model_state_dict.items(): - if model_key not in weights: - error_dict["Not Found"].add(tuple(model_key.split(".")[:-2])) - continue - if model_weight.shape != weights[model_key].shape: - error_dict["Mismatch"].add(tuple(model_key.split(".")[:-2])) - continue - model_state_dict[model_key] = weights[model_key] - - for error_name, error_set in error_dict.items(): - for weight_name in error_set: - logger.warning(f":warning: Weight {error_name} for key: {'.'.join(weight_name)}") - - self.model.load_state_dict(model_state_dict) + else: + # TODO1: autoload old version weight + # TODO2: weight transform if num_class difference + + model_state_dict = self.model.state_dict() + + CHECK_FOR_WEIGHT_MUNGING = True + if CHECK_FOR_WEIGHT_MUNGING: + # Handle the simple case of weight munging ourselves + src_keys = list(weights.keys()) + dst_keys = list(model_state_dict.keys()) + src_roots = {p.split('.')[0] for p in src_keys} + dst_roots = {p.split('.')[0] for p in dst_keys} + if len(src_roots & dst_roots) == 0: + src_prefixes = {tuple(p.split('.')[0:2]) for p in src_keys} + if src_prefixes == {('ema', 'model'), ('model', 'model')}: + logger.warning(":warning: Munging weights") + munged_weights = {} + for key in list(weights.keys()): + prefix = 'model.model.' + if key.startswith(prefix): + new_key = key[len(prefix):] + munged_weights[new_key] = weights[key] + logger.warning(f":warning: Munged {len(munged_weights)} / {len(weights)} tensors") + weights = munged_weights + + error_dict = {"Mismatch": set(), "Not Found": set()} + for model_key, model_weight in model_state_dict.items(): + if model_key not in weights: + error_dict["Not Found"].add(tuple(model_key.split(".")[:-2])) + continue + if model_weight.shape != weights[model_key].shape: + error_dict["Mismatch"].add(tuple(model_key.split(".")[:-2])) + continue + model_state_dict[model_key] = weights[model_key] + + for error_name, error_set in error_dict.items(): + for weight_name in error_set: + logger.warning(f":warning: Weight {error_name} for key: {'.'.join(weight_name)}") + + for error_name, error_set in error_dict.items(): + if len(error_set) == 0: + logger.info(f":white_check_mark: Num: weight {error_name}: {len(error_set)}") + else: + logger.warning(f":warning: Num: weight {error_name}: {len(error_set)}") + + self.model.load_state_dict(model_state_dict) def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True, class_num: int = 80) -> YOLO: @@ -164,6 +236,7 @@ def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True, OmegaConf.set_struct(model_cfg, False) model = YOLO(model_cfg, class_num) if weight_path: + logger.info('🏋 Initializing weights') if weight_path == True: weight_path = Path("weights") / f"{model_cfg.name}.pt" elif isinstance(weight_path, str): @@ -173,8 +246,9 @@ def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True, logger.info(f"🌐 Weight {weight_path} not found, try downloading") prepare_weight(weight_path=weight_path) if weight_path.exists(): + logger.info(f'🏋 Loading weights from {weight_path}') model.save_load_weights(weight_path) logger.info(":white_check_mark: Success load model & weight") else: - logger.info(":white_check_mark: Success load model") + logger.info(":white_check_mark: Success load model without weights") return model diff --git a/yolo/tools/data_loader.py b/yolo/tools/data_loader.py index c44f00c68..dde3030dd 100644 --- a/yolo/tools/data_loader.py +++ b/yolo/tools/data_loader.py @@ -12,7 +12,7 @@ from torch.utils.data import DataLoader, Dataset from yolo.config.config import DataConfig, DatasetConfig -from yolo.tools.data_augmentation import * +from yolo.tools import data_augmentation from yolo.tools.data_augmentation import AugmentationComposer from yolo.tools.dataset_preparation import prepare_dataset from yolo.utils.dataset_utils import ( @@ -33,10 +33,12 @@ def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, phase: str self.dynamic_shape = getattr(data_cfg, "dynamic_shape", False) self.base_size = mean(self.image_size) - transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()] + transforms = [getattr(data_augmentation, aug)(prob) for aug, prob in augment_cfg.items()] self.transform = AugmentationComposer(transforms, self.image_size, self.base_size) self.transform.get_more_data = self.get_more_data - self.img_paths, self.bboxes, self.ratios = tensorlize(self.load_data(Path(dataset_cfg.path), phase_name)) + + data = self.load_data(Path(dataset_cfg.path), phase_name) + self.img_paths, self.bboxes, self.ratios = tensorlize(data) def load_data(self, dataset_path: Path, phase_name: str): """ @@ -81,47 +83,111 @@ def filter_data(self, dataset_path: Path, phase_name: str, sort_image: bool = Fa list: A list of tuples, each containing the path to an image file and its associated segmentation as a tensor. """ images_path = dataset_path / "images" / phase_name + labels_path, data_type = locate_label_paths(dataset_path, phase_name) - images_list = sorted([p.name for p in Path(images_path).iterdir() if p.is_file()]) - if data_type == "json": - annotations_index, image_info_dict = create_image_metadata(labels_path) - data = [] - valid_inputs = 0 - for image_name in track(images_list, description="Filtering data"): - if not image_name.lower().endswith((".jpg", ".jpeg", ".png")): - continue - image_id = Path(image_name).stem + if data_type == 'kwcoco': + """ + More robust data handling that only depends on paths within the + specified manifest file. - if data_type == "json": - image_info = image_info_dict.get(image_id, None) - if image_info is None: - continue - annotations = annotations_index.get(image_info["id"], []) - image_seg_annotations = scale_segmentation(annotations, image_info) - elif data_type == "txt": - label_path = labels_path / f"{image_id}.txt" - if not label_path.is_file(): - continue - with open(label_path, "r") as file: - image_seg_annotations = [list(map(float, line.strip().split())) for line in file] - else: - image_seg_annotations = [] + Principles: - labels = self.load_valid_labels(image_id, image_seg_annotations) + * Dont glob for the images, let the dataset tell you where they are. + + * A Dataset should be referenced as a single URI to a manifest. + The manifest should either contain relevant data or point to + paths for everything. + """ + import kwcoco + coco_dset = kwcoco.CocoDataset(labels_path) + + from yolo.tools.data_conversion import discretize_categories + id_to_idx = discretize_categories(coco_dset.dataset.get("categories", [])) if "categories" in coco_dset.dataset else None + + total_images = coco_dset.n_images - img_path = images_path / image_name if sort_image: - with Image.open(img_path) as img: - width, height = img.size - else: - width, height = 0, 1 - data.append((img_path, labels, width / height)) - valid_inputs += 1 + # Ensure all images have populated sizes + coco_dset._ensure_imgsize() + + ALLOW_EMPTY_IMAGES = 0 + + # Build the expected output + data = [] + valid_inputs = 0 + for coco_img in coco_dset.images().coco_images_iter(): + image_info = coco_img.img + img_path = coco_img.primary_image_filepath() + + if sort_image: + width, height = coco_img['width'], coco_img['height'] + else: + width, height = 0, 1 + + annotations = coco_img.annots().objs + + # Handle filtering as done in + # :func:`dataset_utils.organize_annotations_by_image` + modified_annotations = [] + for anno in annotations: + if id_to_idx: + anno["category_id"] = id_to_idx[anno["category_id"]] + if anno.get("iscrowd", False): # TODO: make configurable + continue + modified_annotations.append(anno) + annotations = modified_annotations + + if ALLOW_EMPTY_IMAGES or len(annotations): + + image_seg_annotations = scale_segmentation(annotations, image_info) + labels = self.load_valid_labels(None, image_seg_annotations) + + data.append((img_path, labels, width / height)) + valid_inputs += 1 + + else: + images_list = sorted([p.name for p in Path(images_path).iterdir() if p.is_file()]) + if data_type == "json": + annotations_index, image_info_dict = create_image_metadata(labels_path) + + data = [] + valid_inputs = 0 + for image_name in track(images_list, description="Filtering data"): + if not image_name.lower().endswith((".jpg", ".jpeg", ".png")): + continue + image_id = Path(image_name).stem + + if data_type == "json": + image_info = image_info_dict.get(image_id, None) + if image_info is None: + continue + annotations = annotations_index.get(image_info["id"], []) + image_seg_annotations = scale_segmentation(annotations, image_info) + elif data_type == "txt": + label_path = labels_path / f"{image_id}.txt" + if not label_path.is_file(): + continue + with open(label_path, "r") as file: + image_seg_annotations = [list(map(float, line.strip().split())) for line in file] + else: + image_seg_annotations = [] + + labels = self.load_valid_labels(image_id, image_seg_annotations) + + img_path = images_path / image_name + if sort_image: + with Image.open(img_path) as img: + width, height = img.size + else: + width, height = 0, 1 + data.append((img_path, labels, width / height)) + valid_inputs += 1 + total_images = len(images_list) data = sorted(data, key=lambda x: x[2], reverse=True) - logger.info(f"Recorded {valid_inputs}/{len(images_list)} valid inputs") + logger.info(f"Recorded {valid_inputs}/{total_images} valid inputs") return data def load_valid_labels(self, label_path: str, seg_data_one_img: list) -> Union[Tensor, None]: @@ -139,8 +205,18 @@ def load_valid_labels(self, label_path: str, seg_data_one_img: list) -> Union[Te bboxes = [] for seg_data in seg_data_one_img: cls = seg_data[0] - points = np.array(seg_data[1:]).reshape(-1, 2) - valid_points = points[(points >= 0) & (points <= 1)].reshape(-1, 2) + # This seems like an incorrect check. Putting my fix inside an if + # in case I don't understand why it is this way. + FIX_INCORRECT_CHECK = 1 + if FIX_INCORRECT_CHECK: + points = np.array(seg_data[1:]).reshape(-1, 2) + # This probably should just be a clamp / clip operation + # but I'm keeping it similar to the original + flags = (points >= 0).all(axis=1) & (points <= 1).all(axis=1) + valid_points = points[flags] + else: + points = np.array(seg_data[1:]).reshape(-1, 2) + valid_points = points[(points >= 0) & (points <= 1)].reshape(-1, 2) if valid_points.size > 1: bbox = torch.tensor([cls, *valid_points.min(axis=0), *valid_points.max(axis=0)]) bboxes.append(bbox) diff --git a/yolo/tools/loss_functions.py b/yolo/tools/loss_functions.py index 79fe1cf96..379a42ce1 100644 --- a/yolo/tools/loss_functions.py +++ b/yolo/tools/loss_functions.py @@ -107,6 +107,21 @@ def __call__(self, predicts: List[Tensor], targets: Tensor) -> Tuple[Tensor, Ten class DualLoss: + """ + Example: + >>> import torch + >>> from yolo.tools.loss_functions import DualLoss + >>> from yolo.utils.bounding_box_utils import Vec2Box + >>> from yolo.utils.config_utils import build_config + >>> cfg = build_config(overrides=['task=train']) + >>> device = 'cpu' + >>> vec2box = Vec2Box(model=None, anchor_cfg=cfg.model.anchor, image_size=cfg.image_size, device=device) + >>> self = DualLoss(cfg, vec2box) + >>> targets = torch.zeros(1, 20, 5, device=device) + >>> aux_predicts = [torch.zeros(1, 8400, *cn, device=device) for cn in [(80,), (4, 16), (4,)]] + >>> main_predicts = [torch.zeros(1, 8400, *cn, device=device) for cn in [(80,), (4, 16), (4,)]] + >>> loss, loss_dict = self(aux_predicts, main_predicts, targets) + """ def __init__(self, cfg: Config, vec2box) -> None: loss_cfg = cfg.task.loss self.loss = YOLOLoss(loss_cfg, vec2box, class_num=cfg.dataset.class_num, reg_max=cfg.model.anchor.reg_max) diff --git a/yolo/tools/solver.py b/yolo/tools/solver.py index 8246a66d2..8590d3aeb 100644 --- a/yolo/tools/solver.py +++ b/yolo/tools/solver.py @@ -55,7 +55,7 @@ def validation_step(self, batch, batch_idx): def on_validation_epoch_end(self): epoch_metrics = self.metric.compute() - del epoch_metrics["classes"] + epoch_metrics.pop("classes", None) self.log_dict(epoch_metrics, prog_bar=True, sync_dist=True, rank_zero_only=True) self.log_dict( {"PyCOCO/AP @ .5:.95": epoch_metrics["map"], "PyCOCO/AP @ .5": epoch_metrics["map_50"]}, @@ -69,6 +69,10 @@ class TrainModel(ValidateModel): def __init__(self, cfg: Config): super().__init__(cfg) self.cfg = cfg + + # TODO: if we defer creating the model until the dataset is loaded, we + # can introspect the number of categories and other things to make user + # configuration have less interdependencies and thus be more robust. self.train_loader = create_dataloader(self.cfg.task.data, self.cfg.dataset, self.cfg.task.task) def setup(self, stage): diff --git a/yolo/utils/config_utils.py b/yolo/utils/config_utils.py new file mode 100644 index 000000000..465373c4e --- /dev/null +++ b/yolo/utils/config_utils.py @@ -0,0 +1,36 @@ +import omegaconf +from typing import List + + +def build_config(overrides: List[str] = []) -> omegaconf.DictConfig: + """ + Creates an explicit config for testing. + + Example: + >>> from yolo.utils.config_utils import build_config + >>> cfg = build_config(overrides=['task=train']) + >>> cfg = build_config(overrides=['task=validation']) + >>> cfg = build_config(overrides=['task=inference']) + """ + import yolo + import os + import pathlib + from hydra import compose, initialize + + # This is annoying that we cant just specify an absolute path when it is + # robustly built. Furthermore, the relative path seems like it isn't even + # from the cwd, but the module that is currently being run. + + # Find the path that we need to be relative to in a somewhat portable + # manner (i.e. will work in a Jupyter snippet). + try: + path_base = pathlib.Path(__file__).parent + except NameError: + path_base = pathlib.Path.cwd() + yolo_path = pathlib.Path(yolo.__file__).parent + rel_yolo_path = pathlib.Path(os.path.relpath(yolo_path, path_base)) + # rel_yolo_path = yolo_path.relative_to(path_base, walk_up=True) # requires Python 3.12 + config_path = os.fspath(rel_yolo_path / 'config') + with initialize(config_path=config_path, version_base=None): + cfg = compose(config_name="config", overrides=overrides) + return cfg diff --git a/yolo/utils/dataset_utils.py b/yolo/utils/dataset_utils.py index dd9a66abc..b794b8c55 100644 --- a/yolo/utils/dataset_utils.py +++ b/yolo/utils/dataset_utils.py @@ -34,6 +34,14 @@ def locate_label_paths(dataset_path: Path, phase_name: Path) -> Tuple[Path, Path if txt_files: return txt_labels_path, "txt" + HANDLE_KWCOCO_FILES = 1 + if HANDLE_KWCOCO_FILES: + candidate = dataset_path / phase_name + if candidate.is_file(): + labels_path = dataset_path / phase_name + data_type = 'kwcoco' + return labels_path, data_type + logger.warning("No labels found in the specified dataset path and phase name.") return [], None @@ -97,12 +105,22 @@ def scale_segmentation( if annotations is None: return None + try: + import kwimage + except ImportError: + kwimage = None + seg_array_with_cat = [] h, w = image_dimensions["height"], image_dimensions["width"] for anno in annotations: category_id = anno["category_id"] if "segmentation" in anno: - seg_list = [item for sublist in anno["segmentation"] for item in sublist] + if kwimage is None: + # original fallback code + seg_list = [item for sublist in anno["segmentation"] for item in sublist] + else: + # Convert to original coco representation + seg_list = kwimage.MultiPolygon.coerce(anno["segmentation"]).to_coco('orig') elif "bbox" in anno: x, y, width, height = anno["bbox"] seg_list = [x, y, x + width, y, x + width, y + height, x, y + height] diff --git a/yolo/utils/logger.py b/yolo/utils/logger.py index 286027205..4880cf07c 100644 --- a/yolo/utils/logger.py +++ b/yolo/utils/logger.py @@ -3,9 +3,16 @@ from lightning.pytorch.utilities.rank_zero import rank_zero_only from rich.console import Console from rich.logging import RichHandler +import os logger = logging.getLogger("yolo") logger.setLevel(logging.DEBUG) logger.propagate = False -if rank_zero_only.rank == 0 and not logger.hasHandlers(): - logger.addHandler(RichHandler(console=Console(), show_level=True, show_path=True, show_time=True, markup=True)) + +# allow the user to get a simpler output +# TODO: needs to be better integrated +DISABLE_RICH_HANDLER = bool(os.environ.get('DISABLE_RICH_HANDLER', '')) + +if not DISABLE_RICH_HANDLER: + if rank_zero_only.rank == 0 and not logger.hasHandlers(): + logger.addHandler(RichHandler(console=Console(), show_level=True, show_path=True, show_time=True, markup=True)) diff --git a/yolo/utils/logging_utils.py b/yolo/utils/logging_utils.py index 28a536222..9ca14d25e 100644 --- a/yolo/utils/logging_utils.py +++ b/yolo/utils/logging_utils.py @@ -48,7 +48,7 @@ def set_seed(seed): if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) # if you are using multi-GPU. - torch.backends.cudnn.deterministic = True + # torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False @@ -107,7 +107,7 @@ def on_train_batch_end(self, trainer, pl_module, outputs, batch: Any, batch_idx: epoch_descript = "[cyan]Train [white]|" batch_descript = "[green]Train [white]|" metrics = self.get_metrics(trainer, pl_module) - metrics.pop("v_num") + metrics.pop("v_num", None) for metrics_name, metrics_val in metrics.items(): if "Loss_step" in metrics_name: epoch_descript += f"{metrics_name.removesuffix('_step').split('/')[1]: ^9}|" @@ -216,12 +216,92 @@ def on_validation_batch_end(self, trainer: Trainer, pl_module, outputs, batch, b pred_boxes = outputs[0] if isinstance(outputs, list) else outputs images = [images[0]] step = trainer.current_epoch + for logger in trainer.loggers: if isinstance(logger, WandbLogger): + # FIXME: not robust to configured image sizes, need to know + # that info. logger.log_image("Input Image", images, step=step) logger.log_image("Ground Truth", images, step=step, boxes=[log_bbox(gt_boxes)]) logger.log_image("Prediction", images, step=step, boxes=[log_bbox(pred_boxes)]) + # TODO: better config + import os + LOG_BATCH_VIZ_TO_DISK = bool(os.environ.get('LOG_BATCH_VIZ_TO_DISK', '')) + if LOG_BATCH_VIZ_TO_DISK: + import einops + import kwimage + + # TODO: + # get a batter output path + import pathlib + root_dpath = pathlib.Path(trainer.default_root_dir) + out_dpath = root_dpath / 'debug_images' / trainer.state.stage.name + out_dpath.mkdir(exist_ok=True, parents=True) + epoch = trainer.current_epoch + + for bx in range(len(images)): + image_chw = images[bx].data.cpu().numpy() + image_hwc = einops.rearrange(image_chw, 'c h w -> h w c') + image_hwc = kwimage.ensure_uint255(image_hwc) + + assert bx == 0, 'not handling multiple per batch' + true_dets = tensor_to_kwimage(gt_boxes).numpy() + pred_dets = tensor_to_kwimage(pred_boxes).numpy() + pred_dets = pred_dets.non_max_supress(thresh=0.3) + # pred_dets = pred_dets.compress(pred_dets.scores > 0.1) + + raw_canvas = image_hwc.copy() + true_canvas = true_dets.draw_on(raw_canvas.copy(), color='green') + pred_canvas = pred_dets.draw_on(raw_canvas.copy(), color='blue') + + raw_canvas = kwimage.draw_header_text(raw_canvas, 'raw') + true_canvas = kwimage.draw_header_text(true_canvas, f'true, n={len(true_dets)}') + pred_canvas = kwimage.draw_header_text(pred_canvas, f'pred, n={len(pred_dets)}') + canvas = kwimage.stack_images([ + raw_canvas, true_canvas, pred_canvas + ], axis=1, pad=3) + + fname = f'img_{epoch:04d}_{batch_idx:04d}.jpg' + fpath = out_dpath / fname + kwimage.imwrite(fpath, canvas) + + +def tensor_to_kwimage(yolo_annot_tensor): + import kwimage + class_idxs = yolo_annot_tensor[:, 0].int() + boxes = kwimage.Boxes(yolo_annot_tensor[:, 1:5], format='xyxy') + dets = kwimage.Detections( + boxes=boxes, + class_idxs=class_idxs + ) + + if yolo_annot_tensor.shape[1] > 5: + scores = yolo_annot_tensor[:, 5] + dets.data['scores'] = scores + return dets + + +def wandb_to_kwimage(wand_annots): + import numpy as np + import kwimage + box_list = [] + class_idxs = [] + for row in wand_annots['predictions']['box_data']: + pos = row['position'] + class_idx = row['class_id'] + xyxy = [pos['minX'], pos['minY'], pos['maxX'], pos['maxY']] + box_list.append(xyxy) + class_idxs.append(class_idx) + + boxes = kwimage.Boxes(np.array(box_list), format='xyxy') + dets = kwimage.Detections( + boxes=boxes, + class_idxs=np.array(class_idxs) + ) + dets = dets.compress(dets.class_idxs > -1) + return dets + def setup_logger(logger_name, quite=False): class EmojiFormatter(logging.Formatter): @@ -264,9 +344,13 @@ def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=Tru logger.setLevel(logging.ERROR) return progress, loggers, save_path - progress.append(YOLORichProgressBar()) - progress.append(YOLORichModelSummary()) + from yolo.utils.logger import DISABLE_RICH_HANDLER + if not DISABLE_RICH_HANDLER: + progress.append(YOLORichProgressBar()) + progress.append(YOLORichModelSummary()) + progress.append(ImageLogger()) + if cfg.use_tensorboard: loggers.append(TensorBoardLogger(log_graph="all", save_dir=save_path)) if cfg.use_wandb: diff --git a/yolo/utils/model_utils.py b/yolo/utils/model_utils.py index 9d6c0ce59..bb78096eb 100644 --- a/yolo/utils/model_utils.py +++ b/yolo/utils/model_utils.py @@ -57,6 +57,14 @@ def on_validation_start(self, trainer: "Trainer", pl_module: "LightningModule"): self.ema_state_dict = deepcopy(pl_module.model.state_dict()) pl_module.ema.load_state_dict(self.ema_state_dict) + @no_grad() + def on_train_batch_start(self, trainer: "Trainer", pl_module: "LightningModule", *args, **kwargs) -> None: + if self.ema_state_dict is None: + # If validation sanity checks are disabled, then we need to + # initialize the ema state before training starts. + self.ema_state_dict = deepcopy(pl_module.model.state_dict()) + pl_module.ema.load_state_dict(self.ema_state_dict) + @no_grad() def on_train_batch_end(self, trainer: "Trainer", pl_module: "LightningModule", *args, **kwargs) -> None: self.step += 1