From cf8c3c82633346a511d4080fc854a67a7077ee0c Mon Sep 17 00:00:00 2001
From: localhorst <localhorst@mosad.xyz>
Date: Tue, 25 Jul 2023 19:49:32 +0200
Subject: [PATCH] error handling and converting extracted photos to avif

---
 extract_photos.py | 97 +++++++++++++++++++++++++++++------------------
 1 file changed, 60 insertions(+), 37 deletions(-)

diff --git a/extract_photos.py b/extract_photos.py
index 9be5c04..8062932 100644
--- a/extract_photos.py
+++ b/extract_photos.py
@@ -6,36 +6,50 @@
 import cv2
 import sys
 import os
+import subprocess
 
-def convert_to_avif(input, output):
+number_of_photos_in_scan = 3
+
+def convert_to_avif(input_dir, output_dir):
     # print("input: " + str(input))
     # print("output: " + str(output))
-    temp = subprocess.Popen(
-        ["avifenc", "--jobs", "all", input, output], stdout=subprocess.PIPE
-    )
-    # get the output as a string
-    # output = str(temp.communicate())
-    # print(output)
-    temp = subprocess.Popen(
-        [
-            "sync",
-        ],
-        stdout=subprocess.PIPE,
-    )
 
+    for path in os.listdir(input_dir):
+        if os.path.isfile(os.path.join(input_dir, path)):
+            input = os.path.join(input_dir, path)
+            output = os.path.join(output_dir,(os.path.splitext(os.path.basename(os.path.normpath(input)))[0] + ".avif"))
+            #print("Found extracted photo: " + str(input))
+            #print("Convert to: " + str(output))
+
+            cmd_handle = subprocess.Popen(["avifenc", "--jobs", "all", input, output], stdout=subprocess.PIPE)
+            cmd_handle.communicate()[0]
+
+            if cmd_handle.returncode != 0:
+                raise Exception("Unable to convert via avifenc! Return code: " + str(cmd_handle.returncode))
+                
+            if os.path.exists(input):
+                os.remove(input)
+
+            cmd_handle = subprocess.Popen(["sync",],stdout=subprocess.PIPE,)
+            cmd_handle.communicate()[0]
+
+            if cmd_handle.returncode != 0:
+                raise Exception("Unable to sync! Return code: " + str(cmd_handle.returncode))
+    
 def auto_crop_scan(src_path, output_dir):
+    areas = list()
+    areas.clear()
+
     # read the input image
     img_src = cv2.imread(src_path)
 
-    if not img_src:
-        print("unable to read image")
-        sys.exit(-1)
-
-    print("read done")
+    if img_src is None:
+        raise Exception("unable to read image " + str(input))
 
     # convert the image to grayscale
     img_gray = cv2.cvtColor(img_src, cv2.COLOR_BGR2GRAY)
-    print("convert to gray done")
+    if img_gray is None:
+        raise Exception("unable to convert image to grayscale")
 
     # apply thresholding on the gray image to create a binary image
     ret, thresh = cv2.threshold(img_gray, 127, 255, 0)
@@ -45,31 +59,24 @@ def auto_crop_scan(src_path, output_dir):
         thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
     )
 
-    areas = list()
-
     for cnt in contours:
-        # print(cv2.contourArea(cnt))
         areas.append(cv2.contourArea(cnt))
 
     areas.sort(reverse=True)
 
-    # print(areas)
-
-    if len(areas) >= 5:
+    if len(areas) >= (number_of_photos_in_scan+1):
         outer = areas[0]
-        inter_min = areas[4]
-        print("Outer area: " + str(outer))
-        print("Inner area: " + str(inter_min))
+        inter_min = areas[number_of_photos_in_scan]
+        #print("Outer area: " + str(outer))
+        #print("Inner area: " + str(inter_min))
 
-        index = 0
+        index = 0 # used for exported file name
 
         for cnt in contours:
             area = cv2.contourArea(cnt)
             if (area < outer) and (area >= inter_min):
                 # compute the bounding rectangle of the contour
                 x, y, w, h = cv2.boundingRect(cnt)
-                # draw the bounding rectangle
-                # imgView = cv2.rectangle(img,(x,y),(x+w,y+h),(0,255,0),2)
 
                 x = x + 5  # TODO
                 y = y + 5  # TODO
@@ -77,7 +84,9 @@ def auto_crop_scan(src_path, output_dir):
                 h = h - 10  # TODO
 
                 img_crop = img_src[y : y + h, x : x + w]
-                # cv2.imshow("cropped", crop_img)
+
+                if img_crop is None:
+                    raise Exception("unable to crop image")
 
                 export_file_name = (
                     os.path.splitext(os.path.basename(os.path.normpath(src_path)))[0]
@@ -85,15 +94,13 @@ def auto_crop_scan(src_path, output_dir):
                     + str(index)
                     + ".png"
                 )
-                print(export_file_name)
 
                 export_file_path = os.path.join(output_dir, export_file_name)
-                print(export_file_path)
 
                 cv2.imwrite(export_file_path, img_crop)
                 index = index + 1
-                # cv2.waitKey(0)
-
+    else:
+        raise Exception("unable to find all photos in scan")
 
 def main():
     print("starting ...")
@@ -118,10 +125,26 @@ def main():
         )
         sys.exit(-1)
 
+    failed_extractions = 0
+
     for path in os.listdir(source_dir_path):
         if os.path.isfile(os.path.join(source_dir_path, path)):
-            print(path)
+            print("Found scan: " + str(os.path.join(source_dir_path, path)), end="")
+            try:
+                auto_crop_scan(os.path.join(source_dir_path, path), tmp_dir_path)
+                convert_to_avif(tmp_dir_path, destination_dir_path)
+                print("  →  Successfully extracted photos from scan")
+            except Exception as e:
+                print(str(e))
+                failed_extractions = failed_extractions + 1
 
+    if(failed_extractions == 0):
+        print("Exiting without error")
+        sys.exit(0)
+    else:
+        print("Unable to extract " + str(failed_extractions) + " scans.")
+        print("Exiting with error")
+        sys.exit(-1)
 
 if __name__ == "__main__":
     main()