From 32d5ac74d6ac4b7124397eaf021c9471a8798a56 Mon Sep 17 00:00:00 2001
From: mhby1g21 <mhby1g21@soton.ac.uk>
Date: Thu, 14 Nov 2024 13:57:58 +0000
Subject: [PATCH] added progress bar, removed auto generate tickbox,
reimplemented pipeline worker and steps
---
scripts/simple_tab.py | 116 +++++++++++++++++++++++-------------------
1 file changed, 65 insertions(+), 51 deletions(-)
diff --git a/scripts/simple_tab.py b/scripts/simple_tab.py
index e415700..88c3924 100644
--- a/scripts/simple_tab.py
+++ b/scripts/simple_tab.py
@@ -1,5 +1,6 @@
from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout,
- QGroupBox, QCheckBox, QMessageBox, QPushButton)
+ QGroupBox, QCheckBox, QMessageBox, QPushButton,
+ QProgressBar)
from PyQt6.QtCore import Qt, QThread, pyqtSignal
import os
import sys
@@ -31,19 +32,35 @@ class PipelineWorker(QThread):
except Exception as e:
self.finished.emit(False, f"Pipeline failed: {str(e)}")
- def run_pipeline(self):
- steps = [
- (self.tab.shift_image, "Shifting input image..."),
- (self.tab.run_depth_estimation, "Running depth estimation..."),
- (self.tab.run_material_recognition, "Running material recognition..."),
- (self.tab.run_edge_net, "Running EdgeNet..."),
- (self.tab.run_mesh_splitting, "Running mesh splitting..."),
- (self.tab.run_blender_flip, "Running Blender flip...")
- ]
+ def shift_image(self):
+ 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()
- for step_func, message in steps:
- self.progress.emit(message)
- step_func()
+ def run_depth_estimation(self):
+ self.progress.emit("Running depth estimation...")
+ self.tab.depth.clean_input_dir()
+ self.tab.depth.remove_rgb()
+ self.tab.depth.copy_file()
+ self.tab.depth.run_depth_estimation()
+
+ def run_material_recognition(self):
+ self.progress.emit("Running material recognition...")
+ self.tab.material.run_all_steps()
+
+ def run_edge_net(self):
+ self.progress.emit("Running EdgeNet...")
+ self.tab.edge_net.include_top = self.tab.include_top_check.isChecked()
+ self.tab.edge_net.clean_output_directory()
+ self.tab.edge_net.run_all_steps()
+
+ def run_pipeline(self):
+ self.shift_image()
+ self.run_depth_estimation()
+ self.run_material_recognition()
+ self.run_edge_net()
class SimpleTab(QWidget):
def __init__(self, config_reader):
@@ -84,10 +101,6 @@ class SimpleTab(QWidget):
# Options
options_layout = QHBoxLayout()
- self.auto_depth_check = QCheckBox("Auto Generate Depth Map")
- self.auto_depth_check.setChecked(True)
- options_layout.addWidget(self.auto_depth_check)
-
self.include_top_check = QCheckBox("Include Top in Mesh")
options_layout.addWidget(self.include_top_check)
@@ -96,12 +109,25 @@ class SimpleTab(QWidget):
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
+ controls_layout.addWidget(self.progress_bar)
+
# Buttons
- buttons = [
- ("Select Input Image", self.handle_file_select, 'left'),
- ("Run Pipeline", self.run_full_pipeline, 'right')
- ]
- controls_layout.addLayout(create_button_layout(*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
+
+ buttons_layout = QHBoxLayout()
+ select_btn = QPushButton("Select Input Image")
+ select_btn.clicked.connect(self.handle_file_select)
+
+ buttons_layout.addWidget(select_btn)
+ buttons_layout.addWidget(self.run_pipeline_btn)
+ controls_layout.addLayout(buttons_layout)
layout.addWidget(controls_group)
@@ -134,6 +160,9 @@ class SimpleTab(QWidget):
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.depth_input_path = file_path
self.depth.depth_input_path = file_path
@@ -149,52 +178,37 @@ class SimpleTab(QWidget):
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)
- self.pipeline_thread.start()
# 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())
-
- # Pipeline steps using existing module implementations
- def shift_image(self):
- if not self.shift_image_check.isChecked():
- self.update_status("Skipping image shift...")
- return
- self.shifter.run_shifter()
-
- def run_depth_estimation(self):
- if not self.auto_depth_check.isChecked():
- self.update_status("Skipping depth estimation - manual depth selected")
- return
- self.depth.run_depth_estimation()
-
- def run_material_recognition(self):
- self.material.run_material_recognition()
-
- def run_edge_net(self):
- self.edge_net.include_top = self.include_top_check.isChecked()
- self.edge_net.run_edge_net()
-
- def run_mesh_splitting(self):
- self.edge_net.run_mesh_splitting()
-
- def run_blender_flip(self):
- self.edge_net.run_blender_flip()
\ No newline at end of file
+ scrollbar.setValue(scrollbar.maximum())
\ No newline at end of file
--
GitLab