robinkrens.nl - git repository - haptic-extractor/blob - haptic-extractor.py

   1 # /usr/bin/python3
   2 #
   3 # Usage:
   4 # haptic-extractor.py [option] input_image output_image
   5 #
   6 # Options
   7 #    --transpose     Generate a vertical profile through the image
   8 #    --color=COLOR  Draw the graph in a different color
   9 #
  10 ####
  11
  12 from PIL import Image
  13 import matplotlib.pyplot as plt
  14 import numpy as np
  15 import signal
  16 import wave
  17 import sys
  18 import os
  19
  20 PROGNAME = os.path.basename(sys.argv[0])
  21 PROGDIR = os.path.dirname(os.path.abspath(sys.argv[0]))
  22
  23 def usage(message):
  24     """Output the script comments as docs"""
  25     print(f"{PROGNAME}: {' '.join(sys.argv[1:])}")
  26     with open(os.path.join(PROGDIR, PROGNAME), 'r') as script_file:
  27         lines = script_file.readlines()[1:]
  28         for line in lines:
  29             if line.startswith('###'):
  30                 break
  31             line = line.strip().lstrip('#').lstrip()
  32             print(line)
  33
  34     sys.exit(1)
  35
  36 def error(message):
  37     """Output an error message and exit"""
  38     print(f"{PROGNAME}: {message}")
  39     sys.exit(1)
  40
  41 COLOR = 'red'
  42 TRANSPOSE = False
  43 DEBUG = False
  44
  45 args = iter(sys.argv[1:])
  46 for arg in args:
  47     if arg in ('--help', '-h'):
  48         usage()
  49     elif arg in ('--transpose', '-t'):
  50         TRANSPOSE = True
  51     elif arg in ('--verbose'):
  52         DEBUG = True
  53     elif arg in ('-', '--'):
  54         break
  55     elif arg.startswith('-'):
  56         usage(f"Unknown option \"{arg}\"")
  57     else:
  58         break
  59
  60 if len(sys.argv) == 1:
  61     usage("Missing input_image")
  62
  63 if len(sys.argv) > 3:
  64     usage("Too many arguments")
  65
  66 def cleanup_handler(signum, frame):
  67     """Cleanup function to be executed on exit"""
  68     os.remove(input_image)
  69     sys.exit(0)
  70
  71 # Register the cleanup function
  72 #for sig in [0, 1, 2, 3]:
  73 #    signal.signal(sig, cleanup_handler)
  74
  75 try:
  76     input_image_path = sys.argv[1]
  77     if len(sys.argv) == 3:
  78         output_image_path = sys.argv[2]
  79     else:
  80         base_name, ext = os.path.splitext(os.path.basename(input_image_path))
  81         output_image_path = f"{base_name}_output{ext}"
  82
  83     image = Image.open(input_image_path)
  84
  85     # Apply the specified vertical transformation
  86     if TRANSPOSE:
  87         image = image.transpose(Image.Transpose.ROTATE_90)
  88
  89     # Crop the image to a 1xwidth size at the center of the height
  90     width, height = image.size
  91     strip_height = 1
  92     top = (height - strip_height) // 2
  93     bottom = top + strip_height
  94     image = image.crop((0, top, width, bottom))
  95
  96     # Save the result
  97     #image.save(output_image_path, format="PNG")
  98
  99 except Exception as e:
 100     print(f"Error: {e}")
 101     usage(e)
 102     sys.exit(1)
 103
 104 #image = Image.open("./im_profile_strip.png")
 105 image = image.convert('L')
 106
 107 # Convert the image to a NumPy array
 108 image_array = np.array(image)
 109 image_flat = image_array.flatten()
 110
 111 # Check the bit depth of the image
 112 bit_depth = image_array.dtype.itemsize * 8
 113
 114 # Print the values in a format similar to the original script
 115 if DEBUG:
 116     for value in image_array.flatten():
 117         print(value, end=' ')
 118
 119     print(f"\nBit Depth: {bit_depth} bits")
 120
 121 fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(3, 3), gridspec_kw={'height_ratios': [3, 1]})
 122 ax1.plot(image_flat, linestyle='-', linewidth=1, color='gray')
 123
 124 # Customize x and y ticks
 125 ax1.set_xticks([])
 126 ax1.set_ylim(0, 255)
 127 ax1.set_yticks([75, 150, 225])
 128 ax1.set_xticklabels([])
 129 ax1.set_yticklabels([])
 130
 131 # resize picture a bit and attach at bottom
 132 width, height = image.size
 133 image = image.resize((width, 5), resample=Image.BICUBIC)
 134 ax2.imshow(image, cmap="gray")
 135 ax2.axis('off')  # Hide axis for the image subplot
 136
 137 plt.subplots_adjust(hspace=0)
 138 fig.tight_layout(pad=0)
 139 fig.tight_layout(rect=[0, 0, 1, 0.90])
 140
 141 print(f"Written output file {output_image_path}")
 142 plt.savefig(output_image_path, dpi=100, bbox_inches='tight')