refactoring

lab2_1/lab2_1.py

@@ -1,161 +1,162 @@
 import os
 import sys
+import traceback
+
 import cv2
 import numpy as np
 from matplotlib import pyplot as plt
 
 
-def menu():
-    print("Выберите пункт меню:\n1.Linear filter\n2.Blur\n3.Median blur\n4.GaussianBlur\n5.Erode")
-    print("6.Dilate\n7.MORPH_OPERATIONS\n8.Sobel\n9.Laplacian\n10.Canny\n11.CalcHist\n12.EqualizeHist")
-    print("13.Save\n14.Add Figure\n15.Del Figure\n16.Exit")
+SOURCE_IMAGE = "test.jpg"
+
+
+def filter_linear(image):
+    kernel = np.array([[-0.1, 0.2, -0.1], [0.2, 3.0, 0.2], [-0.1, 0.2, -0.1]])
+    cv2.imshow("Исходное изображение", image)
+    cv2.imshow("Результат", cv2.filter2D(image.copy(), -1, kernel))
+
+
+def filter_blur(image):
+    cv2.imshow("Исходное изображение", image)
+    cv2.imshow("Результат", cv2.blur(image.copy(), (5, 5)))
+
+
+def filter_median_blur(image):
+    cv2.imshow("Исходное изображение", image)
+    cv2.imshow("Результат", cv2.medianBlur(image.copy(), 5))
+
+
+def filter_gauss_blur(image):
+    cv2.imshow("Исходное изображение", image)
+    cv2.imshow("Результат", cv2.GaussianBlur(image.copy(), (9, 9), cv2.BORDER_DEFAULT))
+
+
+def filter_erode(image):
+    cv2.imshow("Исходное изображение", image)
+    cv2.imshow("Результат", cv2.erode(image.copy(), np.ones((11, 11))))
+
+
+def filter_dilate(image):
+    cv2.imshow("Исходное изображение", image)
+    cv2.imshow("Результат", cv2.dilate(image.copy(), np.ones((11, 11))))
+
+
+def filter_morph(image):
+    # TODO переписать
+    image2 = image.copy()
+    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
+    cv2.imshow("Image_load", image)
+    kernel = np.ones((6, 6), np.uint8)
+    image3 = cv2.morphologyEx(image3, cv2.MORPH_OPEN, kernel, iterations=1)
+    cv2.imshow("MORPH_OPEN", image3)
+    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
+    image3 = cv2.morphologyEx(image3, cv2.MORPH_CLOSE, kernel, iterations=1)
+    cv2.imshow("MORPH_CLOSE", image3)
+    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
+    image3 = cv2.morphologyEx(image3, cv2.MORPH_GRADIENT, kernel, iterations=1)
+    cv2.imshow("MORPH_GRADIENT", image3)
+    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
+    image3 = cv2.morphologyEx(image3, cv2.MORPH_TOPHAT, kernel, iterations=1)
+    cv2.imshow("MORPH_TOPHAT", image3)
+    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
+    image3 = cv2.morphologyEx(image3, cv2.MORPH_BLACKHAT, kernel, iterations=1)
+    cv2.imshow("MORPH_BLACKHAT", image3)
+    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
+
+
+def filter_sobel(image):
+    image2 = image.copy()
+    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
+    image3 = cv2.GaussianBlur(image3, (3, 3), 0)
+    im3 = cv2.Sobel(image3, cv2.CV_64F, 1, 0, ksize=5)
+    im4 = cv2.Sobel(image3, cv2.CV_64F, 0, 1, ksize=5)
+    im5 = cv2.Sobel(image3, cv2.CV_8UC1, 0, 1, ksize=5)
+    cv2.imshow("Image_load", image)
+    cv2.imshow("Result_Image_X", im3)  # x
+    cv2.imshow("Result_Image_Y", im4)  # y
+    cv2.imshow("Result_Gradient", im5)  # gradient
+
+
+def filter_laplacian(image):
+    image2 = image.copy()
+    image3 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY)
+    image3 = cv2.GaussianBlur(image3, (3, 3), 0)
+    im3 = cv2.Laplacian(image3, cv2.CV_64F)
+    cv2.imshow("Image_load", image)
+    cv2.imshow("Result_Image", im3)
+
+
+def filter_canny(image):
+    image2 = image.copy()
+    image3 = cv2.blur(image2, (5, 5))
+    image3 = cv2.Canny(image3, 100, 100)
+    cv2.imshow("Image_load", image)
+    cv2.imshow("Result_Image", image3)
+
+
+def filter_calc_hist(image):
+    image2 = image.copy()
+    color = ('b', 'g', 'r')
+    cv2.imshow("Image_load", image)
+    for i, col in enumerate(color):
+        histr = cv2.calcHist([image2], [i], None, [256], [0, 256])
+        plt.plot(histr, color=col)
+        plt.xlim([0, 256])
+    plt.show()
+    plt.close()
+
+
+def filter_equalize_hist(image):
+    cv2.imshow("Исходное изображение", image)
+    cv2.imshow("Результат", cv2.equalizeHist(cv2.cvtColor(image.copy(), cv2.COLOR_RGB2GRAY)))
+
+
+functions_list = {
+    "lf": {"func": filter_linear, "help": "Linear filter"},
+    "bl": {"func": filter_blur, "help": "Blur"},
+    "mb": {"func": filter_median_blur, "help": "Median blur"},
+    "gb": {"func": filter_gauss_blur, "help": "GaussianBlur"},
+    "er": {"func": filter_erode, "help": "Erode"},
+    "di": {"func": filter_dilate, "help": "Dilate"},
+    "mo": {"func": filter_morph, "help": "MORPH_OPERATIONS"},
+    "so": {"func": filter_sobel, "help": "Sobel"},
+    "la": {"func": filter_laplacian, "help": "Laplacian"},
+    "ca": {"func": filter_canny, "help": "Canny"},
+    "cl": {"func": filter_calc_hist, "help": "CalcHist"},
+    "eh": {"func": filter_equalize_hist, "help": "EqualizeHist"},
+    "q": {"func": lambda image: sys.exit(0), "help": "Exit"},
+}
+
+
+def print_help():
+    print("Доступные функции:")
+    for k in functions_list:
+        print(f"    {k} - {functions_list[k]['help']}")
 
 
 def main():
