diff --git a/codes/ecg-image-generator/extract_leads.py b/codes/ecg-image-generator/extract_leads.py index 9516ca4..d73214b 100644 --- a/codes/ecg-image-generator/extract_leads.py +++ b/codes/ecg-image-generator/extract_leads.py @@ -2,6 +2,7 @@ # Load libraries. import os, sys, argparse import json +import qrcode import numpy as np from scipy.io import savemat, loadmat import matplotlib.pyplot as plt @@ -16,7 +17,7 @@ import random # Run script. -def get_paper_ecg(input_file,header_file,output_directory, seed, add_dc_pulse,add_bw,show_grid, add_print, configs, mask_unplotted_samples = False, start_index = -1, store_configs=False, store_text_bbox=True,key='val',resolution=100,units='inches',papersize='',add_lead_names=True,pad_inches=1,template_file=os.path.join('TemplateFiles','TextFile1.txt'),font_type=os.path.join('Fonts','Times_New_Roman.ttf'),standard_colours=5,full_mode='II',bbox = False,columns=-1): +def get_paper_ecg(input_file,header_file,output_directory, seed, add_dc_pulse,add_bw,show_grid, add_print, configs, encoding, mask_unplotted_samples = False, start_index = -1, store_configs=False, store_text_bbox=True,key='val',resolution=100,units='inches',papersize='',add_lead_names=True,pad_inches=1,template_file=os.path.join('TemplateFiles','TextFile1.txt'),font_type=os.path.join('Fonts','Times_New_Roman.ttf'),standard_colours=5,full_mode='II',bbox = False,columns=-1, add_qr_code=False): # Extract a reduced-lead set from each pair of full-lead header and recording files. full_header_file = header_file @@ -270,7 +271,7 @@ def get_paper_ecg(input_file,header_file,output_directory, seed, add_dc_pulse,ad x_grid,y_grid = ecg_plot(ecg_frame[i], configs=configs, full_header_file=full_header_file, style=grid_colour, sample_rate = rate,columns=columns,rec_file_name = rec_file, output_dir = output_directory, resolution = resolution, pad_inches = pad_inches, lead_index=full_leads, full_mode = full_mode, store_text_bbox = store_text_bbox, show_lead_name=add_lead_names,show_dc_pulse=dc,papersize=papersize,show_grid=(grid),standard_colours=standard_colours,bbox=bbox, print_txt=print_txt, json_dict=json_dict, start_index=start, store_configs=store_configs, lead_length_in_seconds=lead_length_in_seconds) rec_head, rec_tail = os.path.split(rec_file) - + json_dict["x_grid"] = round(x_grid, 3) json_dict["y_grid"] = round(y_grid, 3) json_dict["resolution"] =resolution @@ -285,9 +286,30 @@ def get_paper_ecg(input_file,header_file,output_directory, seed, add_dc_pulse,ad json_dict["full_mode_lead"] =full_mode outfile = os.path.join(output_directory,rec_tail+'.png') - + + if add_qr_code: + img = np.array(Image.open(outfile)) + qr = qrcode.QRCode( + version=1, + error_correction=qrcode.constants.ERROR_CORRECT_L, + box_size=5, + border=4, + ) + qr.add_data(encoding) + qr.make(fit=True) + + qr_img = np.array(qr.make_image(fill_color="black", back_color="white")) + qr_img_color = np.zeros((qr_img.shape[0], qr_img.shape[1], 3)) + qr_img_color[:,:,0] = qr_img*255. + qr_img_color[:,:,1] = qr_img*255. + qr_img_color[:,:,2] = qr_img*255. + + img[:qr_img.shape[0], -qr_img.shape[1]:, :3] = qr_img_color + img = Image.fromarray(img) + img.save(outfile) + json_object = json.dumps(json_dict, indent=4) - + # Writing to sample.json if store_configs: with open(os.path.join(output_directory,rec_tail+'.json'), "w") as f: diff --git a/codes/ecg-image-generator/gen_ecg_image_from_data.py b/codes/ecg-image-generator/gen_ecg_image_from_data.py index d56b3f2..279a5d4 100644 --- a/codes/ecg-image-generator/gen_ecg_image_from_data.py +++ b/codes/ecg-image-generator/gen_ecg_image_from_data.py @@ -30,8 +30,9 @@ def get_parser(): parser.add_argument('--num_columns',type=int,default = -1) parser.add_argument('--full_mode', type=str,default='II') parser.add_argument('--mask_unplotted_samples', action="store_true", default=False) + parser.add_argument('--add_qr_code', action="store_true", default=False) - parser.add_argument('-l', '--link', type=str, required=False,default='https://www.physionet.org/content/ptbdb/1.0.0/') + parser.add_argument('-l', '--link', type=str, required=False,default='') parser.add_argument('-n','--num_words',type=int,required=False,default=5) parser.add_argument('--x_offset',dest='x_offset',type=int,default = 30) parser.add_argument('--y_offset',dest='y_offset',type=int,default = 30) @@ -97,6 +98,7 @@ def writeCSV(args): def run_single_file(args): if hasattr(args, 'st') == True: random.seed(args.seed) + args.encoding = args.input_file filename = args.input_file header = args.header_file @@ -126,7 +128,7 @@ def run_single_file(args): configs = read_config_file(os.path.join(os.getcwd(), args.config_file)) - out_array = get_paper_ecg(input_file=filename,header_file=header, configs=configs, mask_unplotted_samples=args.mask_unplotted_samples, start_index=args.start_index, store_configs=args.store_config, store_text_bbox=args.lead_name_bbox, output_directory=args.output_directory,resolution=resolution,papersize=papersize,add_lead_names=lead,add_dc_pulse=bernoulli_dc,add_bw=bernoulli_bw,show_grid=bernoulli_grid,add_print=bernoulli_add_print,pad_inches=padding,font_type=font,standard_colours=standard_colours,full_mode=args.full_mode,bbox = args.lead_bbox, columns = args.num_columns, seed=args.seed) + out_array = get_paper_ecg(input_file=filename,header_file=header, configs=configs, encoding=args.encoding, mask_unplotted_samples=args.mask_unplotted_samples, start_index=args.start_index, store_configs=args.store_config, store_text_bbox=args.lead_name_bbox, output_directory=args.output_directory,resolution=resolution,papersize=papersize,add_lead_names=lead,add_dc_pulse=bernoulli_dc,add_bw=bernoulli_bw,show_grid=bernoulli_grid,add_print=bernoulli_add_print,pad_inches=padding,font_type=font,standard_colours=standard_colours,full_mode=args.full_mode,bbox = args.lead_bbox, columns = args.num_columns, seed=args.seed, add_qr_code=args.add_qr_code) for out in out_array: if args.store_config: diff --git a/codes/ecg-image-generator/gen_ecg_images_from_data_batch.py b/codes/ecg-image-generator/gen_ecg_images_from_data_batch.py index 6b9fd32..03c2738 100644 --- a/codes/ecg-image-generator/gen_ecg_images_from_data_batch.py +++ b/codes/ecg-image-generator/gen_ecg_images_from_data_batch.py @@ -27,6 +27,7 @@ def get_parser(): parser.add_argument('--num_columns',type=int,default = -1) parser.add_argument('--full_mode', type=str,default='II') parser.add_argument('--mask_unplotted_samples', action="store_true", default=False) + parser.add_argument('--add_qr_code', action="store_true", default=False) parser.add_argument('-l', '--link', type=str, required=False,default='') parser.add_argument('-n','--num_words',type=int,required=False,default=5) @@ -101,9 +102,11 @@ def run(args): args.input_file = os.path.join(args.input_directory, filename) args.header_file = os.path.join(args.input_directory, header) args.start_index = -1 - + folder_struct_list = full_header_file.split('/')[:-1] args.output_directory = os.path.join(original_output_dir, '/'.join(folder_struct_list)) + args.encoding = os.path.splitext(filename)[0] + i += run_single_file(args) if(args.max_num_images != -1 and i >= args.max_num_images):