# -*- coding: utf-8 -*- """ image_mosaic.py -- A set of classes to extend widgets from pyqtgraph and pyqt for annotation purposes Copyright 2020 Monterey Bay Aquarium Research Institute Distributed under MIT license. See license.txt for more infomation. """ import os import cv2 import glob import shutil from config import settings from libs.logger import LOG from libs.annotation import Annotation from libs.widgets import RectWidget from pyqtgraph.Qt import QtCore, QtGui, QtWidgets class ImageMosaic: def __init__(self, graphicsView, data_path, rect_slot, zoom=1.0): self.data_path = data_path self.thumbs_path = os.path.join(data_path, 'rois') self.ann_path = os.path.join(data_path, 'xml') self.graphicsView = graphicsView self.layout = None self.scene = None self.panel = None self.thumbs = [] self.annotations = [] self.roi_map = {} self.sort_key = '' self.hide_labeled = True self.hide_discarded = True self.hide_to_review = True self.layouts = [] # The scene and view for image mosaic self.scene = QtGui.QGraphicsScene() self.graphicsView.setScene(self.scene) self.panel = QtGui.QGraphicsWidget() self.scene.addItem(self.panel) count = 0 # load annotations. for ann_path in sorted(glob.glob(os.path.join(data_path, '*.xml'))): ann = Annotation(xml_path=ann_path) # check that image exists as well image_path = os.path.join(data_path, ann.xml_tree.find('filename').text) if not os.path.exists(image_path): LOG.warn('missing image for annotation: ' + os.path.basename(ann_path)) continue self.annotations.append(ann) self.roi_map[os.path.basename(ann_path)] = [] full_img = None #print(ann.xml_tree.find('filename').text) count += 1 for ind, ob in enumerate(ann.xml_tree.findall('object')): class_name = ob.find('name').text.replace(' ', '-') xmin, ymin, xmax, ymax = ann.get_box(ind) roi_path = os.path.join(data_path, settings.ROI_SUBDIR, os.path.basename(ann_path)[ :-4] + "-" + str(xmin) + "-" + str(ymin) + "-" + str(xmax) + "-" + str(ymax) + ".jpg") if xmax - xmin <= 0 or ymax - ymin <= 0: LOG.warn('bad box, skipping') continue else: LOG.info('loading roi: ' + roi_path) # check if roi exists and if not load full_img and extract if not os.path.exists(roi_path): if full_img is None: full_img = self.load_image(os.path.join(data_path, ann.xml_tree.find('filename').text)) roi = full_img[ymin:ymax, xmin:xmax, :] if not os.path.exists(os.path.join(data_path, settings.ROI_SUBDIR)): os.makedirs(os.path.join(data_path, settings.ROI_SUBDIR)) cv2.imwrite(roi_path, roi) else: roi = cv2.imread(roi_path) # create the widgets rw = RectWidget(ann, roi.copy(), ind) rw.text_label = ob.find('name').text rw.update_zoom(zoom) rw.rectHover.connect(rect_slot) self.thumbs.append(rw) self.roi_map[os.path.basename(ann_path)].append(rw) def load_image(self, path): return cv2.imread(path) def render_mosaic(self, sort_key='Class Name', rows=5, cols=5): add_thumbs = False if self.layout is None: add_thumbs = True self.panel.setLayout(None) self.layout = QtWidgets.QGraphicsLinearLayout(QtCore.Qt.Vertical) self.layout.setContentsMargins(50, 100, 50, 50) self.panel.setLayout(self.layout) if sort_key != self.sort_key: add_thumbs = True for ind, l in enumerate(self.layouts): while l.count(): l.removeItem(l.itemAt(0)) self.layout.removeItem(l) self.layouts = [] if sort_key == 'Class Name': self.thumbs.sort(key=lambda x: x.text_label.lower().strip(' '), reverse=False) if sort_key == 'Timestamp': self.thumbs.sort(key=lambda x: x.annotation_path, reverse=False) if sort_key == 'Height': self.thumbs.sort(key=lambda x: x.rectHeight, reverse=True) self.sort_key = sort_key #TODO: set col width to modulus of screen width COLUMNS = 8 ROWS = int(len(self.thumbs) / COLUMNS) i = 0 while i < len(self.thumbs): row_layout = QtWidgets.QGraphicsLinearLayout(QtCore.Qt.Horizontal) j = 0 while j < COLUMNS and i < len(self.thumbs): hide_thumb = False if self.hide_to_review: hide_thumb = hide_thumb | self.thumbs[i].forReview if self.hide_discarded: hide_thumb = hide_thumb | self.thumbs[i].toDiscard if self.hide_labeled: hide_thumb = hide_thumb | self.thumbs[i].isAnnotated if hide_thumb: self.thumbs[i].hide() else: self.thumbs[i].show() if add_thumbs: row_layout.addItem(self.thumbs[i]) self.thumbs[i].show() j += 1 i += 1 if add_thumbs: self.layout.addItem(row_layout) self.layouts.append(row_layout) def apply_label(self, label): for rect in self.thumbs: if rect.isSelected: rect.text_label = label rect.isSelected = False rect.isAnnotated = True rect.update() # update the label in the annotation and save it along rect.annotation.xml_tree.findall('object')[rect.annotation_object_index].find('name').text = label rect.annotation.save_xml(rect.annotation.xml_path) self.render_mosaic(sort_key=self.sort_key) def move_selected(self, src, dst, review=False, discard=False): move_dir = os.path.join(self.data_path, dst) for ind in range(0, len(self.thumbs)): if self.thumbs[ind].isSelected: if not os.path.exists(move_dir): os.makedirs(move_dir) # annotation and image paths image_path = os.path.join(src, self.thumbs[ind].annotation.xml_tree.find('filename').text) xml_path = os.path.join(src, os.path.basename(self.thumbs[ind].annotation.xml_path)) # check that the sources exist if not os.path.exists(image_path): continue if not os.path.exists(xml_path): continue # Move image file dst_path = os.path.join(move_dir, os.path.basename(image_path)) try: shutil.move(image_path, dst_path) except FileNotFoundError: LOG.warn('Missing image file, nothing to move') # Move xml file dst_path = os.path.join(move_dir, os.path.basename(xml_path)) try: shutil.move(xml_path, dst_path) except FileNotFoundError: LOG.warn('Missing xml file, nothing to move') # mark all thumbs associated with this annotation for review or discard for thumb in self.roi_map[os.path.basename(self.thumbs[ind].annotation.xml_path)]: LOG.info(thumb.annotation.xml_path) thumb.forReview = review thumb.toDiscard = discard thumb.update() self.render_mosaic(sort_key=self.sort_key) def clear_selected(self): for ind in range(0, len(self.thumbs)): self.thumbs[ind].isSelected = False self.thumbs[ind].update() #self.render_mosaic(sort_key=self.sort_key) def update_zoom(self, zoom): for rect in self.thumbs: rect.update_zoom(zoom)