Commit f6894f24 authored by ferrari's avatar ferrari
Browse files

initial commit

parent 63511b45
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
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment