1+ import os
2+ from queue import Queue
3+ from threading import Thread
4+ import time
5+ import numpy as np
6+ import cv2
7+ from models import runmodel
8+ from util import mosaic ,util ,ffmpeg ,filt
9+ from util import image_processing as impro
10+ from .init import video_init
11+
12+
13+ '''
14+ ---------------------Add Mosaic---------------------
15+ '''
16+ def addmosaic_img (opt ,netS ):
17+ path = opt .media_path
18+ print ('Add Mosaic:' ,path )
19+ img = impro .imread (path )
20+ mask = runmodel .get_ROI_position (img ,netS ,opt )[0 ]
21+ img = mosaic .addmosaic (img ,mask ,opt )
22+ impro .imwrite (os .path .join (opt .result_dir ,os .path .splitext (os .path .basename (path ))[0 ]+ '_add.jpg' ),img )
23+
24+ def get_roi_positions (opt ,netS ,imagepaths ,savemask = True ):
25+ # resume
26+ continue_flag = False
27+ if os .path .isfile (os .path .join (opt .temp_dir ,'step.json' )):
28+ step = util .loadjson (os .path .join (opt .temp_dir ,'step.json' ))
29+ resume_frame = int (step ['frame' ])
30+ if int (step ['step' ])> 2 :
31+ mask_index = np .load (os .path .join (opt .temp_dir ,'mask_index.npy' ))
32+ return mask_index
33+ if int (step ['step' ])>= 2 and resume_frame > 0 :
34+ pre_positions = np .load (os .path .join (opt .temp_dir ,'roi_positions.npy' ))
35+ continue_flag = True
36+ imagepaths = imagepaths [resume_frame :]
37+
38+ positions = []
39+ t1 = time .time ()
40+ if not opt .no_preview :
41+ cv2 .namedWindow ('mask' , cv2 .WINDOW_NORMAL )
42+ print ('Step:2/4 -- Find mosaic location' )
43+
44+ img_read_pool = Queue (4 )
45+ def loader (imagepaths ):
46+ for imagepath in imagepaths :
47+ img_origin = impro .imread (os .path .join (opt .temp_dir + '/video2image' ,imagepath ))
48+ img_read_pool .put (img_origin )
49+ t = Thread (target = loader ,args = (imagepaths ,))
50+ t .daemon = True
51+ t .start ()
52+
53+ for i ,imagepath in enumerate (imagepaths ,1 ):
54+ img_origin = img_read_pool .get ()
55+ mask ,x ,y ,size ,area = runmodel .get_ROI_position (img_origin ,netS ,opt )
56+ positions .append ([x ,y ,area ])
57+ if savemask :
58+ t = Thread (target = cv2 .imwrite ,args = (os .path .join (opt .temp_dir + '/ROI_mask' ,imagepath ), mask ,))
59+ t .start ()
60+ if i % 1000 == 0 :
61+ save_positions = np .array (positions )
62+ if continue_flag :
63+ save_positions = np .concatenate ((pre_positions ,save_positions ),axis = 0 )
64+ np .save (os .path .join (opt .temp_dir ,'roi_positions.npy' ),save_positions )
65+ step = {'step' :2 ,'frame' :i + resume_frame }
66+ util .savejson (os .path .join (opt .temp_dir ,'step.json' ),step )
67+
68+ #preview result and print
69+ if not opt .no_preview :
70+ cv2 .imshow ('mask' ,mask )
71+ cv2 .waitKey (1 ) & 0xFF
72+ t2 = time .time ()
73+ print ('\r ' ,str (i )+ '/' + str (len (imagepaths )),util .get_bar (100 * i / len (imagepaths ),num = 35 ),util .counttime (t1 ,t2 ,i ,len (imagepaths )),end = '' )
74+
75+ if not opt .no_preview :
76+ cv2 .destroyAllWindows ()
77+
78+ print ('\n Optimize ROI locations...' )
79+ if continue_flag :
80+ positions = np .concatenate ((pre_positions ,positions ),axis = 0 )
81+ mask_index = filt .position_medfilt (np .array (positions ), 7 )
82+ step = {'step' :3 ,'frame' :0 }
83+ util .savejson (os .path .join (opt .temp_dir ,'step.json' ),step )
84+ np .save (os .path .join (opt .temp_dir ,'roi_positions.npy' ),positions )
85+ np .save (os .path .join (opt .temp_dir ,'mask_index.npy' ),np .array (mask_index ))
86+
87+ return mask_index
88+
89+ def addmosaic_video (opt ,netS ):
90+ path = opt .media_path
91+ fps ,imagepaths = video_init (opt ,path )[:2 ]
92+ length = len (imagepaths )
93+ start_frame = int (imagepaths [0 ][7 :13 ])
94+ mask_index = get_roi_positions (opt ,netS ,imagepaths )[(start_frame - 1 ):]
95+
96+ t1 = time .time ()
97+ if not opt .no_preview :
98+ cv2 .namedWindow ('preview' , cv2 .WINDOW_NORMAL )
99+
100+ # add mosaic
101+ print ('Step:3/4 -- Add Mosaic:' )
102+ t1 = time .time ()
103+ # print(mask_index)
104+ for i ,imagepath in enumerate (imagepaths ,1 ):
105+ mask = impro .imread (os .path .join (opt .temp_dir + '/ROI_mask' ,imagepaths [np .clip (mask_index [i - 1 ]- start_frame ,0 ,1000000 )]),'gray' )
106+ img = impro .imread (os .path .join (opt .temp_dir + '/video2image' ,imagepath ))
107+ if impro .mask_area (mask )> 100 :
108+ try :#Avoid unknown errors
109+ img = mosaic .addmosaic (img , mask , opt )
110+ except Exception as e :
111+ print ('Warning:' ,e )
112+ t = Thread (target = cv2 .imwrite ,args = (os .path .join (opt .temp_dir + '/addmosaic_image' ,imagepath ),img ))
113+ t .start ()
114+ os .remove (os .path .join (opt .temp_dir + '/video2image' ,imagepath ))
115+
116+ #preview result and print
117+ if not opt .no_preview :
118+ cv2 .imshow ('preview' ,img )
119+ cv2 .waitKey (1 ) & 0xFF
120+ t2 = time .time ()
121+ print ('\r ' ,str (i )+ '/' + str (length ),util .get_bar (100 * i / length ,num = 35 ),util .counttime (t1 ,t2 ,i ,length ),end = '' )
122+
123+ print ()
124+ if not opt .no_preview :
125+ cv2 .destroyAllWindows ()
126+ print ('Step:4/4 -- Convert images to video' )
127+ ffmpeg .image2video ( fps ,
128+ opt .temp_dir + '/addmosaic_image/output_%06d.' + opt .tempimage_type ,
129+ opt .temp_dir + '/voice_tmp.mp3' ,
130+ os .path .join (opt .result_dir ,os .path .splitext (os .path .basename (path ))[0 ]+ '_add.mp4' ))
0 commit comments