From 05cadf701f1f0c68ae28d5e55584647f5380b17a Mon Sep 17 00:00:00 2001 From: amitjans <amitjans@ethz.ch> Date: Tue, 1 Jun 2021 16:40:50 +0200 Subject: [PATCH] reorder --- app.py | 3 +- html_settings.py | 13 --- html_tools.py | 36 -------- main.py | 191 ------------------------------------------ settings.py | 3 + tools.py | 211 ++++++++++++++++++++++++++++++++++++++++++++++- 6 files changed, 214 insertions(+), 243 deletions(-) delete mode 100644 html_settings.py delete mode 100644 html_tools.py delete mode 100644 main.py diff --git a/app.py b/app.py index 1b4a597..b1c5f47 100644 --- a/app.py +++ b/app.py @@ -2,8 +2,7 @@ import threading import cv2 from flask import Flask, render_template, Response -from html_tools import plot_detection -from main import Main +from tools import plot_detection, Main app = Flask(__name__) sm = threading.Semaphore(10) diff --git a/html_settings.py b/html_settings.py deleted file mode 100644 index 148b2b0..0000000 --- a/html_settings.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Settings file""" - -# Shared memory name (ion images) -SHM_NAME = 'shm://zwoioncameraone' - -# Semaphore name -SEM_NAME = '/zwocamdataone' - -# Radius of circle that captures the intensity of each ion (plotted on html) -RADIUS = 15 - -# Transparency of the circle plotted in the html -ALPHA = 0.15 diff --git a/html_tools.py b/html_tools.py deleted file mode 100644 index 9aef80e..0000000 --- a/html_tools.py +++ /dev/null @@ -1,36 +0,0 @@ -"""This file contains all the needed functions""" -import logging -import numpy as np -import cv2 -import html_settings - -log = logging.getLogger(__name__) - -def plot_detection(image, points): - """ - Use OpenCV to plot the detected ions on the image, following this color code: - - Red: Bright Ion - - Green: Bright Ion - - Blue: Dark Ion - Code the image into bytes and return it. - """ - image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) - for point in points: - if point[2] != 2: - overlay = image.copy() - cv2.circle(overlay, tuple(np.flip(point[:2])), radius=html_settings.RADIUS, - color=(255, 255, 255), thickness=1) - alpha = html_settings.ALPHA - image = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0) - if point[2] == 1: - image = cv2.circle(image, tuple(np.flip(point[:2])), radius=3, color=(0, 255, 255), - thickness=-1) - else: - image = cv2.circle(image, tuple(np.flip(point[:2])), radius=3, color=(0, 255, 0), - thickness=-1) - else: - image = cv2.circle(image, tuple(np.flip(point[:2])), radius=3, color=(0, 0, 255), - thickness=-1) - - - return image \ No newline at end of file diff --git a/main.py b/main.py deleted file mode 100644 index 2ce87c9..0000000 --- a/main.py +++ /dev/null @@ -1,191 +0,0 @@ -"""MAIN file""" -# import matplotlib -# matplotlib.use('tkagg') -import time -import threading -import warnings -from multiprocessing import Queue -import redis -import numpy as np -import influxdb -import settings -from tools import Memory, MovingAverage, ion_positions, dark_ions, ion_linearity -from tools import detect_ion_cloud, read_factor, compensate_gradient, acquire_semaphore_write, release_semaphore_write -from tools import TrapState, max_ion_intensity, detect_dim_ions, load_image, ion_roundness - - -class Main: - """ - Main class. Does the processing of the images and stores all the outputs. - """ - def __init__(self, sem): - """ - Connect to redis and influxdb. Instantiate all the needed classes and variables. - """ - # Define Semaphore - self.sm = sem - - # Connect to redis and influxdb - self.conn = redis.Redis('localhost', decode_responses=True) - - self.db = influxdb.InfluxDBClient( - host=settings.DB_HOST, - port=settings.DB_PORT, - username=settings.DB_USER, - password=settings.DB_PWD, - database=settings.DB_DATASET, - ssl=False, - ) - - # Instantiate classes - self.memory = Memory() - self.count_avg = MovingAverage() - self.trap_state = TrapState() - - # Define class attributes - self.init_time = time.time() - self.num = 0 - self.output = {} - - def start_thread(self): - """ - Start background thread. - """ - t = threading.Thread(target=self._background_thread, daemon=True) - t.start() - - def _background_thread(self): - """ - Apply the processing on the images. - """ - warnings.simplefilter("once", UserWarning) - while True: - dist_factor = read_factor(self.conn) - image = load_image() - ion_cloud = detect_ion_cloud(image) - - if ion_cloud: - output = { - 'image': image, - 'coord': np.array([[-1, -1, -1]]), - 'trap_state': self.trap_state.trap_state, - 'num_ions': -1, - 'bright_ions': -1, - 'dim_ions': -1, - 'dark_ions': -1, - 'linearity': 'cloud', - 'max_int': -1, - 'num_pixels': -1, - 'noise_mean': self.trap_state.noise[0], - 'noise_std': self.trap_state.noise[1], - 'reorder_avg': self.count_avg.moving_avg, - 'reorder_flag': -1, - 'max_spot_size': -1, - 'roundness': -1 - } - - elif not ion_cloud: - output = self._processing(image, dist_factor) - - # Update running time - running_time = time.time() - self.init_time - output['time'] = 60 - int(running_time) - if running_time >= 60: - self.init_time = time.time() - - # Save variables in self.memory - self._save_in_memory(output) - self.num += 1 - - # Delete unnecessary variables for html and put output in queue - del output['num_ions'] - del output['reorder_flag'] - sem = acquire_semaphore_write(self.sm, num_sem=10) - self.output = output - release_semaphore_write(sem, num_sem=10) - - def _processing(self, image, dist_factor): - coord, max_spot_size = ion_positions(image) - if len(coord) > 0: - roundness = ion_roundness(image, coord[0][:2]) - new_img = compensate_gradient(image, coord) - max_int, num_pix = max_ion_intensity(new_img, coord) - coord = detect_dim_ions(new_img, coord, max_int) - linearity = ion_linearity(coord) - - if len(coord) > 1 and dist_factor is not None: - coord = dark_ions(image, coord, dist_factor) - - _, old_idx = self.memory.get_value() - self.memory.save_value(coord) - coord, new_idx = self.memory.get_value() - # If idx changes, it means that ion reordered - reorder_flag = 0 - if new_idx != old_idx: - reorder_flag = 1 - self.count_avg.update_count() - - # Update trap state - self.trap_state.update_coord(coord) - - num_ions = len(coord) - num_bright_ions = np.sum(coord[:, 2] == 0).item() - num_dim_ions = np.sum(coord[:, 2] == 1).item() - num_dark_ions = np.sum(coord[:, 2] == 2).item() - - # Reset self.memory every NUM_ITERATIONS - if self.num > settings.NUM_ITERATIONS: - self.memory.reset_memory(coord) - self.num = 0 - - return { - 'image': image, - 'coord': coord, - 'trap_state': self.trap_state.trap_state, - 'num_ions': num_ions, - 'bright_ions': num_bright_ions, - 'dim_ions': num_dim_ions, - 'dark_ions': num_dark_ions, - 'linearity': linearity, - 'max_int': int(max_int), - 'num_pixels': int(num_pix), - 'noise_mean': self.trap_state.noise[0], - 'noise_std': self.trap_state.noise[1], - 'reorder_avg': self.count_avg.moving_avg, - 'reorder_flag': reorder_flag, - 'max_spot_size': int(max_spot_size), - 'roundness': int(roundness) - } - - def _save_in_memory(self, output): - """ - Save measured variables in shared memory - """ - self.conn.set(name='/image_analysis/trap_state', value=output['trap_state'], ex=10) - self.conn.set(name='/image_analysis/num_ions/bright_ions', value=output['bright_ions'], ex=10) - self.conn.set(name='/image_analysis/num_ions/dim_ions', value=output['dim_ions'], ex=10) - self.conn.set(name='/image_analysis/num_ions/dark_ions', value=output['dark_ions'], ex=10) - self.conn.set(name='/image_analysis/linearity', value=output['linearity'], ex=10) - self.conn.set(name='/image_analysis/ion_intensity/max_int', value=int(output['max_int']), ex=10) - self.conn.set(name='/image_analysis/ion_intensity/num_pixels', value=int(output['num_pixels']), ex=10) - self.conn.set(name='/image_analysis/noise/mean', value=output['noise_mean']) - self.conn.set(name='/image_analysis/noise/std', value=output['noise_std']) - self.conn.set(name='/image_analysis/reorder/flag', value=output['reorder_flag'], ex=10) - self.conn.set(name='/image_analysis/reorder/avg', value=output['reorder_avg'], ex=10) - self.conn.set(name='/image_analysis/time', value=output['time'], ex=10) - self.conn.set(name='/image_analysis/max_spot_size', value=int(output['max_spot_size']), ex=10) - self.conn.set(name='/image_analysis/ion_roundness', value=int(output['roundness']), ex=10) - - self.db.write_points([{ - 'measurement': 'ion_image', - 'fields': { - 'num_ions': output['num_ions'], - 'num_bright_ions': output['bright_ions'], - 'num_dim_ions': output['dim_ions'], - 'num_dark_ions': output['dark_ions'], - 'reorder_rate': output['reorder_avg'], - 'max_int': output['max_int'], - 'max_spot_size': output['max_spot_size'], - 'ion_roundness': output['roundness'] - } - }]) \ No newline at end of file diff --git a/settings.py b/settings.py index 7b58371..0579433 100644 --- a/settings.py +++ b/settings.py @@ -43,3 +43,6 @@ RADIUS = 15 # Area threshold of biggest contour in image (for ion cloud detection) AREA_THS = 10e3 + +# Transparency of the circle plotted in the html +ALPHA = 0.15 diff --git a/tools.py b/tools.py index f4e0ad2..806ac12 100644 --- a/tools.py +++ b/tools.py @@ -1,9 +1,10 @@ """This file contains all the needed functions""" -import struct import time import warnings import logging import threading +import redis +import influxdb import cv2 import requests import SharedArray @@ -20,6 +21,184 @@ import settings log = logging.getLogger(__name__) + +class Main: + """ + Main class. Does the processing of the images and stores all the outputs. + """ + def __init__(self, sem): + """ + Connect to redis and influxdb. Instantiate all the needed classes and variables. + """ + # Define Semaphore + self.sm = sem + + # Connect to redis and influxdb + self.conn = redis.Redis('localhost', decode_responses=True) + + self.db = influxdb.InfluxDBClient( + host=settings.DB_HOST, + port=settings.DB_PORT, + username=settings.DB_USER, + password=settings.DB_PWD, + database=settings.DB_DATASET, + ssl=False, + ) + + # Instantiate classes + self.memory = Memory() + self.count_avg = MovingAverage() + self.trap_state = TrapState() + + # Define class attributes + self.init_time = time.time() + self.num = 0 + self.output = {} + + def start_thread(self): + """ + Start background thread. + """ + t = threading.Thread(target=self._background_thread, daemon=True) + t.start() + + def _background_thread(self): + """ + Apply the processing on the images. + """ + warnings.simplefilter("once", UserWarning) + while True: + dist_factor = read_factor(self.conn) + image = load_image() + ion_cloud = detect_ion_cloud(image) + + if ion_cloud: + output = { + 'image': image, + 'coord': np.array([[-1, -1, -1]]), + 'trap_state': self.trap_state.trap_state, + 'num_ions': -1, + 'bright_ions': -1, + 'dim_ions': -1, + 'dark_ions': -1, + 'linearity': 'cloud', + 'max_int': -1, + 'num_pixels': -1, + 'noise_mean': self.trap_state.noise[0], + 'noise_std': self.trap_state.noise[1], + 'reorder_avg': self.count_avg.moving_avg, + 'reorder_flag': -1, + 'max_spot_size': -1, + 'roundness': -1 + } + + elif not ion_cloud: + output = self._processing(image, dist_factor) + + # Update running time + running_time = time.time() - self.init_time + output['time'] = 60 - int(running_time) + if running_time >= 60: + self.init_time = time.time() + + # Save variables in self.memory + self._save_in_memory(output) + self.num += 1 + + # Delete unnecessary variables for html and put output in queue + del output['num_ions'] + del output['reorder_flag'] + sem = acquire_semaphore_write(self.sm, num_sem=10) + self.output = output + release_semaphore_write(sem, num_sem=10) + + def _processing(self, image, dist_factor): + coord, max_spot_size = ion_positions(image) + if len(coord) > 0: + roundness = ion_roundness(image, coord[0][:2]) + new_img = compensate_gradient(image, coord) + max_int, num_pix = max_ion_intensity(new_img, coord) + coord = detect_dim_ions(new_img, coord, max_int) + linearity = ion_linearity(coord) + + if len(coord) > 1 and dist_factor is not None: + coord = dark_ions(image, coord, dist_factor) + + _, old_idx = self.memory.get_value() + self.memory.save_value(coord) + coord, new_idx = self.memory.get_value() + # If idx changes, it means that ion reordered + reorder_flag = 0 + if new_idx != old_idx: + reorder_flag = 1 + self.count_avg.update_count() + + # Update trap state + self.trap_state.update_coord(coord) + + num_ions = len(coord) + num_bright_ions = np.sum(coord[:, 2] == 0).item() + num_dim_ions = np.sum(coord[:, 2] == 1).item() + num_dark_ions = np.sum(coord[:, 2] == 2).item() + + # Reset self.memory every NUM_ITERATIONS + if self.num > settings.NUM_ITERATIONS: + self.memory.reset_memory(coord) + self.num = 0 + + return { + 'image': image, + 'coord': coord, + 'trap_state': self.trap_state.trap_state, + 'num_ions': num_ions, + 'bright_ions': num_bright_ions, + 'dim_ions': num_dim_ions, + 'dark_ions': num_dark_ions, + 'linearity': linearity, + 'max_int': int(max_int), + 'num_pixels': int(num_pix), + 'noise_mean': self.trap_state.noise[0], + 'noise_std': self.trap_state.noise[1], + 'reorder_avg': self.count_avg.moving_avg, + 'reorder_flag': reorder_flag, + 'max_spot_size': int(max_spot_size), + 'roundness': int(roundness) + } + + def _save_in_memory(self, output): + """ + Save measured variables in shared memory + """ + self.conn.set(name='/image_analysis/trap_state', value=output['trap_state'], ex=10) + self.conn.set(name='/image_analysis/num_ions/bright_ions', value=output['bright_ions'], ex=10) + self.conn.set(name='/image_analysis/num_ions/dim_ions', value=output['dim_ions'], ex=10) + self.conn.set(name='/image_analysis/num_ions/dark_ions', value=output['dark_ions'], ex=10) + self.conn.set(name='/image_analysis/linearity', value=output['linearity'], ex=10) + self.conn.set(name='/image_analysis/ion_intensity/max_int', value=int(output['max_int']), ex=10) + self.conn.set(name='/image_analysis/ion_intensity/num_pixels', value=int(output['num_pixels']), ex=10) + self.conn.set(name='/image_analysis/noise/mean', value=output['noise_mean']) + self.conn.set(name='/image_analysis/noise/std', value=output['noise_std']) + self.conn.set(name='/image_analysis/reorder/flag', value=output['reorder_flag'], ex=10) + self.conn.set(name='/image_analysis/reorder/avg', value=output['reorder_avg'], ex=10) + self.conn.set(name='/image_analysis/time', value=output['time'], ex=10) + self.conn.set(name='/image_analysis/max_spot_size', value=int(output['max_spot_size']), ex=10) + self.conn.set(name='/image_analysis/ion_roundness', value=int(output['roundness']), ex=10) + + self.db.write_points([{ + 'measurement': 'ion_image', + 'fields': { + 'num_ions': output['num_ions'], + 'num_bright_ions': output['bright_ions'], + 'num_dim_ions': output['dim_ions'], + 'num_dark_ions': output['dark_ions'], + 'reorder_rate': output['reorder_avg'], + 'max_int': output['max_int'], + 'max_spot_size': output['max_spot_size'], + 'ion_roundness': output['roundness'] + } + }]) + + class Memory: """ Stores values in memory and counts how many times each value was stored. @@ -505,6 +684,36 @@ def pos_solver(num): return solution +def plot_detection(image, points): + """ + Use OpenCV to plot the detected ions on the image, following this color code: + - Red: Bright Ion + - Green: Bright Ion + - Blue: Dark Ion + Code the image into bytes and return it. + """ + image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) + for point in points: + if point[2] != 2: + overlay = image.copy() + cv2.circle(overlay, tuple(np.flip(point[:2])), radius=settings.RADIUS, + color=(255, 255, 255), thickness=1) + alpha = settings.ALPHA + image = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0) + if point[2] == 1: + image = cv2.circle(image, tuple(np.flip(point[:2])), radius=3, color=(0, 255, 255), + thickness=-1) + else: + image = cv2.circle(image, tuple(np.flip(point[:2])), radius=3, color=(0, 255, 0), + thickness=-1) + else: + image = cv2.circle(image, tuple(np.flip(point[:2])), radius=3, color=(0, 0, 255), + thickness=-1) + + + return image + + def _laser_switch(boolean): """ Switch on (True) or off (False) the repumping laser. -- GitLab