diff --git a/Install_and_Tutorial.txt b/Install_and_Tutorial.txt index 13e2750..7d15ba0 100644 --- a/Install_and_Tutorial.txt +++ b/Install_and_Tutorial.txt @@ -28,8 +28,6 @@ The model packaged here is model 161. The weights are used to run the detections --- Tutorial for hentAI executable -** If you have an Nvidia gpu and plan on doing video uncensoring, I reccomend getting the code to work instead** - 1. You should use the provided input_images folder to place the images you want to decensor. Remember that you must decensor content with bar censors and mosaic censors separately. Images must be in .png format. Also, make sure to clear out the input folder between runs so you dont decensor the same things. @@ -42,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. @@ -52,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. @@ -95,7 +90,9 @@ NOTE Jpg images can now be processed, but they will be soft-converted to .png fo 5. Install requirements - pip install -r requirements.txt + pip install -r requirements-cpu.txt + OR if you have a CUDA compatible gpu with CUDA 9.0: + pip install -r requirements-gpu.txt 5.5. NOTE: You might not be able to install torch 0.4.1 from pip, like on Windows. In that case, run the following for CPU @@ -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 @@ -154,24 +154,25 @@ you improve on the latest detections. An NVIDIA CUDA compatible card is required send it to me! --------NVIDIA users----------- + You may encounter an error like this: ImportError: libcublas.so.9.0: cannot open shared object file: No such file or directory This means your gpu is being detected, and tensorflow wants to use it. This means you could get 6-30x better performance. 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. Uninstall normal tensorflow, then uninstall protobuf - - pip uninstall tensorflow - - pip uninstall protobuf +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. Reinstall tensorflow-gpu +3. Make sure to install gpu requirements, not the cpu. + + pip install -r requirements-gpu.txt - pip uninstall tensorflow-gpu +4. NOTE: It should be possible to have CUDA 9.0 and CUDA 10 coexist. - pip install tensorflow-gpu==1.9.0 +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/ + \ No newline at end of file diff --git a/README.md b/README.md index 7a10ec6..c873940 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,8 @@ Here is an example of a screentoned image, and what it looks like when removed b * For full video decensoring via ESRGAN, you will need to download Twittman's model [here](https://de-next.owncube.com/index.php/s/mDGmi7NgdyyQRXL) and place it inside the ColabESRGAN/models folder. +* Nvidia GPU owners should install CUDA 9.0, and cuDNN 7.6.4. Note that there are issues with RTX cards and ESRGAN, so if you want to use that I again reccomend the colab notebook instead. + ## Important Notes (READ BEFORE USING) * I highly reccommend running hent-AI on batches, for example one doujin or a few doujins at once. The slowest part of hent-AI is the initialization, so the first inference takes time but the rest will be quicker. @@ -111,6 +113,7 @@ Here is an example of a screentoned image, and what it looks like when removed b * Do not put entire clips through the video detection, it is a very slow task. If you can, edit in only the short clips with visible mosaics, get the decensored output, then edit them in the rest of the video. +* The compiled exe release does not support ESRGAN. If you want to use this, refer to the colab notebook. ## Versions and Downloads @@ -132,6 +135,8 @@ Here is an example of a screentoned image, and what it looks like when removed b * [1.6.5](): Added adaptive mosaic granularity checking via GMP by rekaXua. Added colab file for free cloud-based ESRGAN video decensoring. +* [1.6.7](): Changed ESRGAN processs to run in 2 phases: ESRGAN resize, then mask detection. Slower but more memory forgiving. Added mask blurring for less seams on ESRGAN. For non-ESRGAN, added custom dilation to expand masks. Removed option for jpg, it will be used automatically. Improved file cleaning. + ## Installation directions diff --git a/detector.py b/detector.py index 935a0dd..e048d36 100644 --- a/detector.py +++ b/detector.py @@ -13,6 +13,7 @@ import skimage.draw from skimage.filters import unsharp_mask import imgaug # should augment this improt as well haha +import time # from PIL import Image # Root directory of project @@ -25,7 +26,7 @@ from mrcnn import model as modellib, utils # sys.path.insert(1, 'samples/hentai/') # from hentai import HentaiConfig -from cv2 import VideoCapture, 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 @@ -48,7 +49,6 @@ class HentaiConfig(Config): # Number of classes (including background) NUM_CLASSES = 1 + 1 + 1 - # NOTE: Enable the following and disable above if on Canny edge detector model # Number of training steps per epoch, equal to dataset train size STEPS_PER_EPOCH = 1490 @@ -84,20 +84,21 @@ class InferenceConfig(HentaiConfig): return # Create esrgan instance for detector instance try: - esr_model_path = os.path.join(os.path.abspath('.'), "4x_FatalPixels_340000_G.pth") + self.esr_model_path = os.path.join(os.path.abspath('.'), "4x_FatalPixels_340000_G.pth") except: print("ERROR in Detector init: ESRGAN model not found, make sure you have 4x_FatalPixels_340000_G.pth in this directory") return # Scan for cuda compatible GPU for ESRGAN. Mask-RCNN *should* automatically use a GPU if available. + self.hardware = 'cpu' if self.model.check_cuda_gpu()==True: print("CUDA-compatible GPU located!") - self.esrgan_instance = ColabESRGAN.test.esrgan(model_path=esr_model_path, hw='cuda') - else: - print("No CUDA-compatible GPU located. Using CPU") - self.esrgan_instance = ColabESRGAN.test.esrgan(model_path=esr_model_path, hw='cpu') + self.hardware = 'cuda' + # 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): @@ -112,7 +113,9 @@ def clean_work_dirs(self): # Make sure this is called before using model weights def load_weights(self): - print('Loading weights...', end=' ') + print('Creating model, Loading weights...', end=' ') + self.model = modellib.MaskRCNN(mode="inference", config=self.config, + model_dir=DEFAULT_LOGS_DIR) try: self.model.load_weights(self.weights_path, by_name=True) print("Weights loaded") @@ -125,24 +128,54 @@ 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] + 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) - cover = np.where(mask, image, green).astype(np.uint8) + # dilate mask to ensure proper coverage + mimg = mask.astype('uint8')*255 + 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') + mask_img[:,:,2] = mimg.astype('bool') + + cover = np.where(mask_img.astype('bool'), image, green).astype(np.uint8) else: # error case, return image cover = image - return cover, mask + return cover, mask # Similar to above function, except it places the decensored image over the original image. def splice(self, image, mask, gan_out): if mask.shape[-1] > 0: mask = (np.sum(mask, -1, keepdims=True) < 1) - cover = np.where(mask, image, gan_out).astype(np.uint8) + mask = 1 - mask # invert mask for blending + mask = mask.astype('uint8')*255 + mask = GaussianBlur(mask, (29,29), 0) + # mask_img = np.zeros([mask.shape[0], mask.shape[1],3]).astype('uint8') + # for i in range(3): + # mask_img[:,:,i] = mask + mask_img = mask.astype(float) / 255 + # proper blending courtesy of https://www.learnopencv.com/alpha-blending-using-opencv-cpp-python/ + fg_o = gan_out.astype(float) + bg_o = image.astype(float) + fg = np.zeros([mask.shape[0], mask.shape[1],3]).astype(float) + bg = np.zeros([mask.shape[0], mask.shape[1],3]).astype(float) # create foreground and background images with proper rgb channels + cover = image + for i in range(3): + # Multiply the fg with the mask matte + fg[:,:,i] = multiply(mask_img, fg_o[:,:,i]) + # Multiply the bg with ( 1 - mask_img ) + bg[:,:,i] = multiply(1.0 - mask_img, bg_o[:,:,i]) + # Add the masked fg and bg. + cover[:,:,i] = add(fg[:,:,i], bg[:,:,i]) else: #error case, return image cover=image @@ -152,7 +185,80 @@ def splice(self, image, mask, gan_out): def get_non_png(self): return self.dcp_compat - # Runs hent-AI detection, and ESRGAN on image. Mosaic only. + # function to handle all of the esrgan stuff + def resize_GAN(self, img_path, img_name, is_video=False): + # non-video, standard image + if is_video is False: + # Attempt to obtain image + try: + image = skimage.io.imread(img_path) # problems with strange shapes + if image.ndim != 3: + image = skimage.color.gray2rgb(image) # convert to rgb if greyscale + if image.shape[-1] == 4: + image = image[..., :3] # strip alpha channel + except Exception as e: + print("ERROR in resize_GAN: Image read. Skipping. image_path=", img_path) + print(e) + return + # Calculate mosaic granularity. + granularity = get_mosaic_res(np.array(image)) + if granularity < 10: #TODO: implement adaptive granularity by weighted changes + print("Granularity of image was less than threshold at ", granularity) + granularity = 10 + # Resize image down + try: + mini_img = resize(image, (int(image.shape[1]/granularity), int(image.shape[0]/granularity)), interpolation=INTER_AREA) # TODO: experiment with interpolations + # After resize, run bilateral filter to keep colors coherent + file_name = self.temp_path + img_name[:-4] + '.png' + skimage.io.imsave(file_name, mini_img) + except Exception as e: + print("ERROR in resize_GAN: resize. Skipping. image_path=",img_path, e) + return + # Now run ESRGAN inference + gan_img_path = self.out_path + img_name[:-4] + '.png' + self.esrgan_instance.run_esrgan(test_img_folder=file_name, out_filename=gan_img_path, mosaic_res=granularity) + else: + try: + video_path = img_path + vcapture = VideoCapture(video_path) + width = int(vcapture.get(CAP_PROP_FRAME_WIDTH)) + height = int(vcapture.get(CAP_PROP_FRAME_HEIGHT)) + fps = vcapture.get(CAP_PROP_FPS) + print("Detected fps:", fps) + except Exception as e: + print("ERROR in resize_GAN: video read and init.", e) + return + count = 0 + success = True + print("Video read complete. Starting video phase 1 : resize + GAN") + while success: + print("frame: ", count) + # Read next image + success, image = vcapture.read() + if success: + # OpenCV returns images as BGR, convert to RGB + image = image[..., ::-1] + + granularity = get_mosaic_res(np.array(image)) # pass np array of image as ref to gmp function + if granularity < 10: #TODO: implement adaptive granularity by weighted changes + print('Granularity was less than threshold at ',granularity) + granularity = 10 + + # initial resize frame + mini_img = resize(image, (int(image.shape[1]/granularity), int(image.shape[0]/granularity)), interpolation=INTER_AREA) # downscale to 1/16 + # bil2 = bilateralFilter(mini_img, 3, 70, 70) + file_name = self.temp_path + img_name[:-4] + '.png' # need to save a sequence of pngs for TGAN operation + skimage.io.imsave(file_name, mini_img) # save resized images to temp path. Not used in main ESRGAN function below. + + # run ESRGAN algorithms + gan_img_path = self.out_path + img_name[:-4] + str(count).zfill(6) + '.png' + self.esrgan_instance.run_esrgan(test_img_folder=file_name, out_filename=gan_img_path, mosaic_res=granularity) + + gan_image = skimage.io.imread(gan_img_path) + gan_image = resize(gan_image, (image.shape[1], image.shape[0])) + count += 1 + print('Video: Phase 1 complete!') + # Runs hent-AI detection and splice. Mosaic only. def ESRGAN(self, img_path, img_name, is_video=False): # Image reads if is_video == False: @@ -171,48 +277,17 @@ def ESRGAN(self, img_path, img_name, is_video=False): # Remove bars from detection; class 1 if len(r["scores"]) == 0: - print("Skipping frame with no detection") + print("Skipping image with no detection") return remove_indices = np.where(r['class_ids'] != 2) new_masks = np.delete(r['masks'], remove_indices, axis=2) - # Calculate mosaic granularity. Then, apply pre-sharpen - granularity = get_mosaic_res(np.array(image)) - if granularity < 12: #TODO: implement adaptive granularity by weighted changes - print("Granularity of image was ", granularity) - granularity = 12 - - - # find way to skip frames with no detection - - # Now we have the mask from detection, begin ESRGAN by first resizing img into temp folder. - try: - mini_img = resize(image, (int(image.shape[1]/granularity), int(image.shape[0]/granularity)), interpolation=INTER_AREA) # downscale to 1/16 - # After resize, run bilateral filter to keep colors coherent - # bil2 = bilateralFilter(mini_img, 3, 60, 60) - file_name = self.temp_path + img_name[:-4] + '.png' - skimage.io.imsave(file_name, mini_img) - except Exception as e: - print("ERROR in detector.ESRGAN: resize. Skipping. image_path=",img_path, e) - return - # Now run ESRGAN inference + # load image from esrgan gan_img_path = self.out_path + img_name[:-4] + '.png' - self.esrgan_instance.run_esrgan(test_img_folder=file_name, out_filename=gan_img_path) - # load output from esrgan, will still be 1/4 size of original image - gan_image = skimage.io.imread(gan_img_path) - # Resize to 1/3. Run ESRGAN again. - gan_image = resize(gan_image, (int(gan_image.shape[1]/2), int(gan_image.shape[0]/2))) - file_name = self.temp_path2 + img_name[:-4] + '.png' - skimage.io.imsave(file_name, gan_image) - gan_img_path = self.out_path2 + img_name[:-4] + '.png' - self.esrgan_instance.run_esrgan(test_img_folder=file_name, out_filename=gan_img_path) - gan_image = skimage.io.imread(gan_img_path) gan_image = resize(gan_image, (image.shape[1], image.shape[0])) # Splice newly enhanced mosaic area over original image fin_img = self.splice(image, new_masks, gan_image) - # bilateral filter to soften output image (NOTE: make this optional?) - fin_img = bilateralFilter(fin_img, 9, 70, 70) try: # Save output, now force save as png file_name = self.fin_path + img_name[:-4] + '.png' @@ -230,16 +305,16 @@ def ESRGAN(self, img_path, img_name, is_video=False): print("Detected fps:", fps) # Define codec and create video writer, video output is purely for debugging and educational purpose. Not used in decensoring. - file_name = img_name[:-4] + "_decensored.avi" + file_name = img_name[:-4] + "_decensored.mp4" vwriter = VideoWriter(file_name, - VideoWriter_fourcc(*'MJPG'), + VideoWriter_fourcc(*'mp4v'), fps, (width, height)) except Exception as e: - print("ERROR in TGAN: video read and init.", e) + print("ERROR in ESRGAN: video read and init.", e) return count = 0 success = True - print("Video read complete, starting video detection. NOTE: frame 0 may take up to 1 minute") + print("Video read complete. Starting video phase 2: detection + splice") while success: print("frame: ", count) # Read next image @@ -257,31 +332,16 @@ def ESRGAN(self, img_path, img_name, is_video=False): vwriter.write(image) count += 1 continue - - granularity = get_mosaic_res(np.array(image)) # pass np array of image as ref to gmp function - if granularity < 10: #TODO: implement adaptive granularity by weighted changes - print('Granularity was',granularity) - granularity = 10 # Remove unwanted class, code from https://github.com/matterport/Mask_RCNN/issues/1666 remove_indices = np.where(r['class_ids'] != 2) # remove bars: class 1 new_masks = np.delete(r['masks'], remove_indices, axis=2) - # initial resize frame - mini_img = resize(image, (int(image.shape[1]/granularity), int(image.shape[0]/granularity)), interpolation=INTER_AREA) # downscale to 1/16 - bil2 = bilateralFilter(mini_img, 3, 70, 70) - file_name = self.temp_path + img_name[:-4] + '.png' # need to save a sequence of pngs for TGAN operation - skimage.io.imsave(file_name, bil2) - - # run ESRGAN algorithms - gan_img_path = self.out_path + img_name[:-4] + '.png' - self.esrgan_instance.run_esrgan(test_img_folder=file_name, out_filename=gan_img_path) - + gan_img_path = self.out_path + img_name[:-4] + str(count).zfill(6) + '.png' gan_image = skimage.io.imread(gan_img_path) gan_image = resize(gan_image, (image.shape[1], image.shape[0])) fin_img = self.splice(image, new_masks, gan_image) - fin_img = bilateralFilter(fin_img, 7, 70, 70) # quick bilateral filter to soften splice fin_img = fin_img[..., ::-1] # reverse RGB to BGR for video writing # Add image to video writer vwriter.write(fin_img) @@ -290,30 +350,37 @@ def ESRGAN(self, img_path, img_name, is_video=False): count += 1 vwriter.release() - print('Video 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 + print('Video: Phase 2 complete!') + # 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 - # ColabESRGAN.test.esrgan_warmup(model_path = os.path.join(os.path.abspath('.'), "ColabESRGAN/models/4x_FatalPixels_340000_G.pth")) - # similar to run_on_folder + # Parse directory for files. img_list = [] for file in os.listdir(in_path): - # TODO: check what other filetpyes supported try: if file.endswith('.png') or file.endswith('.PNG') or file.endswith(".jpg") or file.endswith(".JPG") or file.endswith(".mp4") or file.endswith(".avi"): img_list.append((in_path + '/' + file, file)) except Exception as e: print("ERROR in run_ESRGAN: File parsing. file=", file, e) - # begin ESRGAN on every image - file_counter=0 + # begin ESRGAN on every image. Create esrgan instance too. + star = time.perf_counter() + self.esrgan_instance = ColabESRGAN.test.esrgan(model_path=self.esr_model_path, hw=self.hardware) + for img_path, img_name in img_list: + self.resize_GAN(img_path=img_path, img_name=img_name, is_video=is_video) + # destroy esrgan model. Create hent-AI model. + # self.esrgan_instance = [] + del self.esrgan_instance + self.load_weights() for img_path, img_name in img_list: self.ESRGAN(img_path=img_path, img_name=img_name, is_video=is_video) - print('ESRGAN on image', file_counter, 'is complete') - file_counter += 1 + fin = time.perf_counter() + total_time = 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=''): assert image_path @@ -335,9 +402,9 @@ def video_create(self, image_path=None, dcp_path=''): fps = vcapture.get(CAP_PROP_FPS) # Define codec and create video writer, video output is purely for debugging and educational purpose. Not used in decensoring. - file_name = str(file) + '_uncensored.avi' + file_name = str(file) + '_uncensored.mp4' vwriter = VideoWriter(file_name, - VideoWriter_fourcc(*'MJPG'), + VideoWriter_fourcc(*'mp4v'), fps, (width, height)) count = 0 print("Beginning build. Do ensure only relevant images are in source directory") @@ -345,7 +412,6 @@ def video_create(self, image_path=None, dcp_path=''): img_list = [] for file in os.listdir(input_path): - # TODO: check what other filetpyes supported file_s = str(file) if file_s.endswith('.png') or file_s.endswith('.PNG'): img_list.append(input_path + file_s) @@ -365,11 +431,11 @@ 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? - if is_video: # TODO: video capabilities will finalize later + if is_video: # Video capture video_path = image_path vcapture = VideoCapture(video_path) @@ -378,9 +444,9 @@ def detect_and_cover(self, image_path=None, fname=None, save_path='', is_video=F fps = vcapture.get(CAP_PROP_FPS) # Define codec and create video writer, video output is purely for debugging and educational purpose. Not used in decensoring. - file_name = fname + "_with_censor_masks.avi" + file_name = fname + "_with_censor_masks.mp4" vwriter = VideoWriter(file_name, - VideoWriter_fourcc(*'MJPG'), + VideoWriter_fourcc(*'mp4v'), fps, (width, height)) count = 0 success = True @@ -406,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' @@ -422,9 +488,7 @@ def detect_and_cover(self, image_path=None, fname=None, save_path='', is_video=F vwriter.release() print('video complete') else: - # print("Running on ", end='') - # print(image_path) - # Read image + # Run on Image try: image = skimage.io.imread(image_path) # problems with strange shapes if image.ndim != 3: @@ -435,8 +499,16 @@ def detect_and_cover(self, image_path=None, fname=None, save_path='', is_video=F print("ERROR in detect_and_cover: Image read. Skipping. image_path=", image_path) return # Detect objects - # try: - r = self.model.detect([image], verbose=0)[0] + # image_ced =Canny(image=image, threshold1=10, threshold2=42) + # image_ced = 255 - image_ced + # image_ced = cvtColor(image_ced,COLOR_GRAY2RGB) + # skimage.io.imsave(save_path + fname[:-4] + '_ced' + '.png', image_ced) + try: + # r = self.model.detect([image_ced], verbose=0)[0] + r = self.model.detect([image], verbose=0)[0] + except Exception as e: + print("ERROR in detect_and_cover: Model detection.",e) + return # Remove unwanted class, code from https://github.com/matterport/Mask_RCNN/issues/1666 if is_mosaic==True or is_video==True: remove_indices = np.where(r['class_ids'] != 2) # remove bars: class 2 @@ -446,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' @@ -456,12 +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): @@ -474,33 +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)): - # TODO: check what other filetpyes supported 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 diff --git a/hent_AI_COLAB_1.ipynb b/hent_AI_COLAB_1.ipynb index ead156a..b407991 100644 --- a/hent_AI_COLAB_1.ipynb +++ b/hent_AI_COLAB_1.ipynb @@ -185,7 +185,8 @@ "source": [ "# Create directories, you'll only need to do this if you dont already have them in your drive\n", "!mkdir /content/drive/My\\ Drive/hent-AI/\n", - "!mkdir /content/drive/My\\ Drive/hent-AI/videos" + "!mkdir /content/drive/My\\ Drive/hent-AI/videos\n", + "!mkdir /content/drive/My\\ Drive/hent-AI/images" ], "execution_count": 0, "outputs": [] @@ -216,7 +217,7 @@ "source": [ "# Get requirements. This will take some time and lots of disk space. MAKE SURE TO PRESS THE \"RESTART RUNTIME\" BUTTON AT THE BOTTOM OF THE OUTPUT HERE\n", "%cd /content/hent-AI/\n", - "!pip install -r requirements.txt" + "!pip install -r requirements-gpu.txt" ], "execution_count": 0, "outputs": [] @@ -303,11 +304,12 @@ "colab": {} }, "source": [ + "# Ignore this cell \n", "# Remove tensorflow normal to operate on GPU only? NOTE: You will need to authorize both uninstalls. MAKE SURE TO PRESS THE \"RESTART RUNTIME\" BUTTON AT THE BOTTOM OF THE OUTPUT HERE\n", - "!pip uninstall tensorflow \n", - "!pip uninstall protobuf\n", + "# !pip uninstall tensorflow \n", + "# !pip uninstall protobuf\n", "# !pip install tensorflow==1.8.0\n", - "!pip install --force-reinstall tensorflow-gpu==1.9.0 " + "# !pip install --force-reinstall tensorflow-gpu==1.9.0 " ], "execution_count": 0, "outputs": [] @@ -349,7 +351,9 @@ }, "source": [ "# Make sure videos are in the videos folder inside hent-AI\n", - "!python samples/hentai/hentai.py inference --weights=weights.h5 --sources=/content/drive/My\\ Drive/hent-AI/videos/ --dtype=esrgan" + "!python samples/hentai/hentai.py inference --weights=weights.h5 --sources=/content/drive/My\\ Drive/hent-AI/videos/ --dtype=esrgan\n", + "# !python samples/hentai/hentai.py inference --weights=weights.h5 --sources=/content/drive/My\\ Drive/hent-AI/videos/ --dtype=bar --dcpdir=/path/to/dcpdir \n", + "# !python samples/hentai/hentai.py inference --weights=weights.h5 --sources=/content/drive/My\\ Drive/hent-AI/videos/ --dtype=mosaic --dcpdir=/path/to/dcpdir " ], "execution_count": 0, "outputs": [] diff --git a/main.py b/main.py index 3c47800..5c80309 100644 --- a/main.py +++ b/main.py @@ -13,7 +13,7 @@ import shutil from detector import Detector -versionNumber = '1.6.3' +versionNumber = '1.6.7' weights_path = 'weights.h5' # should call it weights.h5 in main dir # tkinter UI globals for window tracking. Sourced from https://stackoverflow.com/a/35486067 @@ -87,7 +87,7 @@ def hentAI_video_create(video_path=None, dcp_dir=None): okbutton.pack() popup.mainloop() -def hentAI_detection(dcp_dir=None, in_path=None, is_mosaic=False, is_video=False, force_jpg=False): +def hentAI_detection(dcp_dir=None, in_path=None, is_mosaic=False, is_video=False, force_jpg=False, dilation=0): # TODO: Create new window? Can show loading bar # hent_win = new_window() # info_label = Label(hent_win, text="Beginning detection") @@ -98,22 +98,21 @@ def hentAI_detection(dcp_dir=None, in_path=None, is_mosaic=False, is_video=False error(5) if in_path==None: error(2) + + dilation = (dilation) * 2 # Dilation performed via kernel, so dimension is doubled if(is_mosaic == True and is_video==False): # Copy input folder to decensor_input_original. NAMES MUST MATCH for DCP print('copying inputs into input_original dcp folder') # print(in_path) # print(listdir(in_path)) - for file in listdir(in_path): - # kinda dumb but check if same file - if force_jpg==True: + for fil in listdir(in_path): + if fil.endswith('jpg') or fil.endswith('png') or fil.endswith('jpeg') or fil.endswith('JPG') or fil.endswith('PNG') or fil.endswith('JPEG'): try: - shutil.copy(in_path + '/' + file, dcp_dir + '/decensor_input_original/' + file) # DCP is compatible with original jpg input. - except: - print("ERROR in hentAI_detection: Mosaic copy + png conversion to decensor_input_original failed!", file) + shutil.copy(in_path + '/' + fil, dcp_dir + '/decensor_input_original/' + fil) # DCP is compatible with original jpg input. + except Exception as e: + print("ERROR in hentAI_detection: Mosaic copy to decensor_input_original failed!", fil, e) return - else: - shutil.copy(in_path + '/' + file, dcp_dir + '/decensor_input_original/') # Run detection if(is_video==True): @@ -123,7 +122,7 @@ def hentAI_detection(dcp_dir=None, in_path=None, is_mosaic=False, is_video=False load_label = Label(loader, text='Now running detections. This can take around a minute or so per image. Please wait') load_label.pack(side=TOP, fill=X, pady=10, padx=20) loader.update() - detect_instance.run_on_folder(input_folder=in_path, output_folder=dcp_dir+'/decensor_input/', is_video=True, orig_video_folder=dcp_dir + '/decensor_input_original/') #no jpg for video detect + detect_instance.run_on_folder(input_folder=in_path, output_folder=dcp_dir+'/decensor_input/', is_video=True, orig_video_folder=dcp_dir + '/decensor_input_original/', dilation=dilation) #no jpg for video detect loader.destroy() else: print('Running detection, outputting to dcp input') @@ -132,7 +131,7 @@ def hentAI_detection(dcp_dir=None, in_path=None, is_mosaic=False, is_video=False load_label = Label(loader, text='Now running detections. This can take around a minute or so per image. Please wait') load_label.pack(side=TOP, fill=X, pady=10, padx=20) loader.update() - detect_instance.run_on_folder(input_folder=in_path, output_folder=dcp_dir+'/decensor_input/', is_video=False, force_jpg=force_jpg, is_mosaic=is_mosaic) + detect_instance.run_on_folder(input_folder=in_path, output_folder=dcp_dir+'/decensor_input/', is_video=False, is_mosaic=is_mosaic, dilation=dilation) loader.destroy() @@ -143,10 +142,6 @@ def hentAI_detection(dcp_dir=None, in_path=None, is_mosaic=False, is_video=False label = Label(popup, text='Process executed successfully! Now you can run DeepCreamPy.') label.pack(side=TOP, fill=X, pady=20, padx=10) num_jpgs = detect_instance.get_non_png() - # Popup for unprocessed jpgs - if(num_jpgs > 0 and force_jpg==False): - label2 = Label(popup, text= str(num_jpgs) + " files are NOT in .png format, and were not processed.\nPlease convert jpgs to pngs.") - label2.pack(side=TOP, fill=X, pady=10, padx=5) # dcprun = Button(popup, text='Run DCP (Only if you have the .exe)', command= lambda: run_dcp(dcp_dir)) # dcprun.pack(pady=10) okbutton = Button(popup, text='Ok', command=popup.destroy) @@ -154,7 +149,7 @@ def hentAI_detection(dcp_dir=None, in_path=None, is_mosaic=False, is_video=False popup.mainloop() # helper function to call TGAN folder function. -def hentAI_TGAN(in_path=None, is_video=False, force_jpg=False): +def hentAI_TGAN(in_path=None, is_video=False, force_jpg=True): print("Starting ESRGAN detection and decensor") loader = Tk() loader.title('Running TecoGAN') @@ -202,20 +197,23 @@ def bar_detect(): out_button.grid(row=1, column=2) # Entry for DCP installation - d_label = Label(bar_win, text = 'DeepCreamPy install folder (usually called dist1): ') + d_label = Label(bar_win, text = 'DeepCreamPy install folder: ') d_label.grid(row=2, padx=20, pady=10) d_entry = Entry(bar_win, textvariable = dvar) d_entry.grid(row=2, column=1, padx=20) dir_button = Button(bar_win, text="Browse", command=dcp_newdir) dir_button.grid(row=2, column=2, padx=20) - boolv = BooleanVar() - cb = Checkbutton(bar_win, text='Force use jpg (will save as png)?', variable = boolv) - cb.grid(row=3,column=2, padx=5) - go_button = Button(bar_win, text="Go!", command = lambda: hentAI_detection(dcp_dir=d_entry.get(), in_path=o_entry.get(), is_mosaic=False, is_video=False, force_jpg=boolv.get())) - go_button.grid(row=3, column=1, pady=10) + dil_label = Label(bar_win, text='Grow detected mask amount (0 to 10)') + dil_label.grid(row=3, padx=10, pady=10) + dil_entry = Entry(bar_win) + dil_entry.grid(row=3, column=1, padx=20) + dil_entry.insert(0, '4') + + go_button = Button(bar_win, text="Go!", command = lambda: hentAI_detection(dcp_dir=d_entry.get(), in_path=o_entry.get(), is_mosaic=False, is_video=False, force_jpg=True, dilation=int(dil_entry.get()) )) + go_button.grid(row=4, column=1, pady=10) back_button = Button(bar_win, text="Back", command = backMain) - back_button.grid(row=3,column=0, padx=10) + back_button.grid(row=4,column=0, padx=10) bar_win.mainloop() @@ -232,20 +230,26 @@ def mosaic_detect(): out_button.grid(row=1, column=2) # Entry for DCP installation - d_label = Label(mos_win, text = 'DeepCreamPy install folder (usually called dist1): ') + d_label = Label(mos_win, text = 'DeepCreamPy install folder: ') d_label.grid(row=2, padx=20, pady=20) d_entry = Entry(mos_win, textvariable = dvar) d_entry.grid(row=2, column=1, padx=20) dir_button = Button(mos_win, text="Browse", command=dcp_newdir) dir_button.grid(row=2, column=2, padx=20) - boolv = BooleanVar() - cb = Checkbutton(mos_win, text='Force use jpg (will save as png)?', variable = boolv) - cb.grid(row=3,column=2, padx=5) - go_button = Button(mos_win, text="Go!", command = lambda: hentAI_detection(dcp_dir=d_entry.get(), in_path=o_entry.get(), is_mosaic=True, is_video=False, force_jpg=boolv.get())) - go_button.grid(row=3,column=1, pady=10) + dil_label = Label(mos_win, text='Grow detected mask amount (0 to 10)') + dil_label.grid(row=3, padx=10, pady=10) + dil_entry = Entry(mos_win) + dil_entry.grid(row=3, column=1, padx=20) + dil_entry.insert(0, '4') + + # boolv = BooleanVar() + # cb = Checkbutton(mos_win, text='Force use jpg (will save as png)?', variable = boolv) + # cb.grid(row=3,column=2, padx=5) # Removing Force jpg option because jpg always works + go_button = Button(mos_win, text="Go!", command = lambda: hentAI_detection(dcp_dir=d_entry.get(), in_path=o_entry.get(), is_mosaic=True, is_video=False,dilation=int(dil_entry.get()), force_jpg=True)) + go_button.grid(row=4,column=1, pady=10) back_button = Button(mos_win, text="Back", command = backMain) - back_button.grid(row=3,column=0, padx=10) + back_button.grid(row=4,column=0, padx=10) mos_win.mainloop() @@ -302,22 +306,28 @@ def video_detect(): out_button.grid(row=1, column=2) # Entry for DCP installation - d_label = Label(vid_win, text = 'DeepCreamPy install folder (usually called dist1): ') + d_label = Label(vid_win, text = 'DeepCreamPy install folder: ') d_label.grid(row=2, padx=20, pady=20) d_entry = Entry(vid_win, textvariable = dvar) d_entry.grid(row=2, column=1, padx=20) dir_button = Button(vid_win, text="Browse", command=dcp_newdir) dir_button.grid(row=2, column=2, padx=20) - go_button = Button(vid_win, text="Begin Detection!", command = lambda: hentAI_detection(dcp_dir=d_entry.get(), in_path=o_entry.get(), is_mosaic=True, is_video=True)) - go_button.grid(row=3, columnspan=2, pady=5) + dil_label = Label(vid_win, text='Grow detected mask amount (0 to 10)') + dil_label.grid(row=3, padx=10, pady=10) + dil_entry = Entry(vid_win) + dil_entry.grid(row=3, column=1, padx=20) + dil_entry.insert(0, '4') + + go_button = Button(vid_win, text="Begin Detection!", command = lambda: hentAI_detection(dcp_dir=d_entry.get(), in_path=o_entry.get(), is_mosaic=True,dilation=int(dil_entry.get()), is_video=True)) + go_button.grid(row=4, columnspan=2, pady=5) vid_label = Label(vid_win, text= 'If you finished the video uncensoring, make images from DCP output back into video format. Check README for usage.') - vid_label.grid(row=4, pady=5, padx=4) + vid_label.grid(row=5, pady=5, padx=4) vid_button = Button(vid_win, text='Begin Video Maker!', command = lambda: hentAI_video_create(dcp_dir=d_entry.get(), video_path=o_entry.get())) - vid_button.grid(row=5, pady=5, padx=10, column=1) + vid_button.grid(row=6, pady=5, padx=10, column=1) back_button = Button(vid_win, text="Back", command = backMain) - back_button.grid(row=5, padx=10, column=0) + back_button.grid(row=6, padx=10, column=0) vid_win.mainloop() @@ -362,8 +372,6 @@ def backMain(): video_button.pack(pady=10, padx=10) video_TG_button = Button(title_window, text="Video (ESRGAN)", command=video_detect_TGAN) # separate window for future functionality changes video_TG_button.pack(pady=10, padx=10) - # detect_instance = Detector(weights_path=weights_path) # Detect instance is the same - # detect_instance.load_weights() title_window.geometry("300x300") title_window.mainloop() @@ -387,7 +395,7 @@ def backMain(): video_TG_button = Button(title_window, text="Video (ESRGAN)", command=video_detect_TGAN) # separate window for future functionality changes video_TG_button.pack(pady=10, padx=10) detect_instance = Detector(weights_path=weights_path) - detect_instance.load_weights() + # detect_instance.load_weights() # instance will load weights on its own title_window.geometry("300x300") title_window.mainloop() diff --git a/requirements-cpu.txt b/requirements-cpu.txt new file mode 100644 index 0000000..1fff81b --- /dev/null +++ b/requirements-cpu.txt @@ -0,0 +1,15 @@ +numpy +scipy +Pillow +cython +matplotlib +scikit-image +tensorflow==1.8.0 +keras==2.2.0 +opencv-python +h5py +imgaug +IPython[all] +pandas>=0.23.1 +torch==0.4.1 +torchvision>=0.2.1 diff --git a/requirements.txt b/requirements-gpu.txt similarity index 90% rename from requirements.txt rename to requirements-gpu.txt index a348081..35aefe2 100644 --- a/requirements.txt +++ b/requirements-gpu.txt @@ -4,7 +4,6 @@ Pillow cython matplotlib scikit-image -tensorflow==1.8.0 tensorflow-gpu==1.9 keras==2.2.0 opencv-python diff --git a/samples/hentai/hentai.py b/samples/hentai/hentai.py index 19954b4..2cc5dae 100644 --- a/samples/hentai/hentai.py +++ b/samples/hentai/hentai.py @@ -12,8 +12,10 @@ """ import os +from os import listdir, system import sys import json +import shutil import datetime import numpy as np import skimage.draw @@ -258,8 +260,11 @@ def train(model): metavar="path to images folder or video folder", help='Source folder to run on') parser.add_argument('--dtype', required=False, - metavar="esrgan", - help='Type of detection: Only esrgan supported now') + metavar="esrgan, bar, mosaic", + help='Type of detection. esrgan for video, bar or mosaic for images') + parser.add_argument('--dcpdir', required=False, + metavar="/path/to/DeepCreamPy installation", + help='Enter path to your DeepCreamPy folder') args = parser.parse_args() # Validate arguments @@ -299,6 +304,7 @@ class InferenceConfig(HentaiConfig): weights_path = args.weights src_path = args.sources + out_path = args.dcpdir # Train or evaluate if args.command == "train": @@ -307,9 +313,20 @@ class InferenceConfig(HentaiConfig): elif args.command == "inference": print("Starting inference") detect_instance = Detector(weights_path=weights_path) # declare instance and load weights - detect_instance.load_weights() + # detect_instance.load_weights() if args.dtype == "esrgan": detect_instance.run_ESRGAN(in_path = src_path, is_video = True, force_jpg = True) + elif args.dtype == "bar": + detect_instance.run_on_folder(input_folder=src_path, output_folder=out_path + '/decensor_input/', is_video=False, force_jpg=True, is_mosaic=False) + elif args.dtype == "mosaic": + # First copy over all original files into input_original folder + for fil in listdir(src_path): + if fil.endswith('jpg') or fil.endswith('png') or fil.endswith('jpeg') or fil.endswith('JPG') or fil.endswith('PNG') or fil.endswith('JPEG'): + try: + shutil.copy(src_path + '/' + fil, out_path + '/decensor_input_original/' + fil) # DCP is compatible with original jpg input. + except Exception as e: + print("ERROR in hentAI_detection: Mosaic copy to decensor_input_original failed!", fil, e) + detect_instance.run_on_folder(input_folder=src_path, output_folder=out_path + '/decensor_input/', is_video=False, force_jpg=True, is_mosaic=True) else: print("'{}' is not recognized. " "Use 'train'".format(args.command))