🔀 [Merge] pull request #14 from LucyTuan/DATASET

[Merge][Add] json annotation data readable in dataloader.py

utils/dataloader.py

@@ -1,18 +1,116 @@
+import json
+import os
+from itertools import chain
 from os import listdir, path
-from typing import List, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import diskcache as dc
 import hydra
 import numpy as np
 import torch
+from data_augment import Compose, HorizontalFlip, MixUp, Mosaic, VerticalFlip
+from drawer import draw_bboxes
 from loguru import logger
 from PIL import Image
 from torch.utils.data import DataLoader, Dataset
 from torchvision.transforms import functional as TF
 from tqdm.rich import tqdm
 
-from utils.data_augment import Compose, HorizontalFlip, MixUp, Mosaic, VerticalFlip
-from utils.drawer import draw_bboxes
+
+def find_labels_path(dataset_path: str, phase_name: str):
+    """
+    Find the path to label files for a specified dataset and phase(e.g. training).
+
+    Args:
+        dataset_path (str): The path to the root directory of the dataset.
+        phase_name (str): The name of the phase for which labels are being searched (e.g., "train", "val", "test").
+
+    Returns:
+        Tuple[str, str]: A tuple containing the path to the labels file and the file format ("json" or "txt").
+    """
+    json_labels_path = path.join(dataset_path, "annotations", f"instances_{phase_name}.json")
+
+    txt_labels_path = path.join(dataset_path, "label", phase_name)
+
+    if path.isfile(json_labels_path):
+        return json_labels_path, "json"
+
+    elif path.isdir(txt_labels_path):
+        txt_files = [f for f in os.listdir(txt_labels_path) if f.endswith(".txt")]
+        if txt_files:
+            return txt_labels_path, "txt"
+
+    raise FileNotFoundError("No labels found in the specified dataset path and phase name.")
+
+
+def create_image_info_dict(labels_path: str) -> Tuple[Dict[str, List], Dict[str, Dict]]:
+    """
+    Create a dictionary containing image information and annotations indexed by image ID.
+
+    Args:
+        labels_path (str): The path to the annotation json file.
+
+    Returns:
+        - annotations_index: A dictionary where keys are image IDs and values are lists of annotations.
+        - image_info_dict: A dictionary where keys are image file names without extension and values are image information dictionaries.
+    """
+    with open(labels_path, "r") as file:
+        labels_data = json.load(file)
+        annotations_index = index_annotations_by_image(labels_data)  # check lookup is a good name?
+        image_info_dict = {path.splitext(img["file_name"])[0]: img for img in labels_data["images"]}
+        return annotations_index, image_info_dict
+
+
+def index_annotations_by_image(data: Dict[str, Any]):
+    """
+    Use image index to lookup every annotations
+    Args:
+        data (Dict[str, Any]): A dictionary containing annotation data.
+
+    Returns:
+        Dict[int, List[Dict[str, Any]]]: A dictionary where keys are image IDs and values are lists of annotations.
+        Annotations with "iscrowd" set to True are excluded from the index.
+
+    """
+    annotation_lookup = {}
+    for anno in data["annotations"]:
+        if anno["iscrowd"]:
+            continue
+        image_id = anno["image_id"]
+        if image_id not in annotation_lookup:
+            annotation_lookup[image_id] = []
+        annotation_lookup[image_id].append(anno)
+    return annotation_lookup
+
+
+def get_scaled_segmentation(
+    annotations: List[Dict[str, Any]], image_dimensions: Dict[str, int]
+) -> Optional[List[List[float]]]:
+    """
+    Scale the segmentation data based on image dimensions and return a list of scaled segmentation data.
+
+    Args:
+        annotations (List[Dict[str, Any]]): A list of annotation dictionaries.
+        image_dimensions (Dict[str, int]): A dictionary containing image dimensions (height and width).
+
+    Returns:
+        Optional[List[List[float]]]: A list of scaled segmentation data, where each sublist contains category_id followed by scaled (x, y) coordinates.
+    """
+    if annotations is None:
+        return None
+
+    seg_array_with_cat = []
+    h, w = image_dimensions["height"], image_dimensions["width"]
+    for anno in annotations:
+        category_id = anno["category_id"]
+        seg_list = [item for sublist in anno["segmentation"] for item in sublist]
+        scaled_seg_data = (
+            np.array(seg_list).reshape(-1, 2) / [w, h]
+        ).tolist()  # make the list group in x, y pairs and scaled with image width, height
+        scaled_flat_seg_data = [category_id] + list(chain(*scaled_seg_data))  # flatten the scaled_seg_data list
+        seg_array_with_cat.append(scaled_flat_seg_data)
+
+    return seg_array_with_cat
 
 
 class YoloDataset(Dataset):
