update make_datasets

make_datasets/csv/video_used_time.csv

+010412_249-1pon-whole1_hd.avi,00:12:00,00:13:33,00:14:26,00:15:06,00:19:35,00:24:30,00:25:53,00:29:29,00:29:55,00:30:30,00:31:43,00:32:54,00:33:39,00:35:55,00:38:30,00:38:49,00:39:47,00:41:15,00:42:35,00:43:15,00:43:50,00:45:30,00:46:33,00:47:35,00:49:10,00:49:20,00:51:04,00:51:20,00:53:10,00:55:05
+011013_511-1pon-whole1_hd.avi,00:16:09,00:16:43,00:19:12,00:19:54,00:24:52,00:26:23,00:29:20,00:31:40,00:32:16,00:36:45,00:37:15,00:37:35,00:38:00,00:38:40,00:41:40,00:46:09,00:57:50,00:58:10
+012514_744-1pon-whole1_hd.mp4,00:08:12,00:12:00,00:12:30,00:17:40,00:19:35,00:20:50,00:21:50,00:24:35,00:29:10,00:30:25,00:33:10,00:39:35,00:40:35,00:42:25,00:42:35,00:57:05,00:58:25,00:59:15
+020916_242-1pon-1080p.mp4,00:13:35,00:15:10,00:18:20,00:26:50,00:31:25,00:33:15,00:34:55,00:37:15,00:38:25,00:39:35,00:41:05,00:41:55,00:42:10,00:43:10,00:43:20,00:45:15,00:45:20,00:46:10,00:47:50,00:49:10,00:50:00,00:50:20,00:52:10,00:56:55,00:57:05,00:57:35,00:59:15,00:59:30
+031516_262-1pon-1080p.mp4,00:09:30,00:13:00,00:13:50,00:14:50,00:16:00,00:20:55,00:25:50,00:26:35,00:30:30,00:32:40,00:38:20,00:38:30,00:39:55,00:42:10,00:43:45,00:45:40,00:46:20,00:47:50,00:48:05,00:49:50,00:51:45,00:51:50,00:57:00
+031716_001-1pon-1080p.mp4,00:02:30,00:02:40,00:02:55,00:04:00,00:04:20,00:05:40,00:06:05,00:06:50,00:08:10,00:08:20,00:08:30,00:08:47,00:10:00,00:10:05,00:10:20,00:10:30,00:11:50,00:12:00,00:12:35,00:13:20,00:14:20,00:15:35
+032113_554-1pon-whole1_hd.avi,00:13:20,00:21:20,00:23:15,00:23:35,00:24:00,00:25:10,0:25:30,00:25:50,00:26:35,00:26:50,00:31:40,00:35:15,00:35:25,00:37:10,00:45:35,00:46:05,00:48:00,00:49:50,00:50:30,00:51:50,00:52:30,00:52:40,00:52:50,00:58:20,00:58:30,00:59:30,00:59:45,01:01:45,01:02:00,01:03:50,01:04:05,01:04:20,01:04:30,01:05:35,01:07:40
+032313_556-1pon-whole1_hd.avi,00:04:05,00:05:00,00:06:40,00:06:50,00:07:50,00:09:10,00:10:30,00:13:15,00:16:05,00:17:35,00:18:20,00:20:25,00:20:30,00:22:30,00:26:50,00:27:30,00:35:30,00:42:40,00:44:09,00:50:00,00:52:50,00:53:40,00:54:15,00:58:00,00:58:25,01:04:05,01:05:05,01:06:15,01:06:50,01:07:51,01:08:10
+032715_001-1pon-1080p.mp4,00:09:20,00:10:35,00:10:45,00:13:25,00:21:20,00:24:50,00:28:10,00:29:26,00:29:52,00:30:55,00:31:10,00:31:55,00:32:20,00:32:40,00:33:10,00:34:30,00:35:40,00:35:50,00:48:30,00:48:50,00:49:45,00:50:15,00:53:55,00:57:13,00:57:20,00:59:00,00:59:55
+032715_004-1pon-1080p.mp4,00:22:30,00:22:55,00:24:44,00:26:15,00:28:00,00:28:40,00:30:40,00:35:40,00:38:20,00:38:50,00:39:50,00:41:30,00:42:10,00:42:30,00:43:40,00:44:05,00:44:35,00:45:17,00:45:36,00:46:23,00:46:55,00:47:20,00:47:40,00:48:05,00:48:30,00:50:50,00:52:00,00:53:30,00:53:45,00:54:25,00:54:45,00:57:40,00:58:00,00:58:40,00:58:50
+040111_063-1pon-whole1_hd.avi,00:08:25,00:08:45,00:09:00,00:10:55,00:16:40,00:17:05,00:17:35,00:19:10,00:27:00,00:28:05,00:29:05,00:31:40,00:36:00,00:37:50,00:45:30,00:46:15,00:47:45,00:50:15,00:52:50,00:53:47,00:53:58,00:55:05,00:56:15,00:58:40,00:59:00,00:59:20,00:59:45
+040814_786-1pon-whole1_hd.avi,00:04:40,00:05:00,00:06:50,00:10:20,00:21:00,00:23:35,00:24:10,00:26:40,00:28:35,00:29:15,00:29:20,00:31:15,00:32:50,00:36:10,00:39:40,00:42:00,00:42:50,00:44:00,00:44:15,00:44:36,00:45:00,00:45:20,00:47:20,00:48:10,00:48:30,00:53:50,00:54:43,00:55:20,00:59:15,00:59:30
+050915_077-1pon-1080p,00:11:00,00:12:30,00:19:20,00:19:50,00:21:00,00:22:00,00:23:40,00:24:30,00:28:20,00:33:50,00:36:00,00:37:30,00:38:50,00:39:30,00:41:50,00:44:20,00:48:45,00:49:25,00:50:45,00:51:00,00:53:05,00:54:00,00:54:27,00:57:30,00:59:10,01:00:30,01:04:10,01:04:20,01:04:30,01:04:50,01:05:20
+052215_084-1pon-1080p.mp4,00:26:50,00:27:15,00:30:20,00:33:20,00:34:00,00:37:00,00:41:00,00:43:00,00:44:30,00:47:40,00:50:35,00:50:40,00:51:40,00:55:20,00:55:50,00:55:55,00:56:20,00:57:30,00:57:40,00:59:10,00:59:15,01:00:05
+062015_101-1pon-1080p.mp4,00:11:00,00:12:47,00:13:10,00:14:20,00:15:20,00:16:20,00:17:10,00:17:25,00:19:45,00:21:05,00:23:40,00:27:40,00:28:10,00:37:15,00:41:30,00:43:20,00:44:25,00:46:51,00:47:20,00:49:00,00:50:40,00:51:50,00:52:50,00:55:00,00:56:20,00:58:10,00:59:00,10:00:00
+062715_105-1pon-1080p.mp4,00:11:30,00:11:55,00:12:00,00:12:30,00:13:45,00:16:50,00:18:25,00:19:20,00:20:40,00:25:15,00:36:20,00:36:40,00:37:25,00:39:05,00:39:50,00:40:55,00:41:55,00:45:40,00:43:30,00:44:15,00:45:30,00:47:40,00:50:05,00:50:10,00:50:20,00:50:30,00:55:10,00:56:35,00:58:40,01:00:15,01:05:05,01:05:15,01:05:30,01:05:50
+1pondo_070315_108_1080p.mp4,00:11:10,00:11:50,00:13:50,00:14:20,00:14:35,00:15:50,00:17:20,00:18:35,00:20:45,00:24:35,00:25:05,00:29:15,00:30:40,00:31:55,00:35:20,00:42:55,00:43:05,00:46:15,00:48:00,00:51:45,00:52:33,00:54:20,00:59:25,00:59:40,01:00:05
+071114_842-1pon-whole1_hd.mp4,00:09:50,00:11:25,00:16:35,00:18:20,00:22:10,00:25:25,00:26:35,00:33:50,00:35:40,00:43:10
+071715_116-1pon-1080p.mp4,00:10:50,00:11:30,00:12:50,00:15:10,00:16:45,00:17:05,00:25:20,00:26:45,00:28:30,00:30:20,00:32:55,00:34:30,00:37:40,00:38:40,00:40:20,00:41:20,00:44:10,00:47:15,00:55:00,00:59:40,00:59:50
+071815_117-1pon-1080p.mp4,00:14:50,00:15:10,00:18:05,00:14:50,00:25:55,00:26:25,00:32:45,00:33:40,00:43:15,00:45:05,00:45:45,00:48:40,00:48:50,00:55:45,10:00:20,01:00:35,01:01:00,01:01:10
\ No newline at end of file

