merged (#110)

* Upload ./stage2_diffusion folder

* 1. Update color: border
2. The white edges of septas have been optimized
3. Fixed keep_memory bug
4. Simple README is provided(not completed)

* format fix

* fixed a float bug

* Merged. Implement sharpening & post-processing

* fixed a type bug

* Correctly Merged

* Fixed type

* 这会总该是对的了吧

* Figure fixed

* 这个真是对的了

---------

Co-authored-by: Drosrin <[email protected]>

Author: Drosrin

stage2_diffusion/figure.py

@@ -251,22 +251,19 @@ def __handle_curve(
         return curve_points[:, 0], curve_points[:, 1], center
 
     def __get_essential_info(self, shape, width, color, trans):
-        if width is None:
-            width = self.__get_width(shape)
-        if color is None:
-            color = self.__get_color(shape)
-        color = np.clip(color, 0, 1)
-        if trans is None:
-            trans = self.__get_transparency(shape)
-        return width, color, trans
+        w = self.__get_width(shape, width)
+        c = self.__get_color(shape, color)
+        c = np.clip(c, 0, 1)
+        t = self.__get_transparency(shape, trans)
+        return w, c, t
 
-    def __get_width(self, shape):
+    def __get_width(self, shape, width: float):
         try:
             return shape["width"]
         except Exception:
-            return 1.8
+            return 1.8 if width is None else width
 
-    def __get_color(self, shape):
+    def __get_color(self, shape, color):
         try:
             return shape["color"]
         except Exception:
@@ -276,19 +273,22 @@ def __get_color(self, shape):
                 else:
                     return (1, 1, 1, 1)
             except Exception:
-                return (random.random(), random.random(), random.random())
+                return (random.random(), random.random(), random.random()) if color is None else color
 
-    def __get_transparency(self, shape):
+    def __get_transparency(self, shape, trans):
         try:
             if shape["fill_mode"] == "no":
-                trans = (0, 0, 0, 0)
+                return (0, 0, 0, 0)
             elif shape["fill_mode"] == "white":
-                trans = (1, 1, 1, 1)
+                return (1, 1, 1, 1)
             elif shape["fill_mode"] == "black":
-                trans = (0, 0, 0, 1)
+                return (0, 0, 0, 1)
+            elif shape["fill_mode"] == "border":
+                return (0, 0, 0, 0)
+            else:
+                raise Exception("Invalid fill_mode arg")
         except Exception:
-            trans = (0, 0, 0, 0)  # no
-        return trans
+            return trans if trans is not None else (0, 0, 0, 0)
 
     def transfer_to_cv2_wrapper(self):
         buf2 = BytesIO()

stage2_diffusion/process.py

@@ -208,6 +208,14 @@ def fill_septas(xs, ys, centers, cv2_fig: np.ndarray, alpha_channel: np.ndarray)
     return cv2_fig, alpha_channel
 
 
+def post_processing(img: np.ndarray):
+    img = cv2.convertScaleAbs(img, alpha=1.5, beta=0)
+    laplacian = cv2.Laplacian(img, cv2.CV_64F)
+    laplacian = np.uint8(np.absolute(laplacian))
+    sharpened = cv2.addWeighted(img, 1, laplacian, 1, 0)
+    return sharpened
+
+
 def generate_one_img(
     idx,
     sample,
@@ -259,10 +267,11 @@ def generate_one_img(
     blended_img = np.where(alpha_mask == 255, septa_overlayer, diffused_img)
 
     blended_img = cv2.cvtColor(blended_img, cv2.COLOR_BGRA2GRAY)
+    final_img = post_processing(blended_img)
     img_path = f"{keyword}/{img_path}"
     if not os.path.exists(f"{keyword}"):
         os.mkdir(f"{keyword}")
-    cv2.imwrite(img_path, blended_img)
+    cv2.imwrite(img_path, final_img)
 
 
 def main():