From c08ede1e769f3f44e473a34fc9775c6d780379fc Mon Sep 17 00:00:00 2001 From: Till Raab Date: Wed, 8 Nov 2023 15:31:17 +0100 Subject: [PATCH] working on bbox corrector --- correct_bboxes.py | 72 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 correct_bboxes.py diff --git a/correct_bboxes.py b/correct_bboxes.py new file mode 100644 index 0000000..ae89da3 --- /dev/null +++ b/correct_bboxes.py @@ -0,0 +1,72 @@ +import os +import sys +from PyQt5.QtWidgets import * +from PyQt5.QtCore import * +import pyqtgraph as pg + +from PIL import Image, ImageOps +import numpy as np +import torch +import torchvision.transforms.functional as F + +import pathlib + +class Bbox_correct_UI(QMainWindow): + def __init__(self, img_path, parent=None): + super(Bbox_correct_UI, self).__init__(parent) + label_path = (pathlib.Path(img_path).parent.parent / 'labels' / pathlib.Path(img_path).name).with_suffix('.txt') + + rec = QApplication.desktop().screenGeometry() + height = rec.height() + width = rec.width() + self.resize(int(.8 * width), int(.8 * height)) + self.setWindowTitle('efishSignalTracker') # set window title + + self.central_widget = QWidget(self) + self.gridLayout = QGridLayout() + self.central_widget.setLayout(self.gridLayout) + self.setCentralWidget(self.central_widget) + + self.plot_handels = [] + self.plot_widgets = [] + self.win = pg.GraphicsLayoutWidget() + + + self.plot_handels.append(pg.ImageItem(ColorMap='viridis')) + self.plot_widgets.append(self.win.addPlot(title="")) + # xxx.setMouseMode(xxx.RectMode) + + self.plot_widgets[0].addItem(self.plot_handels[0], colorMap='viridis') + # self.plot_widgets[0].setLabel('left', 'frequency [Hz]') + # self.plot_widgets[0].setLabel('bottom', 'time [s]') + + self.gridLayout.addWidget(self.win, 0, 0) + + self.plot_handels[0].getViewBox().invertY(True) + img = Image.open(img_path) + img_gray = ImageOps.grayscale(img) + img_array = np.array(img_gray).T + + + labels = np.loadtxt(label_path, delimiter=' ') + + self.plot_handels[0].setImage(np.array(img_array)) + + for l in labels: + x_center, y_center, width, height = l[1:] * img_array.shape[1] + x0, y0, = x_center-width/2, y_center-height/2 + ROI = pg.RectROI((x0, y0), size=(width, height)) + self.plot_widgets[0].addItem(ROI) + + self.plot_widgets[0].setYRange(0, img_array.shape[0], padding=0) + self.plot_widgets[0].setXRange(0, img_array.shape[1], padding=0) + +def main_UI(): + app = QApplication(sys.argv) # create application + img_path = sys.argv[1] + w = Bbox_correct_UI(img_path) # create window + w.show() + sys.exit(app.exec_()) + +if __name__ == '__main__': + main_UI() \ No newline at end of file