datasets/lvis.py

import cv2
import numpy as np
import os
from .data_utils import * 
from .base import BaseDataset
from lvis import LVIS
from pathlib import Path
from util.box_ops import compute_iou_matrix, draw_bboxes
import shutil

IS_VERIFY = False

class LVISDataset(BaseDataset):
    def __init__(self, construct_dataset_dir, obj_thr=20, area_ratio=0.02):
        self.obj_thr = obj_thr
        self.construct_dataset_dir = construct_dataset_dir
        os.makedirs(Path(self.construct_dataset_dir), exist_ok=True)
        self.area_ratio = area_ratio
        self.sample_list = os.listdir(self.construct_dataset_dir)

    def _get_image_path(self, file_name):
        for img_dir in self.image_dir:
            path = img_dir / file_name
            if path.exists():
                return str(path)
        raise FileNotFoundError(f"File {file_name} not found in any of the image_dir.")

    def _intersect_2_obj(self, image_dir, lvis_api, imgs_info, annos, idx):
        self.image_dir = image_dir
        image_name = imgs_info[idx]['coco_url'].split('/')[-1]
        image_path = self._get_image_path(image_name)
        image = cv2.imread(image_path)

        h, w = image.shape[0:2]
        image_area = h*w

        anno = annos[idx]

        # filter by area
        obj_ids = []
        obj_areas = []
        obj_bbox = []
        for i in range(len(anno)):
            obj = anno[i]
            area = obj['area']
            bbox = obj['bbox'] # xyhw
            if area > image_area * self.area_ratio:
                obj_ids.append(i)
                obj_areas.append(area)
                obj_bbox.append(bbox)

        if len(obj_bbox) < 2:
            print(f"[Info] Skip image index {image_name[:-4]} due to insufficient bbox.")
            return

        # filter by IOU
        bbox_xyxy = []
        for box in obj_bbox:
            x, y, w, h = box
            bbox_xyxy.append([x, y, x + w, y + h])
        bbox_xyxy = np.array(bbox_xyxy)  # shape: [N, 4]

        if IS_VERIFY:
            os.makedirs(Path(self.construct_dataset_dir) / image_name[:-4], exist_ok=True)
            image_with_boxes = draw_bboxes(image, bbox_xyxy)
            cv2.imwrite(str(Path(self.construct_dataset_dir) / image_name[:-4] / "bboxes_image.png"), image_with_boxes)
        
        iou_matrix = compute_iou_matrix(bbox_xyxy)
        np.fill_diagonal(iou_matrix, -1) # Exclude self-comparisons (i.e., each box with itself)

        max_index = np.unravel_index(np.argmax(iou_matrix), iou_matrix.shape)
        index0, index1 = max_index[0], max_index[1]
        max_iou = iou_matrix[index0, index1]

        if max_iou <= 0:
            print(f"[Info] Skip image index {image_name[:-4]} due to no overlapping bboxes.")
            return

        os.makedirs(Path(self.construct_dataset_dir) / image_name[:-4], exist_ok=True)
        dst = Path(self.construct_dataset_dir) / image_name[:-4] / "image.jpg"
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(image_path, dst)

        anno_id = anno[obj_ids[index0]]
        mask = lvis_api.ann_to_mask(anno_id)
        cv2.imwrite(str(Path(self.construct_dataset_dir) / image_name[:-4] / "object_0_mask.png"), 255*mask)
        patch = self.get_patch(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), mask)
        patch = cv2.cvtColor(patch, cv2.COLOR_RGB2BGR)
        cv2.imwrite(str(Path(self.construct_dataset_dir) / image_name[:-4] / "object_0.png"), patch)

        if IS_VERIFY:
            mask_color = np.stack([mask * 255]*3, axis=-1).astype(np.uint8)
            highlight = np.zeros_like(image)
            highlight[:, :, 2] = 255  # red channel
            alpha = 0.5
            image_with_boxes = np.where(mask_color == 255, cv2.addWeighted(image_with_boxes, 1 - alpha, highlight, alpha, 0), image_with_boxes)

        anno_id = anno[obj_ids[index1]]
        mask = lvis_api.ann_to_mask(anno_id)
        cv2.imwrite(str(Path(self.construct_dataset_dir) / image_name[:-4] / "object_1_mask.png"), 255*mask)
        patch = self.get_patch(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), mask)
        patch = cv2.cvtColor(patch, cv2.COLOR_RGB2BGR)
        cv2.imwrite(str(Path(self.construct_dataset_dir) / image_name[:-4] / "object_1.png"), patch)

        if IS_VERIFY:
            mask_color = np.stack([mask * 255]*3, axis=-1).astype(np.uint8)
            highlight = np.zeros_like(image)
            highlight[:, :, 0] = 255  # blue channel
            alpha = 0.5
            image_with_boxes = np.where(mask_color == 255, cv2.addWeighted(image_with_boxes, 1 - alpha, highlight, alpha, 0), image_with_boxes)
            cv2.imwrite(str(Path(self.construct_dataset_dir) / image_name[:-4] / "highlighted_image.png"), image_with_boxes)

    def _intersect_3_obj(self, image_dir, lvis_api, imgs_info, annos, idx):
        self.image_dir = image_dir
        image_name = imgs_info[idx]['coco_url'].split('/')[-1]
        image_path = self._get_image_path(image_name)
        image = cv2.imread(image_path)

        h, w = image.shape[0:2]
        image_area = h * w

        anno = annos[idx]

        # filter by area
        obj_ids = []
        obj_areas = []
        obj_bbox = []
        for i, obj in enumerate(anno):
            area = obj['area']
            bbox = obj['bbox']  # xywh
            if area > image_area * self.area_ratio:
                obj_ids.append(i)
                obj_areas.append(area)
                obj_bbox.append(bbox)

        if len(obj_bbox) < 3:
            print(f"[Info] Skip image index {image_name[:-4]} due to insufficient bbox (need >=3, got {len(obj_bbox)}).")
            return

        # calculate IOU matrix
        bbox_xyxy = []
        for box in obj_bbox:
            x, y, w_box, h_box = box
            bbox_xyxy.append([x, y, x + w_box, y + h_box])
        bbox_xyxy = np.array(bbox_xyxy)  # shape: [N, 4]

        if IS_VERIFY:
            os.makedirs(Path(self.construct_dataset_dir) / image_name[:-4], exist_ok=True)
            image_with_boxes = draw_bboxes(image, bbox_xyxy)
            cv2.imwrite(str(Path(self.construct_dataset_dir) / image_name[:-4] / "bboxes_image.png"), image_with_boxes)

        iou_matrix = compute_iou_matrix(bbox_xyxy)
        np.fill_diagonal(iou_matrix, -1)  # Exclude self-comparisons

        # find 3 overlapped objects
        positive_iou = np.where(iou_matrix > 0, iou_matrix, 0.0)
        row_sums = positive_iou.sum(axis=1)
        anchor = int(np.argmax(row_sums))

        partner_candidates = np.argsort(iou_matrix[anchor])[::-1]
        partners = [int(p) for p in partner_candidates if iou_matrix[anchor, p] > 0]

        if len(partners) < 2:
            print(f"[Info] Skip image index {image_name[:-4]} due to not enough overlapping bboxes for 3 objects.")
            return

        index0 = anchor
        index1 = partners[0]
        index2 = partners[1]

        max_iou_pair = max(iou_matrix[index0, index1], iou_matrix[index0, index2], iou_matrix[index1, index2])
        if max_iou_pair <= 0:
            print(f"[Info] Skip image index {image_name[:-4]} due to no overlapping bboxes.")
            return

        # copy original image
        out_dir = Path(self.construct_dataset_dir) / image_name[:-4]
        out_dir.mkdir(parents=True, exist_ok=True)
        dst = out_dir / "image.jpg"
        shutil.copy(image_path, dst)

        # first object
        anno_id = anno[obj_ids[index0]]
        mask0 = lvis_api.ann_to_mask(anno_id)
        cv2.imwrite(str(out_dir / "object_0_mask.png"), 255 * mask0)
        patch0 = self.get_patch(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), mask0)
        patch0 = cv2.cvtColor(patch0, cv2.COLOR_RGB2BGR)
        cv2.imwrite(str(out_dir / "object_0.png"), patch0)

        if IS_VERIFY:
            mask_color = np.stack([mask0 * 255] * 3, axis=-1).astype(np.uint8)
            highlight = np.zeros_like(image)
            highlight[:, :, 2] = 255  # red channel
            alpha = 0.5
            image_with_boxes = np.where(
                mask_color == 255,
                cv2.addWeighted(image_with_boxes, 1 - alpha, highlight, alpha, 0),
                image_with_boxes
            )

        # second object
        anno_id = anno[obj_ids[index1]]
        mask1 = lvis_api.ann_to_mask(anno_id)
        cv2.imwrite(str(out_dir / "object_1_mask.png"), 255 * mask1)
        patch1 = self.get_patch(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), mask1)
        patch1 = cv2.cvtColor(patch1, cv2.COLOR_RGB2BGR)
        cv2.imwrite(str(out_dir / "object_1.png"), patch1)

        if IS_VERIFY:
            mask_color = np.stack([mask1 * 255] * 3, axis=-1).astype(np.uint8)
            highlight = np.zeros_like(image)
            highlight[:, :, 0] = 255  # blue channel
            alpha = 0.5
            image_with_boxes = np.where(
                mask_color == 255,
                cv2.addWeighted(image_with_boxes, 1 - alpha, highlight, alpha, 0),
                image_with_boxes
            )

        # third object
        anno_id = anno[obj_ids[index2]]
        mask2 = lvis_api.ann_to_mask(anno_id)
        cv2.imwrite(str(out_dir / "object_2_mask.png"), 255 * mask2)
        patch2 = self.get_patch(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), mask2)
        patch2 = cv2.cvtColor(patch2, cv2.COLOR_RGB2BGR)
        cv2.imwrite(str(out_dir / "object_2.png"), patch2)

        if IS_VERIFY:
            mask_color = np.stack([mask2 * 255] * 3, axis=-1).astype(np.uint8)
            highlight = np.zeros_like(image)
            highlight[:, :, 1] = 255  # green channel
            alpha = 0.5
            image_with_boxes = np.where(
                mask_color == 255,
                cv2.addWeighted(image_with_boxes, 1 - alpha, highlight, alpha, 0),
                image_with_boxes
            )
            cv2.imwrite(str(out_dir / "highlighted_image.png"), image_with_boxes)


    def _get_sample(self, idx):
        sample_path = os.path.join(self.construct_dataset_dir, self.sample_list[idx])
        image = cv2.cvtColor(cv2.imread(os.path.join(sample_path, "image.jpg")), cv2.COLOR_BGR2RGB)
        object_0 = cv2.cvtColor(cv2.imread(os.path.join(sample_path, "object_0.png")), cv2.COLOR_BGR2RGB)
        object_1 = cv2.cvtColor(cv2.imread(os.path.join(sample_path, "object_1.png")), cv2.COLOR_BGR2RGB)
        mask_0 = cv2.imread(os.path.join(sample_path, "object_0_mask.png"), cv2.IMREAD_GRAYSCALE)
        mask_1 = cv2.imread(os.path.join(sample_path, "object_1_mask.png"), cv2.IMREAD_GRAYSCALE)
        collage = self._construct_collage(image, object_0, object_1, mask_0, mask_1)
        return collage

    def __len__(self):
        return len(os.listdir(self.construct_dataset_dir))


