From fbefd891430a940e11bce42492f651dab399e738 Mon Sep 17 00:00:00 2001 From: Ethan Suleman <70169323+esuleman@users.noreply.github.com> Date: Fri, 15 Oct 2021 12:06:50 -0400 Subject: [PATCH 01/14] Create shapenet_tfrecord_gen.py --- .../mesh_rcnn/data/shapenet_tfrecord_gen.py | 125 ++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py new file mode 100644 index 00000000000..49452f46218 --- /dev/null +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py @@ -0,0 +1,125 @@ +from absl import app # pylint:disable=unused-import +from absl import flags +import numpy as np + +import tensorflow as tf + +import multiprocessing as mp +from official.vision.beta.data import tfrecord_lib +from research.object_detection.utils import dataset_util +from official.vision.beta.data.tfrecord_lib import convert_to_feature + +flags.DEFINE_multi_string('shapenet_dir', '', 'Directory containing ShapeNet.') +flags.DEFINE_string('output_file_prefix', '/tmp/train', 'Path to output file') +flags.DEFINE_integer('num_shards', 32, 'Number of shards for output file.') + +FLAGS = flags.FLAGS + +logger = tf.get_logger() +logger.setLevel(logging.INFO) + +def create_tf_example(image): + with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["img"]), 'rb') as fid: + encoded_img_jpg = fid.read() + + img_width, img_height = image["img_size"] + img_filename = image["img"] + img_category = image["category"] + keypoints_2d = image["2d_keypoints"] + + feature_dict = {"img/height": convert_to_feature(img_height), + "img/width": convert_to_feature(img_width), + "img/category": convert_to_feature(img_category), + "img/filename": convert_to_feature(img_filename), + "img/encoded": convert_to_feature(encoded_img_jpg), + "img/2d_keypoints": convert_to_feature(keypoints_2d)} + + with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["mask"]), 'rb') as fid: + encoded_mask_jpg = fid.read() + + feature_dict.update({"mask": convert_to_feature(encoded_mask_jpg)}) + + with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["model"]), 'rb') as fid: + encoded_model = fid.read() + + with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["3d_keypoints"]), 'rb') as fid: + keypoints_3d = fid.read() + + model_raw = image["model_raw"] + model_source = image["model_source"] + + feature_dict.update({"model": convert_to_feature(encoded_model), + "model/raw": convert_to_feature(model_raw), + "model/source": convert_to_feature(model_source), + "model/3d_keypoints": convert_to_feature(keypoints_3d)}) + + with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["voxel"]), 'rb') as fid: + encoded_voxel = fid.read() + + rot_mat = image["rot_mat"] + trans_mat = image["trans_mat"] + focal_length = image["focal_length"] + cam_position = image["cam_positon"] + inplane_rotation = image["inplane_rotation"] + truncated = image["truncated"] + occluded = image["occluded"] + slightly_occluded = image["slightly_occluded"] + bbox = image["bbox"] + + feature_dict.update({"voxel": convert_to_feature(encoded_voxel), + "rot_mat": convert_to_feature(rot_mat), + "trans_mat": convert_to_feature(trans_mat), + "focal_length": convert_to_feature(focal_length), + "cam_position": convert_to_feature(cam_position), + "inplane_rotation": convert_to_feature(inplane_rotation), + "truncated": convert_to_feature(truncated), + "occluded": convert_to_feature(occluded), + "slightly_occluded": convert_to_feature(slightly_occluded), + "bbox": convert_to_feature(bbox)}) + + example = tf.train.Example( + features=tf.train.Features(feature=feature_dict)) + + return example, 0 + +def generate_annotations(images, shapenet_dir): + for image in images: + yield {"img": image["img"], "category": image["category"], "img_size": image["img_size"], "2d_keypoints": image["2d_keypoints"], + "mask": image["mask"], "img_source": image["img_source"], "model": image["model"], "model_raw": image["model_raw"], + "model_source": image["model_source"], "3d_keypoints": image["3d_keypoints"], "voxel": image["voxel"], "rot_mat": image["rot_mat"], + "trans_mat": image["trans_mat"], "focal_length": image["focal_length"], "cam_position": image["cam_position"], + "inplane_rotation": image["inplane_rotation"], "truncated": image["truncated"], "occluded": image["occluded"], + "slightly_ocluded": image["slightly_occluded"], "bbox": image["bbox"], "shapenet_dir": shapenet_dir} + +def _create_tf_record_from_shapenet_dir(shapenet_dir, + output_path, + num_shards): + """Loads Shapenet json files and converts to tf.Record format. + Args: + images_info_file: shapenet_dir download directory + output_path: Path to output tf.Record file. + num_shards: Number of output files to create. + """ + + logging.info('writing to output path: %s', output_path) + + images = json.load(open(os.join(shapenet_dir, "taxonomy.json"))) + + shapenet_annotations_iter = generate_annotations( + images=images, shapenet_dir=shapenet_dir) + + num_skipped = tfrecord_lib.write_tf_record_dataset( + output_path, shapenet_annotations_iter, create_tf_example, num_shards) + + logging.info('Finished writing, skipped %d annotations.', num_skipped) + +def main(_): + assert FLAGS.shapenet_dir, '`shapenet_dir` missing.' + + directory = os.path.dirname(FLAGS.output_file_prefix) + if not tf.io.gfile.isdir(directory): + tf.io.gfile.makedirs(directory) + + +if __name__ == '__main__': + app.run(main) From f52e78eb8ab3ab284ecf5df6ec5ab40240ae57dd Mon Sep 17 00:00:00 2001 From: Abhirakshak <37345811+Abhirakshak@users.noreply.github.com> Date: Fri, 29 Oct 2021 14:48:17 -0400 Subject: [PATCH 02/14] updated with shapenet annotations --- .../mesh_rcnn/data/shapenet_tfrecord_gen.py | 103 +++++++----------- 1 file changed, 37 insertions(+), 66 deletions(-) diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py index 49452f46218..a165050910d 100644 --- a/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py @@ -1,13 +1,15 @@ from absl import app # pylint:disable=unused-import from absl import flags import numpy as np - +import json import tensorflow as tf +import os import multiprocessing as mp from official.vision.beta.data import tfrecord_lib from research.object_detection.utils import dataset_util from official.vision.beta.data.tfrecord_lib import convert_to_feature +from create_pix3d_tf_record import parse_obj_file flags.DEFINE_multi_string('shapenet_dir', '', 'Directory containing ShapeNet.') flags.DEFINE_string('output_file_prefix', '/tmp/train', 'Path to output file') @@ -18,65 +20,21 @@ logger = tf.get_logger() logger.setLevel(logging.INFO) + def create_tf_example(image): - with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["img"]), 'rb') as fid: - encoded_img_jpg = fid.read() - - img_width, img_height = image["img_size"] - img_filename = image["img"] - img_category = image["category"] - keypoints_2d = image["2d_keypoints"] - - feature_dict = {"img/height": convert_to_feature(img_height), - "img/width": convert_to_feature(img_width), - "img/category": convert_to_feature(img_category), - "img/filename": convert_to_feature(img_filename), - "img/encoded": convert_to_feature(encoded_img_jpg), - "img/2d_keypoints": convert_to_feature(keypoints_2d)} - - with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["mask"]), 'rb') as fid: - encoded_mask_jpg = fid.read() - - feature_dict.update({"mask": convert_to_feature(encoded_mask_jpg)}) - - with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["model"]), 'rb') as fid: - encoded_model = fid.read() - - with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["3d_keypoints"]), 'rb') as fid: - keypoints_3d = fid.read() - - model_raw = image["model_raw"] - model_source = image["model_source"] - - feature_dict.update({"model": convert_to_feature(encoded_model), - "model/raw": convert_to_feature(model_raw), - "model/source": convert_to_feature(model_source), - "model/3d_keypoints": convert_to_feature(keypoints_3d)}) - - with tf.io.gfile.GFile(os.join(image["shapenet_dir"], image["voxel"]), 'rb') as fid: - encoded_voxel = fid.read() - - rot_mat = image["rot_mat"] - trans_mat = image["trans_mat"] - focal_length = image["focal_length"] - cam_position = image["cam_positon"] - inplane_rotation = image["inplane_rotation"] - truncated = image["truncated"] - occluded = image["occluded"] - slightly_occluded = image["slightly_occluded"] - bbox = image["bbox"] - - feature_dict.update({"voxel": convert_to_feature(encoded_voxel), - "rot_mat": convert_to_feature(rot_mat), - "trans_mat": convert_to_feature(trans_mat), - "focal_length": convert_to_feature(focal_length), - "cam_position": convert_to_feature(cam_position), - "inplane_rotation": convert_to_feature(inplane_rotation), - "truncated": convert_to_feature(truncated), - "occluded": convert_to_feature(occluded), - "slightly_occluded": convert_to_feature(slightly_occluded), - "bbox": convert_to_feature(bbox)}) - + + model_id = image["model_id"] + label = image["label"] + + temp_file_dir = os.join(image["shapenet_dir"], image["synset_id"]) + model_vertices, model_faces = parse_obj_file(os.join(temp_file_dir, image["model_id"])) + + feature_dict = {"model_id": convert_to_feature(model_id), + "label": convert_to_feature(label), + "vertices": convert_to_feature(model_vertices), + "faces": convert_to_feature(model_faces)} + + example = tf.train.Example( features=tf.train.Features(feature=feature_dict)) @@ -84,12 +42,10 @@ def create_tf_example(image): def generate_annotations(images, shapenet_dir): for image in images: - yield {"img": image["img"], "category": image["category"], "img_size": image["img_size"], "2d_keypoints": image["2d_keypoints"], - "mask": image["mask"], "img_source": image["img_source"], "model": image["model"], "model_raw": image["model_raw"], - "model_source": image["model_source"], "3d_keypoints": image["3d_keypoints"], "voxel": image["voxel"], "rot_mat": image["rot_mat"], - "trans_mat": image["trans_mat"], "focal_length": image["focal_length"], "cam_position": image["cam_position"], - "inplane_rotation": image["inplane_rotation"], "truncated": image["truncated"], "occluded": image["occluded"], - "slightly_ocluded": image["slightly_occluded"], "bbox": image["bbox"], "shapenet_dir": shapenet_dir} + yield {"shapenet_dir": shapenet_dir, + "label": image["label"], + "model_id": image["model_id"], + "synset_id": image["synset_id"]} def _create_tf_record_from_shapenet_dir(shapenet_dir, output_path, @@ -103,7 +59,22 @@ def _create_tf_record_from_shapenet_dir(shapenet_dir, logging.info('writing to output path: %s', output_path) - images = json.load(open(os.join(shapenet_dir, "taxonomy.json"))) + #create synset ID to label mapping dictionary + with open(shapenet_synset_dict.json, "r") as dict_file: + synset_dict = json.load(dict_file) + + #images list + images = [] + + for _, synset_directories, _ in os.walk(shapenet_dir): + for synset_directory in synset_directories: + for _, object_directories, _ in os.walk(os.path.join(shapenet_dir, synset_directory)): + for object_directory in object_directories: + image = {"model_id" : object_directory, + "label" : synset_dict[synset_directory], + "shapenet_dir": shapenet_dir, + "synset_id": synset_directory} + images.append(image) shapenet_annotations_iter = generate_annotations( images=images, shapenet_dir=shapenet_dir) From 4c99fe49fc0c56b8d3129256f6a008e24da248de Mon Sep 17 00:00:00 2001 From: Abhirakshak <37345811+Abhirakshak@users.noreply.github.com> Date: Fri, 5 Nov 2021 13:16:31 -0400 Subject: [PATCH 03/14] Delete create_pix3d_tf_record.py --- .../mesh_rcnn/data/create_pix3d_tf_record.py | 179 ------------------ 1 file changed, 179 deletions(-) delete mode 100644 official/vision/beta/projects/mesh_rcnn/data/create_pix3d_tf_record.py diff --git a/official/vision/beta/projects/mesh_rcnn/data/create_pix3d_tf_record.py b/official/vision/beta/projects/mesh_rcnn/data/create_pix3d_tf_record.py deleted file mode 100644 index 249de1bca8d..00000000000 --- a/official/vision/beta/projects/mesh_rcnn/data/create_pix3d_tf_record.py +++ /dev/null @@ -1,179 +0,0 @@ -# Copyright 2021 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# Author: Jacob Zietek - -# https://github.com/PurdueDualityLab/tf-models/blob/master/official/vision/beta/data/create_coco_tf_record.py reference - -r"""Convert raw Pix3D dataset to TFRecord format. -This scripts follows the label map decoder format and supports detection -boxes, instance masks and captions. -Example usage: - python create_pix3d_tf_record.py --logtostderr \ - --pix3d_dir="${TRAIN_IMAGE_DIR}" \ - --output_file_prefix="${OUTPUT_DIR/FILE_PREFIX}" \ - --num_shards=100 -""" - -import collections -import json -import logging -import os -import json -from re import I - -from absl import app # pylint:disable=unused-import -from absl import flags -import numpy as np - -import tensorflow as tf - -import multiprocessing as mp -from official.vision.beta.data import tfrecord_lib -from research.object_detection.utils import dataset_util -from official.vision.beta.data.tfrecord_lib import convert_to_feature - -flags.DEFINE_multi_string('pix3d_dir', '', 'Directory containing Pix3d.') -flags.DEFINE_string('output_file_prefix', '/tmp/train', 'Path to output file') -flags.DEFINE_integer('num_shards', 32, 'Number of shards for output file.') - -FLAGS = flags.FLAGS - -logger = tf.get_logger() -logger.setLevel(logging.INFO) - - -def create_tf_example(image): - """Converts image and annotations to a tf.Example proto. - Args: - image: dict with keys: [img, category, img_size, 2d_keypoints, mask, img_source, model, - model_raw, model_source, 3d_keypoints, voxel, rot_mat, trans_mat, - focal_length, cam_position, inplane_rotation, truncated, occluded, - slightly_occluded, bbox] - Returns: - example: The converted tf.Example - num_annotations_skipped: Number of (invalid) annotations that were ignored. - Raises: - ValueError: if the image is not able to be found. This indicates the file structure - of the Pix3D folder is incorrect. - """ - - with tf.io.gfile.GFile(os.join(image["pix3d_dir"], image["img"]), 'rb') as fid: - encoded_img_jpg = fid.read() - - img_width, img_height = image["img_size"] - img_filename = image["img"] - img_category = image["category"] - keypoints_2d = image["2d_keypoints"] - - feature_dict = {"img/height": convert_to_feature(img_height), - "img/width": convert_to_feature(img_width), - "img/category": convert_to_feature(img_category), - "img/filename": convert_to_feature(img_filename), - "img/encoded": convert_to_feature(encoded_img_jpg), - "img/2d_keypoints": convert_to_feature(keypoints_2d)} - - with tf.io.gfile.GFile(os.join(image["pix3d_dir"], image["mask"]), 'rb') as fid: - encoded_mask_jpg = fid.read() - - feature_dict.update({"mask": convert_to_feature(encoded_mask_jpg)}) - - with tf.io.gfile.GFile(os.join(image["pix3d_dir"], image["model"]), 'rb') as fid: - encoded_model = fid.read() - - with tf.io.gfile.GFile(os.join(image["pix3d_dir"], image["3d_keypoints"]), 'rb') as fid: - keypoints_3d = fid.read() - - model_raw = image["model_raw"] - model_source = image["model_source"] - - feature_dict.update({"model": convert_to_feature(encoded_model), - "model/raw": convert_to_feature(model_raw), - "model/source": convert_to_feature(model_source), - "model/3d_keypoints": convert_to_feature(keypoints_3d)}) - - with tf.io.gfile.GFile(os.join(image["pix3d_dir"], image["voxel"]), 'rb') as fid: - encoded_voxel = fid.read() - - rot_mat = image["rot_mat"] - trans_mat = image["trans_mat"] - focal_length = image["focal_length"] - cam_position = image["cam_positon"] - inplane_rotation = image["inplane_rotation"] - truncated = image["truncated"] - occluded = image["occluded"] - slightly_occluded = image["slightly_occluded"] - bbox = image["bbox"] - - # Where are these supposed to be categorized??? - feature_dict.update({"voxel": convert_to_feature(encoded_voxel), - "rot_mat": convert_to_feature(rot_mat), - "trans_mat": convert_to_feature(trans_mat), - "focal_length": convert_to_feature(focal_length), - "cam_position": convert_to_feature(cam_position), - "inplane_rotation": convert_to_feature(inplane_rotation), - "truncated": convert_to_feature(truncated), - "occluded": convert_to_feature(occluded), - "slightly_occluded": convert_to_feature(slightly_occluded), - "bbox": convert_to_feature(bbox)}) - - example = tf.train.Example( - features=tf.train.Features(feature=feature_dict)) - - return example, 0 - - -def generate_annotations(images, pix3d_dir): - """Generator for Pix3D annotations.""" - for image in images: - yield {"img": image["img"], "category": image["category"], "img_size": image["img_size"], "2d_keypoints": image["2d_keypoints"], - "mask": image["mask"], "img_source": image["img_source"], "model": image["model"], "model_raw": image["model_raw"], - "model_source": image["model_source"], "3d_keypoints": image["3d_keypoints"], "voxel": image["voxel"], "rot_mat": image["rot_mat"], - "trans_mat": image["trans_mat"], "focal_length": image["focal_length"], "cam_position": image["cam_position"], - "inplane_rotation": image["inplane_rotation"], "truncated": image["truncated"], "occluded": image["occluded"], - "slightly_ocluded": image["slightly_occluded"], "bbox": image["bbox"], "pix3d_dir": pix3d_dir} - - -def _create_tf_record_from_pix3d_dir(pix3d_dir, - output_path, - num_shards): - """Loads Pix3D json files and converts to tf.Record format. - Args: - images_info_file: pix3d_dir download directory - output_path: Path to output tf.Record file. - num_shards: Number of output files to create. - """ - - logging.info('writing to output path: %s', output_path) - - images = json.load(open(os.join(pix3d_dir, "pix3d.json"))) - - pix3d_annotations_iter = generate_annotations( - images=images, pix3d_dir=pix3d_dir) - - num_skipped = tfrecord_lib.write_tf_record_dataset( - output_path, pix3d_annotations_iter, create_tf_example, num_shards) - - logging.info('Finished writing, skipped %d annotations.', num_skipped) - - -def main(_): - assert FLAGS.pix3d_dir, '`pix3d_dir` missing.' - - directory = os.path.dirname(FLAGS.output_file_prefix) - if not tf.io.gfile.isdir(directory): - tf.io.gfile.makedirs(directory) - - -if __name__ == '__main__': - app.run(main) From 43d5e806ba7dab9bfcd6e59eff492f23bc0c46d4 Mon Sep 17 00:00:00 2001 From: Abhirakshak <37345811+Abhirakshak@users.noreply.github.com> Date: Fri, 5 Nov 2021 13:18:03 -0400 Subject: [PATCH 04/14] adding parse_obj function --- .../mesh_rcnn/data/shapenet_tfrecord_gen.py | 37 ++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py index a165050910d..4d8dad035c6 100644 --- a/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py @@ -9,7 +9,6 @@ from official.vision.beta.data import tfrecord_lib from research.object_detection.utils import dataset_util from official.vision.beta.data.tfrecord_lib import convert_to_feature -from create_pix3d_tf_record import parse_obj_file flags.DEFINE_multi_string('shapenet_dir', '', 'Directory containing ShapeNet.') flags.DEFINE_string('output_file_prefix', '/tmp/train', 'Path to output file') @@ -20,6 +19,42 @@ logger = tf.get_logger() logger.setLevel(logging.INFO) +def parse_obj_file(file): + """ + Parses relevant data out of a .obj file. This contains all of the model information. + Args: + file: file path to .obj file + Return: + vertices: vertices of object + faces: faces of object + """ + vertices = [] + faces = [] + + obj_file = open(file, 'r') + lines = obj_file.readlines() + + for line in lines: + lineID = line[0:2] + + if lineID == "v ": + vertex = line[2:].split(" ") + + for i, v in enumerate(vertex): + vertex[i] = float(v) + + vertices.append(vertex) + + if lineID == "f ": + + face = line[2:].split(" ") + + for i, f in enumerate(face): + face[i] = [int(x) - 1 for x in f.split("/")] + + faces.append(face) + + return vertices, faces def create_tf_example(image): From f555f861ffead2793020648a139789272a052eb6 Mon Sep 17 00:00:00 2001 From: Abhirakshak <37345811+Abhirakshak@users.noreply.github.com> Date: Fri, 5 Nov 2021 13:24:15 -0400 Subject: [PATCH 05/14] adding labels dictionary file --- .../mesh_rcnn/data/shapenet_synset_dict.json | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 official/vision/beta/projects/mesh_rcnn/data/Users/abhirakshak10/Documents/Purdue/VIP/tfmodels_repo/tf-models/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json diff --git a/official/vision/beta/projects/mesh_rcnn/data/Users/abhirakshak10/Documents/Purdue/VIP/tfmodels_repo/tf-models/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json b/official/vision/beta/projects/mesh_rcnn/data/Users/abhirakshak10/Documents/Purdue/VIP/tfmodels_repo/tf-models/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json new file mode 100644 index 00000000000..b2fc62ae621 --- /dev/null +++ b/official/vision/beta/projects/mesh_rcnn/data/Users/abhirakshak10/Documents/Purdue/VIP/tfmodels_repo/tf-models/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json @@ -0,0 +1,59 @@ +{ + "04379243": "table", + "02958343": "car", + "03001627": "chair", + "02691156": "airplane", + "04256520": "sofa", + "04090263": "rifle", + "03636649": "lamp", + "04530566": "watercraft", + "02828884": "bench", + "03691459": "loudspeaker", + "02933112": "cabinet", + "03211117": "display", + "04401088": "telephone", + "02924116": "bus", + "02808440": "bathtub", + "03467517": "guitar", + "03325088": "faucet", + "03046257": "clock", + "03991062": "flowerpot", + "03593526": "jar", + "02876657": "bottle", + "02871439": "bookshelf", + "03642806": "laptop", + "03624134": "knife", + "04468005": "train", + "02747177": "trash bin", + "03790512": "motorbike", + "03948459": "pistol", + "03337140": "file cabinet", + "02818832": "bed", + "03928116": "piano", + "04330267": "stove", + "03797390": "mug", + "02880940": "bowl", + "04554684": "washer", + "04004475": "printer", + "03513137": "helmet", + "03761084": "microwaves", + "04225987": "skateboard", + "04460130": "tower", + "02942699": "camera", + "02801938": "basket", + "02946921": "can", + "03938244": "pillow", + "03710193": "mailbox", + "03207941": "dishwasher", + "04099429": "rocket", + "02773838": "bag", + "02843684": "birdhouse", + "03261776": "earphone", + "03759954": "microphone", + "04074963": "remote", + "03085013": "keyboard", + "02834778": "bicycle", + "02954340": "cap", + "02858304": "boat", + "02992529": "mobile phone" +} From 192032e14af8fc5548e86874d137ef1621cde2ac Mon Sep 17 00:00:00 2001 From: Abhirakshak <37345811+Abhirakshak@users.noreply.github.com> Date: Fri, 5 Nov 2021 13:25:26 -0400 Subject: [PATCH 06/14] Delete official/vision/beta/projects/mesh_rcnn/data/Users/abhirakshak10/Documents/Purdue/VIP/tfmodels_repo/tf-models/official/vision/beta/projects/mesh_rcnn/data directory --- .../mesh_rcnn/data/shapenet_synset_dict.json | 59 ------------------- 1 file changed, 59 deletions(-) delete mode 100644 official/vision/beta/projects/mesh_rcnn/data/Users/abhirakshak10/Documents/Purdue/VIP/tfmodels_repo/tf-models/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json diff --git a/official/vision/beta/projects/mesh_rcnn/data/Users/abhirakshak10/Documents/Purdue/VIP/tfmodels_repo/tf-models/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json b/official/vision/beta/projects/mesh_rcnn/data/Users/abhirakshak10/Documents/Purdue/VIP/tfmodels_repo/tf-models/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json deleted file mode 100644 index b2fc62ae621..00000000000 --- a/official/vision/beta/projects/mesh_rcnn/data/Users/abhirakshak10/Documents/Purdue/VIP/tfmodels_repo/tf-models/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json +++ /dev/null @@ -1,59 +0,0 @@ -{ - "04379243": "table", - "02958343": "car", - "03001627": "chair", - "02691156": "airplane", - "04256520": "sofa", - "04090263": "rifle", - "03636649": "lamp", - "04530566": "watercraft", - "02828884": "bench", - "03691459": "loudspeaker", - "02933112": "cabinet", - "03211117": "display", - "04401088": "telephone", - "02924116": "bus", - "02808440": "bathtub", - "03467517": "guitar", - "03325088": "faucet", - "03046257": "clock", - "03991062": "flowerpot", - "03593526": "jar", - "02876657": "bottle", - "02871439": "bookshelf", - "03642806": "laptop", - "03624134": "knife", - "04468005": "train", - "02747177": "trash bin", - "03790512": "motorbike", - "03948459": "pistol", - "03337140": "file cabinet", - "02818832": "bed", - "03928116": "piano", - "04330267": "stove", - "03797390": "mug", - "02880940": "bowl", - "04554684": "washer", - "04004475": "printer", - "03513137": "helmet", - "03761084": "microwaves", - "04225987": "skateboard", - "04460130": "tower", - "02942699": "camera", - "02801938": "basket", - "02946921": "can", - "03938244": "pillow", - "03710193": "mailbox", - "03207941": "dishwasher", - "04099429": "rocket", - "02773838": "bag", - "02843684": "birdhouse", - "03261776": "earphone", - "03759954": "microphone", - "04074963": "remote", - "03085013": "keyboard", - "02834778": "bicycle", - "02954340": "cap", - "02858304": "boat", - "02992529": "mobile phone" -} From 38903c8c263f079106b947ac77b1ef5efa96c131 Mon Sep 17 00:00:00 2001 From: Abhirakshak <37345811+Abhirakshak@users.noreply.github.com> Date: Fri, 5 Nov 2021 13:26:40 -0400 Subject: [PATCH 07/14] Adding labels dictionary .json file --- .../mesh_rcnn/data/shapenet_synset_dict.json | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json b/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json new file mode 100644 index 00000000000..b2fc62ae621 --- /dev/null +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_synset_dict.json @@ -0,0 +1,59 @@ +{ + "04379243": "table", + "02958343": "car", + "03001627": "chair", + "02691156": "airplane", + "04256520": "sofa", + "04090263": "rifle", + "03636649": "lamp", + "04530566": "watercraft", + "02828884": "bench", + "03691459": "loudspeaker", + "02933112": "cabinet", + "03211117": "display", + "04401088": "telephone", + "02924116": "bus", + "02808440": "bathtub", + "03467517": "guitar", + "03325088": "faucet", + "03046257": "clock", + "03991062": "flowerpot", + "03593526": "jar", + "02876657": "bottle", + "02871439": "bookshelf", + "03642806": "laptop", + "03624134": "knife", + "04468005": "train", + "02747177": "trash bin", + "03790512": "motorbike", + "03948459": "pistol", + "03337140": "file cabinet", + "02818832": "bed", + "03928116": "piano", + "04330267": "stove", + "03797390": "mug", + "02880940": "bowl", + "04554684": "washer", + "04004475": "printer", + "03513137": "helmet", + "03761084": "microwaves", + "04225987": "skateboard", + "04460130": "tower", + "02942699": "camera", + "02801938": "basket", + "02946921": "can", + "03938244": "pillow", + "03710193": "mailbox", + "03207941": "dishwasher", + "04099429": "rocket", + "02773838": "bag", + "02843684": "birdhouse", + "03261776": "earphone", + "03759954": "microphone", + "04074963": "remote", + "03085013": "keyboard", + "02834778": "bicycle", + "02954340": "cap", + "02858304": "boat", + "02992529": "mobile phone" +} From b4819d0a1da063e8289ee6b7df5cc315b108f3e6 Mon Sep 17 00:00:00 2001 From: Ethan Suleman <70169323+esuleman@users.noreply.github.com> Date: Tue, 9 Nov 2021 02:19:20 -0500 Subject: [PATCH 08/14] added shapenet visualizer --- .../mesh_rcnn/data/shapenet_visualizer.py | 141 ++++++++++++++++++ 1 file changed, 141 insertions(+) create mode 100644 official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py new file mode 100644 index 00000000000..11a82e2b3fe --- /dev/null +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py @@ -0,0 +1,141 @@ +# Copyright 2021 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Author: Jacob Zietek + +r"""Takes TFRecord containing ShapeNet dataset and visualizes it. +Example usage: + python shapenet_visualizer.py --logtostderr \ + --num_models=2 \ + --shapenet_dir="${SHAPENET_TFRECORD_DIR}" \ + --output_file_prefix="${OUTPUT_DIR/FILE_PREFIX}" +""" + +import json +import logging +import os +import json + +from absl import app # pylint:disable=unused-import +from absl import flags +import numpy as np +from numpy.core.defchararray import encode +from numpy.lib.arraysetops import isin +from numpy.lib.type_check import imag +import scipy.io as sio +import cv2 + +import tensorflow as tf +import itertools +import multiprocessing as mp + +flags.DEFINE_multi_string('shapenet_dir', '', 'Directory containing ShapeNet.') +flags.DEFINE_string('output_file_prefix', '/tmp/output', 'Path to output files') +flags.DEFINE_integer('num_models', 2, 'Number of models rebuilt from TFRecord.') + +FLAGS = flags.FLAGS + +logger = tf.get_logger() +logger.setLevel(logging.INFO) + + +def write_obj_file(vertices, + faces, + filename): + """Writes a new .obj file from data. + Args: + verticies: List of vertices + faces: List of faces + filename: Filename to write .obj file to + """ + logging.info(f"Logging file {filename}") + with open(filename, 'w') as f: + for vertex in vertices: + print("v " + ' '.join(map(str, vertex)), file=f) + + for face in faces: + ret = "f " + + for coordinate in face: + ret += str(coordinate[0]) + " " + + print(ret, file=f) + + +def write_masked_image(mask, + image, + filename): + """Writes a new .png file from data. + Args: + mask: 2d array of mask data + image: 3d array of image data + filename: Filename to write .png file to + """ + logging.info(f"Logging file {filename}") + dim1 = len(mask) + dim2 = len(mask[0]) + + for i in range(dim1): + for j in range(dim2): + if mask[i][j] > 0: # add green mask to image + image[i][j][1] += 50 + + cv2.imwrite(filename, np.array(image)) + + +def visualize_tf_record(shapenet_dir, + output_path, + num_models): + """Visualizes shapenet data in TFRecord format. + Args: + shapenet_dir: shapenet_dir for TFRecords + output_path: Path to output .obj files. + num_models: Number of output files to create. + """ + + logging.info( + f"Starting to visualize {num_models} models from the TFRecords in directory {shapenet_dir} into {output_path}.") + + filenames = [os.path.join(shapenet_dir, x) for x in os.listdir(shapenet_dir)] + + raw_dataset = tf.data.TFRecordDataset(filenames) + + for raw_record in raw_dataset.take(num_models): + example = tf.train.Example() + example.ParseFromString(raw_record.numpy()) + features = example.features.feature + vertices = tf.io.parse_tensor(features["model/vertices"].bytes_list.value[0], tf.float32).numpy().tolist() + faces = tf.io.parse_tensor(features["model/faces"].bytes_list.value[0], tf.int32).numpy().tolist() + mask = tf.io.parse_tensor(features["mask"].bytes_list.value[0], tf.int32).numpy().tolist() + image = tf.io.parse_tensor(features["img/encoded"].bytes_list.value[0], tf.int32).numpy().tolist() + + filename = str(features["img/filename"].bytes_list.value[0]).split("/")[2][:-1] + filename = filename.split(".")[0] + filename = os.path.join(output_path, filename) + + write_obj_file(vertices, faces, filename + ".obj") + write_masked_image(mask, image, filename + ".png") + + +def main(_): + assert FLAGS.shapenet_dir, '`shapenet_dir` missing.' + + directory = os.path.dirname(FLAGS.output_file_prefix) + if not tf.io.gfile.isdir(directory): + tf.io.gfile.makedirs(directory) + + visualize_tf_record(FLAGS.shapenet_dir[0], FLAGS.output_file_prefix, FLAGS.num_models) + + +if __name__ == '__main__': + app.run(main) From f56176bff2d697989450237ce0930c96188e962c Mon Sep 17 00:00:00 2001 From: Ethan Suleman <70169323+esuleman@users.noreply.github.com> Date: Tue, 9 Nov 2021 02:20:54 -0500 Subject: [PATCH 09/14] Update shapenet_visualizer.py --- .../mesh_rcnn/data/shapenet_visualizer.py | 23 +------------------ 1 file changed, 1 insertion(+), 22 deletions(-) diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py index 11a82e2b3fe..4296fe71c47 100644 --- a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py @@ -70,28 +70,7 @@ def write_obj_file(vertices, ret += str(coordinate[0]) + " " print(ret, file=f) - - -def write_masked_image(mask, - image, - filename): - """Writes a new .png file from data. - Args: - mask: 2d array of mask data - image: 3d array of image data - filename: Filename to write .png file to - """ - logging.info(f"Logging file {filename}") - dim1 = len(mask) - dim2 = len(mask[0]) - - for i in range(dim1): - for j in range(dim2): - if mask[i][j] > 0: # add green mask to image - image[i][j][1] += 50 - - cv2.imwrite(filename, np.array(image)) - + def visualize_tf_record(shapenet_dir, output_path, From 3e105eb4ca50320874ce67772b19cc016f7c102b Mon Sep 17 00:00:00 2001 From: Ethan Suleman <70169323+esuleman@users.noreply.github.com> Date: Tue, 9 Nov 2021 02:21:11 -0500 Subject: [PATCH 10/14] Update shapenet_visualizer.py --- .../vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py index 4296fe71c47..c994b676d04 100644 --- a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# Author: Jacob Zietek +# Author: JEthan Suleman & Abhi Raja r"""Takes TFRecord containing ShapeNet dataset and visualizes it. Example usage: From 5fe4e326e9e2c56fe5174aecab60441035f72180 Mon Sep 17 00:00:00 2001 From: Ethan Suleman <70169323+esuleman@users.noreply.github.com> Date: Tue, 9 Nov 2021 02:21:21 -0500 Subject: [PATCH 11/14] Update shapenet_visualizer.py --- .../vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py index c994b676d04..6c4143cef5c 100644 --- a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# Author: JEthan Suleman & Abhi Raja +# Author: Ethan Suleman & Abhi Raja r"""Takes TFRecord containing ShapeNet dataset and visualizes it. Example usage: From d15c1f4fdf6ab860c164d0dc2efe64e64d4cd24b Mon Sep 17 00:00:00 2001 From: Abhirakshak <37345811+Abhirakshak@users.noreply.github.com> Date: Tue, 9 Nov 2021 02:23:11 -0500 Subject: [PATCH 12/14] Update shapenet_visualizer.py --- .../beta/projects/mesh_rcnn/data/shapenet_visualizer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py index 6c4143cef5c..8a2f9d1988d 100644 --- a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py @@ -93,8 +93,8 @@ def visualize_tf_record(shapenet_dir, example = tf.train.Example() example.ParseFromString(raw_record.numpy()) features = example.features.feature - vertices = tf.io.parse_tensor(features["model/vertices"].bytes_list.value[0], tf.float32).numpy().tolist() - faces = tf.io.parse_tensor(features["model/faces"].bytes_list.value[0], tf.int32).numpy().tolist() + vertices = tf.io.parse_tensor(features["vertices"].bytes_list.value[0], tf.float32).numpy().tolist() + faces = tf.io.parse_tensor(features["faces"].bytes_list.value[0], tf.int32).numpy().tolist() mask = tf.io.parse_tensor(features["mask"].bytes_list.value[0], tf.int32).numpy().tolist() image = tf.io.parse_tensor(features["img/encoded"].bytes_list.value[0], tf.int32).numpy().tolist() From 0c00a3f6b0f503a21c8a3e7e7e2b73a86dd60531 Mon Sep 17 00:00:00 2001 From: Abhirakshak <37345811+Abhirakshak@users.noreply.github.com> Date: Tue, 9 Nov 2021 02:28:11 -0500 Subject: [PATCH 13/14] Update shapenet_visualizer.py --- .../vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py index 8a2f9d1988d..bfc2d05428f 100644 --- a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py @@ -95,8 +95,6 @@ def visualize_tf_record(shapenet_dir, features = example.features.feature vertices = tf.io.parse_tensor(features["vertices"].bytes_list.value[0], tf.float32).numpy().tolist() faces = tf.io.parse_tensor(features["faces"].bytes_list.value[0], tf.int32).numpy().tolist() - mask = tf.io.parse_tensor(features["mask"].bytes_list.value[0], tf.int32).numpy().tolist() - image = tf.io.parse_tensor(features["img/encoded"].bytes_list.value[0], tf.int32).numpy().tolist() filename = str(features["img/filename"].bytes_list.value[0]).split("/")[2][:-1] filename = filename.split(".")[0] From ad499e67835a4dc64d68e2faf9d3a8cb38cd63ab Mon Sep 17 00:00:00 2001 From: Ethan Suleman Date: Tue, 9 Nov 2021 21:09:29 -0500 Subject: [PATCH 14/14] BROKEN --- .../mesh_rcnn/data/shapenet_tfrecord_gen.py | 62 ++++++++++-------- .../mesh_rcnn/data/shapenet_visualizer.py | 64 ++++++++++--------- 2 files changed, 69 insertions(+), 57 deletions(-) diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py index 4d8dad035c6..e09d4a9bec2 100644 --- a/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_tfrecord_gen.py @@ -1,17 +1,17 @@ -from absl import app # pylint:disable=unused-import -from absl import flags -import numpy as np import json -import tensorflow as tf +import logging import os -import multiprocessing as mp +import tensorflow as tf +from absl import app # pylint:disable=unused-import +from absl import flags + from official.vision.beta.data import tfrecord_lib -from research.object_detection.utils import dataset_util from official.vision.beta.data.tfrecord_lib import convert_to_feature -flags.DEFINE_multi_string('shapenet_dir', '', 'Directory containing ShapeNet.') -flags.DEFINE_string('output_file_prefix', '/tmp/train', 'Path to output file') +flags.DEFINE_multi_string('shapenet_dir', '', 'Directory containing ' + 'ShapeNet.') +flags.DEFINE_string('output_file_prefix', '', 'Path to output file') flags.DEFINE_integer('num_shards', 32, 'Number of shards for output file.') FLAGS = flags.FLAGS @@ -19,6 +19,7 @@ logger = tf.get_logger() logger.setLevel(logging.INFO) + def parse_obj_file(file): """ Parses relevant data out of a .obj file. This contains all of the model information. @@ -36,10 +37,10 @@ def parse_obj_file(file): for line in lines: lineID = line[0:2] - + if lineID == "v ": vertex = line[2:].split(" ") - + for i, v in enumerate(vertex): vertex[i] = float(v) @@ -48,7 +49,7 @@ def parse_obj_file(file): if lineID == "f ": face = line[2:].split(" ") - + for i, f in enumerate(face): face[i] = [int(x) - 1 for x in f.split("/")] @@ -56,8 +57,8 @@ def parse_obj_file(file): return vertices, faces -def create_tf_example(image): +def create_tf_example(image): model_id = image["model_id"] label = image["label"] @@ -69,22 +70,23 @@ def create_tf_example(image): "vertices": convert_to_feature(model_vertices), "faces": convert_to_feature(model_faces)} - example = tf.train.Example( features=tf.train.Features(feature=feature_dict)) return example, 0 + def generate_annotations(images, shapenet_dir): for image in images: - yield {"shapenet_dir": shapenet_dir, - "label": image["label"], - "model_id": image["model_id"], - "synset_id": image["synset_id"]} + yield {"shapenet_dir": shapenet_dir, + "label": image["label"], + "model_id": image["model_id"], + "synset_id": image["synset_id"]} + def _create_tf_record_from_shapenet_dir(shapenet_dir, - output_path, - num_shards): + output_path, + num_shards): """Loads Shapenet json files and converts to tf.Record format. Args: images_info_file: shapenet_dir download directory @@ -94,21 +96,22 @@ def _create_tf_record_from_shapenet_dir(shapenet_dir, logging.info('writing to output path: %s', output_path) - #create synset ID to label mapping dictionary - with open(shapenet_synset_dict.json, "r") as dict_file: + # create synset ID to label mapping dictionary + with open('C:/Users/Ethan/PycharmProjects/tf-models/official/vision/beta/projects/mesh_rcnn/data' + '/shapenet_synset_dict.json', "r") as dict_file: synset_dict = json.load(dict_file) - #images list + # images list images = [] - for _, synset_directories, _ in os.walk(shapenet_dir): + for _, synset_directories, _ in os.walk(shapenet_dir[0]): for synset_directory in synset_directories: - for _, object_directories, _ in os.walk(os.path.join(shapenet_dir, synset_directory)): + for _, object_directories, _ in os.walk(os.path.join(shapenet_dir[0], synset_directory)): for object_directory in object_directories: - image = {"model_id" : object_directory, - "label" : synset_dict[synset_directory], - "shapenet_dir": shapenet_dir, - "synset_id": synset_directory} + image = {"model_id": object_directory, + "label": synset_dict[synset_directory], + "shapenet_dir": shapenet_dir, + "synset_id": synset_directory} images.append(image) shapenet_annotations_iter = generate_annotations( @@ -119,6 +122,7 @@ def _create_tf_record_from_shapenet_dir(shapenet_dir, logging.info('Finished writing, skipped %d annotations.', num_skipped) + def main(_): assert FLAGS.shapenet_dir, '`shapenet_dir` missing.' @@ -126,6 +130,8 @@ def main(_): if not tf.io.gfile.isdir(directory): tf.io.gfile.makedirs(directory) + _create_tf_record_from_shapenet_dir('shapenet_dir', 'tmp', 32) + if __name__ == '__main__': app.run(main) diff --git a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py index bfc2d05428f..6b0be09eb3b 100644 --- a/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py +++ b/official/vision/beta/projects/mesh_rcnn/data/shapenet_visualizer.py @@ -11,37 +11,30 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# Author: Ethan Suleman & Abhi Raja -r"""Takes TFRecord containing ShapeNet dataset and visualizes it. +r"""Takes TFRecord containing Pix3D dataset and visualizes it. Example usage: - python shapenet_visualizer.py --logtostderr \ + python pix3d_visualizer.py --logtostderr \ --num_models=2 \ - --shapenet_dir="${SHAPENET_TFRECORD_DIR}" \ - --output_file_prefix="${OUTPUT_DIR/FILE_PREFIX}" + --pix3d_dir="${PIX3D_TFRECORD_DIR}" \ + --output_folder="${OUTPUT_DIR}" """ -import json import logging import os -import json from absl import app # pylint:disable=unused-import from absl import flags import numpy as np -from numpy.core.defchararray import encode -from numpy.lib.arraysetops import isin -from numpy.lib.type_check import imag -import scipy.io as sio import cv2 import tensorflow as tf -import itertools -import multiprocessing as mp -flags.DEFINE_multi_string('shapenet_dir', '', 'Directory containing ShapeNet.') -flags.DEFINE_string('output_file_prefix', '/tmp/output', 'Path to output files') -flags.DEFINE_integer('num_models', 2, 'Number of models rebuilt from TFRecord.') +flags.DEFINE_multi_string('pix3d_dir', 'C:/Users/Ethan/tmp', 'Directory containing Pix3d.') +flags.DEFINE_string('output_folder', 'C:/Users/Ethan/tmp/output', + 'Path to output files') +flags.DEFINE_integer( + 'num_models', 2, 'Number of models rebuilt from TFRecord.') FLAGS = flags.FLAGS @@ -60,6 +53,7 @@ def write_obj_file(vertices, """ logging.info(f"Logging file {filename}") with open(filename, 'w') as f: + for vertex in vertices: print("v " + ' '.join(map(str, vertex)), file=f) @@ -70,22 +64,22 @@ def write_obj_file(vertices, ret += str(coordinate[0]) + " " print(ret, file=f) - -def visualize_tf_record(shapenet_dir, + +def visualize_tf_record(pix3d_dir, output_path, num_models): - """Visualizes shapenet data in TFRecord format. + """Visualizes pix3d data in TFRecord format. Args: - shapenet_dir: shapenet_dir for TFRecords + pix3d_dir: pix3d_dir for TFRecords output_path: Path to output .obj files. num_models: Number of output files to create. """ logging.info( - f"Starting to visualize {num_models} models from the TFRecords in directory {shapenet_dir} into {output_path}.") + f"Starting to visualize {num_models} models from the TFRecords in directory {pix3d_dir} into {output_path}.") - filenames = [os.path.join(shapenet_dir, x) for x in os.listdir(shapenet_dir)] + filenames = [os.path.join(pix3d_dir, x) for x in os.listdir(pix3d_dir)] raw_dataset = tf.data.TFRecordDataset(filenames) @@ -93,10 +87,21 @@ def visualize_tf_record(shapenet_dir, example = tf.train.Example() example.ParseFromString(raw_record.numpy()) features = example.features.feature - vertices = tf.io.parse_tensor(features["vertices"].bytes_list.value[0], tf.float32).numpy().tolist() - faces = tf.io.parse_tensor(features["faces"].bytes_list.value[0], tf.int32).numpy().tolist() - filename = str(features["img/filename"].bytes_list.value[0]).split("/")[2][:-1] + vertices = tf.io.parse_tensor( + features["model/vertices"].bytes_list.value[0], tf.float32).numpy().tolist() + faces = tf.io.parse_tensor( + features["model/faces"].bytes_list.value[0], tf.int32).numpy().tolist() + mask = cv2.imdecode(np.fromstring( + features["mask"].bytes_list.value[0], np.uint8), cv2.IMREAD_GRAYSCALE).tolist() + image = cv2.imdecode(np.fromstring( + features["img/encoded"].bytes_list.value[0], np.uint8), flags=1).tolist() + + # mask = tf.io.parse_tensor(features["mask"].bytes_list.value[0], tf.int32).numpy().tolist() + # image = tf.io.parse_tensor(features["img/encoded"].bytes_list.value[0], tf.int32).numpy().tolist() + + filename = str( + features["img/filename"].bytes_list.value[0]).split("/")[2][:-1] filename = filename.split(".")[0] filename = os.path.join(output_path, filename) @@ -105,14 +110,15 @@ def visualize_tf_record(shapenet_dir, def main(_): - assert FLAGS.shapenet_dir, '`shapenet_dir` missing.' + assert FLAGS.pix3d_dir, '`pix3d_dir` missing.' - directory = os.path.dirname(FLAGS.output_file_prefix) + directory = os.path.dirname(FLAGS.output_folder) if not tf.io.gfile.isdir(directory): tf.io.gfile.makedirs(directory) - visualize_tf_record(FLAGS.shapenet_dir[0], FLAGS.output_file_prefix, FLAGS.num_models) + visualize_tf_record( + FLAGS.pix3d_dir[0], FLAGS.output_folder, FLAGS.num_models) if __name__ == '__main__': - app.run(main) + app.run(main) \ No newline at end of file