def __getitem__(self, item):
image_id = self.img_list[item]
if self.is_training:
# Read annotation
annotation_id = "{}/gt_{}".format(
"/".join(image_id.split("/")[0:-1]),
image_id.split("/")[-1].replace(".jpg", ''))
annotation_path = os.path.join(self.data_root, annotation_id)
polygons = self.parse_txt(annotation_path)
else:
polygons = None
# Read image data
image_path = os.path.join(self.data_root, image_id)
image = pil_load_img(image_path)
try:
h, w, c = image.shape
assert (c == 3)
except:
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = np.array(image)
print("MMMMMMMMMMMMMMMMMM")
return self.get_training_data(image,
polygons,
image_id=image_id,
image_path=image_path)
def __getitem__(self, item):
# Read image data
image_id = self.image_list[item]
image_path = os.path.join(self.image_root, image_id)
image = pil_load_img(image_path)
return self.get_test_data(image, image_id=image_id, image_path=image_path)
def __getitem__(self, item):
# Read annotation
annotation_id = self.annotation_list[item]
annotation_path = os.path.join(self.annotation_root, annotation_id)
image_id, polygons = self.parse_txt(annotation_path)
# Read image data
image_path = os.path.join(self.image_root, image_id)
image = pil_load_img(image_path)
return self.get_training_data(image,
polygons,
image_id=image_id,
image_path=image_path)
def __getitem__(self, item):
image_id = self.image_list[item]
image_path = os.path.join(self.image_root, image_id)
# Read image data
image = pil_load_img(image_path)
try:
# Read annotation
annotation_id = self.annotation_list[item]
annotation_path = os.path.join(self.annotation_root, annotation_id)
polygons = self.parse_txt(annotation_path)
except Exception as e:
print(e)
polygons = None
return self.get_training_data(image, polygons, image_id=image_id, image_path=image_path)
def __getitem__(self, item):
image_id = self.image_list[item]
image_path = os.path.join(self.image_root, image_id)
image = pil_load_img(image_path)
annotation_id = self.annotation_list[item]
annotation_path = os.path.join(self.annotation_root, annotation_id)
polygons = self.parse_txt(annotation_path)
for i, polygon in enumerate(polygons):
polygon.find_bottom_and_sideline()
return self.get_training_data(image,
polygons,
image_id=image_id,
image_path=image_path)
def __getitem__(self, item):
image_id = self.image_list[item]
image_path = os.path.join(self.image_root, image_id)
# print(image_path)
annotation_id = self.annotation_list[item]
annotation_path = os.path.join(self.annotation_root, annotation_id)
polygons = self.parse_mat(annotation_path)
for i, polygon in enumerate(polygons):
if polygon.text != '#':
polygon.find_bottom_and_sideline()
# print(image_path, annotation_path)
# Read image data
image = pil_load_img(image_path)
H, W, _ = image.shape
if self.transform:
image, polygons = self.transform(image, copy.copy(polygons))
tcl_mask = np.zeros(image.shape[:2], np.uint8)
radius_map = np.zeros(image.shape[:2], np.float32)
sin_map = np.zeros(image.shape[:2], np.float32)
cos_map = np.zeros(image.shape[:2], np.float32)
for i, polygon in enumerate(polygons):
if polygon.text != '#':
sideline1, sideline2, center_points, radius = polygon.disk_cover(
n_disk=cfg.n_disk)
self.make_text_center_line(sideline1, sideline2, center_points,
radius, tcl_mask, radius_map,
sin_map, cos_map)
tr_mask, train_mask = self.make_text_region(image, polygons)
# to pytorch channel sequence
image = image.transpose(2, 0, 1)
meta = {
'image_id': image_id,
'image_path': image_path,
'Height': H,
'Width': W
}
return image, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map, meta
def __getitem__(self, item):
image_id = self.image_list[item]
image_path = os.path.join(self.image_root, image_id)
# Read image data
image = pil_load_img(image_path)
try:
h, w, c = image.shape
assert (c == 3)
except:
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = np.array(image)
return self.get_test_data(image,
image_id=image_id,
image_path=image_path)
def __getitem__(self, item):
image_id = self.image_list[item]
image_path = os.path.join(self.image_root, image_id)
# Read image data
image = pil_load_img(image_path)
# Read annotation
polygons = self.load_json(os.path.join(self.annotation_root, image_id.replace('jpg', 'json')))
for i, polygon in enumerate(polygons):
if polygon.text != '#':
polygon.find_bottom_and_sideline()
try:
return self.get_training_data(image, polygons, image_id=image_id, image_path=image_path)
except Exception as e:
print(e)
def __getitem__(self, item):
# Read annotation
annotation_id = self.annotation_list[item]
annotation_path = os.path.join(self.annotation_root, annotation_id)
image_id, polygons = self.parse_txt(annotation_path)
# Read image data
image_path = os.path.join(self.image_root, image_id)
# print(image_path)
image = pil_load_img(image_path)
# cv2.imwrite('image.jpg', image)
for i, polygon in enumerate(polygons):
if polygon.text != '#':
polygon.find_bottom_and_sideline()
return self.get_training_data(image,
polygons,
image_id=image_id,
image_path=image_path)
def __getitem__(self, item):
image_id = self.image_list[item]
image_path = os.path.join(self.image_root, image_id)
# Read image data
image = pil_load_img(image_path)
try:
h, w, c = image.shape
assert (c == 3)
except:
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = np.array(image)
# Read annotation
annotation_id = self.annotation_list[item]
annotation_path = os.path.join(self.annotation_root, annotation_id)
polygons = self.parse_carve_txt(annotation_path)
return self.get_training_data(image,
polygons,
image_id=image_id,
image_path=image_path)
def __getitem__(self, item):
image_id = self.image_list[item]
image_path = os.path.join(self.image_root, image_id)
# Read image data
# PIL读取图片后转为numpy
image = pil_load_img(image_path)
# Read annotation
annotation_id = self.annotation_list[item]
annotation_path = os.path.join(self.annotation_root, annotation_id)
#读取mat文件
polygons = self.parse_mat(annotation_path)
for i, polygon in enumerate(polygons):
if polygon.text != '#':
# 计算首位俩边和上下俩边
polygon.find_bottom_and_sideline()
return self.get_training_data(image,
polygons,
image_id=image_id,
image_path=image_path)
def Predict(self,
image_path,
output_img_path="output.jpg",
output_txt_path="output.txt",
tr_thresh=0.4,
tcl_thresh=0.4):
cfg = self.system_dict["local"]["cfg"]
model = self.system_dict["local"]["model"]
start = time.time()
image = pil_load_img(image_path)
transform = BaseTransform(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds)
H, W, _ = image.shape
image, polygons = transform(image)
# to pytorch channel sequence
image = image.transpose(2, 0, 1)
meta = {
'image_id': 0,
'image_path': image_path,
'Height': H,
'Width': W
}
image = torch.from_numpy(np.expand_dims(image, axis=0))
image = to_device(image)
if (self.system_dict["local"]["cfg"].cuda):
torch.cuda.synchronize()
end = time.time()
print("Image loading time: {}".format(end - start))
start = time.time()
detector = TextDetector(model,
tr_thresh=tr_thresh,
tcl_thresh=tcl_thresh)
# get detection result
contours, output = detector.detect(image)
torch.cuda.synchronize()
end = time.time()
print("Inference time - {}".format(end - start))
start = time.time()
tr_pred, tcl_pred = output['tr'], output['tcl']
img_show = image[0].permute(1, 2, 0).cpu().numpy()
img_show = ((img_show * cfg.stds + cfg.means) * 255).astype(np.uint8)
img_show, contours = rescale_result(img_show, contours, H, W)
pred_vis = visualize_detection(img_show, contours)
cv2.imwrite(output_img_path, pred_vis)
# write to file
self.write_to_file(contours, output_txt_path)
end = time.time()
print("Writing output time - {}".format(end - start))
