From 1b04ccfcc6e264d2dcd4f345bcc85a7bbed887e9 Mon Sep 17 00:00:00 2001 From: Patrick Esser Date: Fri, 16 Mar 2018 11:39:46 +0100 Subject: [PATCH] clean --- batches_pg2.py | 188 ++----------------------------------------------- main.py | 145 ++------------------------------------ 2 files changed, 10 insertions(+), 323 deletions(-) diff --git a/batches_pg2.py b/batches_pg2.py index def02516..21278797 100644 --- a/batches_pg2.py +++ b/batches_pg2.py @@ -60,14 +60,6 @@ def preprocess(x): return np.cast[np.float32](x / 127.5 - 1.0) -def preprocess_mask(x): - """From uint8 mask to [0,1].""" - mask = np.cast[np.float32](x / 255.0) - if mask.shape[-1] == 3: - mask = np.amax(mask, axis = -1, keepdims = True) - return mask - - def postprocess(x): """[-1,1] to uint8.""" x = (x + 1.0) / 2.0 @@ -112,8 +104,6 @@ def make_joint_img(img_shape, jo, joints): for i in range(3): imgs.append(np.zeros(img_shape[:2], dtype = "uint8")) - assert("cnose" in jo) - # MSCOCO body = ["lhip", "lshoulder", "rshoulder", "rhip"] body_pts = np.array([[joints[jo.index(part),:] for part in body]]) if np.min(body_pts) >= 0: @@ -174,24 +164,6 @@ def valid_joints(*joints): return (j >= 0).all() -def zoom(img, factor, center = None): - shape = img.shape[:2] - if center is None or not valid_joints(center): - center = np.array(shape) / 2 - e1 = np.array([1,0]) - e2 = np.array([0,1]) - - dst_center = np.array(center) - dst_e1 = e1 * factor - dst_e2 = e2 * factor - - src = np.float32([center, center+e1, center+e2]) - dst = np.float32([dst_center, dst_center+dst_e1, dst_center+dst_e2]) - M = cv2.getAffineTransform(src, dst) - - return cv2.warpAffine(img, M, shape, flags = cv2.INTER_AREA, borderMode = cv2.BORDER_REPLICATE) - - def get_crop(bpart, joints, jo, wh, o_w, o_h, ar = 1.0): bpart_indices = [jo.index(b) for b in bpart] part_src = np.float32(joints[bpart_indices]) @@ -272,7 +244,6 @@ def get_crop(bpart, joints, jo, wh, o_w, o_h, ar = 1.0): def normalize(imgs, coords, stickmen, jo): - out_imgs = list() out_stickmen = list() @@ -317,58 +288,6 @@ def normalize(imgs, coords, stickmen, jo): img = np.concatenate(part_imgs, axis = 2) stickman = np.concatenate(part_stickmen, axis = 2) - """ - bpart = ["lshoulder","lhip","rhip","rshoulder"] - dst = np.float32([[0.0,0.0],[0.0,1.0],[1.0,1.0],[1.0,0.0]]) - bpart_indices = [jo.index(b) for b in bpart] - part_src = np.float32(joints[bpart_indices]) - part_dst = np.float32(wh * dst) - - M = cv2.getPerspectiveTransform(part_src, part_dst) - img = cv2.warpPerspective(img, M, (h,w), borderMode = cv2.BORDER_REPLICATE) - stickman = cv2.warpPerspective(stickman, M, (h,w), borderMode = cv2.BORDER_REPLICATE) - """ - - """ - # center of possible rescaling - c = joints[jo.index("cneck")] - - # find valid body part for scale estimation - a = joints[jo.index("lshoulder")] - b = joints[jo.index("lhip")] - target_length = 33.0 - if not valid_joints(a,b): - a = joints[jo.index("rshoulder")] - b = joints[jo.index("rhip")] - target_length = 33.0 - if not valid_joints(a,b): - a = joints[jo.index("rshoulder")] - b = joints[jo.index("relbow")] - target_length = 33.0 / 2 - if not valid_joints(a,b): - a = joints[jo.index("lshoulder")] - b = joints[jo.index("lelbow")] - target_length = 33.0 / 2 - if not valid_joints(a,b): - a = joints[jo.index("lwrist")] - b = joints[jo.index("lelbow")] - target_length = 33.0 / 2 - if not valid_joints(a,b): - a = joints[jo.index("rwrist")] - b = joints[jo.index("relbow")] - target_length = 33.0 / 2 - - if valid_joints(a,b): - body_length = np.linalg.norm(b - a) - factor = target_length / body_length - img = zoom(img, factor, center = c) - stickman = zoom(stickman, factor, center = c) - else: - factor = 0.25 - img = zoom(img, factor, center = c) - stickman = zoom(stickman, factor, center = c) - """ - out_imgs.append(img) out_stickmen.append(stickman) out_imgs = np.stack(out_imgs) @@ -376,71 +295,16 @@ def normalize(imgs, coords, stickmen, jo): return out_imgs, out_stickmen -def make_mask_img(img_shape, jo, joints): - scale_factor = img_shape[1] / 128 - masks = 3*[None] - for i in range(3): - masks[i] = np.zeros(img_shape[:2], dtype = "uint8") - - body = ["lhip", "lshoulder", "rshoulder", "rhip"] - body_pts = np.array([[joints[jo.index(part),:] for part in body]], dtype = np.int32) - cv2.fillPoly(masks[1], body_pts, 255) - - head = ["lshoulder", "chead", "rshoulder"] - head_pts = np.array([[joints[jo.index(part),:] for part in head]], dtype = np.int32) - cv2.fillPoly(masks[2], head_pts, 255) - - thickness = int(15 * scale_factor) - lines = [[ - ("rankle", "rknee"), - ("rknee", "rhip"), - ("rhip", "lhip"), - ("lhip", "lknee"), - ("lknee", "lankle") ], [ - ("rhip", "rshoulder"), - ("rshoulder", "relbow"), - ("relbow", "rwrist"), - ("rhip", "lhip"), - ("rshoulder", "lshoulder"), - ("lhip", "lshoulder"), - ("lshoulder", "lelbow"), - ("lelbow", "lwrist")], [ - ("rshoulder", "chead"), - ("rshoulder", "lshoulder"), - ("lshoulder", "chead")]] - for i in range(len(lines)): - for j in range(len(lines[i])): - line = [jo.index(lines[i][j][0]), jo.index(lines[i][j][1])] - a = tuple(np.int_(joints[line[0]])) - b = tuple(np.int_(joints[line[1]])) - cv2.line(masks[i], a, b, color = 255, thickness = thickness) - - for i in range(3): - r = int(11 * scale_factor) - if r % 2 == 0: - r = r + 1 - masks[i] = cv2.GaussianBlur(masks[i], (r,r), 0) - maxmask = np.max(masks[i]) - if maxmask > 0: - masks[i] = masks[i] / maxmask - mask = np.stack(masks, axis = -1) - mask = np.uint8(255 * mask) - - return mask - - class IndexFlow(object): """Batches from index file.""" - def __init__( self, shape, index_path, train, - mask = True, fill_batches = True, shuffle = True, - return_keys = ["imgs", "joints"]): + return_keys = ["imgs", "joints", "norm_imgs", "norm_joints"]): self.shape = shape self.batch_size = self.shape[0] self.img_shape = self.shape[1:] @@ -448,7 +312,6 @@ def __init__( self.index = pickle.load(f) self.basepath = os.path.dirname(index_path) self.train = train - self.mask = mask self.fill_batches = fill_batches self.shuffle_ = shuffle self.return_keys = return_keys @@ -487,7 +350,7 @@ def __next__(self): if self.fill_batches and batch_indices.shape[0] != self.batch_size: n_missing = self.batch_size - batch_indices.shape[0] batch_indices = np.concatenate([batch_indices, self.indices[:n_missing]], axis = 0) - assert(batch_indices.shape[0] == self.batch_size) + assert batch_indices.shape[0] == self.batch_size batch_indices = np.array(batch_indices) batch["indices"] = batch_indices @@ -502,6 +365,7 @@ def __next__(self): batch["imgs"] = list() for i in batch_indices: fname = self.index["imgs"][i] + # TODO this should be more uniform traintest = "train" if self.train else "test" path = os.path.join(self.basepath, "..", "original", "filted_up_{}".format(traintest), fname) batch["imgs"].append(load_img(path, target_size = self.img_shape)) @@ -519,25 +383,6 @@ def __next__(self): batch["joints"] = np.stack(batch["joints"]) batch["joints"] = preprocess(batch["joints"]) - if False and self.mask: - if "masks" in self.index: - batch_masks = list() - for i in batch_indices: - fname = self.index["masks"][i] - path = os.path.join(self.basepath, fname) - batch_masks.append(load_img(path, target_size = self.img_shape)) - else: - # generate mask based on joint coordinates - batch_masks = list() - for joints in batch["joints_coordinates"]: - mask = make_mask_img(self.img_shape, self.jo, joints) - batch_masks.append(mask) - batch["masks"] = np.stack(batch_masks) - batch["masks"] = preprocess_mask(batch["masks"]) - # apply mask to images - batch["imgs"] = batch["imgs"] * batch["masks"] - - imgs, joints = normalize(batch["imgs"], batch["joints_coordinates"], batch["joints"], self.jo) batch["norm_imgs"] = imgs batch["norm_joints"] = joints @@ -556,12 +401,11 @@ def get_batches( shape, index_path, train, - mask, fill_batches = True, shuffle = True, return_keys = ["imgs", "joints", "norm_imgs", "norm_joints"]): """Buffered IndexFlow.""" - flow = IndexFlow(shape, index_path, train, mask, fill_batches, shuffle, return_keys) + flow = IndexFlow(shape, index_path, train, fill_batches, shuffle, return_keys) return BufferedWrapper(flow) @@ -575,27 +419,7 @@ def get_batches( shape = (16, 128, 128, 3), index_path = sys.argv[1], train = True, - mask = False, shuffle = True) - X, C = next(batches) - plot_batch(X, "unmasked.png") + X, C, XN, CN = next(batches) + plot_batch(X, "images.png") plot_batch(C, "joints.png") - - """ - batches = get_batches( - shape = (16, 128, 128, 3), - index_path = sys.argv[1], - train = True, - mask = True) - X, C = next(batches) - plot_batch(X, "masked.png") - - batches = get_batches( - shape = (16, 32, 32, 3), - index_path = sys.argv[1], - train = True, - mask = True) - X, C = next(batches) - plot_batch(X, "masked32.png") - plot_batch(C, "joints32.png") - """ diff --git a/main.py b/main.py index 958e683f..b1d8346c 100644 --- a/main.py +++ b/main.py @@ -40,22 +40,6 @@ def init_logging(out_base_dir): return out_dir, logger -class Timer(object): - def __init__(self): - self.tick() - - - def tick(self): - self.start_time = time.time() - - - def tock(self): - self.end_time = time.time() - time_since_tick = self.end_time - self.start_time - self.tick() - return time_since_tick - - class Model(object): def __init__(self, opt, out_dir, logger): self.batch_size = opt.batch_size @@ -159,7 +143,6 @@ def likelihood_loss(self, x, params): def define_graph(self): # pretrained net for perceptual loss - #self.vgg19 = deeploss.JigsawFeatures(session) self.vgg19 = deeploss.VGG19Features(session) global_step = tf.Variable(0, trainable = False, name = "global_step") @@ -253,8 +236,6 @@ def define_graph(self): self.img_ops["test_sample"] = test_sample self.img_ops["x"] = self.x self.img_ops["c"] = self.c - for i in range(N_BOXES): - self.img_ops["xn{}".format(i)] = self.xn[:,:,:,i*3:(i+1)*3] for i, l in enumerate(self.vgg19.losses): self.log_ops["vgg_loss_{}".format(i)] = l @@ -416,21 +397,6 @@ def test(self, c_batch): return results - def mcmc(self, c_batch, n_iters = 10): - results = dict() - results["cond"] = c_batch - sample = session.run( - self.img_ops["test_sample"], {self.c: c_batch}) - results["sample_{}".format(0)] = sample - for i in range(n_iters - 1): - sample = session.run( - self.img_ops["sample"], { - self.x: sample, - self.c: c_batch}) - results["sample_{:03}".format(i+1)] = sample - return results - - def reconstruct(self, x_batch, c_batch): return session.run( self.reconstruction, @@ -477,9 +443,6 @@ def transfer(self, x_encode, c_encode, c_decode): parser.add_argument("--ckpt_freq", default = 1000, type = int, help = "frequency to checkpoint") parser.add_argument("--test_freq", default = 1000, type = int, help = "frequency to test") parser.add_argument("--drop_prob", default = 0.1, type = float, help = "Dropout probability") - parser.add_argument("--mask", dest = "mask", action = "store_true", help = "Use masked data") - parser.add_argument("--no-mask", dest = "mask", action = "store_false", help = "Do not use mask") - parser.set_defaults(mask = True) opt = parser.parse_args() if not os.path.exists(opt.data_index): @@ -494,9 +457,9 @@ def transfer(self, x_encode, c_encode, c_decode): data_shape = [batch_size] + img_shape init_shape = [opt.init_batches * batch_size] + img_shape - batches = get_batches(data_shape, opt.data_index, mask = opt.mask, train = True) - init_batches = get_batches(init_shape, opt.data_index, mask = opt.mask, train = True) - valid_batches = get_batches(data_shape, opt.data_index, mask = opt.mask, train = False) + batches = get_batches(data_shape, opt.data_index, train = True) + init_batches = get_batches(init_shape, opt.data_index, train = True) + valid_batches = get_batches(data_shape, opt.data_index, train = False) logger.info("Number of training samples: {}".format(batches.n)) logger.info("Number of validation samples: {}".format(valid_batches.n)) if valid_batches.n == 0: @@ -514,7 +477,7 @@ def transfer(self, x_encode, c_encode, c_decode): batch_size = opt.batch_size img_shape = 2*[opt.spatial_size] + [3] data_shape = [batch_size] + img_shape - valid_batches = get_batches(data_shape, opt.data_index, mask = opt.mask, train = False) + valid_batches = get_batches(data_shape, opt.data_index, train = False) model = Model(opt, out_dir, logger) model.restore_graph(opt.checkpoint) @@ -525,105 +488,5 @@ def transfer(self, x_encode, c_encode, c_decode): plot_batch(x_gen[k], os.path.join( out_dir, "testing_{}_{:07}.png".format(k, i))) - elif opt.mode == "add_reconstructions": - if not opt.checkpoint: - raise Exception("Testing requires --checkpoint") - batch_size = opt.batch_size - img_shape = 2*[opt.spatial_size] + [3] - data_shape = [batch_size] + img_shape - batches = get_batches(data_shape, opt.data_index, mask = opt.mask, - train = True, return_index_id = True) - valid_batches = get_batches(data_shape, opt.data_index, - mask = opt.mask, train = False, return_index_id = True) - model = Model(opt, out_dir, logger) - model.restore_graph(opt.checkpoint) - - # open index file to get image filenames and update with - # reconstruction data - with open(opt.data_index, "rb") as f: - index = pickle.load(f) - index_dir = os.path.dirname(opt.data_index) - index["reconstruction"] = len(index["imgs"]) * [None] - index["sample"] = len(index["imgs"]) * [None] - - def process_batches(batches): - for i in trange(math.ceil(batches.n / batch_size)): - X_batch, C_batch, I_batch = next(batches) - # reconstructions - R_batch = model.reconstruct(X_batch, C_batch) - R_batch = postprocess(R_batch) # to uint8 for saving - # samples from pose - S_batch = model.test(C_batch)["test_sample"] - S_batch = postprocess(S_batch) # to uint8 for saving - for batch_i, i in enumerate(I_batch): - original_fname = index["imgs"][i] - reconstr_fname = original_fname.rsplit(".", 1)[0] + "_reconstruction.png" - reconstr_path = os.path.join(index_dir, reconstr_fname) - sample_fname = original_fname.rsplit(".", 1)[0] + "_sample.png" - sample_path = os.path.join(index_dir, sample_fname) - index["reconstruction"][i] = reconstr_path - index["sample"][i] = sample_path - PIL.Image.fromarray(R_batch[batch_i,...]).save(reconstr_path) - PIL.Image.fromarray(S_batch[batch_i,...]).save(sample_path) - process_batches(batches) - process_batches(valid_batches) - - # write updated index - with open(opt.data_index, "wb") as f: - pickle.dump(index, f) - logger.info("Wrote {}".format(opt.data_index)) - - elif opt.mode == "transfer": - if not opt.checkpoint: - opt.checkpoint = "log/2017-10-24T16:34:09/checkpoints/model.ckpt-100000" - batch_size = opt.batch_size - img_shape = 2*[opt.spatial_size] + [3] - data_shape = [batch_size] + img_shape - valid_batches = get_batches(data_shape, opt.data_index, - mask = opt.mask, train = False) - model = Model(opt, out_dir, logger) - model.restore_graph(opt.checkpoint) - - ids = ["00038", "00281", "01166", "x", "06909", "y", "07586", "07607", "z", "09874"] - for step in trange(10): - X_batch, C_batch, XN_batch, CN_batch = next(valid_batches) - bs = X_batch.shape[0] - imgs = list() - imgs.append(np.zeros_like(X_batch[0,...])) - for r in range(bs): - imgs.append(C_batch[r,...]) - for i in range(bs): - x_infer = XN_batch[i,...] - c_infer = CN_batch[i,...] - imgs.append(X_batch[i,...]) - - x_infer_batch = x_infer[None,...].repeat(bs, axis = 0) - c_infer_batch = c_infer[None,...].repeat(bs, axis = 0) - c_generate_batch = C_batch - results = model.transfer(x_infer_batch, c_infer_batch, c_generate_batch) - for j in range(bs): - imgs.append(results[j,...]) - imgs = np.stack(imgs, axis = 0) - plot_batch(imgs, os.path.join( - out_dir, - "transfer_{}.png".format(ids[step]))) - - elif opt.mode == "mcmc": - if not opt.checkpoint: - raise Exception("Testing requires --checkpoint") - batch_size = opt.batch_size - img_shape = 2*[opt.spatial_size] + [3] - data_shape = [batch_size] + img_shape - valid_batches = get_batches(data_shape, opt.data_index, mask = opt.mask, train = False) - model = Model(opt, out_dir, logger) - model.restore_graph(opt.checkpoint) - - for i in trange(valid_batches.n // batch_size): - X_batch, C_batch = next(valid_batches) - x_gen = model.mcmc(C_batch) - for k in x_gen: - plot_batch(x_gen[k], os.path.join( - out_dir, - "mcmc_{}_{:07}.png".format(k, i))) else: raise NotImplemented()