plot_sourece_labels.py 3.54 KB
import os

import cv2
import numpy as np
import pandas as pd
import tqdm


def get_source_image_det(crop_position, predict_positions):
    result = []
    x1, y1, x2, y2 = crop_position
    for p in predict_positions:
        px1, py1, px2, py2, score = p
        w, h = px2 - px1, py2 - py1
        result.append([x1 + px1, y1 + py1, x1 + px1 + w, y1 + py1 + h, score])
    return result


def decode_label(image, label_path):
    data = open(label_path).readlines()
    h, w, c = image.shape
    result = []
    for d in data:
        d = [float(i) for i in d.strip().split(' ')]
        cls, cx, cy, cw, ch, score = d
        cx, cy, cw, ch = cx * w, cy * h, cw * w, ch * h
        result.append([int(cx - cw // 2), int(cy - ch // 2), int(cx + cw // 2), int(cy + ch // 2), score])
    return result


if __name__ == '__main__':
    source_image_path = '/data/situ_invoice_bill_data/new_data/qfs_bank_bill_data/gongshang/authentic/images/val'
    val_image_path = '/home/situ/qfs/invoice_tamper/09_project/project/tamper_det/data/images/crop_img'
    predict_label_path = '/home/situ/qfs/invoice_tamper/09_project/project/tamper_det/runs/detect/exp/labels'
    crop_csv_path = '/data/situ_invoice_bill_data/new_data/qfs_bank_bill_data/gongshang/croped_merge.csv'
    predict_labels = os.listdir(predict_label_path)
    source_images = os.listdir(source_image_path)
    data = pd.read_csv(crop_csv_path)
    img_name = data.loc[:, 'img_name'].tolist()
    crop_position1 = data.loc[:, 'name_crop_coord'].tolist()
    crop_position2 = data.loc[:, 'number_crop_coord'].tolist()
    cc = '/data/situ_invoice_bill_data/new_data/qfs_bank_bill_data/gongshang/tampered/images/val/ps3'
    for im in os.listdir(cc):
        print(im)
        img = cv2.imread(os.path.join(cc, im))
        img_ = img.copy()
        id = img_name.index(im)
        name_crop_position = [int(i) for i in crop_position1[id].split(',')]
        number_crop_position = [int(i) for i in crop_position2[id].split(',')]
        nx1, ny1, nx2, ny2 = name_crop_position
        nux1, nuy1, nux2, nuy2 = number_crop_position
        if im[:-4] + '_hname.txt' in predict_labels:

            h, w, c = img[ny1:ny2, nx1:nx2, :].shape
            data = open(os.path.join(predict_label_path, im[:-4] + '_hname.txt')).readlines()
            for d in data:
                cls, cx, cy, cw, ch, score = [float(i) for i in d.strip().split(' ')]
                cx, cy, cw, ch = int(cx * w), int(cy * h), int(cw * w), int(ch * h)
                cx1, cy1 = cx - cw // 2, cy - ch // 2
                x1, y1, x2, y2 = nx1 + cx1, ny1 + cy1, nx1 + cx1 + cw, ny1 + cy1 + ch
                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
                cv2.putText(img, f'tampered:{score}', (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
        if im[:-4] + '_hnumber.txt' in predict_labels:
            h, w, c = img[nuy1:nuy2, nux1:nux2, :].shape
            data = open(os.path.join(predict_label_path, im[:-4] + '_hname.txt')).readlines()
            for d in data:
                cls, cx, cy, cw, ch, score = [float(i) for i in d.strip().split(' ')]
                cx, cy, cw, ch = int(cx * w), int(cy * h), int(cw * w), int(ch * h)
                cx1, cy1 = cx - cw // 2, cy - ch // 2
                x1, y1, x2, y2 = nux1 + cx1, nuy1 + cy1, nux1 + cx1 + cw, nuy1 + cy1 + ch
                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
                cv2.putText(img, f'tampered:{score}', (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
        result = np.vstack((img_, img))
        cv2.imwrite(f'z/{im}', result)