make_datasets/cut_video.py

+import os
+import numpy as np
+import cv2
+import random
+import csv
+
+import sys
+sys.path.append("..")
+from util import util,ffmpeg
+from util import image_processing as impro
+
+files = util.Traversal('/media/hypo/Media/download')
+videos = util.is_videos(files)
+
+
+video_times = []
+useable_videos = []
+video_dict = {}
+reader = csv.reader(open('./csv/video_used_time.csv'))
+for line in reader:
+    useable_videos.append(line[0])
+    video_times.append(line[1:])
+    video_dict[line[0]]=line[1:]
+
+in_cnt = 0
+out_cnt = 502
+for video in videos:
+    if os.path.basename(video) in useable_videos:
+        # print(video)
+        for i in range(len(video_times[in_cnt])):
+            ffmpeg.cut_video(video, video_dict[os.path.basename(video)][i], '00:00:05', './video/'+'%04d'%out_cnt+'.mp4')
+            out_cnt +=1
+        in_cnt += 1          

make_datasets/draw_mask.py

+import cv2
+import numpy as np
+import datetime
+import os
+import random
+
+def resize(img,size):
+    h, w = img.shape[:2]
+    if w >= h:
+        res = cv2.resize(img,(int(size*w/h), size))
+    else:
+        res = cv2.resize(img,(size, int(size*h/w)))
+   
+    return res
+
+# mouse callback function
+drawing = False # true if mouse is pressed
+ix,iy = -1,-1
+brushsize = 20
+def draw_circle(event,x,y,flags,param):
+    global ix,iy,drawing,brushsize
+
+    if event == cv2.EVENT_LBUTTONDOWN:
+        drawing = True
+        ix,iy = x,y
+
+    elif event == cv2.EVENT_MOUSEMOVE:
+        if drawing == True:
+            cv2.circle(img,(x,y),brushsize,(0,255,0),-1)
+
+    elif event == cv2.EVENT_LBUTTONUP:
+        drawing = False
+        cv2.circle(img,(x,y),brushsize,(0,255,0),-1)
+
+def makemask(img):
+    # starttime = datetime.datetime.now()
+    mask = np.zeros(img.shape, np.uint8)
+    for row in range(img.shape[0]):
+        for col in range(img.shape[1]):
+            # if (img[row,col,:] == [0,255,0]).all(): #too slow
+            if img[row,col,0] == 0:
+                if img[row,col,1] == 255:
+                    if img[row,col,2] == 0:
+                        mask[row,col,:] = [255,255,255]
+    # endtime = datetime.datetime.now()
+    # print('Cost time:',(endtime-starttime))
+    return mask
+
+files = os.listdir('./origin_image')
+files_new =files.copy()
+print('find image:',len(files))
+masks = os.listdir('./mask')
+print('mask:',len(masks))
+for i in range(len(masks)):
+    masks[i]=masks[i].replace('.png','.jpg')
+for file in files:
+    if file  in masks:
+        files_new.remove(file)
+files = files_new
+# files = list(set(files)) #Distinct 
+print('remain:',len(files))
+random.shuffle (files)
+# files.sort()
+cnt = 0
+
+for file in files:
+    cnt += 1
+    img = cv2.imread('./origin_image/'+file)
+    img = resize(img,512)
+    cv2.namedWindow('image')
+    cv2.setMouseCallback('image',draw_circle) #MouseCallback
+    while(1):
+
+        cv2.imshow('image',img)
+        k = cv2.waitKey(1) & 0xFF
+        if k == ord(' '):
+            img = resize(img,256)
+            mask = makemask(img)
+            cv2.imwrite('./mask/'+os.path.splitext(file)[0]+'.png',mask)
+            print('./mask/'+os.path.splitext(file)[0]+'.png')
+            # cv2.destroyAllWindows()
+            print('remain:',len(files)-cnt)
+            brushsize = 20
+            break
+        elif k == ord('a'):
+            brushsize -= 5
+            if brushsize<5:
+                brushsize = 5
+            print('brushsize:',brushsize)
+        elif k == ord('d'):
+            brushsize += 5
+            print('brushsize:',brushsize)
+        elif k == ord('w'):
+            print('remain:',len(files)-cnt)
+            break
+
+# cv2.destroyAllWindows()
\ No newline at end of file

