Select Git revision
simple_tab.py
simple_tab.py 18.78 KiB
from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout,
QGroupBox, QCheckBox, QMessageBox, QPushButton,
QProgressBar,QComboBox,QLabel)
from PyQt6.QtCore import Qt, QThread, pyqtSignal, QTimer
import os
import sys
# Import utilities from debug_tool
from debug_tool.utils.qt_widgets import (create_group_with_text, create_button_layout,
create_info_group, create_preview_group)
from debug_tool.utils.file_handlers import select_file, clean_directory, copy_file, run_command
from debug_tool.utils.image_handlers import update_preview
# Import existing module implementations
from debug_tool.tabs.shifter_tab import ShifterTab
from debug_tool.tabs.depth_tab import DepthTab
from debug_tool.tabs.material_tab import MaterialTab
from debug_tool.tabs.edge_net_tab import EdgeNetTab
class PipelineWorker(QThread):
progress = pyqtSignal(str)
finished = pyqtSignal(bool, str)
def __init__(self, tab_instance):
super().__init__()
self.tab = tab_instance
def run(self):
try:
self.run_pipeline()
self.finished.emit(True, "Pipeline completed successfully!")
except Exception as e:
print(f"Pipeline failed with error: {str(e)}") # Debug print
self.finished.emit(False, f"Pipeline failed: {str(e)}")
def clean_temp_files(self):
self.tab.depth.clean_input_dir() # This one doesn't have dialogs
self.progress.emit("Cleaning input directory...")
self.tab.depth.remove_rgb()
self.progress.emit("Removing RGB image from 360monodepthexecution...")
self.tab.material.clean_working_dir(silent=True)
self.progress.emit("Cleaning Dynamic-Backward-Attention-Transformer temp directory...")
self.tab.edge_net.clean_output_directory(silent=True)
self.progress.emit("Cleaning EdgeNet output directory...")
def copy_file(self):
self.progress.emit("Copying input file to scripts/360monodepthexecution...")
self.tab.depth.copy_file()
def shift_image(self):
print("Starting shift_image") # Debug print
if not self.tab.shift_image_check.isChecked():
self.progress.emit("Skipping image shift...")
return
self.progress.emit("Shifting input image...")
self.tab.shifter.run_shifter()
print("Completed shift_image") # Debug print
def run_depth_estimation(self):
print("Starting depth_estimation") # Debug print
self.progress.emit("Running depth estimation...")
self.tab.depth.run_depth_estimation()
print("Completed depth_estimation") # Debug print
def run_material_recognition(self):
print("Starting material_recognition")
self.progress.emit("Running material recognition...")
try:
print(f"Input file path: {self.tab.material.input_file_path}")
print("Running split 360...")
success = self.tab.material.run_split_360()
if not success:
raise Exception("Split 360 failed")
print("Running material recognition...")
success = self.tab.material.run_material_recognition()
if not success:
raise Exception("Material recognition step failed")
print("Starting combine step...")
print(f"Current working directory: {os.getcwd()}")
print(f"Material recognition directory: {self.tab.material.material_recog_dir}")
print(f"Checking if cubemap directory exists: {os.path.exists(self.tab.material.cubemap_dir)}")
print(f"Checking if material output directory exists: {os.path.exists(self.tab.material.material_output_dir)}")
print("Files in cubemap directory:")
if os.path.exists(self.tab.material.cubemap_dir):
print("\n".join(os.listdir(self.tab.material.cubemap_dir)))
print("Files in material output directory:")
if os.path.exists(self.tab.material.material_output_dir):
print("\n".join(os.listdir(self.tab.material.material_output_dir)))
success = self.tab.material.run_combine()
if not success:
raise Exception("Combine step failed")
except Exception as e:
print(f"Material recognition error: {str(e)}")
raise
print("Completed material_recognition")
def run_edge_net(self):
print("Starting edge_net")
self.progress.emit("Running EdgeNet enhance360.py and infer360.py...")
try:
self.tab.edge_net.include_top = self.tab.include_top_check.isChecked()
self.tab.edge_net._run_edge_net_process()
print("Completed edge_net")
except Exception as e:
print(f"EdgeNet failed: {str(e)}")
raise
def run_post_processing(self):
self.progress.emit("Running post-processing...")
self.tab.edge_net._run_mesh_split_process()
self.progress.emit("Completed mesh split (replace.py)")
self.tab.edge_net._run_blender_flip_process()
self.progress.emit("Completed blender flip (blenderFlip.py)")
self.progress.emit("Post-processing completed!")
def run_pipeline(self):
self.clean_temp_files()
self.copy_file()
self.shift_image()
self.run_depth_estimation()
self.run_material_recognition()
self.run_edge_net()
self.run_post_processing()
self.progress.emit("Pipeline completed!")
class SimpleTab(QWidget):
def __init__(self, config_reader):
super().__init__()
self.config_reader = config_reader
self.input_path = None
self.pipeline_thread = None
# Initialize module instances
self.shifter = ShifterTab(self.config_reader)
self.depth = DepthTab(self.config_reader)
self.material = MaterialTab(self.config_reader)
self.edge_net = EdgeNetTab(self.config_reader)
# Hide their UIs as we'll use our own
self.shifter.hide()
self.depth.hide()
self.material.hide()
self.edge_net.hide()
self.file_selected = False
self.flash_timer = QTimer(self)
self.flash_timer.timeout.connect(self.toggle_flash)
self.setup_ui()
def setup_ui(self):
layout = QVBoxLayout(self)
# Controls section
controls_group = QGroupBox("Pipeline Controls")
controls_group.setStyleSheet("""
QGroupBox {
font-weight: bold;
border: 2px solid grey;
border-radius: 10px;
margin-top: 10px;
background-color: #3e3e3e;
padding: 15px;
}
QGroupBox::title {
margin: 10px;
background-color: transparent;
color: white;
}
""")
controls_layout = QVBoxLayout(controls_group)
# Info display
info_rows = [
("Input Image:", "No file selected"),
("Status:", "Ready - Waiting for input"),
]
self.info_group, self.info_labels = create_info_group("Information", info_rows)
self.info_group.setStyleSheet("""
QGroupBox {
font-weight: bold;
border: 2px solid grey;
border-radius: 10px;
margin-top: 10px;
background-color: #3e3e3e;
padding: 20px;
}
QGroupBox::title {
margin: 10px;
background-color: transparent;
color: white;
}
QLabel{
margin: 5px;
background-color: #3e3e3e;
color: white;
}
""")
for label in self.info_labels.values():
label.setStyleSheet("""
QLabel{
margin: 5px;
background-color: #3e3e3e;
color: white;
}
""")
controls_layout.addWidget(self.info_group)
# Options
options_layout = QHBoxLayout()
self.include_top_check = QCheckBox("Include Top in Mesh")
self.include_top_check.setStyleSheet("""
QCheckBox {
margin: 5px;
padding: 5px;
background-color: #3e3e3e;
color: white;
border: none;
border-radius: 5px;
}
""")
options_layout.addWidget(self.include_top_check)
self.shift_image_check = QCheckBox("Shift Input Image")
self.shift_image_check.setStyleSheet("QCheckBox { margin: 5px; background-color: #3e3e3e;}")
options_layout.addWidget(self.shift_image_check)
# SSC Model selection
ssc_model_layout = QHBoxLayout()
ssc_model_label = QLabel("SSC Model:")
#ssc_model_label.setAlignment(Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter)
ssc_model_label.setStyleSheet("""
QLabel {
margin: 5px;
background-color: transparent;
color: white;
font-weight: bold;
}
""")
ssc_model_label.setFixedWidth(ssc_model_label.sizeHint().width())
ssc_model_layout.addWidget(ssc_model_label)
self.ssc_model_combo = QComboBox()
self.ssc_model_combo.addItems(["EdgeNet360", "MDBNet"])
self.ssc_model_combo.setStyleSheet("""
QComboBox {
margin: 5px;
padding: 5px;
background-color: #1e1e1e;
color: white;
border: none;
border-radius: 5px;
}
QComboBox QAbstractItemView {
background-color: #1e1e1e;
color: white;
selection-background-color: #5e5e5e;
border-radius: 5px;
}
""")
self.ssc_model_combo.setFixedWidth(150)
ssc_model_layout.addWidget(self.ssc_model_combo)
ssc_model_layout.setAlignment(Qt.AlignmentFlag.AlignLeft)
# Add the horizontal layout to the options layout
options_layout.addLayout(ssc_model_layout)
controls_layout.addLayout(options_layout)
# Progress Bar
self.progress_bar = QProgressBar()
self.progress_bar.setMinimum(0)
self.progress_bar.setMaximum(0) # Makes it an indefinite progress bar
self.progress_bar.hide() # Hidden by default
self.progress_bar.setStyleSheet("""
QProgressBar {
border: 2px solid grey;
border-radius: 5px;
text-align: center;
}
QProgressBar::chunk {
background-color: #05B8CC;
width: 20px;
}
""")
controls_layout.addWidget(self.progress_bar)
# Buttons
self.run_pipeline_btn = QPushButton("Run Pipeline")
self.run_pipeline_btn.clicked.connect(self.run_full_pipeline)
self.run_pipeline_btn.setEnabled(False) # Disabled by default
self.run_pipeline_btn.setStyleSheet("""
QPushButton {
margin: 5px;
padding: 5px;
border-radius: 10px;
}
QPushButton:enabled {
background-color: green;
color: white;
}
QPushButton:disabled {
background-color: red;
color: white;
}
""")
self.run_pipeline_btn.setFixedSize(600, 40) # Explicit size
buttons_layout = QHBoxLayout()
self.select_btn = QPushButton("Select Input Image")
self.select_btn.clicked.connect(self.handle_file_select)
self.select_btn.setStyleSheet("""
QPushButton {
margin: 5px;
padding: 5px;
border-radius: 10px;
}
""")
self.select_btn.setFixedSize(600, 40) # Explicit size
buttons_layout.addWidget(self.select_btn)
buttons_layout.addWidget(self.run_pipeline_btn)
controls_layout.addLayout(buttons_layout)
layout.addWidget(controls_group)
# Status section
status_group, self.status_text = create_group_with_text("Pipeline Status", 300)
status_group.setStyleSheet("""
QGroupBox {
font-weight: bold;
border: 2px solid grey;
border-radius: 10px;
margin-top: 10px;
background-color: #3e3e3e;
padding: 20px;
}
QGroupBox::title {
margin: 10px;
background-color: transparent;
color: white;
}
""")
self.status_text.setStyleSheet("""
QTextEdit {
background-color: #1e1e1e;
color: white;
border: 2px solid grey;
border-radius: 10px;
padding: 10px;
font-size: 14px;
}
QTextEdit:focus {
border: 2px solid #05B8CC;
}
""")
layout.addWidget(status_group)
# Preview section
preview_group = QGroupBox("Preview")
preview_group.setStyleSheet("""
QGroupBox {
font-weight: bold;
border: 2px solid grey;
border-radius: 10px;
margin-top: 10px;
background-color: #3e3e3e;
padding: 20px;
}
QGroupBox::title {
margin: 10px;
background-color: transparent;
color: white;
}
""")
preview_layout = QHBoxLayout(preview_group)
input_group, self.input_preview = create_preview_group("Input Image")
output_group, self.output_preview = create_preview_group("Current Output")
input_group.setStyleSheet("""
QGroupBox {
font-weight: bold;
border: 2px solid grey;
border-radius: 10px;
margin-top: 10px;
background-color: #3e3e3e;
padding: 20px;
}
QGroupBox::title {
margin: 10px;
background-color: transparent;
color: white;
}
QLabel {
margin: 5px;
background-color: #3e3e3e;
color: white;
}
""")
output_group.setStyleSheet("""
QGroupBox {
font-weight: bold;
border: 2px solid grey;
border-radius: 10px;
margin-top: 10px;
background-color: #3e3e3e;
padding: 20px;
}
QGroupBox::title {
margin: 10px;
background-color: transparent;
color: white;
}
QLabel {
margin: 5px;
background-color: #3e3e3e;
color: white;
}
""")
preview_layout.addWidget(input_group)
preview_layout.addWidget(output_group)
layout.addWidget(preview_group)
# Start flashing if no file is selected
self.start_flashing()
def handle_file_select(self):
file_path = select_file(
self,
"Select Input Image",
"Images (*.png *.jpg *.jpeg)",
initial_dir=self.config_reader.directories['edgeNetDir'] + '/Data'
)
if file_path:
self.input_path = file_path
self.info_labels["Input Image:"].setText(os.path.basename(file_path))
self.update_status(f"Selected input file: {file_path}")
update_preview(self.input_preview, file_path,
error_callback=self.update_status)
# Enable the run pipeline button
self.run_pipeline_btn.setEnabled(True)
# Provide input path to all modules
self.shifter.input_file_path = file_path
self.depth.depth_input_path = file_path
self.material.input_file_path = file_path
# self.edge_net.input_path = file_path # edgenet have default input path
self.file_selected = True
self.flash_timer.stop()
self.select_btn.setStyleSheet("QPushButton { margin: 5px; padding: 5px; border-radius: 10px;}")
def start_flashing(self):
if not self.file_selected:
self.flash_timer.start(1000) # Flash every 1000 milliseconds
def toggle_flash(self):
current_style = self.select_btn.styleSheet()
if "background-color: DarkOrange;" in current_style:
self.select_btn.setStyleSheet("QPushButton { margin: 5px; padding: 5px; border-radius: 10px;}")
else:
self.select_btn.setStyleSheet("QPushButton { margin: 5px; padding: 5px; background-color: DarkOrange; border-radius: 10px;}")
def run_full_pipeline(self):
if not self.input_path:
QMessageBox.warning(self, "Warning", "Please select an input file first")
return
if self.pipeline_thread and self.pipeline_thread.isRunning():
QMessageBox.warning(self, "Warning", "Pipeline is already running")
return
# Show progress bar and update status
self.progress_bar.show()
self.run_pipeline_btn.setEnabled(False)
self.pipeline_thread = PipelineWorker(self)
# Connect signals
self.pipeline_thread.progress.connect(self.update_status)
self.pipeline_thread.finished.connect(self.pipeline_completed)
# Disable controls while running
self.setEnabled(False)
self.progress_bar.setEnabled(True) # Keep progress bar enabled
# Start the pipeline
self.pipeline_thread.start()
def pipeline_completed(self, success, message):
self.setEnabled(True)
self.progress_bar.hide()
self.run_pipeline_btn.setEnabled(True)
self.update_status(message)
if success:
QMessageBox.information(self, "Success", "Pipeline completed successfully!")
else:
QMessageBox.critical(self, "Error", f"Pipeline failed: {message}")
def update_status(self, message):
self.status_text.append(message)
self.info_labels["Status:"].setText(message.split("...")[-1] if "..." in message else message)
# Scroll to bottom
scrollbar = self.status_text.verticalScrollBar()
scrollbar.setValue(scrollbar.maximum())