if __name__ == "__main__":
    '''
    two-object case: train/test: 34610/8859
    '''
    import argparse

    parser = argparse.ArgumentParser(description="LVISDataset Analysis")
    parser.add_argument("--dataset_dir", type=str, required=True, help="Path to the dataset directory.")
    parser.add_argument("--construct_dataset_dir", type=str, default='bin', help="Path to the debug bin directory.")
    parser.add_argument("--dataset_name", type=str, default='COCO', help="Dataset name.")
    parser.add_argument('--is_train', action='store_true', help="Train/Test")
    parser.add_argument('--is_build_data', action='store_true', help="Build data")
    parser.add_argument('--is_multiple', action='store_true', help="Multiple/Two objects")
    parser.add_argument("--area_ratio", type=float, default=0.01171, help="Area ratio for filtering out small objects.")
    parser.add_argument("--obj_thr", type=int, default=20, help="Object threshold for filtering.")
    parser.add_argument("--index", type=int, default=0, help="Index of the sample to test.")
    args = parser.parse_args()

    image_dirs = [
        Path(args.dataset_dir) / args.dataset_name / "train2017",
        Path(args.dataset_dir) / args.dataset_name / "val2017",
    ]

    if args.is_train: 
        json_path = Path(args.dataset_dir) / args.dataset_name / "lvis_v1/lvis_v1_train.json"
        max_num = 2000000
    else:
        json_path = Path(args.dataset_dir) / args.dataset_name / "lvis_v1/lvis_v1_val.json"
        max_num = 30000

    dataset = LVISDataset(
        construct_dataset_dir = args.construct_dataset_dir,
        obj_thr = args.obj_thr, 
        area_ratio = args.area_ratio, 
    )
    
    lvis_api = LVIS(json_path)
    img_ids = sorted(lvis_api.imgs.keys())
    imgs_info = lvis_api.load_imgs(img_ids)
    annos = [lvis_api.img_ann_map[img_id] for img_id in img_ids]

    if args.is_build_data: 
        if not args.is_multiple:
            for index in range(max_num):
                dataset._intersect_2_obj(image_dirs, lvis_api, imgs_info, annos, index)
                # dataset._intersect_3_obj(image_dirs, lvis_api, imgs_info, annos, index)
    else:
        for index in range(len(os.listdir(args.construct_dataset_dir))):
            collage = dataset._get_sample(index)