make_datasets/use_addmosaic_model_make_dataset.py

+import sys
+import os
+import random
+import datetime
+
+import numpy as np
+import cv2
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+from torch import optim
+
+from unet import UNet
+from mosaic import random_mosaic
+import image_processing as impro
+
+
+
+def runmodel(img,net):
+    img=impro.image2folat(img,3)
+    img=img.reshape(1,3,128,128)
+    img = torch.from_numpy(img)
+    img=img.cuda()
+    pred = net(img)
+    pred = (pred.cpu().detach().numpy()*255)
+    pred = pred.reshape(128,128).astype('uint8')
+    return pred
+
+
+
+dir_img = './origin_image/'
+dir_mosaic = './mosaic/'
+dir_mask = './mask/'
+dir_dataset = './dataset/'
+dir_checkpoint = 'checkpoints/'
+
+net = UNet(n_channels = 3, n_classes = 1)
+net.load_state_dict(torch.load(dir_checkpoint+'mosaic_position.pth'))
+net.cuda()
+net.eval()
+# cudnn.benchmark = True
+files = os.listdir(dir_mosaic)
+
+for i,file in enumerate(files,1):
+    orgin_image = cv2.imread(dir_img+file)
+    mosaic_image = cv2.imread(dir_mosaic+file)
+    img = impro.resize(mosaic_image,128)
+    img1,img2 = impro.spiltimage(img)
+    mask1 =runmodel(img1,net)
+    mask2 =runmodel(img2,net)
+    mask = impro.mergeimage(mask1,mask2,img)
+
+    # test_mask = mask.copy()
+
+    mask = impro.mask_threshold(mask,blur=5,threshold=128)
+    if impro.mask_area(mask) > 1:
+        h,w = orgin_image.shape[:2]
+        mosaic_image = cv2.resize(mosaic_image,(w,h))
+        # test_mask  = cv2.resize(test_mask,(w,h))
+        # test_mask  = impro.ch_one2three(test_mask)
+
+        x,y,size,area = impro.boundingSquare(mask,Ex_mul=1.5)
+        rat = min(orgin_image.shape[:2])/128.0
+        x,y,size = int(rat*x),int(rat*y),int(rat*size)
+        orgin_crop = orgin_image[y-size:y+size,x-size:x+size]
+        mosaic_crop = mosaic_image[y-size:y+size,x-size:x+size]
+        # mosaic_crop = test_mask[y-size:y+size,x-size:x+size]
+
+        result = impro.makedataset(mosaic_crop,orgin_crop)
+        cv2.imwrite(dir_dataset+file,result)
+    if i%1000==0:
+        print(i,'image finished.')

