Merge pull request #89 from RUTILEA/feature/fix_for_broken_image

Feature/fix for broken image

src/main/python/model/dataset.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from model.project import Project
 from model.supporting_model import TrimmingData
-import os, cv2
+import os, shutil, imageio
 
 
 class Dataset:
@@ -12,6 +12,7 @@ class Category(Enum):
         TRAINING_OK = auto()
         TEST_OK = auto()
         TEST_NG = auto()
+        TRUNCATED = auto()
 
     @classmethod
     def _root_path(cls) -> str:
@@ -25,6 +26,8 @@ def images_path(cls, category: Category) -> Path:
             return Path(cls._root_path() + '/test/OK')
         elif category is cls.Category.TEST_NG:
             return Path(cls._root_path() + '/test/NG')
+        elif category is cls.Category.TRUNCATED:
+            return Path(cls._root_path() + '/truncated')
         else:
             assert False
 
@@ -46,11 +49,15 @@ def generate_image_path(cls, category: Category, cam_number: int, file_extension
         return cls.images_path(category).joinpath(file_name)
 
     @classmethod
-    def trim_image(cls, path: Path, save_path: Path, data: TrimmingData):
-        img = cv2.imread(path)
+    def trim_image(cls, path: Path, save_path: Path, data: TrimmingData) -> Path:
+        try:
+            img = imageio.imread(path)
+        except:
+            return path
+        file_name = os.path.basename(path)
         position = data.position
         size = data.size
         rect = img[int(position[1]):int(position[1])+size[1], int(position[0]):int(position[0])+size[0]]
-        file_name = os.path.basename(path)
-        cv2.imwrite(os.path.join(save_path, file_name), rect)
+        imageio.imwrite(os.path.join(save_path, file_name), rect)
+        return
 

src/main/python/model/learning_model.py

@@ -2,7 +2,7 @@
 from module.novelty_detector import NoveltyDetector
 from model.dataset import Dataset
 from model.project import Project
-import threading, os, numpy as np
+import threading, os, shutil, numpy as np
 from statistics import stdev, mean
 from math import sqrt
 
@@ -169,10 +169,13 @@ def load_weights(self):
     def start_predict(self, image_paths):
         image_path = image_paths[0]
         trimming_data = Project.latest_trimming_data()
-        Dataset.trim_image(image_path, os.path.dirname(image_path), trimming_data)
+        truncated_image_path = Dataset.trim_image(image_path, os.path.dirname(image_path), trimming_data)
+        if truncated_image_path:
+            return truncated_image_path
         self.predicting_start.emit()
         predict_thread = threading.Thread(target=self.predict, args=([image_paths]))
         predict_thread.start()
+        return
 
     def predict(self, image_paths):
         scores = self.__model.predict_paths(image_paths)

src/main/python/view/dataset.py

@@ -60,6 +60,7 @@ def __init__(self):
         self.watcher.directoryChanged.connect(self.on_dataset_directory_changed)
 
         self.select_area_dialog = None
+        self.msgBox = None
 
         LearningModel.default().training_finished.connect(self.on_finished_training)
 
@@ -175,6 +176,7 @@ def on_clicked_train_button(self):
 
     def on_finished_selecting_area(self, data: TrimmingData):
         categories = [Dataset.Category.TRAINING_OK, Dataset.Category.TEST_OK, Dataset.Category.TEST_NG]
+        truncated_image_paths = []
         for category in categories:
             dir_path = Dataset.images_path(category)
             save_path = Dataset.trimmed_path(category)
@@ -188,8 +190,26 @@ def on_finished_selecting_area(self, data: TrimmingData):
                 file_list = [img for img in file_list if
                                   Path(img).suffix in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']]
                 for file_name in file_list:
-                    Dataset.trim_image(os.path.join(dir_path, file_name), save_path, data)
-        Project.save_latest_trimming_data(data)
+                    truncated_image_path = Dataset.trim_image(os.path.join(dir_path, file_name), save_path, data)
+                    if truncated_image_path:
+                        file_name = os.path.basename(truncated_image_path)
+                        shutil.move(truncated_image_path,
+                                    os.path.join(Dataset.images_path(Dataset.Category.TRUNCATED), file_name))
+                        truncated_image_paths.append(truncated_image_path)
+            Project.save_latest_trimming_data(data)
+
+        # alert for moving truncated images
+        if truncated_image_paths:
+            self.msgBox = QMessageBox()
+            self.msgBox.setText(str(len(truncated_image_paths))+'枚の画像を読み込めませんでした. これらの画像はtruncatedフォルダに移動されました.\n\n'\
+                                + 'このままトレーニングを開始しますか？')
+            self.msgBox.setStandardButtons(self.msgBox.Yes | self.msgBox.No)
+            self.msgBox.setDefaultButton(self.msgBox.Yes)
+            reply = self.msgBox.exec()
+            if reply == self.msgBox.No:
+                return
+
+        # start training
         LearningModel.default().start_training()
 
     def on_dataset_directory_changed(self, directory: str):

src/main/python/view/inspection.py

@@ -56,6 +56,8 @@ def __init__(self):
         self.ui.OK_counter_label.setText(str(self.ok_counter))
         self.ui.NG_counter_label.setText(str(self.ng_counter))
 
+        self.msgBox = None
+
     @property
     def ok_counter(self):
         return self.__ok_counter
@@ -97,7 +99,9 @@ def on_clicked_inspect_button(self):
     def on_image_saved(self, image_path):
         # FIXME: refactor the structure of camera model class not to call this function from camera_model.capture
         if os.path.basename(os.path.dirname(image_path)) == 'tmp':
-            self.learning_model.start_predict([image_path])
+            path = self.learning_model.start_predict([image_path])
+            if path:
+                self.alert_for_truncated_image(path)
 
     def on_finished_predicting(self, result):
         image_path = result['image_paths'][0]
@@ -145,7 +149,15 @@ def on_clicked_inspection_existing_image_button(self):
             file_name = f'camera_0_{timestamp}.{ext}'
             copied_image_path = Project.project_path() + '/tmp/' + file_name
             copy2(original_image_path, copied_image_path)
-            self.learning_model.start_predict([copied_image_path])
+            path = self.learning_model.start_predict([copied_image_path])
+            if path:
+                self.alert_for_truncated_image(original_image_path)
+                return
             self.ui.inspect_button.setDisabled(True)
             self.ui.inspect_existing_image_button.setDisabled(True)
 
+    def alert_for_truncated_image(self, path):
+        self.msgBox = QMessageBox()
+        self.msgBox.setText('この画像は破損している可能性があります. 検品はスキップされます.\n\n'
+                            + '検品を続けるにはOKを押してください.')
+        self.msgBox.exec()

src/main/python/view/new_project.py

@@ -15,6 +15,7 @@
         - NG
     - train 
         - OK
+    - truncated
 - inspection_results
     - images 
 - tmp
@@ -70,6 +71,7 @@ def on_clicked_create_button(self):
             os.path.join(save_location_path, 'dataset/test/OK'),
             os.path.join(save_location_path, 'dataset/test/NG'),
             os.path.join(save_location_path, 'dataset/train/OK'),
+            os.path.join(save_location_path, 'dataset/truncated'),
             os.path.join(save_location_path, 'inspection_results'),
             os.path.join(save_location_path, 'inspection_results/images'),
             os.path.join(save_location_path, 'tmp'),

src/main/python/view/ui/dataset.py

@@ -2,12 +2,14 @@
 
 # Form implementation generated from reading ui file 'dataset.ui'
 #
-# Created by: PyQt5 UI code generator 5.9.2
+# Created by: PyQt5 UI code generator 5.13.0
 #
 # WARNING! All changes made in this file will be lost!
 
+
 from PyQt5 import QtCore, QtGui, QtWidgets
 
+
 class Ui_Dataset(object):
     def setupUi(self, Dataset):
         Dataset.setObjectName("Dataset")
@@ -171,5 +173,4 @@ def retranslateUi(self, Dataset):
         self.select_images_action.setToolTip(_translate("Dataset", "既存の画像"))
         self.camera_action.setText(_translate("Dataset", "カメラで撮影"))
         self.camera_action.setToolTip(_translate("Dataset", "カメラで撮影"))
-
 from qrc import icon_rc

src/main/python/view/ui/inspection.py

@@ -2,12 +2,14 @@
 
 # Form implementation generated from reading ui file 'inspection.ui'
 #
-# Created by: PyQt5 UI code generator 5.9.2
+# Created by: PyQt5 UI code generator 5.13.0
 #
 # WARNING! All changes made in this file will be lost!
 
+
 from PyQt5 import QtCore, QtGui, QtWidgets
 
+
 class Ui_inspection(object):
     def setupUi(self, inspection):
         inspection.setObjectName("inspection")
@@ -299,6 +301,5 @@ def retranslateUi(self, inspection):
         self.NG_message_label.setText(_translate("inspection", "この製品は不良品です"))
         self.ng_score.setText(_translate("inspection", "<html><head/><body><p align=\"center\"><br/></p></body></html>"))
         self.total_header_label.setText(_translate("inspection", "累計"))
-
 from view.q_camera_view_finder_with_guide import QCameraViewFinderWithGuide
 from qrc import icon_rc

src/main/python/view/ui/main_window.py

@@ -2,12 +2,14 @@
 
 # Form implementation generated from reading ui file 'main_window.ui'
 #
-# Created by: PyQt5 UI code generator 5.9.2
+# Created by: PyQt5 UI code generator 5.13.0
 #
 # WARNING! All changes made in this file will be lost!
 
+
 from PyQt5 import QtCore, QtGui, QtWidgets
 
+
 class Ui_MainWindow(object):
     def setupUi(self, MainWindow):
         MainWindow.setObjectName("MainWindow")
@@ -116,5 +118,4 @@ def retranslateUi(self, MainWindow):
         self.action_about_SDT.setText(_translate("MainWindow", "SDTestについて"))
         self.action_about_SDT.setToolTip(_translate("MainWindow", "SDTestについて"))
         self.action_quit_SDTest.setText(_translate("MainWindow", "SDTestを終了"))
-
 from qrc import icon_rc

src/main/python/view/ui/startup.py

@@ -2,12 +2,14 @@
 
 # Form implementation generated from reading ui file 'startup.ui'
 #
-# Created by: PyQt5 UI code generator 5.9.2
+# Created by: PyQt5 UI code generator 5.13.0
 #
 # WARNING! All changes made in this file will be lost!
 
+
 from PyQt5 import QtCore, QtGui, QtWidgets
 
+
 class Ui_StartupWidget(object):
     def setupUi(self, StartupWidget):
         StartupWidget.setObjectName("StartupWidget")
@@ -166,6 +168,5 @@ def retranslateUi(self, StartupWidget):
         self.logo_title_label.setText(_translate("StartupWidget", "Software-Defined Test"))
         self.new_project_button.setText(_translate("StartupWidget", "新規プロジェク ト"))
         self.open_project_button.setText(_translate("StartupWidget", "開く"))
-
 from qrc import icon_rc
 from qrc import logo_rc

src/main/python/view/ui/test.py

@@ -2,12 +2,14 @@
 
 # Form implementation generated from reading ui file 'test.ui'
 #
-# Created by: PyQt5 UI code generator 5.9.2
+# Created by: PyQt5 UI code generator 5.13.0
 #
 # WARNING! All changes made in this file will be lost!
 
+
 from PyQt5 import QtCore, QtGui, QtWidgets
 
+
 class Ui_Test(object):
     def setupUi(self, Test):
         Test.setObjectName("Test")
@@ -236,5 +238,4 @@ def retranslateUi(self, Test):
         self.false_negative_rate_title_label.setText(_translate("Test", "見逃し率"))
         self.false_negative_rate_label.setText(_translate("Test", "1%"))
         self.details_button.setText(_translate("Test", "詳細"))
-
 from qrc import icon_rc