@@ -44,9 +142,7 @@ class YoloDataset(Dataset):
 
         if data is None:
             logger.info("Generating {} cache", phase_name)
-            images_path = path.join(dataset_path, "images", phase_name)
-            labels_path = path.join(dataset_path, "labels", phase_name)
-            data = self.filter_data(images_path, labels_path)
+            data = self.filter_data(dataset_path, phase_name)
             cache[phase_name] = data
 
         cache.close()
@@ -54,7 +150,7 @@ class YoloDataset(Dataset):
         data = cache[phase_name]
         return data
 
-    def filter_data(self, images_path: str, labels_path: str) -> list:
+    def filter_data(self, dataset_path: str, phase_name: str) -> list:
         """
         Filters and collects dataset information by pairing images with their corresponding labels.
 
@@ -63,29 +159,49 @@ class YoloDataset(Dataset):
             labels_path (str): Path to the directory containing label files.
 
         Returns:
-            list: A list of tuples, each containing the path to an image file and its associated labels as a tensor.
+            list: A list of tuples, each containing the path to an image file and its associated segmentation as a tensor.
         """
+        images_path = path.join(dataset_path, "images", phase_name)
+        labels_path, data_type = find_labels_path(dataset_path, phase_name)
+        images_list = sorted(os.listdir(images_path))
+        if data_type == "json":
+            annotations_index, image_info_dict = create_image_info_dict(labels_path)
+
         data = []
         valid_inputs = 0
-        images_list = sorted(listdir(images_path))
         for image_name in tqdm(images_list, desc="Filtering data"):
             if not image_name.lower().endswith((".jpg", ".jpeg", ".png")):
                 continue
-
-            img_path = path.join(images_path, image_name)
-            base_name, _ = path.splitext(image_name)
-            label_path = path.join(labels_path, f"{base_name}.txt")
-
-            if path.isfile(label_path):
-                labels = self.load_valid_labels(label_path)
-                if labels is not None:
-                    data.append((img_path, labels))
-                    valid_inputs += 1
+            image_id, _ = path.splitext(image_name)
+
+            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 = get_scaled_segmentation(annotations, image_info)
+                if not image_seg_annotations:
+                    continue
+
+            elif data_type == "txt":
+                label_path = path.join(labels_path, f"{image_id}.txt")
+                if not path.isfile(label_path):
+                    continue
+                with open(label_path, "r") as file:
+                    image_seg_annotations = [
+                        list(map(float, line.strip().split())) for line in file
+                    ]  # add a comment for this line, complicated, do you need "list", im not sure
+
+            labels = self.load_valid_labels(image_id, image_seg_annotations)
+            if labels is not None:
+                img_path = path.join(images_path, image_name)
+                data.append((img_path, labels))
+                valid_inputs += 1
 
         logger.info("Recorded {}/{} valid inputs", valid_inputs, len(images_list))
         return data
 
-    def load_valid_labels(self, label_path: str) -> Union[torch.Tensor, None]:
+    def load_valid_labels(self, label_path, seg_data_one_img) -> Union[torch.Tensor, None]:
         """
         Loads and validates bounding box data is [0, 1] from a label file.
 
@@ -96,15 +212,13 @@ class YoloDataset(Dataset):
             torch.Tensor or None: A tensor of all valid bounding boxes if any are found; otherwise, None.
         """
         bboxes = []
-        with open(label_path, "r") as file:
-            for line in file:
-                parts = list(map(float, line.strip().split()))
-                cls = parts[0]
-                points = np.array(parts[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)
+        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)
+            if valid_points.size > 1:
+                bbox = torch.tensor([cls, *valid_points.min(axis=0), *valid_points.max(axis=0)])
+                bboxes.append(bbox)
 
         if bboxes:
             return torch.stack(bboxes)