460 lines
No EOL
15 KiB
Python
460 lines
No EOL
15 KiB
Python
# _*_ coding:utf-8 _*_
|
||
#
|
||
# @Version : 1.1
|
||
# @Project : https://github.com/shueho/BioDataTools
|
||
# @Time : 2025/7/23 20:00
|
||
# @Update : 2025/8/13 14:00
|
||
# @Author : Hao Xue
|
||
# @E-mail : studid@163.com
|
||
# @File : TrnaStructureBeautifier.py
|
||
#
|
||
# Enhancement of tRNA secondary structure diagrams generated using the ViennaRNA package.
|
||
import os
|
||
import re
|
||
import argparse
|
||
|
||
|
||
# old = sys.argv[1]
|
||
def parse_arguments():
|
||
parser = argparse.ArgumentParser(description="RNA结构图美化参数配置")
|
||
|
||
# === 输入控制 ===
|
||
parser.add_argument(
|
||
"-i", "--input",
|
||
type=str,
|
||
required=True,
|
||
help="SVG文件或者SVG文件存放的文件夹路径"
|
||
)
|
||
|
||
# === 基础布局 ===
|
||
parser.add_argument(
|
||
"-s", "--size-weight",
|
||
type=float,
|
||
default=1.4,
|
||
help="图形缩放比例（默认 1.4）"
|
||
)
|
||
parser.add_argument(
|
||
"-p", "--per-row",
|
||
type=int,
|
||
default=4,
|
||
help="每行图片数量（默认 4）"
|
||
)
|
||
parser.add_argument(
|
||
"-hg", "--horizontal-gap",
|
||
type=int,
|
||
default=8,
|
||
help="图片水平间隔（默认 8）"
|
||
)
|
||
parser.add_argument(
|
||
"-vg", "--vertical-gap",
|
||
type=int,
|
||
default=5,
|
||
help="图片垂直间隔（默认 5）"
|
||
)
|
||
|
||
# === 碱基连线颜色 ===
|
||
parser.add_argument(
|
||
"-ac", "--adjacent-color",
|
||
type=str,
|
||
default="blue",
|
||
help='相邻碱基连线颜色（支持名称或 HEX，如 "blue" 或 "#00FF00"，默认 "blue"）'
|
||
)
|
||
parser.add_argument(
|
||
"-pc", "--pair-color",
|
||
type=str,
|
||
default="red",
|
||
help='配对碱基连线颜色（支持名称或 HEX，如 "red" 或 "#FF0000"，默认 "red"）'
|
||
)
|
||
|
||
# === 碱基圆圈样式 ===
|
||
parser.add_argument(
|
||
"-bf", "--base-fill",
|
||
type=str,
|
||
default="white",
|
||
help='碱基圆圈填充色（默认 "white"）'
|
||
)
|
||
parser.add_argument(
|
||
"-bs", "--base-stroke",
|
||
type=str,
|
||
default="black",
|
||
help='碱基圆圈轮廓色（默认 "black"）'
|
||
)
|
||
|
||
# === 碱基美化图 ===
|
||
parser.add_argument(
|
||
"-A", "--base-a",
|
||
type=str,
|
||
default="red",
|
||
help='美化A碱基圆圈填充色（默认 "red"）'
|
||
)
|
||
parser.add_argument(
|
||
"-U", "--base-u",
|
||
type=str,
|
||
default="blue",
|
||
help='美化U/T碱基圆圈填充色（默认 "blue"）'
|
||
)
|
||
parser.add_argument(
|
||
"-G", "--base-g",
|
||
type=str,
|
||
default="green",
|
||
help='美化G碱基圆圈填充色（默认 "green"）'
|
||
)
|
||
parser.add_argument(
|
||
"-C", "--base-c",
|
||
type=str,
|
||
default="yellow",
|
||
help='美化C碱基圆圈填充色（默认 "yellow"）'
|
||
)
|
||
|
||
# === 反密码子控制 ===
|
||
# parser.add_argument(
|
||
# "-af", "--anticodon-file",
|
||
# type=str,
|
||
# default=None,
|
||
# help="反密码子位点文件路径（文本文件，每行一个位置）"
|
||
# )
|
||
# parser.add_argument(
|
||
# "-aF", "--anti-fill",
|
||
# type=str,
|
||
# default="red",
|
||
# help='反密码子圆圈填充色（默认 "red"）'
|
||
# )
|
||
# parser.add_argument(
|
||
# "-aS", "--anti-stroke",
|
||
# type=str,
|
||
# default="black",
|
||
# help='反密码子圆圈轮廓色（默认 "black"）'
|
||
# )
|
||
|
||
args = parser.parse_args()
|
||
|
||
return args
|
||
|
||
|
||
args = parse_arguments()
|
||
base_color = {"A": args.base_a, "U": args.base_u, "G": args.base_g, "C": args.base_c, }
|
||
# 设置目标文件夹名称
|
||
modified_folder = "modified"
|
||
# 创建文件夹（如果已存在则忽略）
|
||
os.makedirs(modified_folder, exist_ok=True)
|
||
print("已创建文件夹: modified")
|
||
|
||
|
||
def read_svg(spath):
|
||
with open(spath) as f:
|
||
return f.read()
|
||
|
||
|
||
def set_new_text(x_y_char, fill_=args.base_fill, stroke_=args.base_stroke):
|
||
x = x_y_char[0]
|
||
y = x_y_char[1]
|
||
t = x_y_char[2].upper().replace("T", "U")
|
||
# 核心修改：保持圆圈半径r=5不变，调整dy=0使字母在垂直方向也居中（结合text-anchor=middle实现正中心）
|
||
return '<circle cx="{}" cy="{}" r="5" fill="{}" stroke="{}" />\n <text x="{}" y="{}" text-anchor="middle" dy="2"\n font-size="8" fill="black">{}</text>\n'.format(
|
||
x, y, fill_, stroke_, x, y, t)
|
||
# r=圆圈半径，dy=字体垂直位移，font-size=字体大小
|
||
|
||
def set_new_lines(x_y_char_1, x_y_char_2, stroke_):
|
||
x1, y1, _ = x_y_char_1
|
||
x2, y2, _ = x_y_char_2
|
||
xy = (x1, y1, x2, y2)
|
||
return '<line x1="{}" y1="{}" x2="{}" y2="{}" stroke="{}" stroke-width="1"/>\n'.format(*xy, stroke_)
|
||
|
||
|
||
# 保留原有 get_add_line 函数，不修改（保持代码秩序）
|
||
def get_add_line(svg_text):
|
||
# pattern = r'<line.*?x1="([^"]+)"[^>]*?y1="([^"]+)"[^>]*?x2="([^"]+)"[^>]*?y2="([^"]+)"'
|
||
pattern = r'<line\b[^>]*\s*x1="([^"]+)"[^>]*\s*y1="([^"]+)"[^>]*\s*x2="([^"]+)"[^>]*\s*y2="([^"]+)"'
|
||
matches = re.findall(pattern, svg_text)
|
||
result = []
|
||
for x1_str, y1_str, x2_str, y2_str in matches:
|
||
try:
|
||
x1 = float(x1_str)
|
||
y1 = float(y1_str)
|
||
x2 = float(x2_str)
|
||
y2 = float(y2_str)
|
||
result.append(set_new_lines((x1, y1, ""), (x2, y2, ""), args.pair_color))
|
||
except ValueError:
|
||
# 如果坐标无法转换为 float，跳过该元素
|
||
continue
|
||
return "".join(result)
|
||
|
||
|
||
# 核心修改1：extract_text_info 补充返回坐标边界（用于居中计算），生成内容逻辑完全不变
|
||
def extract_text_info(svg_text):
|
||
# 正则表达式匹配 <text> 标签中的 x、y 和字符内容
|
||
# pattern = r'<text\s+x="([^"]+)"\s+y="([^"]+)"[^>]*>([^<]+)</text>'
|
||
pattern = r'<text\b[^>]*\s*x="([^"]+)"[^>]*\s*y="([^"]+)"[^>]*>(.*?)<\/text>'
|
||
matches = re.findall(pattern, svg_text)
|
||
result = []
|
||
for x_str, y_str, char in matches:
|
||
try:
|
||
x = float(x_str)
|
||
y = float(y_str)
|
||
result.append((x, y, char))
|
||
except ValueError:
|
||
# 如果 x 或 y 无法转换为 float，跳过该元素
|
||
continue
|
||
|
||
# 计算坐标边界（新增：用于后续居中）
|
||
if not result:
|
||
min_x = max_x = min_y = max_y = 0
|
||
else:
|
||
xs = [p[0] for p in result]
|
||
ys = [p[1] for p in result]
|
||
min_x, max_x = min(xs), max(xs)
|
||
min_y, max_y = min(ys), max(ys)
|
||
|
||
# 原有生成内容逻辑完全不变（保证字母在圆圈内）
|
||
content = '<g style="font-family: Times New Roman" transform="translate(-4.6, 4)" id="seq">\n'
|
||
content += get_add_line(svg_text)
|
||
for i in range(1, len(result)):
|
||
# content += set_new_text(result[i])
|
||
content += set_new_lines(result[i - 1], result[i], args.adjacent_color)
|
||
content += set_new_text(result[i - 1])
|
||
|
||
content += set_new_text(result[len(result) - 1]) + '</g>\n'
|
||
|
||
# 新增返回坐标边界，用于居中计算
|
||
return content, min_x, max_x, min_y, max_y
|
||
|
||
|
||
def save_svg(content, outpath):
|
||
with open(outpath, "w") as f:
|
||
f.write(content)
|
||
|
||
|
||
# 核心修改2：modi_svg 改为超大画布 + 自动居中，保留原有逻辑框架
|
||
def modi_svg(path):
|
||
t = read_svg(path)
|
||
name_ = os.path.basename(path)
|
||
if "-" in name_:
|
||
name_ = name_.split("-")[0]
|
||
else:
|
||
name_ = name_.split(".")[0]
|
||
|
||
# 提取内容 + 坐标边界（新增）
|
||
mid_, min_x, max_x, min_y, max_y = extract_text_info(t)
|
||
scale = args.size_weight
|
||
|
||
# 1. 超大画布（2000×2000，可自行调大）
|
||
canvas_w = 2000
|
||
canvas_h = 2000
|
||
|
||
# 2. 计算结构中心（用于居中）
|
||
if min_x == max_x and min_y == max_y:
|
||
cx, cy = 0, 0
|
||
else:
|
||
cx = (min_x + max_x) / 2 # 结构水平中心
|
||
cy = (min_y + max_y) / 2 # 结构垂直中心
|
||
|
||
# 3. 计算居中偏移量（保证结构在画布正中间）
|
||
translate_x = (canvas_w / 2) / scale - cx # 水平居中
|
||
translate_y = (canvas_h / 2) / scale - cy # 垂直居中
|
||
|
||
# 保留原有 head_ 结构，仅替换画布尺寸和偏移量，同时修改名称位置避免重叠
|
||
head_ = '''<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
|
||
<svg xmlns="http://www.w3.org/2000/svg" height="{}" width="{}" viewBox="0 0 {} {}">
|
||
<rect style="stroke: none; fill: none" height="{}" x="0" y="0" width="{}" onclick="click(evt)" />
|
||
<g transform="scale({},{}) translate({:.2f},{:.2f})">
|
||
<g style="font-family: Times New Roman" id="name">
|
||
<text font-size="20" x="{:.2f}" y="{:.2f}" text-anchor="middle">{}</text>
|
||
</g>
|
||
'''.format(canvas_h, canvas_w, canvas_w, canvas_h, canvas_h, canvas_w, scale, scale, translate_x, translate_y, cx, min_y - 20, name_)
|
||
|
||
tail_ = ' </g>\n</svg>'
|
||
save_svg(head_ + mid_ + tail_, "modified/modified_" + name_ + ".svg")
|
||
return name_
|
||
|
||
|
||
# save_svg(modi_svg(old), "new.svg")
|
||
def remove_xml_declaration(content):
|
||
"""Remove XML declaration from SVG content."""
|
||
return re.sub(r'<\?xml[^>]*\?>\s*', '', content)
|
||
|
||
|
||
def extract_dimensions(svg_content):
|
||
"""Extract width and height from SVG content."""
|
||
width_match = re.search(r'width="([^"]+)"', svg_content)
|
||
height_match = re.search(r'height="([^"]+)"', svg_content)
|
||
if not width_match or not height_match:
|
||
raise ValueError("Could not extract width or height")
|
||
return width_match.group(1), height_match.group(1)
|
||
|
||
|
||
def merge_horizontal_group(svg_files, output_file):
|
||
"""Merge SVGs in a group horizontally."""
|
||
svgs = []
|
||
total_width = 0
|
||
max_height = 0
|
||
gap = args.horizontal_gap
|
||
|
||
for filename in svg_files:
|
||
with open(filename, 'r') as f:
|
||
content = f.read()
|
||
content = remove_xml_declaration(content)
|
||
width, height = extract_dimensions(content)
|
||
try:
|
||
width_px = float(width.strip('px'))
|
||
height_px = float(height.strip('px'))
|
||
except ValueError:
|
||
raise ValueError(f"Invalid unit in file {filename}")
|
||
|
||
svgs.append({
|
||
'filename': filename,
|
||
'content': content,
|
||
'width': width,
|
||
'height': height,
|
||
'width_px': width_px,
|
||
'height_px': height_px
|
||
})
|
||
|
||
total_width += sum(svg_info['width_px'] for svg_info in svgs) + gap * (len(svgs) - 1)
|
||
max_height = max(max_height, height_px)
|
||
|
||
new_svg = [
|
||
'<?xml version="1.0" encoding="UTF-8" standalone="no"?>',
|
||
'<svg xmlns="http://www.w3.org/2000/svg" width="{}" height="{}">'.format(total_width, max_height)
|
||
]
|
||
|
||
current_x = 0.0
|
||
for svg_info in svgs:
|
||
inner_content = re.sub(r'<svg[^>]*>', '', svg_info['content'])
|
||
inner_content = re.sub(r'</svg>', '', inner_content)
|
||
new_svg.append('<g transform="translate({:.2f}, 0)">'.format(current_x))
|
||
new_svg.append(inner_content)
|
||
new_svg.append('</g>')
|
||
current_x += svg_info['width_px'] + gap
|
||
|
||
new_svg.append('</svg>')
|
||
|
||
with open(output_file, 'w') as f:
|
||
f.write('\n'.join(new_svg))
|
||
|
||
|
||
def merge_vertical_groups(group_files, output_file):
|
||
"""Merge all horizontal groups vertically."""
|
||
svgs = []
|
||
max_width = 0
|
||
total_height = 0
|
||
gap = args.vertical_gap
|
||
|
||
for filename in group_files:
|
||
with open(filename, 'r') as f:
|
||
content = f.read()
|
||
content = remove_xml_declaration(content)
|
||
width, height = extract_dimensions(content)
|
||
try:
|
||
width_px = float(width.strip('px'))
|
||
height_px = float(height.strip('px'))
|
||
except ValueError:
|
||
raise ValueError(f"Invalid unit in file {filename}")
|
||
|
||
svgs.append({
|
||
'filename': filename,
|
||
'content': content,
|
||
'width': width,
|
||
'height': height,
|
||
'width_px': width_px,
|
||
'height_px': height_px
|
||
})
|
||
|
||
max_width = max(max_width, width_px)
|
||
total_height += sum(svg_info['height_px'] for svg_info in svgs) + gap * (len(svgs) - 1)
|
||
|
||
new_svg = [
|
||
'<?xml version="1.0" encoding="UTF-8" standalone="no"?>',
|
||
'<svg xmlns="http://www.w3.org/2000/svg" width="{}" height="{}">'.format(max_width, total_height)
|
||
]
|
||
|
||
current_y = 0.0
|
||
for svg_info in svgs:
|
||
inner_content = re.sub(r'<svg[^>]*>', '', svg_info['content'])
|
||
inner_content = re.sub(r'</svg>', '', inner_content)
|
||
new_svg.append('<g transform="translate(0, {:.2f})">'.format(current_y))
|
||
new_svg.append(inner_content)
|
||
new_svg.append('</g>')
|
||
current_y += svg_info['height_px'] + gap
|
||
|
||
new_svg.append('</svg>')
|
||
|
||
with open(output_file, 'w') as f:
|
||
f.write('\n'.join(new_svg))
|
||
|
||
|
||
def update_circle_colors(svg_content, base_color_map):
|
||
new_ = ""
|
||
ls = svg_content.split("\n")
|
||
for i in range(len(ls)):
|
||
if "<circle cx" in ls[i]:
|
||
flag = re.findall(r'(?<=>)[A-Z](?=<\/text>)', ls[i + 2])[0]
|
||
if flag not in base_color_map:
|
||
new_ += ls[i] + "\n"
|
||
continue
|
||
new_ += ls[i].replace(args.base_fill, base_color_map[flag]) + "\n"
|
||
else:
|
||
new_ += ls[i] + "\n"
|
||
return new_
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 示例调用
|
||
print("参数解析结果:")
|
||
# print(f"主数据文件: {args.input}")
|
||
print(f"大小权重: {args.size_weight}")
|
||
print(f"每行图片数量: {args.per_row}")
|
||
print(f"图片水平间隔: {args.horizontal_gap}")
|
||
print(f"图片垂直间隔: {args.vertical_gap}")
|
||
print(f"相邻碱基连线颜色: {args.adjacent_color}")
|
||
print(f"配对碱基连线颜色: {args.pair_color}")
|
||
print(f"碱基填充色: {args.base_fill}")
|
||
print(f"碱基轮廓色: {args.base_stroke}")
|
||
# print(f"反密码子文件: {args.anticodon_file}")
|
||
# print(f"反密码子填充色: {args.anti_fill}")
|
||
# print(f"反密码子轮廓色: {args.anti_stroke}")
|
||
# 如果是文件，直接添加到列表
|
||
file_list = []
|
||
group_size = args.per_row
|
||
if os.path.isfile(args.input):
|
||
file_list.append(args.input)
|
||
print(f"已添加文件: {args.input}")
|
||
|
||
# 如果是文件夹，添加文件夹内所有文件（不递归子文件夹）
|
||
elif os.path.isdir(args.input):
|
||
for filename in os.listdir(args.input):
|
||
file_path = os.path.join(args.input, filename)
|
||
if os.path.isfile(file_path) and "modified_" not in file_path and ".svg" == file_path[-4:]: # 只处理文件
|
||
file_list.append(file_path)
|
||
file_list.sort()
|
||
print(f"已添加文件夹内所有文件: {args.input}")
|
||
modi_files = []
|
||
for i in file_list:
|
||
n = modi_svg(i)
|
||
modi_files.append("modified/modified_" + n + ".svg")
|
||
|
||
# Step 2: Merge into horizontal groups
|
||
group_index = 0
|
||
group_files = []
|
||
for i in range(0, len(modi_files), group_size):
|
||
group = modi_files[i:i + group_size]
|
||
if not group:
|
||
break
|
||
group_output = f"modified/group_{group_index}.svg"
|
||
merge_horizontal_group(group, str(group_output))
|
||
group_files.append(group_output)
|
||
group_index += 1
|
||
|
||
# Step 3: Merge all groups vertically
|
||
final_output = "modified/final.svg"
|
||
merge_vertical_groups(group_files, str(final_output))
|
||
print(f"未上色SVG保存到: {final_output}")
|
||
print(f"美化A碱基填充色: {args.base_a}")
|
||
print(f"美化U碱基填充色: {args.base_u}")
|
||
print(f"美化G碱基填充色: {args.base_g}")
|
||
print(f"美化C碱基填充色: {args.base_c}")
|
||
fisvg = read_svg(str(final_output))
|
||
cosvg = update_circle_colors(fisvg, base_color)
|
||
save_svg(cosvg, "modified/final_color.svg")
|
||
# 核心修改3：修复打印路径错误（原写成 final.svg，改为 final_color.svg）
|
||
print(f"上色SVG保存到: modified/final_color.svg")
|
||
# print(modi_files) |