demo/simple_dataset.py

import torch
import torch.nn as nn
from torch.utils import data
import os
from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC
import glob
import numpy as np
import open3d as o3d
import cv2
from datasets.taxonomy import category_map as category_ids

classname_map={
    "chair":["chair","stool"],
    "cabinet":["dishwasher","cabinet","oven","refrigerator",'storage'],
    "sofa":["sofa"],
    "table":["table"],
    "bed":["bed"],
    "shelf":["shelf"]
}
classname_remap={ #map small categories to six large categories
    "chair":"chair",
    "stool":"chair",
    "dishwasher":"cabinet",
    "cabinet":"cabinet",
    "oven":"cabinet",
    "refrigerator":"cabinet",
    "storage":"cabinet",
    "sofa":"sofa",
    "table":"table",
    "bed":"bed",
    "shelf":"shelf"
}

def image_transform(n_px):
    return Compose([
        Resize(n_px, interpolation=BICUBIC),
        CenterCrop(n_px),
        ToTensor(),
        Normalize((0.48145466, 0.4578275, 0.40821073),
                 (0.26862954, 0.26130258, 0.27577711)),
    ])
class Simple_InTheWild_dataset(data.Dataset):
    def __init__(self,dataset_dir="/data1/haolin/data/real_scene_process_data",scene_id="letian-310",n_px=224):
        self.dataset_dir=dataset_dir
        self.preprocess = image_transform(n_px)
        self.image_path = []
        if scene_id=="all":
            scene_list=os.listdir(self.dataset_dir)
            for id in scene_list:
                image_folder=os.path.join(self.dataset_dir,id,"6_images")
                self.image_path+=glob.glob(image_folder+"/*/*jpg")
        else:
            image_folder = os.path.join(self.dataset_dir, scene_id, "6_images")
            self.image_path += glob.glob(image_folder + "/*/*jpg")
    def __len__(self):
        return len(self.image_path)

    def __getitem__(self,index):
        path=self.image_path[index]
        basename=os.path.basename(path)[:-4]
        model_id=path.split(os.sep)[-2]
        scene_id=path.split(os.sep)[-4]
        image=Image.open(path)
        image_tensor=self.preprocess(image)

        return {"images":image_tensor,"image_name":basename,"model_id":model_id,"scene_id":scene_id}

class InTheWild_Dataset(data.Dataset):
    def __init__(self,data_dir="/data1/haolin/data/real_scene_process_data/letian-310",scene_id="letian-310",
                 par_pc_size=2048,category="chair",max_n_imgs=5):
        self.par_pc_size=par_pc_size
        self.data_dir=data_dir
        self.category=category
        self.max_n_imgs=max_n_imgs

        self.models=[]
        category_list=classname_map[category]
        modelid_list=[]
        for cat in category_list:
            if scene_id=="all":
                scene_list=os.listdir(self.data_dir)
                for id in scene_list:
                    data_folder=os.path.join(self.data_dir,id)
                    modelid_list+=glob.glob(data_folder+"/6_images/%s*"%(cat))
            else:
                data_folder=os.path.join(self.data_dir,scene_id)
                modelid_list+=glob.glob(data_folder+"/6_images/%s*"%(cat))
        sceneid_list = [item.split("/")[-3] for item in modelid_list]
        modelid_list=[item.split("/")[-1] for item in modelid_list]
        for idx,modelid in enumerate(modelid_list):
            scene_id=sceneid_list[idx]
            image_folder=os.path.join(self.data_dir,scene_id,"6_images",modelid)
            image_list=os.listdir(image_folder)
            if len(image_list)==0:
                continue
            imageid_list=[item[0:-4] for item in image_list]
            imageid_list.sort(key=lambda x:int(x))
            partial_path=os.path.join(self.data_dir,scene_id,"5_partial_points",modelid+".ply")
            if os.path.exists(partial_path)==False: continue
            self.models+=[
                {'model_id':modelid,
                 "scene_id":scene_id,
                 "partial_path":partial_path,
                 "imageid_list":imageid_list,
                 }
            ]
    def __len__(self):
        return len(self.models)

    def __getitem__(self,idx):
        model = self.models[idx]['model_id']
        scene_id=self.models[idx]['scene_id']
        imageid_list = self.models[idx]['imageid_list']
        partial_path=self.models[idx]['partial_path']
        n_frames=min(len(imageid_list),self.max_n_imgs)
        img_indexes=np.linspace(start=0,stop=len(imageid_list)-1,num=n_frames).astype(np.int32)

        '''load partial points'''
        par_point_o3d = o3d.io.read_point_cloud(partial_path)
        par_points = np.asarray(par_point_o3d.points)
        replace = par_points.shape[0] < self.par_pc_size
        ind = np.random.default_rng().choice(par_points.shape[0], self.par_pc_size, replace=replace)
        par_points=par_points[ind]
        par_points=torch.from_numpy(par_points).float()

        '''load image features'''
        image_list=[]
        valid_frames = []
        image_namelist=[]
        for img_index in img_indexes:
            image_name = imageid_list[img_index]
            image_feat_path = os.path.join(self.data_dir,scene_id, "7_img_feature", model,image_name + '.npz')
            image = np.load(image_feat_path)["img_features"]
            image_list.append(torch.from_numpy(image).float())
            image_namelist.append(image_name)
            valid_frames.append(True)
        '''load original image'''
        org_img_list=[]
        for img_index in img_indexes:
            image_name = imageid_list[img_index]
            image_path = os.path.join(self.data_dir,scene_id, "6_images", model,image_name+".jpg")
            org_image = cv2.imread(image_path)
            org_image = cv2.resize(org_image, dsize=(224, 224), interpolation=cv2.INTER_LINEAR)
            org_img_list.append(org_image)

        '''load project matrix'''
        proj_mat_list=[]
        for img_index in img_indexes:
            image_name = imageid_list[img_index]
            proj_mat_path = os.path.join(self.data_dir,scene_id, "8_proj_matrix", model, image_name + ".npy")
            proj_mat = np.load(proj_mat_path)
            proj_mat_list.append(proj_mat)

        '''load transformation matrix'''
        tran_mat_path = os.path.join(self.data_dir,scene_id, "10_tran_matrix", model+".npy")
        tran_mat = np.load(tran_mat_path)

        '''category code, not used for category specific models'''
        category_id = category_ids[self.category]
        one_hot = torch.zeros((6)).float()
        one_hot[category_id] = 1.0

        ret_dict={
            "model_id":model,
            "scene_id":scene_id,
            "par_points":par_points,
            "proj_mat":torch.stack([torch.from_numpy(mat) for mat in proj_mat_list], dim=0),
            "tran_mat":torch.from_numpy(tran_mat).float(),
            "image":torch.stack(image_list,dim=0),
            "org_image":org_img_list,
            "valid_frames":torch.tensor(valid_frames).bool(),
            "category_ids": category_id,
            "category_code":one_hot,
        }
        return ret_dict