ecg gen: code for QR code addition

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:

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:

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):