inference_scripts/generate_mask.py

import re
import math
from tqdm import tqdm
import os
from PIL import Image, ImageDraw
import numpy as np
import json
import argparse
import matplotlib.pyplot as plt
import matplotlib.patches as patches



def calculate_iou(box1, box2):
    """
    快速计算两个水平矩形框的 IoU（Intersection over Union）。

    参数:
    - box1, box2: (x1, y1, x2, y2) 矩形框，表示左上角和右下角的坐标。
    
    返回:
    - IoU (float): 交并比（Intersection over Union）
    """
    
    # 计算交集矩形的左上角和右下角坐标
    xi1 = max(box1[0], box2[0])
    yi1 = max(box1[1], box2[1])
    xi2 = min(box1[2], box2[2])
    yi2 = min(box1[3], box2[3])
    
    # 计算交集的宽和高，如果没有交集则宽或高为0
    inter_width = xi2 - xi1
    inter_height = yi2 - yi1
    
    if inter_width <= 0 or inter_height <= 0:
        return 0.0  # 没有交集
    
    # 交集面积
    inter_area = inter_width * inter_height
    
    # 计算每个矩形框的面积
    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
    
    # 并集面积 = 两个矩形的面积之和 - 交集面积
    union_area = box1_area + box2_area - inter_area
    
    # 计算IoU
    return inter_area / union_area


def decode_string(s, length):
    rows = s.strip().split(';')  # 先分割每一行
    result = []
    
    for row in rows:
        if not row.strip():  # 如果字符串是空的，跳过
            continue
        decoded_row = []
        
        groups = row.split(',')  # 分割每个组
        # print(row)
        # print(groups)
        for group in groups:
            try:
                num, count = group.split('*')  # 分割数字和次数
                decoded_row.extend([int(num.strip())] * int(count.strip()))  # 解码并扩展
            except:
                decoded_row.extend([0] * length)
        # print(len(decoded_row))
        if len(decoded_row) > length:
            decoded_row = decoded_row[:length]
            # print(decoded_row)
        elif len(decoded_row) < length:
            decoded_row.extend([0] * (length - len(decoded_row)))
            # print(decoded_row)
        result.append(decoded_row)  # 将解码的行添加到结果中


    if len(result) < length:
        for _ in range(length - len(result)):
            result.append([0] * length)
        # print(result)

    return result

def extract_bboxes(output, length):
    """
    Extract bounding box coordinates from the given string using regular expressions.
    :param output: String containing bounding box coordinates in the format {<bx_left><by_top><bx_right><by_bottom>|θ}
    :return: List of bounding boxes, each in the format [bx_left, by_top, bx_right, by_bottom, θ]
    """
    # 修改正则表达式，确保最后一个数字和管道符号能够正确匹配
    # segments = re.findall(r'<seg>(.*?)</seg>', output)[0]
    mask = decode_string(output, length)
    mask = np.array(mask)
    # print(mask.shape)
    arr = np.argwhere(mask == 1).tolist()

    # bboxes = [arr[0][1], arr[0][0], arr[-1][1], arr[-1][0]]
    # bboxes = [arr[0][1] + 0.5, arr[0][0] + 0.5, arr[-1][1] + 0.5, arr[-1][0] + 0.5]
    arr = [coord for pair in arr for coord in pair]
    arr = [point + 0.5 for point in arr]


    # 计算奇数索引和偶数索引的最大值和最小值
    if arr == []:
        return None, None
    odd_index_elements = [arr[i] for i in range(1, len(arr), 2)]
    even_index_elements = [arr[i] for i in range(0, len(arr), 2)]
    odd_max = max(odd_index_elements)
    odd_min = min(odd_index_elements)
    even_max = max(even_index_elements)
    even_min = min(even_index_elements)
    if odd_min == odd_max:
        odd_min -= 0.5
        odd_max += 0.5
    if even_min == even_max:
        even_min -= 0.5
        even_max += 0.5
    bboxes = [odd_min, even_min, odd_max, even_max]
    
    
    # quit()
    # pattern = r'\[([0-9, ]+)\]'
    # matches = re.findall(pattern, output)
    # bboxes = [list(map(float, match.split(","))) for match in matches]

    return bboxes, arr

# 读取JSONL文件并将每行解析为Python字典，存入列表
def load_jsonl(filename):
    data = []
    with open(filename, 'r') as jsonl_file:
        for line in jsonl_file:
            data.append(json.loads(line.strip()))
    return data