make_datasets/use_drawn_mask_make_dataset.py

+import numpy as np
+import cv2
+import os
+from torchvision import transforms
+from PIL import Image
+import random
+import sys
+sys.path.append("..")
+import util.image_processing as impro
+from util import util,mosaic
+import datetime
+import shutil
+
+mask_path = '/media/hypo/Porject/Datasets/unet/av/mask'
+img_path ='/media/hypo/Porject/Datasets/unet/av/origin_image'
+output_dir = './datasets_img'
+util.makedirs(output_dir)
+HD = True # if false make dataset for pix2pix, if Ture for pix2pix_HD
+MASK = False # if True, output mask,too
+OUT_SIZE = 256
+FOLD_NUM = 5
+
+if HD:
+    train_A_path = os.path.join(output_dir,'train_A')
+    train_B_path = os.path.join(output_dir,'train_B')
+    util.makedirs(train_A_path)
+    util.makedirs(train_B_path)
+else:
+    train_path = os.path.join(output_dir,'train')
+    util.makedirs(train_path)
+if MASK:
+    mask_path = os.path.join(output_dir,'mask')
+    util.makedirs(mask_path)
+
+mask_names = os.listdir(mask_path)
+img_names = os.listdir(img_path)
+mask_names.sort()
+img_names.sort()
+print('Find images:',len(img_names))
+
+cnt = 0
+for fold in range(FOLD_NUM):
+    for img_name,mask_name in zip(img_names,mask_names):
+        try:
+            img = impro.imread(os.path.join(img_path,img_name))
+            mask = impro.imread(os.path.join(mask_path,mask_name),'gray')
+            mask = impro.resize_like(mask, img)
+            x,y,size,area = impro.boundingSquare(mask, 1.5)
+            if area > 100:
+                img = impro.resize(img[y-size:y+size,x-size:x+size],OUT_SIZE) 
+                mask =  impro.resize(mask[y-size:y+size,x-size:x+size],OUT_SIZE)
+                img_mosaic = mosaic.addmosaic_random(img, mask)
+
+                if HD:
+                    cv2.imwrite(os.path.join(train_A_path,'%05d' % cnt+'.jpg'), img_mosaic)
+                    cv2.imwrite(os.path.join(train_B_path,'%05d' % cnt+'.jpg'), img)
+                else:
+                    merge_img = impro.makedataset(img_mosaic, img)
+                    cv2.imwrite(os.path.join(train_path,'%05d' % cnt+'.jpg'), merge_img)
+                if MASK:
+                    cv2.imwrite(os.path.join(mask_path,'%05d' % cnt+'.png'), mask)
+                print("Processing:",img_name," ","Remain:",len(img_names)*FOLD_NUM-cnt)
+                
+        except Exception as e:
+            print(img_name,e)
+        cnt += 1

