fixed dilation, custom dilation, other fixes and features

Install_and_Tutorial.txt

@@ -40,8 +40,8 @@ The model packaged here is model 161. The weights are used to run the detections
 
 5. For "Your own input image folder", select the folder with your images in it.
    For "DCP install directory", selec the parent directory of DeepCreamPy (usually called dist 1)
+   New for 1.6.7, dilation amount will expand the mask by any positive number of pixels (I reccommend lower integers).
 
-NOTE Jpg images can now be processed, but they will be soft-converted to .png for compatibility with DCP.
 
 6. Now you can hit the Go button. Loading the nueral network will take some time. 
    The image detections will also take some time, up to a minute per image once it gets rolling depending on your computer.
@@ -50,9 +50,6 @@ NOTE Jpg images can now be processed, but they will be soft-converted to .png fo
    corresponding folders in your DeepCreamPy directory.
 
 8. Now you should run DeepCreamPy, and you can close hentAI. Be sure to select the appropriate censor type in DCP.
-
-9. If you choose the ESRGAN options, detection and decensoring will be done together so DeepCreamPy won't be needed.
-   The output of this will be in the ESR_output folder, but videos will be written to the main directory
 
 --- hentAI video detecting (Experimental, mosaic only)
 1. Place the input .mp4 into its own folder. 
@@ -115,15 +112,18 @@ or if you have a Cuda compatible card:
 	For Screentone Remover, instructions are in its own file. Use this if you are using 
 	non-colored images with the printed effect. 
 	NOTE: If the doujin was tagged as [Digital] and does not have the screentone effect, you
-	do not need to use Screentone Remover.
+	probably do not need to use Screentone Remover.
 
 1. First, have an input folder ready with whatever images you want to decensor.
    Do not have bar censors and mosaic censors in the same folder.
 
-2. To detect, simply run the main script
+2. To detect, simply run the main script and follow tutorial for the exe above.
 
 	python main.py
 
+2a. NOTE: Code and Colab versions will have ESRGAN as well. Detection and decensoring will be done together so DeepCreamPy won't be needed.
+   The output of this will be in the ESR_output folder, but videos will be written to the main directory
+
 3. Training:
 **If you are interested in training, please contact me and I may provide you with
 the current dataset. All I ask is that you also send me your trained model should
@@ -161,10 +161,18 @@ This means your gpu is being detected, and tensorflow wants to use it. This mean
 1. Install CUDA 9.0 here: https://developer.nvidia.com/cuda-90-download-archive?target_os=Windows&target_arch=x86_64
 Get the right one for your computer
 
-2. Make sure to install gpu requiremtnes
+4. You also need cuDNN v7.6.4 runtime for CUDA 9.0. You need an Nvidia account to access the download, but it that is free.
+Get it from here: https://developer.nvidia.com/cudnn
+
+3. Make sure to install gpu requirements, not the cpu.
 
 	pip install -r requirements-gpu.txt
+
+4. NOTE: It should be possible to have CUDA 9.0 and CUDA 10 coexist.
+
+5. NOTE FOR RTX owners: There may exist some issue with RTX cards and the torch model in ESRGAN. I don't think it has been resolved yet, 
+so for now if you encounter errors such as CUDNN_STATUS_SUCCESS or CUDDNN_STATUS_ALLOC_FAILED, then change line 95 in detector.py
+to say 'cpu' instead of 'cuda'. This will force only ESRGAN to use CPU. This will be slow, so at this point i'd reccommend using colab instead.
 
-4. (Training only) For training, you may need cudnn 7.0 runtime. You need an Nvidia account to access the download, but it that is free.
-I dont have the exact link, but you can Google search or look for it here: https://developer.nvidia.com/ 
+
 

detector.py

@@ -26,7 +26,7 @@
 from mrcnn import model as modellib, utils
 # sys.path.insert(1, 'samples/hentai/')
 # from hentai import HentaiConfig
-from cv2 import imshow, waitKey, multiply, add, dilate, VideoCapture, Canny, cvtColor,COLOR_GRAY2RGB, imdecode, CAP_PROP_FRAME_HEIGHT, CAP_PROP_FRAME_WIDTH, CAP_PROP_FPS, VideoWriter, VideoWriter_fourcc, resize, INTER_LANCZOS4, INTER_AREA, GaussianBlur, filter2D, bilateralFilter, blur
+from cv2 import imshow, waitKey, multiply, add, erode, VideoCapture, Canny, cvtColor,COLOR_GRAY2RGB, imdecode, CAP_PROP_FRAME_HEIGHT, CAP_PROP_FRAME_WIDTH, CAP_PROP_FPS, VideoWriter, VideoWriter_fourcc, resize, INTER_LANCZOS4, INTER_AREA, GaussianBlur, filter2D, bilateralFilter, blur
 import ColabESRGAN.test
 from green_mask_project_mosaic_resolution import get_mosaic_res
 
