detection.py 文件源码

def main():
    # matplotlib.use('qt5agg')
    import matplotlib.pyplot as plt
    from matplotlib.patches import Rectangle

    init_model()
    MAT_DIR = './mat/test'
    LABEL_DIR = './label/test'
    for dirpath, dirnames, filenames in os.walk(MAT_DIR):
        print(dirpath)
        for filename in filenames:
            if filename == 'full.mat':
                data = sio.loadmat(os.path.join(dirpath, filename))
                img = data['data']
                centers = detection(img)
                img_id = dirpath.split('/')[-1]
                label_file = os.path.join(LABEL_DIR, img_id + '.mat')
                labels = sio.loadmat(label_file)['label']
                distance = (lambda x1, y1, x2, y2: abs(x1 - x2) + abs(y1 - y2))

                centers = cluster(centers)

                TP = 0
                for x, y in labels:
                    for x_, y_ in centers:
                        if distance(x, y, x_, y_) < 36:
                            TP += 1
                            break
                precision = float(TP) / len(centers)
                recall = float(TP) / len(labels)
                f_score = 2 * (precision * recall) / (precision + recall)
                six.print_(precision, recall, f_score)

                f = open(dirpath.split('/')[-1] + '-predict.txt', 'w')
                for x, y in centers:
                    f.write(str(x) + ' ' + str(y) + '\n')
                f.close()
                f = open(dirpath.split('/')[-1] + '-label.txt', 'w')
                for x, y in labels:
                    f.write(str(x) + ' ' + str(y) + '\n')
                f.close()

                # img = img / np.float32(256)
                # plt.imshow(img, cmap=plt.cm.gray)
                # currentAxis = plt.gca()
                # for x, y in labels:
                #     currentAxis.add_patch(Rectangle((y - 90, x - 90),
                #                                     180, 180, fill=None,
                #                                     alpha=1, color='blue'))
                # for x, y in centers:
                #     currentAxis.add_patch(Rectangle((y - 90, x - 90),
                #                                     180, 180, fill=None,
                #                                     alpha=1, color='red'))
                # plt.show()