make_datasets/use_irregular_holes_mask.py → make_datasets/use_irregular_holes_mask_make_dataset.py

@@ -11,18 +11,19 @@ from util import util,mosaic
 import datetime
 
 ir_mask_path = './Irregular_Holes_mask'
-# img_path = 'D:/MyProject_new/face_512'
-img_path ='/media/hypo/Hypoyun/Hypoyun/手机摄影/20190219'
-output_dir = './datasets'
+img_path ='/home/hypo/桌面/FaceRankSample' 
+output_dir = './datasets_img'
 util.makedirs(output_dir)
-HD = True #if false make dataset for pix2pix, if Ture for pix2pix_HD
-MASK = True
-if HD:
+MOD = 'HD' #HD | pix2pix | mosaic
+MASK = False # if True, output mask,too
+BOUNDING = False # if true the mosaic size will be more big
+
+if MOD='HD':
     train_A_path = os.path.join(output_dir,'train_A')
     train_B_path = os.path.join(output_dir,'train_B')
     util.makedirs(train_A_path)
     util.makedirs(train_B_path)
-else:
+elif MOD='pix2pix':
     train_path = os.path.join(output_dir,'train')
     util.makedirs(train_path)
 if MASK:
@@ -51,19 +52,30 @@ for i,img_name in enumerate(img_names,1):
         img = np.array(img)
         img = img[...,::-1]
 
-        mask = Image.open(os.path.join(ir_mask_path,random.choices(mask_names)[0]))
-        mask = transform_mask(mask)
-        mask = np.array(mask)
-
-        mosaic_img = mosaic.addmosaic_random(img, mask)          
-        if HD:
-            cv2.imwrite(os.path.join(train_A_path,'%05d' % i+'.jpg'), mosaic_img)
-            cv2.imwrite(os.path.join(train_B_path,'%05d' % i+'.jpg'), img)
+        if BOUNDING:
+            mosaic_area = 0
+            while mosaic_area < 16384:
+                mask = Image.open(os.path.join(ir_mask_path,random.choices(mask_names)[0]))
+                mask = transform_mask(mask)
+                mask = np.array(mask)
+                mosaic_area = impro.mask_area(mask)
+            mosaic_img = mosaic.addmosaic_random(img, mask,'bounding') 
+            BOUNDING_flag = '_bound'
+        else:
+            mask = Image.open(os.path.join(ir_mask_path,random.choices(mask_names)[0]))
+            mask = transform_mask(mask)
+            mask = np.array(mask)
+            mosaic_img = mosaic.addmosaic_random(img, mask)
+            BOUNDING_flag = ''    
+        
+        if HD:#[128:384,128:384,:] --->256
+            cv2.imwrite(os.path.join(train_A_path,'%05d' % i+BOUNDING_flag+'.jpg'), mosaic_img)
+            cv2.imwrite(os.path.join(train_B_path,'%05d' % i+BOUNDING_flag+'.jpg'), img)
         else:
             merge_img = impro.makedataset(mosaic_img, img)
