class OnFrameAddedEventListener(EventListener):
"""
This Event Listener will:
- Move the frame image file from the input directory (frames_path) to the processing directory.
- Save the frame data to the state of the app (filename, path, ...)
"""
def __init__(self, frames_path: str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/input',
processing_dir: str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/processing',
frame_ext: str = 'png'):
self.frames_path = frames_path
self.processing_dir = processing_dir
self.frame_ext = frame_ext
self.state = DictionaryState()
os.makedirs(os.path.join(self.processing_dir, 'frames'), exist_ok=True)
os.makedirs(os.path.join(self.processing_dir, 'boxes'), exist_ok=True)
def handle(self, frame_id: int) -> None:
"""
:param frame_id: id of the frame
"""
filename = '{}.{}'.format(frame_id, self.frame_ext)
path = os.path.join(self.processing_dir, 'frames', filename)
# os.rename(os.path.join(self.frames_path, filename), path)
shutil.copyfile(os.path.join(self.frames_path, filename), path)
self.state.insert(StateGroups.FRAME, frame_id, {'filename': filename, 'path': path})
class OnEndOfStreamEventListener(EventListener):
"""
This event is fired at the end of the event log.
This Event Listener will:
- Export the state of the application that contains all the information about the frames and boxes to a json file.
- If clear_state = True, it will reset the state to its initial value
- If clear_storage = True, it will remove all the images (frames and boxes) from the processing directory.
"""
def __init__(
self,
output_dir:
str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/output',
processing_dir:
str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/processing',
filename: str = 'annotated_data.json',
clear_state: bool = False,
clear_storage: bool = False):
self.state = DictionaryState()
self.output_dir = output_dir
self.processing_dir = processing_dir
self.filename = filename
self.clear_state = clear_state
self.clear_storage = clear_storage
if not os.path.isdir(self.output_dir):
os.makedirs(self.output_dir, exist_ok=True)
def handle(self):
with open(os.path.join(self.output_dir, self.filename), 'w') as fp:
json.dump(self.state.state, fp)
if self.clear_state:
self.state.reset_state()
if self.clear_storage:
shutil.rmtree(self.processing_dir)
class OnBoxAttributeChangedEventListener(EventListener):
def __init__(self):
self.state = DictionaryState()
def handle(self, box_id: str, attribute_id: str, attribute_value: str):
self.state.update(StateGroups.BOX, box_id, (attribute_id, attribute_value))
if attribute_id == 'label' and attribute_value == 'traffic_light':
pub.sendMessage(Topics.EXTRACT_BOX, box_id=box_id)
def __init__(self, frames_path: str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/input',
processing_dir: str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/processing',
frame_ext: str = 'png'):
self.frames_path = frames_path
self.processing_dir = processing_dir
self.frame_ext = frame_ext
self.state = DictionaryState()
os.makedirs(os.path.join(self.processing_dir, 'frames'), exist_ok=True)
os.makedirs(os.path.join(self.processing_dir, 'boxes'), exist_ok=True)
class OnBoxDeletedEventListener(EventListener):
def __init__(self):
self.state = DictionaryState()
def handle(self, box_id: str):
box = self.state.query(StateGroups.BOX, box_id)
if 'path' in box and os.path.isfile(box['path']):
os.remove(box['path'])
self.state.remove(StateGroups.BOX, box_id)
def __init__(self,
boxes_path: str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/processing/boxes',
boxes_ext: str = 'png',
width_threshold: int = 10,
height_threshold: int = 10,
):
self.state = DictionaryState()
self.boxes_path = boxes_path
self.boxes_ext = boxes_ext
self.width_threshold = width_threshold
self.height_threshold = height_threshold
class OnExtractBoxEventListener(EventListener):
"""
This Event Listener will:
- Assess if the bounding box is contained (fully or partially) in the image or not.
- If the BBox is:
- contained fully: extract the whole BBox and save it
- contained partially: extract the part that is contained in the image and saved it
- not contained: does nothing
- update the label of the box "is_contained" to be either 1 (yes), 0 (no) or 0.5(partly).
- update the label of the box "path" to be the path where the BBox is saved.
- publish an event so the annotator annotates the BBox.
"""
def __init__(self,
boxes_path: str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/processing/boxes',
boxes_ext: str = 'png',
width_threshold: int = 10,
height_threshold: int = 10,
):
self.state = DictionaryState()
self.boxes_path = boxes_path
self.boxes_ext = boxes_ext
self.width_threshold = width_threshold
self.height_threshold = height_threshold
def handle(self, box_id: str):
# querying the box and converting its data to int instead of float.
box = self.state.query(StateGroups.BOX, box_id)
box_data = [box['x'], box['y'], box['width'], box['height']]
x, y, width, height = list(map(int, box_data))
frame = self.state.query(StateGroups.FRAME, box['frame_id'])
# read image and convert colors
image = cv2.imread(frame['path'])
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# check if the box is contained or not and update box attributes
box_iou, (x_in, y_in, width_in, height_in) = iou(image, x, y, width, height)
pub.sendMessage(Topics.BOX_ATTRIBUTE_CHANGED, box_id=box_id, attribute_id='is_contained',
attribute_value=box_iou)
if box_iou == 0:
return
if width_in < self.width_threshold or height_in < self.height_threshold:
return
# extracting the box from the image and saving it.
box_image = image[y_in:y_in + height_in, x_in:x_in + width_in, :]
filename = '{}.{}'.format(box_id, self.boxes_ext)
path = os.path.join(self.boxes_path, filename)
plt.imsave(path, box_image)
pub.sendMessage(Topics.BOX_ATTRIBUTE_CHANGED, box_id=box_id, attribute_id='path', attribute_value=path)
pub.sendMessage(Topics.CLASSIFY_TRAFFIC_LIGHT, box_id=box_id)
class OnBoxCreatedEventListener(EventListener):
def __init__(self):
self.state = DictionaryState()
def handle(self, frame_id: int, box_id: str, x: float, y: float,
width: float, height: float):
self.state.insert(StateGroups.BOX, box_id, {
'x': x,
'y': y,
'width': width,
'height': height,
'frame_id': frame_id
})
def __init__(
self,
output_dir:
str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/output',
processing_dir:
str = '/home/bothmena/Projects/PyCharm/ai_training/UnderstandAI/Homework/var/processing',
filename: str = 'annotated_data.json',
clear_state: bool = False,
clear_storage: bool = False):
self.state = DictionaryState()
self.output_dir = output_dir
self.processing_dir = processing_dir
self.filename = filename
self.clear_state = clear_state
self.clear_storage = clear_storage
if not os.path.isdir(self.output_dir):
os.makedirs(self.output_dir, exist_ok=True)
class OnBoxMovedEventListener(EventListener):
def __init__(self):
self.state = DictionaryState()
def handle(self, box_id: str, x: float, y: float, width: float,
height: float):
for key, value in {
'x': x,
'y': y,
'width': width,
'height': height
}.items():
if value is not None:
self.state.update(StateGroups.BOX, box_id, (key, value))
box = self.state.query(StateGroups.BOX, box_id)
if 'label' in box and box['label'] == 'traffic_light':
pub.sendMessage(Topics.EXTRACT_BOX, box_id=box_id)
class OnClassifyTrafficLightEventListener(EventListener):
def __init__(self):
self.state = DictionaryState()
def handle(self, box_id: str):
# query the box for state and load the bbox image
box = self.state.query(StateGroups.BOX, box_id)
annotation = ['green', 'red', 'orange', 'off'][np.random.choice([0, 1, 2, 3])]
pub.sendMessage(Topics.BOX_ATTRIBUTE_CHANGED, box_id=box_id, attribute_id='color', attribute_value=annotation)
def __init__(self):
self.state = DictionaryState()