working on bbox corrector
This commit is contained in:
parent
5f270c41f2
commit
c08ede1e76
72
correct_bboxes.py
Normal file
72
correct_bboxes.py
Normal file
@ -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()
|
Loading…
Reference in New Issue
Block a user