-    image1 = cv2.imread("D:\\MACH\\LAB_2\\image0.jpg")
-    image2 = image1.copy()
-    im = image1.copy()
+    image1 = cv2.imread(SOURCE_IMAGE)
     while True:
-        menu()
         try:
-            s = int(input())
-            if s == 1:
-                image2 = image1.copy()
-                kernel = np.array([[-0.1, 0.2, -0.1], [0.2, 3.0, 0.2], [-0.1, 0.2, -0.1]])
-                image2 = cv2.filter2D(image2, -1, kernel)
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", image2)
-                cv2.waitKey(0)
-            if s == 2:
-                image2 = image1.copy()
-                image2 = cv2.blur(image2, (5, 5))
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", image2)
-                cv2.waitKey(0)
-            if s == 3:
-                image2 = image1.copy()
-                image2 = cv2.medianBlur(image2, 5)
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", image2)
-                cv2.waitKey(0)
-            if s == 4:
-                image2 = image1.copy()
-                image2 = cv2.GaussianBlur(image2, (9, 9), cv2.BORDER_DEFAULT)
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", image2)
-                cv2.waitKey(0)
-            if s == 5:
-                image2 = image1.copy()
-                image2 = cv2.erode(image2, np.ones((11, 11)))
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", image2)
-                cv2.waitKey(0)
-            if s == 6:
-                image2 = image1.copy()
-                image2 = cv2.dilate(image2, np.ones((11, 11)))
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", image2)
-                cv2.waitKey(0)
-            if s == 7:
-                image2 = image1.copy()
-                image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
-                cv2.imshow("Image_load", image1)
-                kernel = np.ones((6, 6), np.uint8)
-                image3 = cv2.morphologyEx(image3, cv2.MORPH_OPEN, kernel, iterations=1)
-                cv2.imshow("MORPH_OPEN", image3)
-                image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
-                image3 = cv2.morphologyEx(image3, cv2.MORPH_CLOSE, kernel, iterations=1)
-                cv2.imshow("MORPH_CLOSE", image3)
-                image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
-                image3 = cv2.morphologyEx(image3, cv2.MORPH_GRADIENT, kernel, iterations=1)
-                cv2.imshow("MORPH_GRADIENT", image3)
-                image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
-                image3 = cv2.morphologyEx(image3, cv2.MORPH_TOPHAT, kernel, iterations=1)
-                cv2.imshow("MORPH_TOPHAT", image3)
-                image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
-                image3 = cv2.morphologyEx(image3, cv2.MORPH_BLACKHAT, kernel, iterations=1)
-                cv2.imshow("MORPH_BLACKHAT", image3)
-                image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
-                cv2.waitKey(0)
-            if s == 8:
-                image2 = image1.copy()
-                image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
-                image3 = cv2.GaussianBlur(image3, (3, 3), 0)
-                im3 = cv2.Sobel(image3, cv2.CV_64F, 1, 0, ksize=5)
-                im4 = cv2.Sobel(image3, cv2.CV_64F, 0, 1, ksize=5)
-                im5 = cv2.Sobel(image3, cv2.CV_8UC1, 0, 1, ksize=5)
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image_X", im3)  # x
-                cv2.imshow("Result_Image_Y", im4)  # y
-                cv2.imshow("Result_Gradient", im5)  # gradient
-                cv2.waitKey(0)
-            if s == 9:
-                image2 = image1.copy()
-                image3 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY)
-                image3 = cv2.GaussianBlur(image3, (3, 3), 0)
-                im3 = cv2.Laplacian(image3, cv2.CV_64F)
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", im3)  # x
-                cv2.waitKey(0)
-            if s == 10:
-                image2 = image1.copy()
-                image3 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY)
-                image3 = cv2.blur(image2, (5, 5))
-                image3 = cv2.Canny(image3, 100, 100)
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", image3)  # x
-                cv2.waitKey(0)
-            if s == 11:
-                image2 = image1.copy()
-                color = ('b', 'g', 'r')
-                cv2.imshow("Image_load", image1)
-                for i, col in enumerate(color):
-                    histr = cv2.calcHist([image2], [i], None, [256], [0, 256])
-                    plt.plot(histr, color=col)
-                    plt.xlim([0, 256])
-                plt.show()
-                plt.close()
-                cv2.waitKey(0)
-            if s == 12:
-                image2 = image1.copy()
-                image2 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY)
-                image2 = cv2.equalizeHist(image2)
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", image2)
-                cv2.waitKey(0)
-            if s == 13:
-                cv2.imshow("Image_load", image1)
-                cv2.imshow("Result_Image", image2)
-                cv2.imwrite("D:\\im1.jpg", image1)
-                cv2.imwrite("D:\\im2.jpg", image2)
-                cv2.waitKey(0)
-            if s == 14:
-                image2 = image1.copy()
-                cv2.imshow("Image_load", image1)
-                image2 = cv2.line(image2, (15, 15), (270, 270), (76, 187, 23), 10)  # линия
-                image2 = cv2.rectangle(image2, (20, 20), (100, 100), (0, 0, 255), 10)  # нарисуем четырехугольник
-                image2 = cv2.circle(image2, (100, 100), 40, (0, 255, 0), 10)  # нарисуем круг
-                cv2.imshow("Result_Image", image2)
-                image1 = image2
-                cv2.waitKey(0)
-            if s == 15:
-                cv2.imwrite("1.jpg", im)
-                cv2.imwrite("2.jpg", image1)
-                img = cv2.imread("1.jpg")
-                mask = cv2.imread("2.jpg", 0)
-                res = cv2.bitwise_and(img, img, mask=mask)
-                image2 = res.copy()
-                image1 = res.copy()
-                cv2.imshow("Result_Image", res)
+            cmd = input(">> ")
+            if cmd == "":
+                continue
+            elif cmd == "?":
+                print_help()
+            else:
+                if cmd in functions_list:
+                    try:
+                        functions_list[cmd]["func"](image=image1.copy())
+                        cv2.waitKey(0)
+                        cv2.destroyAllWindows()
+                    except Exception:
+                        print("Произошла ошибка при выполнении функции")
+                        traceback.print_exc()
+                else:
+                    print("Функция не найдена! введите '?' для получения справки")
 
-                cv2.waitKey(0)
-                os.remove("1.jpg")
-                os.remove("2.jpg")
-            if s == 16:
-                cv2.destroyAllWindows()
-                sys.exit(0)
-        except ValueError:
-            print("Неверный пункт меню!!!! Выберите другое!")
+        except Exception:
+            traceback.print_exc()
 
 
 if __name__ == "__main__":