def annotations(self, i): annotations = list() tree = ET.parse( os.path.join(self.images[i][0], 'Annotations', self.images[i][1] + '.xml')) for child in tree.getroot(): if not child.tag == 'object': continue if not self.difficult and bool(int(child.find('difficult').text)): continue bndbox = child.find('bndbox') rect = Rect.LTRB(*(float(bndbox.find(t).text) for t in ('xmin', 'ymin', 'xmax', 'ymax'))) cls = self.names.index(child.find('name').text) annotations.append((rect, cls)) return annotations
src = cv2.imread(args.image, cv2.IMREAD_COLOR) x = cv2.resize(src, (size, size)).astype(np.float32) x -= (103.939, 116.779, 123.68) x = x.transpose(2, 0, 1) x = x[np.newaxis] loc, conf = model(x) boxes, conf = multibox_encoder.decode(loc.data[0], conf.data[0]) conf = conf[:, 1:] img = src.copy() selected = set() for i in conf.max(axis=1).argsort()[::-1]: box = Rect.LTRB(*boxes[i]).scale(*img.shape[:2][::-1]) if len(selected) > 0: iou = max(box.iou(s) for s in selected) if iou > 0.45: continue selected.add(box) if conf[i].max() < 0.9: break cv2.rectangle( img, (int(box.left), int(box.top)), (int(box.right), int(box.bottom)), (0, 0, 255), 3)