def __getitem__(self, item): if self.mode == 'Train': sketch_path = self.Train_Sketch[item] vector_x = self.Coordinate[sketch_path] sketch_img, sketch_points = rasterize_Sketch(vector_x) sketch_img = Image.fromarray(sketch_img).convert('RGB') n_flip = random.random() if n_flip > 0.5: sketch_img = F.hflip(sketch_img) sketch_points[:, 0] = -sketch_points[:, 0] + 256. sketch_img = self.train_transform(sketch_img) sample = { 'sketch_img': sketch_img, 'sketch_path': sketch_path, 'sketch_points': sketch_points, 'sketch_label': self.name2num[sketch_path.split('/')[0]] } elif self.mode == 'Test': sketch_path = self.Test_Sketch[item] vector_x = self.Coordinate[sketch_path] sketch_img, sketch_points = rasterize_Sketch(vector_x) sketch_img = self.test_transform( Image.fromarray(sketch_img).convert('RGB')) sample = { 'sketch_img': sketch_img, 'sketch_path': sketch_path, 'sketch_points': sketch_points, 'sketch_label': self.name2num[sketch_path.split('/')[0]] } return sample
def __getitem__(self, item): if self.mode == 'Train': path = self.Train_Sketch[item] anchor_sketch_vector = self.Coordinate[path] anchor_sketch = Image.fromarray( rasterize_Sketch(anchor_sketch_vector)).convert('RGB') class_name = path.split('/')[0] sketch_positive_path = random.choice(self.seen_dict[class_name]) sketch_positive_vector = self.Coordinate[sketch_positive_path] sketch_positive = Image.fromarray( rasterize_Sketch(sketch_positive_vector)).convert('RGB') possible_negative_class = random.choice( list(set(self.Seen_Class) - set(class_name))) sketch_negative_path = random.choice( self.seen_dict[possible_negative_class]) sketch_negative_vector = self.Coordinate[sketch_negative_path] sketch_negative = Image.fromarray( rasterize_Sketch(sketch_negative_vector)).convert('RGB') else: path = self.Test_Sketch[item] anchor_sketch_vector = self.Coordinate[path] anchor_sketch = Image.fromarray( rasterize_Sketch(anchor_sketch_vector)).convert('RGB') class_name = path.split('/')[0] sketch_positive_path = random.choice(self.seen_dict[class_name]) sketch_positive_vector = self.Coordinate[sketch_positive_path] sketch_positive = Image.fromarray( rasterize_Sketch(sketch_positive_vector)).convert('RGB') possible_negative_class = random.choice( list(set(self.Seen_Class) - set(class_name))) sketch_negative_path = random.choice( self.seen_dict[possible_negative_class]) sketch_negative_vector = self.Coordinate[sketch_negative_path] sketch_negative = Image.fromarray( rasterize_Sketch(sketch_negative_vector)).convert('RGB') if self.hp.data_encoding_type == '3point': stroke_wise_split_anchor_list = np.split( anchor_sketch_vector, np.where(anchor_sketch_vector[:, 2])[0] + 1, axis=0)[:-1] stroke_wise_split_positive_list = np.split( sketch_positive_vector, np.where(sketch_positive_vector[:, 2])[0] + 1, axis=0)[:-1] stroke_wise_split_negative_list = np.split( sketch_negative_vector, np.where(sketch_negative_vector[:, 2])[0] + 1, axis=0)[:-1] elif self.hp.data_encoding_type == '5point': anchor_sketch_vector = self.to_delXY(anchor_sketch_vector) stroke_wise_split_anchor_list = np.split( anchor_sketch_vector, np.where(anchor_sketch_vector[:, 3])[0] + 1, axis=0) sketch_positive_vector = self.to_delXY(sketch_positive_vector) stroke_wise_split_positive_list = np.split( sketch_positive_vector, np.where(sketch_positive_vector[:, 3])[0] + 1, axis=0) sketch_negative_vector = self.to_delXY(sketch_negative_vector) stroke_wise_split_negative_list = np.split( sketch_negative_vector, np.where(sketch_negative_vector[:, 3])[0] + 1, axis=0) else: raise ValueError( 'invalid option for --data_encoding_type. Valid options: 3point/5point' ) stroke_wise_split_anchor = [ torch.from_numpy(x) for x in stroke_wise_split_anchor_list ] every_stroke_len_anchor = [ len(stroke) for stroke in stroke_wise_split_anchor ] num_stroke_per_anchor = len(every_stroke_len_anchor) assert sum(every_stroke_len_anchor) == anchor_sketch_vector.shape[0] stroke_wise_split_positive = [ torch.from_numpy(x) for x in stroke_wise_split_positive_list ] every_stroke_len_positive = [ len(stroke) for stroke in stroke_wise_split_positive ] num_stroke_per_positive = len(every_stroke_len_positive) assert sum( every_stroke_len_positive) == sketch_positive_vector.shape[0] stroke_wise_split_negative = [ torch.from_numpy(x) for x in stroke_wise_split_negative_list ] every_stroke_len_negative = [ len(stroke) for stroke in stroke_wise_split_negative ] num_stroke_per_negative = len(every_stroke_len_negative) assert sum( every_stroke_len_negative) == sketch_negative_vector.shape[0] sample = { 'path': path, 'label': self.name2num[class_name], 'anchor_sketch_image': self.train_transform(anchor_sketch), 'anchor_sketch_vector': anchor_sketch_vector, 'num_stroke_per_anchor': num_stroke_per_anchor, 'every_stroke_len_anchor': every_stroke_len_anchor, 'stroke_wise_split_anchor': stroke_wise_split_anchor, 'sketch_positive': self.train_transform(sketch_positive), 'sketch_positive_vector': sketch_positive_vector, 'num_stroke_per_positive': num_stroke_per_positive, 'every_stroke_len_positive': every_stroke_len_positive, 'stroke_wise_split_positive': stroke_wise_split_positive, 'sketch_negative': self.train_transform(sketch_negative), 'sketch_negative_vector': sketch_negative_vector, 'num_stroke_per_negative': num_stroke_per_negative, 'every_stroke_len_negative': every_stroke_len_negative, 'stroke_wise_split_negative': stroke_wise_split_negative } return sample
print(i_rdp, num, key) sketch_points = Coordinate[key] sketch_points_orig = sketch_points sketch_points = sketch_points.astype(np.float) # sketch_points[:, :2] = sketch_points[:, :2] / np.array([800, 800]) # sketch_points[:, :2] = sketch_points[:, :2] * 256 sketch_points = np.round(sketch_points) all_strokes = np.split(sketch_points, np.where(sketch_points[:, 2])[0] + 1, axis=0)[:-1] max_points_old.append(sketch_points_orig.shape) sketch_img_orig = rasterize_Sketch(sketch_points_orig) sketch_img_orig = Image.fromarray(sketch_img_orig).convert('RGB') # sketch_img_orig.show() sketch_img_orig.save(str(i_rdp) + '/' + str(num) + '.jpg') sketch_points_sampled_new = [] for stroke in all_strokes: stroke_new = rdp(stroke[:, :2], epsilon=i_rdp, algo="iter") stroke_new = np.hstack( (stroke_new, np.zeros((stroke_new.shape[0], 1)))) stroke_new[-1, -1] = 1. # print(stroke_new.shape, stroke.shape) sketch_points_sampled_new.append(stroke_new) sketch_points_new = np.vstack(sketch_points_sampled_new) max_points_new.append(sketch_points_new.shape[0])
def show(step, batch, model, save_path): if not os.path.isdir(save_path + '/' + time_id): os.makedirs(save_path + '/' + time_id) output_anc, num_stroke_anc = model.Network(batch, type='anchor') # b,N1,512 output_pos, num_stroke_pos = model.Network(batch, type='positive') # b,N2,512 mask_anc, mask_pos = map(make_mask, [num_stroke_anc, num_stroke_pos]) corr_xpos = model.neighbour(output_anc, output_pos, mask_anc, mask_pos) ''' I need the n1 x n2 matrix For every stroke in N1, the stroke having highest correlation will be painted. ''' i_sample = 0 # range = 0 to batch-size A = corr_xpos[i_sample] # Taking 1 sample out of the batch anc_max = torch.argmax(A, dim=1) # 1 x N1 anc_vec = batch['anchor_sketch_vector'] anc_img_orig = Image.fromarray( 255 - rasterize_Sketch(anc_vec[i_sample].numpy())).convert('RGB') anc_stroke_num = batch['num_stroke_per_anchor'][i_sample] start_index = sum(batch['num_stroke_per_anchor'][:i_sample]) anc_sample = batch['stroke_wise_split_anchor'][start_index:start_index + anc_stroke_num].numpy() pos_vec = batch['sketch_positive_vector'] pos_img_orig = Image.fromarray( 255 - rasterize_Sketch(pos_vec[i_sample].numpy())).convert('RGB') pos_stroke_num = batch['num_stroke_per_positive'][i_sample] start_index = sum(batch['num_stroke_per_positive'][:i_sample]) pos_sample = batch['stroke_wise_split_positive'][start_index:start_index + pos_stroke_num].numpy() im = Image.new('RGB', (266 * anc_stroke_num, 5 + 261 * 2)) # width x height for i_stroke in range(anc_stroke_num): anc_img = anc_img_orig.copy() draw_anc = ImageDraw.Draw(anc_img) anc_stroke_points = np.where( 255 - rasterize_Sketch(anc_sample[i_stroke]) == 0) for i_point in range(len(anc_stroke_points[0])): draw_anc.point( (anc_stroke_points[1][i_point], anc_stroke_points[0][i_point]), fill='red') pos_img = pos_img_orig.copy() draw_pos = ImageDraw.Draw(pos_img) pos_stroke_points = np.where( 255 - rasterize_Sketch(pos_sample[anc_max[i_stroke].item()]) == 0) for i_point in range(len(pos_stroke_points[0])): draw_pos.point( (pos_stroke_points[1][i_point], pos_stroke_points[0][i_point]), fill='blue') im.paste(anc_img, (5 + i_stroke * 266, 5)) #width x height im.paste(pos_img, (5 + i_stroke * 266, 266)) im.save(f'{save_path}/{time_id}/Step_{step}.png') return
def __getitem__(self, item): sample = {} if self.mode == "Train": path = self.Train_List[item] if self.data == "coordinate": vector_x = self.Coordinate[path] sketch_img, sketch_points = rasterize_Sketch(vector_x) sketch_image = Image.fromarray(sketch_img).convert("RGB") sketch_classname = path.split("/")[0] positive_path = self.train_images_classes[ f"{sketch_classname}"][randint( 0, len(self.train_images_classes[f"{sketch_classname}"]) - 1)] positive_classname = sketch_classname possible_list = list(self.train_images_classes.keys()) possible_list.remove(sketch_classname) negative_classname = possible_list[randint( 0, len(possible_list) - 1)] negative_path = self.train_images_classes[ f"{negative_classname}"][randint( 0, len(self.train_images_classes[f"{negative_classname}"]) - 1)] # print(path,positive_path,negative_path) positive_img = Image.open(positive_path).convert("RGB") negative_img = Image.open(negative_path).convert("RGB") n_flip = random.random() if n_flip > 0.5: sketch_image = F.hflip(sketch_image) positive_img = F.hflip(positive_img) negative_img = F.hflip(negative_img) sketch_image = self.train_transform(sketch_image) positive_img = self.train_transform(positive_img) negative_img = self.train_transform(negative_img) sample = { "sketch_img": sketch_image, "sketch_label": self.name2num[sketch_classname], "sketch_points": sketch_points, "positive_img": positive_img, "positive_label": self.name2num[positive_classname], "negative_img": negative_img, "negative_label": self.name2num[negative_classname], } elif self.mode == "Test": path = self.Test_List[item] if self.data == "coordinate": vector_x = self.Coordinate[path] sketch_img, sketch_points = rasterize_Sketch(vector_x) sketch_image = Image.fromarray(sketch_img).convert("RGB") sketch_classname = path.split("/")[0] positive_path = self.test_images_classes[ f"{sketch_classname}"][randint( 0, len(self.test_images_classes[f"{sketch_classname}"]) - 1)] positive_classname = sketch_classname positive_img = Image.open(positive_path).convert("RGB") sketch_image = self.test_transform(sketch_image) positive_image = self.test_transform(positive_img) # sample = { # "sketch_img": sketch_image, # "sketch_label": self.name2num[sketch_classname], # "positive_img": positive_image, # "positive_label": self.name2num[positive_classname], # "test_data":self.test_images_values, # "name2num":self.name2num # } sample = { "sketch_img": sketch_image, "sketch_label": self.name2num[sketch_classname], "sketch_points": sketch_points, "positive_img": positive_image, } elif self.mode == "Test_photo": path = self.test_photo_all[item] positive_img = Image.open(path).convert("RGB") positive_image = self.test_transform(positive_img) positive_classname = path.split('/')[-2] positive_label = self.name2num[positive_classname] sample = { "positive_img": positive_image, "positive_label": positive_label, "path": path } return sample