-            cv2.imwrite(os.path.join(train_path,'%05d' % i+'.jpg'), merge_img)
+            cv2.imwrite(os.path.join(train_path,'%05d' % i+BOUNDING_flag+'.jpg'), merge_img)
         if MASK:
-            cv2.imwrite(os.path.join(mask_path,'%05d' % i+'.png'), mask)
+            cv2.imwrite(os.path.join(mask_path,'%05d' % i+BOUNDING_flag+'.png'), mask)
         print("Processing:",img_name," ","Remain:",len(img_names)-i)
     except Exception as e:
         print(img_name,e)

options.py

+import argparse
+import os
+
+class Options():
+    def __init__(self):
+        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+        self.initialized = False
+
+    def initialize(self):
+
+        #base
+        self.parser.add_argument('--use_gpu', action='store_true', help='if input it, use gpu')
+        self.parser.add_argument('--media_path', type=str, default='./hands_test.mp4',help='your videos or images path')
+        self.parser.add_argument('--mode', type=str, default='auto',help='add or clean mosaic into your media  auto | add | clean')
+        self.parser.add_argument('--model_path', type=str, default='./pretrained_models/add_hands_128.pth',help='pretrained model path')
+        self.parser.add_argument('--result_dir', type=str, default='./result',help='output result will be saved here')
+        self.parser.add_argument('--tempimage_type', type=str, default='png',help='type of temp image, png | jpg, png is better but occupy more storage space')
+
+        #AddMosaic
+        self.parser.add_argument('--mosaic_mod', type=str, default='squa_avg',help='type of mosaic -> squa_avg | squa_random | squa_avg_circle_edge | rect_avg | random')
+        self.parser.add_argument('--mosaic_size', type=int, default=0,help='mosaic size,if 0 auto size')
+        self.parser.add_argument('--mask_extend', type=int, default=10,help='more mosaic area')
+        self.parser.add_argument('--mask_threshold', type=int, default=64,help='threshold of recognize mosaic position 0~255')
+        self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin')
+        
+        #CleanMosaic
+        self.parser.add_argument('--netG', type=str, default='auto',help='select model to use for netG(clean mosaic) -> auto | unet_128 | resnet_9blocks | HD')
+        self.parser.add_argument('--mosaic_position_model_path', type=str, default='auto',help='name of model use to find mosaic position')
+        self.parser.add_argument('--no_feather', action='store_true', help='if true, no edge feather and color correction, but run faster')
+        self.parser.add_argument('--medfilt_num', type=int, default=11,help='medfilt window of mosaic movement in the video')
+        self.initialized = True
+
+
+    def getparse(self):
+        if not self.initialized:
+            self.initialize()
+        self.opt = self.parser.parse_args()
+
+        if self.opt.mode == 'auto':
+            if 'add' in self.opt.model_path:
+                self.opt.mode = 'add'
+            elif 'clean' in self.opt.model_path:
+                self.opt.mode = 'clean'
+            else:
+                print('Please input running mode!')
+
+        if self.opt.netG == 'auto' and self.opt.mode =='clean':
+            if 'unet_128' in self.opt.model_path:
+                self.opt.netG = 'unet_128'
+            elif 'resnet_9blocks' in self.opt.model_path:
+                self.opt.netG = 'resnet_9blocks'
+            elif 'HD' in self.opt.model_path:
+                self.opt.netG = 'HD'
+            else:
+                print('Type of Generator error!')
+
+
+        if self.opt.mosaic_position_model_path == 'auto':
+            _path = os.path.join(os.path.split(self.opt.model_path)[0],'mosaic_position.pth')
+            self.opt.mosaic_position_model_path = _path
+            # print(self.opt.mosaic_position_model_path)
+
+        return self.opt
\ No newline at end of file

