Skip to content
Snippets Groups Projects
Commit f6894f24 authored by ferrari's avatar ferrari
Browse files

initial commit

parent 63511b45
No related branches found
No related tags found
No related merge requests found
import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
import imageio as io
import os
import sys
from matplotlib.widgets import Button, Slider
from matplotlib.backend_bases import MouseButton
import pandas as pd
def wb(channel, perc = 0.05):
mi, ma = (np.percentile(channel, perc, axis=(0,1)), np.percentile(channel,100.0-perc, axis=(0,1)))
channel = np.uint8(np.clip((channel-mi)*255.0/(ma-mi+1e-18), 0, 255))
return channel
def wb2(channel, per=0.05):
hsv = rgb_to_hsv(channel/255)
grey = (channel/255).sum(-1)
mi, ma = (np.percentile(grey, per), np.percentile(grey, 100-per))
if mi > 0.01:
hue = (hsv[grey <= mi][..., 0].mean() + hsv[grey >= ma][..., 0].mean())/2
else:
hue = hsv[grey >= ma][..., 0].mean()
hsv[..., 0] += 0.5 - hue
hsv[..., 0] %= 1
return hsv_to_rgb(hsv)
class Callback(object):
def __init__(self, inp):
self.reader = io.get_reader(inp, format='ffmpeg')
self._len = self.reader.count_frames()
self.fig = plt.figure(inp, figsize=[16, 9], constrained_layout=True)
gs = self.fig.add_gridspec(2, 1, height_ratios=[9, 1])
controls = gs[1].subgridspec(2, 5, height_ratios=[1, 3])
self.screen = self.fig.add_subplot(gs[0])
self.screen.set(xticks=[], yticks=[])
self.idx = 0
self.img = self.screen.imshow(self.reader.get_data(self.idx))
self.bar = Slider(self.fig.add_subplot(controls[0, :], zorder=10),
'frame #', 0, self._len, valinit=0, valstep=1, valfmt='%5d ')
self.bar.valtext.set_family('monospace')
self.bar.on_changed(self.new_frame)
self.resize_b = Button(self.fig.add_subplot(controls[1, 0], zorder=10), 'Resize plot')
self.resize_b.on_clicked(self.resize)
self.normal_b = Button(self.fig.add_subplot(controls[1, 1], zorder=10), 'Auto\nwhite balance')
self.normal_b.on_clicked(self.normalize)
for v in "fullscreen,home,back,forward,pan,zoom,save,quit,grid,yscale,xscale,all_axes".split(','):
plt.rcParams[f'keymap.{v}'] = []
self.cid = self.fig.canvas.mpl_connect('key_press_event', self.key_pressed)
self.cid2 = self.fig.canvas.mpl_connect('button_press_event', self.mouse_press)
self.cid3 = self.fig.canvas.mpl_connect('button_release_event', self.mouse_release)
self.mouse_pos = np.zeros((2,2)) #left/right click, x/y
self.sperm_pos = np.zeros((2,2,2)) #sperm whale, start/stop, x/y
self.pos_select = [0, 0]
self.screen.set_xlim(self.screen.get_xlim())
self.screen.set_ylim(self.screen.get_ylim())
self.line1, self.line2 = self.screen.plot([[-100, -100], [-150, -150]], [[0, 0], [0, 0]], c='r', marker='x')
text_ax = self.fig.add_subplot(controls[1, -2:])
self.text = text_ax.text(0.5, 0.5, 'Sperm whale 1 : ( , ) - ( , ) ratio - \n'
'Sperm whale 2 : ( , ) - ( , ) angle ° ',
horizontalalignment='center',
verticalalignment='center', transform=text_ax.transAxes, family='monospace')
text_ax.axis('off')
self.closed = False
plt.draw()
plt.pause(0.2)
self.fig.set_constrained_layout(False)
def new_frame(self, val):
self.idx = int(val)
self.img.set_data(self.reader.get_data(self.idx))
self.line1.set_data([-100, -150], [0, 0])
self.line2.set_data([-100, -150], [0, 0])
self.sperm_pos[:] = 0
self.text.set_text('Sperm whale 1 : ( , ) - ( , ) ratio - \n'
'Sperm whale 2 : ( , ) - ( , ) angle ° ')
plt.draw()
def resize(self, event):
self.fig.set_constrained_layout(True)
plt.draw()
plt.pause(0.2)
self.fig.set_constrained_layout(False)
def normalize(self, event):
self.img.set_data(wb(wb2(self.reader.get_data(self.idx))))
plt.draw()
def key_pressed(self, event):
key = event.key
move = 0
if key == 'j' or key == 'left':
move = -1
elif key == 'k' or key == 'right':
move = 1
elif key == 'up':
move = -5
elif key == 'pageup':
move = -10
elif key == 'down':
move = 5
elif key == 'pagedown':
move = 10
elif key == 'r':
self.resize(None)
elif key == 'n':
self.normalize(None)
if move:
self.bar.set_val(min(max(0, self.idx + move), self._len - 1))
def mouse_press(self, event):
if event.inaxes != self.screen or self.fig.canvas.manager.toolbar.mode:
return None
if event.button == MouseButton.LEFT:
self.mouse_pos[0] = event.xdata, event.ydata
elif event.button == MouseButton.RIGHT:
self.mouse_pos[1] = event.xdata, event.ydata
def mouse_release(self, event):
if event.inaxes != self.screen or self.fig.canvas.manager.toolbar.mode:
return None
if event.button == MouseButton.LEFT:
if np.sqrt(np.square(self.mouse_pos[0] - np.array([event.xdata, event.ydata])).sum(-1)) < 10:
self.sperm_pos[0, self.pos_select[0]] = event.xdata, event.ydata
self.pos_select[0] ^= 1
else:
self.sperm_pos[0, 0] = self.mouse_pos[0]
self.sperm_pos[0, 1] = event.xdata, event.ydata
self.pos_select[0] = 0
self.line1.set_data(self.sperm_pos[0].T)
elif event.button == MouseButton.RIGHT:
if np.sqrt(np.square(self.mouse_pos[1] - np.array([event.xdata, event.ydata])).sum(-1)) < 10:
self.sperm_pos[1, self.pos_select[1]] = event.xdata, event.ydata
self.pos_select[1] ^= 1
else:
self.sperm_pos[1, 0] = self.mouse_pos[1]
self.sperm_pos[1, 1] = event.xdata, event.ydata
self.pos_select[1] = 0
self.line2.set_data(self.sperm_pos[1].T)
vec1 = self.sperm_pos[0, 1] - self.sperm_pos[0, 0]
vec2 = self.sperm_pos[1, 1] - self.sperm_pos[1, 0]
len1 = np.sqrt(np.square(vec1).sum())
len2 = np.sqrt(np.square(vec2).sum())
if len2:
self.text.set_text(f'Sperm whale 1 : ({self.sperm_pos[0,0,0]:5.0f}, {self.sperm_pos[0,0,1]:5.0f}) '
f'- ({self.sperm_pos[0,1,0]:5.0f}, {self.sperm_pos[0,1,1]:5.0f})'
f' ratio {len1/len2:5.2f} - {len2/len1:5.2f}\n'
f'Sperm whale 2 : ({self.sperm_pos[1,0,0]:5.0f}, {self.sperm_pos[1,0,1]:5.0f}) '
f'- ({self.sperm_pos[1,1,0]:5.0f}, {self.sperm_pos[1,1,1]:5.0f})'
f' angle {np.arccos(np.abs(np.dot(vec1,vec2))/(len1*len2))*180/np.pi:5.2f} ° ')
else:
self.text.set_text(f'Sperm whale 1 : ({self.sperm_pos[0,0,0]:5.0f}, {self.sperm_pos[0,0,1]:5.0f}) '
f'- ({self.sperm_pos[0,1,0]:5.0f}, {self.sperm_pos[0,1,1]:5.0f}) ratio - \n'
f'Sperm whale 2 : ( , ) - ( , ) angle ° ')
plt.draw()
def close(self):
self.reader.close()
del self.resize_b
self.closed = True
def __del__(self):
if not self.closed:
self.close()
def main(args):
if args.out == '':
outpath = args.input.rsplit('.', 1)[0] + '.pred.h5'
else:
outpath = args.out
if os.path.isfile(outpath):
if not (args.erase or args.resume):
print(f'Out file {outpath} already exist and erase or resume option isn\'t set.')
return 1
elif args.resume:
print(f'Out file {outpath} does not already exist and resume option is set.')
return 1
callback = Callback(args.input)
plt.show()
callback.close()
return 0
if __name__ == '__main__':
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("input", type=str, help="Input file")
parser.add_argument("--out", type=str, default='', help="Output file. Default to the input_path'.pred.h5'")
group = parser.add_mutually_exclusive_group()
group.add_argument("--erase", action='store_true', help="If out file exist and this option is not given,"
" the computation will be halted")
group.add_argument("--resume", action='store_true', help="If out file exist and this option is given,"
" the previous annotation file will be loaded")
args = parser.parse_args()
sys.exit(main(args))
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment