def __init__(self,
set_list,
char_to_idx,
augmentation=False,
img_height=32,
random_subset_size=None):
self.img_height = img_height
self.ids = set_list
self.ids.sort()
self.detailed_ids = []
for ids_idx, paths in enumerate(self.ids):
json_path, img_path = paths
d = safe_load.json_state(json_path)
if d is None:
continue
for i in range(len(d)):
if 'hw_path' not in d[i] and 'image_path' not in d[i]:
continue
self.detailed_ids.append((ids_idx, i))
if random_subset_size is not None:
self.detailed_ids = random.sample(
self.detailed_ids,
min(random_subset_size, len(self.detailed_ids)))
print(len(self.detailed_ids))
self.char_to_idx = char_to_idx
self.augmentation = augmentation
self.warning = False
def __init__(self,
set_list,
rescale_range=None,
transform=None,
random_subset_size=None):
self.rescale_range = rescale_range
self.ids = set_list
self.ids.sort()
new_ids = []
for json_path, img_path in self.ids:
gt_json = safe_load.json_state(json_path)
if gt_json is None:
continue
failed = False
for j, gt_item in enumerate(gt_json):
if 'sol' not in gt_item:
failed = True
break
if failed:
continue
new_ids.append([json_path, img_path])
self.ids = new_ids
if random_subset_size is not None:
self.ids = random.sample(self.ids,
min(random_subset_size, len(self.ids)))
print("[SOL count]", len(self.ids))
self.transform = transform
def __getitem__(self, idx):
gt_json_path, img_path = self.ids[idx]
gt_json = safe_load.json_state(gt_json_path)
if gt_json is None:
return None
org_img = cv2.imread(img_path)
target_dim1 = int(
np.random.uniform(self.rescale_range[0], self.rescale_range[1]))
s = target_dim1 / float(org_img.shape[1])
target_dim0 = int(org_img.shape[0] / float(org_img.shape[1]) *
target_dim1)
org_img = cv2.resize(org_img, (target_dim1, target_dim0),
interpolation=cv2.INTER_CUBIC)
gt = np.zeros((1, len(gt_json), 4), dtype=np.float32)
for j, gt_item in enumerate(gt_json):
if 'sol' not in gt_item:
continue
x0 = gt_item['sol']['x0']
x1 = gt_item['sol']['x1']
y0 = gt_item['sol']['y0']
y1 = gt_item['sol']['y1']
gt[:, j, 0] = x0 * s
gt[:, j, 1] = y0 * s
gt[:, j, 2] = x1 * s
gt[:, j, 3] = y1 * s
if self.transform is not None:
out = self.transform({"img": org_img, "sol_gt": gt})
org_img = out['img']
gt = out['sol_gt']
org_img = augmentation.apply_random_color_rotation(org_img)
org_img = augmentation.apply_tensmeyer_brightness(org_img)
img = org_img.transpose([2, 1, 0])[None, ...]
img = img.astype(np.float32)
img = torch.from_numpy(img)
img = img / 128.0 - 1.0
if gt.shape[1] == 0:
gt = None
else:
gt = torch.from_numpy(gt)
return {"img": img, "sol_gt": gt}
def __getitem__(self, idx):
ids_idx, line_idx = self.detailed_ids[idx]
gt_json_path, img_path = self.ids[ids_idx]
gt_json = safe_load.json_state(gt_json_path)
positions = []
positions_xy = []
if 'lf' not in gt_json[line_idx]:
return None
for step in gt_json[line_idx]['lf']:
x0 = step['x0']
x1 = step['x1']
y0 = step['y0']
y1 = step['y1']
positions_xy.append((torch.Tensor([[x1, x0], [y1, y0]])))
dx = x0 - x1
dy = y0 - y1
d = math.sqrt(dx**2 + dy**2)
mx = (x0 + x1) / 2.0
my = (y0 + y1) / 2.0
#Not sure if this is right...
theta = -math.atan2(dx, -dy)
positions.append(torch.Tensor([mx, my, theta, d / 2, 1.0]))
img = cv2.imread(img_path)
if self.augmentation:
img = augmentation.apply_random_color_rotation(img)
img = augmentation.apply_tensmeyer_brightness(img)
img = img.astype(np.float32)
img = img.transpose()
img = img / 128.0 - 1.0
img = torch.from_numpy(img)
gt = gt_json[line_idx]['gt']
result = {
"img": img,
"lf_xyrs": positions,
"lf_xyxy": positions_xy,
"gt": gt
}
return result
def __getitem__(self, idx):
ids_idx, line_idx = self.detailed_ids[idx]
gt_json_path, img_path = self.ids[ids_idx]
gt_json = safe_load.json_state(gt_json_path)
if gt_json is None:
return None
if 'hw_path' not in gt_json[line_idx] and 'image_path' not in gt_json[
line_idx]:
return None
hw_path = gt_json[line_idx]['hw_path'] if "hw_path" in gt_json[
line_idx] else gt_json[line_idx]["image_path"]
hw_path = hw_path.split("/")[-1:]
hw_path = "/".join(hw_path)
hw_folder = os.path.dirname(gt_json_path)
img = cv2.imread(os.path.join(hw_folder, hw_path))
if img is None:
return None
if img.shape[0] != self.img_height:
if img.shape[0] < self.img_height and not self.warning:
self.warning = True
print("WARNING: upsampling image to fit size")
percent = float(self.img_height) / img.shape[0]
img = cv2.resize(img, (0, 0),
fx=percent,
fy=percent,
interpolation=cv2.INTER_CUBIC)
if img is None:
return None
if self.augmentation:
img = augmentation.apply_random_color_rotation(img)
img = augmentation.apply_tensmeyer_brightness(img)
img = warp_image(img)
img = img.astype(np.float32)
img = img / 128.0 - 1.0
gt = gt_json[line_idx]['gt']
if len(gt) == 0:
return None
gt_label = string_utils.str2label_single(gt, self.char_to_idx)
return {"line_img": img, "gt": gt, "gt_label": gt_label}
def __getitem__(self, idx):
gt_json_path, img_path = self.ids[idx]
gt_json = []
if not self.ignore_json:
gt_json = safe_load.json_state(gt_json_path)
if gt_json is None:
return None
org_img = cv2.imread(img_path)
full_img = org_img.astype(np.float32)
full_img = full_img.transpose([2,1,0])[None,...]
full_img = torch.from_numpy(full_img)
full_img = full_img / 128 - 1
target_dim1 = self.resize_width
s = target_dim1 / float(org_img.shape[1])
target_dim0 = int(org_img.shape[0]/float(org_img.shape[1]) * target_dim1)
img = cv2.resize(org_img,(target_dim1, target_dim0), interpolation = cv2.INTER_CUBIC)
img = img.astype(np.float32)
img = img.transpose([2,1,0])[None,...]
img = torch.from_numpy(img)
img = img / 128 - 1
image_key = gt_json_path[:-len('.json')]
return {
"resized_img": img,
"full_img": full_img,
"resize_scale": 1.0/s,
"gt_lines": [x['gt'] for x in gt_json],
"img_key": image_key,
"json_path": gt_json_path,
"gt_json": gt_json
}
def __init__(self, set_list, random_subset_size=None, augmentation=False):
self.augmentation = augmentation
self.ids = set_list
self.ids.sort()
self.detailed_ids = []
for ids_idx, paths in enumerate(self.ids):
json_path, img_path = paths
d = safe_load.json_state(json_path)
if d is None:
continue
for i in range(len(d)):
if 'steps' not in d[i]:
continue
self.detailed_ids.append((ids_idx, i))
if random_subset_size is not None:
self.detailed_ids = random.sample(
self.detailed_ids, min(len(self.ids), random_subset_size))
print(len(self.detailed_ids))
def __getitem__(self, idx):
ids_idx, line_idx = self.detailed_ids[idx]
gt_json_path, img_path = self.ids[ids_idx]
gt_json = safe_load.json_state(gt_json_path)
absolute_steps = []
if 'steps' not in gt_json[line_idx]:
return None
steps = gt_json[line_idx]['steps']
for step_index, step_data in enumerate(
gt_json[line_idx]['steps'][:-1]):
next_step_data = steps[step_index + 1]
step_points = step_to_points(step_data)
next_step_points = step_to_points(next_step_data)
angle_to_next = angle_between_points(
step_points["base_point"], next_step_points["base_point"])
if step_index <= len(steps) - 2:
angle_after_that = angle_to_next
else:
angle_after_that = angle_between_points(
next_step_points["base_point"],
step_to_points(steps[step_index + 2])["base_point"])
absolute_steps.append(
torch.stack([
torch.tensor([
step_points["upper_point"].x,
step_points["upper_point"].y
]),
torch.tensor([
step_points["base_point"].x,
step_points["base_point"].y
]),
torch.tensor([
step_points["lower_point"].x,
step_points["lower_point"].y
]),
torch.stack([torch.tensor(angle_to_next),
torch.tensor(0)]),
torch.stack([
torch.tensor(step_points["stop_confidence"]),
torch.tensor(0)
])
# torch.tensor(angle_after_that)
]))
img = cv2.imread(img_path)
if self.augmentation:
img = augmentation.apply_random_color_rotation(img)
img = augmentation.apply_tensmeyer_brightness(img)
img = img.astype(np.float32)
img = img.transpose()
img = img / 128.0 - 1.0
img = torch.from_numpy(img)
gt = gt_json[line_idx]['gt']
result = {
"img_path": img_path,
"img": img,
"steps": torch.stack(absolute_steps),
"gt": gt
}
return result