@@ -93,11 +93,12 @@ class InferenceConfig(HentaiConfig):
         if self.model.check_cuda_gpu()==True:
             print("CUDA-compatible GPU located!")
             self.hardware = 'cuda'
-        # destroy model. Will re load it during weight load.
+        # destroy model. Will re init during weight load.
         self.model = []
 
     # Clean out temp working images from all directories in ESR_temp. Code from https://stackoverflow.com/questions/185936/how-to-delete-the-contents-of-a-folder
     def clean_work_dirs(self):
+        print("Cleaning work dirs...")
         folders = [self.out_path, self.out_path2, self.temp_path, self.temp_path2]
         for folder in folders:
             for filename in os.listdir(folder):
@@ -127,19 +128,19 @@ def load_weights(self):
     mask: instance segmentation mask [height, width, instance count]
     Returns result covered image.
     """
-    def apply_cover(self, image, mask):
+    def apply_cover(self, image, mask, dilation):
         # Copy color pixels from the original color image where mask is set
         green = np.zeros([image.shape[0], image.shape[1], image.shape[2]], dtype=np.uint8)
         green[:,:] = [0, 255, 0]
-        dilation = 11 #NOTE: Change this to modify dilation amount. Can also change iterations below
+
         if mask.shape[-1] > 0:
             # We're treating all instances as one, so collapse the mask into one layer
             mask = (np.sum(mask, -1, keepdims=True) < 1)
             # dilate mask to ensure proper coverage
             mimg = mask.astype('uint8')*255
-            kernel = np.ones((dilation,dilation), np.uint8) # custom sized kernel.
-            mimg = dilate(src=mask.astype('uint8'), kernel=kernel, iterations=1)
-            # dilation returns image with channels stripped (?!?). Reconstruct image channels (probably only need 1 channel)
+            kernel = np.ones((dilation,dilation), np.uint8)
+            mimg = erode(src=mask.astype('uint8'), kernel=kernel, iterations=1) #
+            # dilation returns image with channels stripped (?!?). Reconstruct image channels
             mask_img = np.zeros([mask.shape[0], mask.shape[1],3]).astype('bool')
             mask_img[:,:,0] = mimg.astype('bool')
             mask_img[:,:,1] = mimg.astype('bool')
@@ -350,11 +351,10 @@ def ESRGAN(self, img_path, img_name, is_video=False):
 
             vwriter.release()
             print('Video: Phase 2 complete!')
-        print("Process complete. Cleaning work directories...")
-        self.clean_work_dirs() #NOTE: DISABLE ME if you want to keep the images in the working dirs
+
 
     # ESRGAN folder running function
-    def run_ESRGAN(self, in_path = None, is_video = False, force_jpg = False):
+    def run_ESRGAN(self, in_path = None, is_video = False, force_jpg = True):
         assert in_path
 
         # Parse directory for files.
@@ -378,7 +378,8 @@ def run_ESRGAN(self, in_path = None, is_video = False, force_jpg = False):
             self.ESRGAN(img_path=img_path, img_name=img_name, is_video=is_video)
         fin = time.perf_counter()
         total_time = fin-star
-        print("Completed ESRGAN detection and decensor in {:.4f} seconds".format(fin, star))
+        print("Completed ESRGAN detection and decensor in {:.4f} seconds".format(total_time))
+        self.clean_work_dirs() #NOTE: DISABLE ME if you want to keep the images in the working dirs
         #TODO: maybe unload hent-AI tf model here
 
     def video_create(self, image_path=None, dcp_path=''):
@@ -430,7 +431,7 @@ def video_create(self, image_path=None, dcp_path=''):
     # save path and orig video folder are both paths, but orig video folder is for original mosaics to be saved.
     # fname = filename.
     # image_path = path of input file, image or video
-    def detect_and_cover(self, image_path=None, fname=None, save_path='', is_video=False, orig_video_folder=None, force_jpg=False, is_mosaic=False):
+    def detect_and_cover(self, image_path=None, fname=None, save_path='', is_video=False, orig_video_folder=None, force_jpg=False, is_mosaic=False, dilation=0):
         assert image_path
         assert fname # replace these with something better?
 
@@ -471,7 +472,7 @@ def detect_and_cover(self, image_path=None, fname=None, save_path='', is_video=F
                     new_masks = np.delete(r['masks'], remove_indices, axis=2)
 
                     # Apply cover
-                    cov, mask = self.apply_cover(image, new_masks)
+                    cov, mask = self.apply_cover(image, new_masks, dilation)
 
                     # save covered frame into input for decensoring path
                     file_name = save_path + im_name + str(count).zfill(6) + '.png'
@@ -517,7 +518,7 @@ def detect_and_cover(self, image_path=None, fname=None, save_path='', is_video=F
             # except:
             #     print("ERROR in detect_and_cover: Model detect")
 
-            cov, mask = self.apply_cover(image, new_masks)
+            cov, mask = self.apply_cover(image, new_masks, dilation)
             try:
                 # Save output, now force save as png
                 file_name = save_path + fname[:-4] + '.png'
@@ -527,14 +528,16 @@ def detect_and_cover(self, image_path=None, fname=None, save_path='', is_video=F
             # print("Saved to ", file_name)
 
     # Function for file parsing, calls the aboven detect_and_cover
-    def run_on_folder(self, input_folder, output_folder, is_video=False, orig_video_folder=None, force_jpg=False, is_mosaic=False):
+    def run_on_folder(self, input_folder, output_folder, is_video=False, orig_video_folder=None, is_mosaic=False, dilation=0):
         assert input_folder
         assert output_folder # replace with catches and popups
 
-        # if force_jpg==True:
-        #     print("WARNING: force_jpg=True. jpg support is not guaranteed, beware.")
         self.esrgan_instance = [] # rare case where esrgan instance not destroyed but new action started, catch it here
         self.load_weights()
+        if dilation < 0:
+            print("ERROR: dilation value < 0")
+            return
+        print("Will expand each mask by {} pixels".format(dilation/2))
 
         file_counter = 0
         if(is_video == True):
@@ -547,32 +550,31 @@ def run_on_folder(self, input_folder, output_folder, is_video=False, orig_video_
 
             for vid_path, vid_name in vid_list:
                 # video will not support separate mask saves
-                self.detect_and_cover(vid_path, vid_name, output_folder, is_video=True, orig_video_folder=orig_video_folder)
-                print('detection on video', file_counter, 'is complete')
+                star = time.perf_counter()
+                self.detect_and_cover(vid_path, vid_name, output_folder, is_video=True, orig_video_folder=orig_video_folder, dilation=dilation)
+                fin = time.perf_counter()
+                total_time = fin-star
+                print('Detection on video', file_counter, 'finished in {:.4f} seconds'.format(total_time))
                 file_counter += 1
         else:
             # obtain inputs from the input folder
             img_list = []
             for file in os.listdir(str(input_folder)):
                 file_s = str(file)
                 try:
-                    if force_jpg == False:
-                        if file_s.endswith('.png') or file_s.endswith('.PNG'):
-                            img_list.append((input_folder + '/' + file_s, file_s))
-                        elif file_s.endswith(".jpg") or file_s.endswith(".JPG"):
-                            # img_list.append((input_folder + '/' + file_s, file_s)) # Do not add jpgs. Conversion to png must happen first
-                            self.dcp_compat += 1
-                    else:
-                        if file_s.endswith('.png') or file_s.endswith('.PNG') or file_s.endswith(".jpg") or file_s.endswith(".JPG"):
-                            img_list.append((input_folder + '/' + file_s, file_s))
+                    if file_s.endswith('.png') or file_s.endswith('.PNG') or file_s.endswith(".jpg") or file_s.endswith(".JPG"):
+                        img_list.append((input_folder + '/' + file_s, file_s))
                 except:
                     print("ERROR in run_on_folder: File parsing. file=", file_s)
 
 
             # save run detection with outputs to output folder
             for img_path, img_name in img_list:
-                self.detect_and_cover(img_path, img_name, output_folder, force_jpg=force_jpg, is_mosaic=is_mosaic)  #sending force_jpg for debugging
-                print('Detection on image', file_counter, 'is complete')
+                star = time.perf_counter()
+                self.detect_and_cover(img_path, img_name, output_folder, is_mosaic=is_mosaic, dilation=dilation)  #sending force_jpg for debugging
+                fin = time.perf_counter()
+                total_time = fin-star
+                print('Detection on image', file_counter, 'finished in {:.4f} seconds'.format(total_time))
                 file_counter += 1