util/data.py

@@ -8,29 +8,40 @@ transform = transforms.Compose([
     ]  
 )  
 
-def tensor2im(image_tensor, imtype=np.uint8, rgb2bgr = True):
+def tensor2im(image_tensor, imtype=np.uint8, gray=False, rgb2bgr = True):
     image_tensor =image_tensor.data
     image_numpy = image_tensor[0].cpu().float().numpy()
+    # if gray:
+    #     image_numpy = (image_numpy+1.0)/2.0 * 255.0
+    # else:
     if image_numpy.shape[0] == 1:
         image_numpy = np.tile(image_numpy, (3, 1, 1))
     image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
-    if rgb2bgr:
+    if rgb2bgr and not gray:
         image_numpy = image_numpy[...,::-1]-np.zeros_like(image_numpy)
     return image_numpy.astype(imtype)
 
 
-def im2tensor(image_numpy, imtype=np.uint8, bgr2rgb = True, reshape = True, use_gpu = True,  use_transform = True):
-    h, w ,ch = image_numpy.shape
-    if bgr2rgb:
-        image_numpy = image_numpy[...,::-1]-np.zeros_like(image_numpy)
-    if use_transform:
-        image_tensor = transform(image_numpy)
-    else:
-        image_numpy = image_numpy/255.0
-        image_numpy = image_numpy.transpose((2, 0, 1))
+def im2tensor(image_numpy, imtype=np.uint8, gray=False,bgr2rgb = True, reshape = True, use_gpu = True,  use_transform = True):
+    
+    if gray:
+        h, w = image_numpy.shape
+        image_numpy = (image_numpy/255.0-0.5)/0.5
         image_tensor = torch.from_numpy(image_numpy).float()
-    if reshape:
-        image_tensor=image_tensor.reshape(1,ch,h,w)
+        if reshape:
+            image_tensor=image_tensor.reshape(1,1,h,w)
+    else:
+        h, w ,ch = image_numpy.shape
+        if bgr2rgb:
+            image_numpy = image_numpy[...,::-1]-np.zeros_like(image_numpy)
+        if use_transform:
+            image_tensor = transform(image_numpy)
+        else:
+            image_numpy = image_numpy/255.0
+            image_numpy = image_numpy.transpose((2, 0, 1))
+            image_tensor = torch.from_numpy(image_numpy).float()
+        if reshape:
+            image_tensor=image_tensor.reshape(1,ch,h,w)
     if use_gpu:
         image_tensor = image_tensor.cuda()
     return image_tensor

util/ffmpeg.py

 import os,json
 
+# ffmpeg 3.4.6
+
 def video2image(videopath,imagepath):
     os.system('ffmpeg -i "'+videopath+'" -f image2 '+imagepath)
 
@@ -15,9 +17,26 @@ def get_video_infos(videopath):
     cmd_str =  'ffprobe -v quiet -print_format json -show_format -show_streams -i "' + videopath + '"'  
     out_string = os.popen(cmd_str).read()
     infos = json.loads(out_string)
-    fps = eval(infos['streams'][0]['avg_frame_rate'])
-    endtime = float(infos['format']['duration'])
-    width = int(infos['streams'][0]['width'])
-    height = int(infos['streams'][0]['height'])
+    try:
+        fps = eval(infos['streams'][0]['avg_frame_rate'])
+        endtime = float(infos['format']['duration'])
+        width = int(infos['streams'][0]['width'])
+        height = int(infos['streams'][0]['height'])
+    except Exception as e:
+        fps = eval(infos['streams'][1]['r_frame_rate'])
+        endtime = float(infos['format']['duration'])
+        width = int(infos['streams'][1]['width'])
+        height = int(infos['streams'][1]['height'])
+
     return fps,endtime,width,height
 
