Skip to content

Commit

Permalink
dataPrepareOK
Browse files Browse the repository at this point in the history
  • Loading branch information
F-jie committed Feb 23, 2022
1 parent fa98e68 commit 4940b5c
Show file tree
Hide file tree
Showing 5 changed files with 122 additions and 46 deletions.
1 change: 1 addition & 0 deletions ARI/ARIAnnotation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ def visBBOX2D(self):
cv2.imshow("ImageWithBBOX2D", imageBOX)
cv2.waitKey()
cv2.destroyAllWindows()
break

def visBBOX3D(self):
image = cv2.imread(self.filePath)
Expand Down
134 changes: 98 additions & 36 deletions ARI/ARICOCO.py
View file
Original file line number Diff line number Diff line change
@@ -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.",
Expand All @@ -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 = []
Expand All @@ -77,27 +116,50 @@ 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'])
cv2.imshow("ImageWithBBOX2D", image)
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)
6 changes: 3 additions & 3 deletions ARI/ARIUtils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions Linemod.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
23 changes: 18 additions & 5 deletions main.py
View file
Original file line number Diff line number Diff line change
@@ -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!")

0 comments on commit 4940b5c

Please sign in to comment.