def dumpTextImg(self,i):
box_num = 0
txt_name = os.path.join(self.output_dir, '{}_{}.txt'.format(self.dump_prefix,str(i).zfill(6)))
fw = open(txt_name, 'w')
for idx, fg_image in enumerate(self.fg_images):
box = self.text_boxes[idx]
if box is None:
continue
self.pil_img.paste(fg_image, (box[0], box[1]), fg_image.split()[3])
coord = self.text_coordinates_in_fg[idx] + [box[0], box[1]]
s = ""
for pt in coord:
x, y = pt
s = s+str(x)+','+str(y)+','
s += '0'
fw.write(s+'\n')
box_num += 1
fw.close()
pasted_img = cv2.cvtColor(np.asarray(self.pil_img),cv2.COLOR_RGB2BGR)
image_name = '{}_{}.jpg'.format(self.dump_prefix,str(i).zfill(6))
self.dumpImgToPath(image_name, pasted_img)
if box_num<=0:
LOG.logW("Image {} paste no text.".format(i))
def earlyIter(self):
start = time.time()
self.sample = self.sample.to(self.device)
self.target = [anno.to(self.device) for anno in self.target]
if not self.is_train:
return
self.data_cpu2gpu_time.update(time.time() - start)
try:
self.addGraph(self.sample)
except:
LOG.logW(
"Tensorboard addGraph failed. You network foward may have more than one parameters?"
)
LOG.logW("Seems you need reimplement preIter function.")
def postIter(self):
preds = self.config.output[
self.config.output[..., 4] >
self.config.conf_thres] # filter with classifier confidence
if not preds.size(0):
LOG.logW("file: {0} >>> non - detect".format(self.config.filepath))
self.config.non_det_num += 1
return
# Compute conf
preds[:, 5:] *= preds[:, 4:5] # conf = obj_conf * cls_conf
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
preds[:, 0] = preds[:, 0] - preds[:, 2] / 2.0
preds[:, 1] = preds[:, 1] - preds[:, 3] / 2.0
preds[:, 2] += preds[:, 0]
preds[:, 3] += preds[:, 1]
# Detections matrix nx6 (xyxy, conf, cls)
conf, idx = preds[:, 5:].max(dim=1, keepdim=True)
preds = torch.cat(
(preds[:, :4], conf, idx),
dim=1)[conf.view(-1) >
self.config.conf_thres] # filter with bbox confidence
if not preds.size(0):
LOG.logW("file: {0} >>> non - detect".format(self.config.filepath))
self.config.non_det_num += 1
return
# nms on per class
max_side = 4096
class_offset = preds[:, 5:6] * max_side
boxes, scores = preds[:, :4] + class_offset, preds[:, 4]
idxs = ops.nms(boxes, scores, self.config.iou_thres)
preds = torch.stack([preds[i] for i in idxs], dim=0)
if not preds.size(0):
LOG.logW("file: {0} >>> non - detect".format(self.config.filepath))
self.config.non_det_num += 1
return
# coords scale
classes = preds[:, -1].long().tolist()
scores = [i.item() for i in preds[:, -2]]
coords = preds[:, :4]
coords -= self.config.pad
coords /= self.config.ratio
coords = coords.long().tolist()
LOG.logI(
"file: {0} >>> class: {1} >>> score: {2} >>> coord: {3}".format(
self.config.filepath, classes, scores, coords))
self.plotRectangle(self.config.filepath, (classes, scores, coords),
self.config.show_output_dir)
def test(self):
self.config.non_det_num = 0
super(Yolov5Test, self).test()
LOG.logW("\nnon - detect: {} ".format(self.config.non_det_num) +
' * ' * 70)