def folder_creat_if_not_exist(folder):
    if not os.path.exists(folder):
        os.makedirs(folder)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process some paths.')
    parser.add_argument('--scale', required=True, help='Normalize scale')
    parser.add_argument('--image-folder', required=True, help='Image directory')
    parser.add_argument('--answers-file', required=True, help='Target jsonl directory')
    parser.add_argument('--vis-dir', default=None, help='Base URL for the API')

    args = parser.parse_args()
    scale = int(args.scale)

    # 从 jsonl 文件中加载数据
    predict = load_jsonl(args.answers_file)
    total_cnt = len(predict)
    correct = 0
    format_error = 0
    i = 0
    for i, predict in tqdm(enumerate(predict), total=total_cnt):
        answer = predict['answer']
        answer = answer.strip()
        gt_bbox = predict['bbox']
  
        answer = answer.replace("others", "0")
        answer = answer.replace("object", "1")
        # answer = answer.replace("<seg>", "")
        # answer = answer.replace("</seg>", "")
        # try:
        predict_boxes, predict_points = extract_bboxes(answer, scale)
        if predict_boxes == None:
            format_error += 1
            continue   
        # except:
        #     format_error += 1
        #     continue            

        ori_img_path = args.image_folder + predict['image_id']
        image = Image.open(ori_img_path)
        width, height = image.size
        scale = int(args.scale)
        vis_dir = args.vis_dir

        if vis_dir:
            folder_creat_if_not_exist(vis_dir)
            draw = ImageDraw.Draw(image)
        
        try:
            pred_bbox = predict_boxes
            pred_bbox[0] = pred_bbox[0] / scale * width
            pred_bbox[1] = pred_bbox[1] / scale * height
            pred_bbox[2] = pred_bbox[2] / scale * width
            pred_bbox[3] = pred_bbox[3] / scale * height

            scaled_predict_points = [x  / scale * width if i % 2 == 0 else x / scale * height for i, x in enumerate(predict_points)]

            # compute  IoU
            iou_score = calculate_iou(gt_bbox, pred_bbox)
            if iou_score >= 0.5:
                correct += 1

            if vis_dir:

                coordinates = [num-0.5 for num in predict_points]
            
                # 定义图像的宽度和高度（单位为英寸）以及目标分辨率（单位为像素）
                width_inch = scale  # 图的逻辑宽度（英寸）
                height_inch = scale  # 图的逻辑高度（英寸）
                output_width = width  # 输出图像的宽度（像素）
                output_height = height  # 输出图像的高度（像素）

                # 计算 DPI（每英寸的像素点数）
                dpi_x = output_width / width_inch
                dpi_y = output_height / height_inch

                # 将坐标列表拆分为 x 和 y 的两部分
                # x_coords = coordinates[::2]  # 偶数索引为 x 坐标
                # y_coords = coordinates[1::2]  # 奇数索引为 y 坐标
                y_coords = coordinates[::2]  # 偶数索引为 x 坐标
                x_coords = coordinates[1::2]  # 奇数索引为 y 坐标

                # 创建图形和坐标轴，设置图的宽高为 width_inch x height_inch，背景为黑色
                fig, ax = plt.subplots(figsize=(width_inch, height_inch))
                ax.set_facecolor('black')  # 设置背景为黑色

                # 设置图的宽高范围为 0 到 24（根据你的坐标范围，也可以调整为其他范围）
                ax.set_xlim(0, width_inch)
                ax.set_ylim(0, height_inch)

                # 将 y 轴的方向翻转，使得图像的 (0, 0) 在左上角（符合常见图像坐标系）
                ax.invert_yaxis()

                # 绘制白色方块
                for x, y in zip(x_coords, y_coords):
                    # 创建一个左上角为 (x, y)，宽度和高度为 1 的矩形
                    rect = patches.Rectangle((x, y), 1.01, 1.01, linewidth=1, edgecolor='white', facecolor='white')
                    ax.add_patch(rect)

                # 隐藏网格线
                ax.grid(False)

                # 隐藏坐标轴
                ax.set_xticks([])
                ax.set_yticks([])

                # 保存为PNG图片，确保输出的宽度和高度为指定的像素值
                mask_save_path = vis_dir + predict['image_id'].split('.')[0] + f'_{i}.png'
                plt.savefig(mask_save_path, dpi=(dpi_x + dpi_y) / 2, bbox_inches='tight', pad_inches=0)

                # mask_save_path = vis_dir + predict['image_id'].split('.')[0] + f'_{i}.png'
                # cv2.imwrite(mask_save_path, predict_mask)
                draw.rectangle(gt_bbox, outline="red", width=5)
                draw.rectangle(pred_bbox, outline="blue", width=5)
                coordinates = [(int(scaled_predict_points[i+1]), int(scaled_predict_points[i])) for i in range(0, len(scaled_predict_points), 2)]
                point_color = (0, 255, 0)  # 红色
                radius = 5
                for point in coordinates:
                    x, y = point
                    # 用椭圆代表点，指定左上角和右下角的坐标范围
                    draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=point_color)
                image.save(vis_dir + predict['image_id'].split('.')[0] + f'_{i}.jpg')
        except:
            format_error += 1
            continue
        
    print(f"Evaluating  ...")
    print(f'Precision @ 0.5: {correct / total_cnt} \n')
    print(f'Format error ratio: {format_error / total_cnt} \n')