From 4940b5c0b36e5c43a4f7176de9df0654ea08f9d2 Mon Sep 17 00:00:00 2001 From: F-jie Date: Thu, 24 Feb 2022 01:07:43 +0800 Subject: [PATCH] dataPrepareOK --- ARI/ARIAnnotation.py | 1 + ARI/ARICOCO.py | 134 +++++++++++++++++++++++++++++++------------ ARI/ARIUtils.py | 6 +- Linemod.py | 4 +- main.py | 23 ++++++-- 5 files changed, 122 insertions(+), 46 deletions(-) diff --git a/ARI/ARIAnnotation.py b/ARI/ARIAnnotation.py index 988060c..79b91a6 100644 --- a/ARI/ARIAnnotation.py +++ b/ARI/ARIAnnotation.py @@ -32,6 +32,7 @@ def visBBOX2D(self): cv2.imshow("ImageWithBBOX2D", imageBOX) cv2.waitKey() cv2.destroyAllWindows() + break def visBBOX3D(self): image = cv2.imread(self.filePath) diff --git a/ARI/ARICOCO.py b/ARI/ARICOCO.py index 6bb1d90..b8c79ff 100644 --- a/ARI/ARICOCO.py +++ b/ARI/ARICOCO.py @@ -1,17 +1,29 @@ import os import cv2 import json +import random +import shutil +import numpy as np from ARI.ARIAnnotation import ARIDataset from ARI.ARIUtils import draw2DBBOX class ARICOCO(object): - def __init__(self, dataset: ARIDataset) -> None: + def __init__(self, dataset: ARIDataset, train_ratio=0.7) -> None: super().__init__() + self.train_ratio = train_ratio self.dataset = dataset + self.split() self.jsonData = self.json() + def split(self): + num_imgs = len(self.dataset.images) + numImgs_train = int(num_imgs*self.train_ratio) + imgIds = list(range(num_imgs)) + random.shuffle(imgIds) + self.trainImgIds = imgIds[:numImgs_train] + def info(self): return { "description":"This is stable 1.0 version of the 2014 MS COCO dataset.", @@ -30,36 +42,63 @@ def licenses(self): }] def images(self): - images = [] + images_train = [] + images_val = [] for id, image in enumerate(self.dataset.images): - images.append({ - "license": 0, - "file_name": "{:04d}.png".format(id), - "coco_url": "", - "height": 480, - "width": 640, - "date_captured": "", - "flickr_url": "", - "id": id - }) - return images + if id in self.trainImgIds: + images_train.append({ + "license": 0, + "file_name": "{:04d}.png".format(id), + "coco_url": "", + "height": 480, + "width": 640, + "date_captured": "", + "flickr_url": "", + "id": id + }) + else: + images_val.append({ + "license": 0, + "file_name": "{:04d}.png".format(id), + "coco_url": "", + "height": 480, + "width": 640, + "date_captured": "", + "flickr_url": "", + "id": id + }) + return images_train, images_val def annotations(self): - annotations = [] + annotations_train = [] + annotations_val = [] numberOfAnno = 0 for id, image in enumerate(self.dataset.images): - for anno in image.annotations: - annotations.append({ - "id": numberOfAnno, - "image_id": id, - "category_id": self.queryIDWithClsID(anno.cls), - "segmentation": [0, 0], - "area": anno.bbox[2] * anno.bbox[3], - "bbox": anno.bbox, # [x, y, width, height] - "iscrowd": 1, - }) - numberOfAnno = numberOfAnno + 1 - return annotations + if id in self.trainImgIds: + for anno in image.annotations: + annotations_train.append({ + "id": numberOfAnno, + "image_id": id, + "category_id": self.queryIDWithClsID(anno.cls), + "segmentation": [0, 0], + "area": anno.bbox[2] * anno.bbox[3], + "bbox": anno.bbox, # [x, y, width, height] + "iscrowd": 1, + }) + numberOfAnno = numberOfAnno + 1 + else: + for anno in image.annotations: + annotations_val.append({ + "id": numberOfAnno, + "image_id": id, + "category_id": self.queryIDWithClsID(anno.cls), + "segmentation": [0, 0], + "area": anno.bbox[2] * anno.bbox[3], + "bbox": anno.bbox, # [x, y, width, height] + "iscrowd": 1, + }) + numberOfAnno = numberOfAnno + 1 + return annotations_train, annotations_val def categories(self): categories = [] @@ -77,19 +116,31 @@ def queryIDWithClsID(self, clsID): return item['id'] def json(self): - return { + images = self.images() + annos = self.annotations() + json_train = { "info": self.info(), "licenses": self.licenses(), - "images": self.images(), - "annotations": self.annotations(), + "images": images[0], + "annotations": annos[0], "categories": self.categories() } + json_val = { + "info": self.info(), + "licenses": self.licenses(), + "images": images[1], + "annotations": annos[1], + "categories": self.categories() + } + return json_train, json_val def visbbox(self, sourceDir, annoIDs): for annoID in annoIDs: - numberOfAnno = len(self.jsonData['annotations']) - assert annoID < numberOfAnno, "annoID shold less than {}".format(numberOfAnno) - anno = self.jsonData['annotations'][annoID] + + if annoID not in list(self.jsonData[0].keys()): + continue + + anno = self.jsonData[0]['annotations'][annoID] imagePath = os.path.join(sourceDir, "{:04d}.png".format(anno['image_id'])) image = cv2.imread(imagePath) image = draw2DBBOX(image.copy(), anno['bbox']) @@ -97,7 +148,18 @@ def visbbox(self, sourceDir, annoIDs): cv2.waitKey() cv2.destroyAllWindows() - def saveCOCO(self, cocoJsonFile): - with open(cocoJsonFile, 'w') as fp: - json.dump(self.jsonData, fp) - + def saveCOCO(self, cocoPath): + annoFile_train = os.path.join(cocoPath, "annotations/instances_train2017.json") + annoFile_val = os.path.join(cocoPath, "annotations/instances_val2017.json") + with open(annoFile_train, 'w') as fp: + json.dump(self.jsonData[0], fp) + with open(annoFile_val, "w") as fp: + json.dump(self.jsonData[1], fp) + + trainDir = os.path.join(cocoPath, "train2017") + valDir = os.path.join(cocoPath, "val2017") + for i in range(len(self.dataset.images)): + if i in self.trainImgIds: + shutil.copy(self.dataset.images[i].filePath, trainDir) + else: + shutil.copy(self.dataset.images[i].filePath, valDir) diff --git a/ARI/ARIUtils.py b/ARI/ARIUtils.py index 963bb91..4122229 100644 --- a/ARI/ARIUtils.py +++ b/ARI/ARIUtils.py @@ -18,8 +18,8 @@ def transRT2Matrix4X4(R: list, t: list): def calculateKeyPoint2D(points3D, intrinsic, RT): RT_np = np.array(RT).reshape(3, 4) - points3D_np = np.array(points3D).reshape(8, 3) - pad = np.ones((8, 1)) + points3D_np = np.array(points3D).reshape(9, 3) + pad = np.ones((9, 1)) points3D_np = np.concatenate((points3D_np, pad), axis=1).T tmp = np.matmul(np.matmul(intrinsic, RT_np), points3D_np) @@ -54,7 +54,7 @@ def draw3DBBOX(image, bbox): point5 = tuple(bbox[8:10]) point6 = tuple(bbox[10:12]) point7 = tuple(bbox[12:14]) - point8 = tuple(bbox[14:]) + point8 = tuple(bbox[14:16]) color = (0, 255, 255) thickness = 1 diff --git a/Linemod.py b/Linemod.py index 91f2f5c..9bda373 100644 --- a/Linemod.py +++ b/Linemod.py @@ -43,11 +43,11 @@ def parseDataset(self, clsId): for imageId in annos.keys(): imageInfo = self.parseToStandardAnno(annos[imageId]) imageInfo.filePath = os.path.join(self.sourceDir, - "data\{:02d}\\rgb\{:04d}.png".format(clsId, imageId)) + "data/{:02d}/rgb/{:04d}.png".format(clsId, imageId)) self.images.append(imageInfo) def parsekeyPoint3D(self): - modelsInfo = load_yaml(os.path.join(self.sourceDir, "models\models_info.yml")) + modelsInfo = load_yaml(os.path.join(self.sourceDir, "models/models_info.yml")) for cls in modelsInfo.keys(): modelInfo = modelsInfo[cls] diff --git a/main.py b/main.py index d5872d2..585dab6 100644 --- a/main.py +++ b/main.py @@ -1,11 +1,24 @@ import os from Linemod import Linemod from ARI.ARICOCO import ARICOCO +import argparse if __name__ == "__main__": - sourceDir = "E:\code\PoseFuse\data\Linemod_preprocessed" - linemod = Linemod(sourceDir, "linemod") - coco = ARICOCO(linemod) + parser = argparse.ArgumentParser("Try use Transforrmer estimate 6DoF Pose", add_help=False) + parser.add_argument("--use", type=str, default="Linemod2COCO", help="Values: Linemod2COCO, ...") + args = parser.parse_args() - coco.visbbox(os.path.join(sourceDir, "data\\02\\rgb"), [190,8990,910,3763]) - coco.saveCOCO(os.path.join(sourceDir, "coco\instances_train2017.json")) + if args.use == "Linemod2COCO": + linemod_root = "/home/hnu/data/Linemod_preprocessed" + linemod_cls_2_dir = os.path.join(linemod_root, "data/02/rgb") + linemod = Linemod(linemod_root, "linemod-2") + coco = ARICOCO(linemod) + coco.visbbox(linemod_cls_2_dir, [190,8990,910,3763]) + + for i, image in enumerate(linemod.images): + image.visBBOX3D() + if i == 1: + break; + coco.saveCOCO("/home/hnu/Documents/data/") + else: + raise ValueError(f"arg {args.use} is not support!")