+def cut_video(in_path,start_time,last_time,out_path,vcodec='h265'):
+    if vcodec == 'copy':
+        os.system('ffmpeg -ss '+start_time+' -t '+last_time+' -i "'+in_path+'" -vcodec copy -acodec copy '+out_path)
+    elif vcodec == 'h264':    
+        os.system('ffmpeg -ss '+start_time+' -t '+last_time+' -i "'+in_path+'" -vcodec libx264 -b 12M '+out_path)
+    elif vcodec == 'h265':
+        os.system('ffmpeg -ss '+start_time+' -t '+last_time+' -i "'+in_path+'" -vcodec libx265 -b 12M '+out_path)
+
+def continuous_screenshot(videopath,savedir,fps):
+    os.system('ffmpeg -i '+videopath+' -vf fps='+str(fps)+' '+savedir+'/'+'%05d.jpg')

util/image_processing.py

 import cv2
 import numpy as np
 
+def imread(file_path,mod = 'normal'):
 
-# imread for chinese path in windows
-def imread(file_path):
-    cv_img = cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),-1)
+    if mod == 'normal':
+        cv_img = cv2.imread(file_path)
+    elif mod == 'gray':
+        cv_img = cv2.imread(file_path,0)
+    elif mod == 'all':
+        cv_img = cv2.imread(file_path,-1)
+
+    # # imread for chinese path in windows but no EXIF
+    # cv_img = cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),-1)
     return cv_img
 
 def resize(img,size):
@@ -17,6 +24,11 @@ def resize(img,size):
         res = cv2.resize(img,(size, int(size*h/w)))
     return res
 
+def resize_like(img,img_like):
+    h, w = img_like.shape[:2]
+    img = cv2.resize(img, (w,h))
+    return img
+
 def ch_one2three(img):
     #zeros = np.zeros(img.shape[:2], dtype = "uint8")
     # ret,thresh = cv2.threshold(img,127,255,cv2.THRESH_BINARY)

util/mosaic.py

@@ -126,9 +126,9 @@ def addmosaic_random(img,mask,area_type = 'normal'):
         area = 0
     area = area/(alpha*alpha)
     if area>50000:
-        img_mosaic = random_mod(img,mask,alpha*random.uniform(16,30))
+        img_mosaic = random_mod(img,mask,alpha*random.uniform(8,30)) #16,30
     elif 20000<area<=50000:
-        img_mosaic = random_mod(img,mask,alpha*random.uniform(12,20))
+        img_mosaic = random_mod(img,mask,alpha*random.uniform(8,20)) #12,20
     elif 5000<area<=20000:
         img_mosaic = random_mod(img,mask,alpha*random.uniform(8,15))
     elif 0<=area<=5000:
@@ -139,11 +139,11 @@ def addmosaic_random(img,mask,area_type = 'normal'):
 
 def random_mod(img,mask,n):
     ran=random.random()
-    if ran < 0.2:
+    if ran < 0.3:
         img = addmosaic_normal(img,mask,n,model = 'squa_mid')
-    if 0.2 <= ran < 0.4:
+    if 0.3 <= ran < 0.5:
         img = addmosaic_normal(img,mask,n,model = 'squa_avg')
-    elif 0.4 <= ran <0.6:
+    elif 0.5 <= ran <0.6:
         img = addmosaic_normal(img,mask,n,model = 'squa_avg_circle_edge')
     else:
         img = addmosaic_normal(img,mask,n,model = 'rect_avg')

util/util.py

@@ -21,11 +21,25 @@ def is_img(path):
 def is_video(path):
     ext = os.path.splitext(path)[1]
     ext = ext.lower()
-    if ext in ['.mp4','.flv','.avi','.mov','.mkv','.wmv','.rmvb']:
+    if ext in ['.mp4','.flv','.avi','.mov','.mkv','.wmv','.rmvb','.mts']:
         return True
     else:
         return False
-        
+
+def is_imgs(paths):
+    tmp = []
+    for path in paths:
+        if is_img(path):
+            tmp.append(path)
+    return tmp
+
+def is_videos(paths):
+    tmp = []
+    for path in paths:
+        if is_video(path):
+            tmp.append(path)
+    return tmp  
+
 def  writelog(path,log):
     f = open(path,'a+')
     f.write(log+'\n')