Skip to content
Hansongli edited this page Feb 20, 2025 · 6 revisions

AiCore - Code Review

AiCore.py :

Main Method:

Initialization of Main Window:

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())

This function is to start the main loop. app is to create an instance of the application. The window is to initialize the Main Window.

MainWindow Class:

class MainWindow(QMainWindow):

MainWindow class inherits the QMainWindow from PyQT library

    dataUpdated = pyqtSignal(str)

the dataUpdated is to communicate with the SegmentAndMap.py`

Initialization:

    def __init__(self):
        global canEnable
        super().__init__()
        self.setWindowFlags(Qt.FramelessWindowHint)
        self.setWindowTitle("AiCore x SpatterSense")
        self.title_bar = TitleBar(self,self)
        self.setGeometry(100, 100, 1200, 800)

The global variable canEnable is used for button toggling later on. FramelessWindowHint is used to customize the title bar later on.

        self.label = QLabel("AI Core Viewer", self)
        self.label.setGeometry(10, 10, 780, 30)

        self.default_size = (10, 10)
        self.textures = {}
        self.image_paths = {}
        self.segments = []
        self.end_points =[]
        self.average_end_point = np.array([0.0, 0.0, 0.0])  
        self.previous_data = None  

        self.main_widget = QWidget()
        self.setCentralWidget(self.main_widget)

        self.main_layout = QVBoxLayout(self.main_widget)
        self.main_layout.addWidget(self.title_bar)

        self.content_layout = QHBoxLayout()        
        
        self.tab_layout = QHBoxLayout()
        self.tabs = QTabWidget()
        self.tab_layout.addWidget(self.tabs)
        
        self.viewer3D = QWidget()
        self.viewer_layout3D = QVBoxLayout(self.viewer3D)
        self.docker3d = QHBoxLayout()
        self.viewer_layout3D.addLayout(self.docker3d)
        self.plotter3D = QtInteractor(self.viewer3D)
        self.export = QPushButton("Export")
        self.export.clicked.connect(self.export_plotter)
        self.add_head_btn = QPushButton("Add Head")
        self.add_head_btn.clicked.connect(self.add_head)
        self.docker3d.addWidget(self.add_head_btn)
        self.docker3d.addWidget(self.export)
        self.viewer_layout3D.addWidget(self.plotter3D.interactor)
        
        self.viewer2D = QWidget()
        self.viewer_layout2D = QHBoxLayout(self.viewer2D)
        
        self.headIcon = QPixmap(self.get_resource_path("images/head.png"))
        self.headscaled = self.headIcon.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)
        self.exportIcon = QPixmap(self.get_resource_path("images/Export.png"))
        self.exportscaled = self.exportIcon.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)
        
        self.add_head_btn.setIcon(QIcon(self.headscaled))
        self.export.setIcon(QIcon(self.exportscaled))
        
        self.tabs.addTab(self.viewer3D, "3D Simulation")
        self.tabs.addTab(self.viewer2D, "2D Analysis")        
        self.content_layout.addLayout(self.tab_layout)
        
        self.sidebar = QWidget()
        self.sidebar.setObjectName("sidebar")
        self.sidebar.setFixedWidth(250)
        self.shadow = QGraphicsDropShadowEffect()
        self.shadow.setBlurRadius(12)
        self.shadow.setXOffset(5)
        self.shadow.setYOffset(3)
        self.shadow.setColor(QColor(0,0,0,100))
        self.sidebar.setGraphicsEffect(self.shadow)
        self.sidebar_layout = QVBoxLayout(self.sidebar)
        
        self.floorIcon = QPixmap(self.get_resource_path("images/floor.png"))
        self.scaled_pixmap1 = self.floorIcon.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)
        self.wallrIcon = QPixmap(self.get_resource_path("images/wallR.png"))
        self.scaled_pixmap2 = self.wallrIcon.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)
        self.walllIcon = QPixmap(self.get_resource_path("images/wallL.png"))
        self.scaled_pixmap3 = self.walllIcon.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)
        self.wallbIcon = QPixmap(self.get_resource_path("images/wallB.png"))
        self.scaled_pixmap4 = self.wallbIcon.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)
        self.wallfIcon = QPixmap(self.get_resource_path("images/wallF.png"))
        self.scaled_pixmap5 = self.wallfIcon.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)
        self.generateicon = QPixmap(self.get_resource_path("images/generateicon.png"))
        self.scaled_pixmap6 = self.generateicon.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)
        self.selectspatter = QPixmap(self.get_resource_path("images/select spatter.png"))
        self.scaled_pixmap7 = self.selectspatter.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)
        self.deletePlaneIcon = QPixmap(self.get_resource_path("images/delete.png"))
        self.scaled_pixmap8 = self.deletePlaneIcon.scaled(500, 500, aspectRatioMode=Qt.KeepAspectRatio)

        self.assets_frame = QFrame()
        self.assets_frame.setFrameStyle(QFrame.Raised)
        self.assets_layout = QGridLayout(self.assets_frame)
        
        self.Header3D = QLabel("3D Assets")
        self.sidebar_layout.addWidget(self.Header3D)
        self.sidebar_layout.addWidget(self.assets_frame)

        self.add_floor_btn = QPushButton("Floor")
        self.add_floor_btn.clicked.connect(lambda: self.add_plane_with_image("floor"))
        self.add_floor_btn.setIcon(QIcon(self.scaled_pixmap1))
        self.add_floor_btn.setIconSize(QSize(20,20))

        self.add_right_wall_btn = QPushButton("Right Wall")
        self.add_right_wall_btn.clicked.connect(lambda: self.add_plane_with_image("right"))
        self.add_right_wall_btn.setIcon(QIcon(self.scaled_pixmap2))
        self.add_right_wall_btn.setIconSize(QSize(20,20))

        self.add_left_wall_btn = QPushButton("Left Wall")
        self.add_left_wall_btn.clicked.connect(lambda: self.add_plane_with_image("left"))
        self.add_left_wall_btn.setIcon(QIcon(self.scaled_pixmap3))
        self.add_left_wall_btn.setIconSize(QSize(20,20))

        self.add_back_wall_btn = QPushButton("Back Wall")
        self.add_back_wall_btn.clicked.connect(lambda: self.add_plane_with_image("back"))
        self.add_back_wall_btn.setIcon(QIcon(self.scaled_pixmap4))
        self.add_back_wall_btn.setIconSize(QSize(20,20))

        self.add_front_wall_btn = QPushButton("Front Wall")
        self.add_front_wall_btn.clicked.connect(lambda: self.add_plane_with_image("front"))
        self.add_front_wall_btn.setIcon(QIcon(self.scaled_pixmap5))
        self.add_front_wall_btn.setIconSize(QSize(20,20))
        
        self.del_floor_btn = QPushButton()
        self.del_floor_btn.clicked.connect(lambda: self.delete_plane(self.plotter3D.renderer.actors["floor_plane"]))
        self.del_right_wall_btn = QPushButton()
        self.del_right_wall_btn.clicked.connect(lambda: self.delete_plane(self.plotter3D.renderer.actors["right_plane"]))
        self.del_left_wall_btn = QPushButton()
        self.del_left_wall_btn.clicked.connect(lambda: self.delete_plane(self.plotter3D.renderer.actors["left_plane"]))
        self.del_back_wall_btn = QPushButton()
        self.del_back_wall_btn.clicked.connect(lambda: self.delete_plane(self.plotter3D.renderer.actors["back_plane"]))
        self.del_front_wall_btn = QPushButton()
        self.del_front_wall_btn.clicked.connect(lambda: self.delete_plane(self.plotter3D.renderer.actors["front_plane"]))
        
        self.asset_buttons = [ self.add_floor_btn, self.del_floor_btn, self.add_right_wall_btn, self.del_right_wall_btn, self.add_left_wall_btn, self.del_left_wall_btn, self.add_back_wall_btn, self.del_back_wall_btn, self.add_front_wall_btn, self.del_front_wall_btn ]

        for i, button in enumerate(self.asset_buttons):
            row = i // 2
            col = i % 2
            if i % 2 != 0:
                button.setFixedSize(QSize(20,20))
                button.setIcon(QIcon(self.deletePlaneIcon))
            self.assets_layout.addWidget(button, row, col)
            
        self.analysis_label = QLabel("Analysis")
        self.sidebar_layout.addWidget(self.analysis_label)
        
        self.analysis_frame = QFrame()
        self.analysis_frame.setFrameShape(QFrame.StyledPanel)
        self.analysis_frame.setFrameShadow(QFrame.Raised)
        self.analysis_layout = QGridLayout(self.analysis_frame)
                
        self.add_points_btn = QPushButton("Add Points")
        self.add_points_btn.setIcon(QIcon(self.scaled_pixmap7))
        self.add_points_btn.setIconSize(QSize(20,20))
        self.add_points_btn.clicked.connect(self.open_image_with_interaction)
        self.sidebar_layout.addWidget(self.add_points_btn)

        self.texture_select = QComboBox()
        self.texture_select.addItem("Floor")
        self.texture_select.addItem("Right")
        self.texture_select.addItem("Left")
        self.texture_select.addItem("Back")
        self.texture_select.addItem("Front")
        self.sidebar_layout.addWidget(self.texture_select)
        
        self.analysis_buttons = [self.add_points_btn, self.texture_select]

        for i, button in enumerate(self.analysis_buttons):
            self.analysis_layout.addWidget(button, 0, i)
            
        self.sidebar_layout.addWidget(self.analysis_frame)

        self.ObjectListLabel = QLabel("Spatters:")
        self.sidebar_layout.addWidget(self.ObjectListLabel)
        self.object_list = QListWidget()
        self.object_list.itemClicked.connect(self.on_object_selected)
        self.sidebar_layout.addWidget(self.object_list)

        self.delete_button = QPushButton("Delete Selected")
        self.delete_button.clicked.connect(self.delete_selected_object)
        self.sidebar_layout.addWidget(self.delete_button)

        self.sidebar_layout.addStretch()

        self.content_layout.addWidget(self.sidebar)

        self.main_layout.addLayout(self.content_layout)

        self.bottom_bar = QWidget()
        self.bottom_bar_layout = QHBoxLayout(self.bottom_bar)
        self.bottom_bar.setFixedHeight(35)
        self.InformationHeader = QLabel("Information")
        self.stainCount = QLabel("Spatter Count: 0")
        self.AngleReport = QLabel("Impact Angle: 0")
        self.HeightReport = QLabel("Point of Origin: 0")
        self.Conclusive = QLabel("")
        self.bottom_bar_layout.addWidget(self.InformationHeader)
        self.bottom_bar_layout.addWidget(self.stainCount)
        self.bottom_bar_layout.addWidget(self.AngleReport)
        self.bottom_bar_layout.addWidget(self.HeightReport)
        self.bottom_bar_layout.addWidget(self.Conclusive)

        self.main_layout.addWidget(self.bottom_bar)

        self.enableUI(canEnable)

        self.setStyleSheet(self.load_stylesheet(self.get_resource_path("style/style.css")))
        if getattr(sys, 'frozen', False):
            pyi_splash.close()
        self.plotter3D.set_background("#3f3f3f")
        self.configure_plotter()
        
    def export_plotter(self):
        opt = QFileDialog.Options()
        file_path, _ = QFileDialog.getSaveFileName(
            self,
            "Save Screenshot",
            "",
            "PNG Files (*.png);;All Files (*)",
            options=opt
        )
        
        if file_path: 
            self.plotter3D.screenshot(file_path)

SegmentAndMap.py :

class SegmentAndMap(QWidget):
    dataUpdated = pyqtSignal(str)