Update gradio_tabs/vid_edit.py

gradio_tabs/vid_edit.py

@@ -21,7 +21,7 @@ labels_k = [
 	'pout',
 	'open->close',
 	'"O" mouth',
-	'apple cheek',
+	'smile',
 
 	'close->open',
 	'eyebrows',
@@ -58,16 +58,24 @@ def img_preprocessing(img_path, size):
 
 def resize(img, size):
 	transform = torchvision.transforms.Compose([
-		torchvision.transforms.Resize(size, antialias=True),
-		torchvision.transforms.CenterCrop(size)
+		torchvision.transforms.Resize((size, size), antialias=True),
 	])
 
 	return transform(img)
 
 
+def resize_back(img, w, h):
+	transform = torchvision.transforms.Compose([
+		torchvision.transforms.Resize((h, w), antialias=True),
+	])
+
+	return transform(img)
+    
+
 def vid_preprocessing(vid_path, size):
 	vid_dict = torchvision.io.read_video(vid_path, pts_unit='sec')
 	vid = vid_dict[0].permute(0, 3, 1, 2).unsqueeze(0)	# btchw
+	_,_,_,h,w = vid.size()
 	fps = vid_dict[2]['video_fps']
 	vid_norm = (vid / 255.0 - 0.5) * 2.0  # [-1, 1]
 
@@ -75,7 +83,7 @@ def vid_preprocessing(vid_path, size):
 		resize(vid_norm[:, i, :, :, :], size).unsqueeze(1) for i in range(vid.size(1))
 	], dim=1)
 
-	return vid_norm, fps
+	return vid_norm, fps, w, h
 
 
 def img_denorm(img):
@@ -92,7 +100,8 @@ def vid_denorm(vid):
 	return vid
 
 
-def img_postprocessing(image):
+def img_postprocessing(image, w, h):
+	image = resize_back(image, w, h)
 	image = image.permute(0, 2, 3, 1)
 	edited_image = img_denorm(image)
 	img_output = (edited_image[0].numpy() * 255).astype(np.uint8)
@@ -102,10 +111,14 @@ def img_postprocessing(image):
 		return temp_file.name
 
 
-def vid_all_save(vid_d, vid_a, fps):
+def vid_all_save(vid_d, vid_a, w, h, fps):
 
-	vid_d = rearrange(vid_d, 'b t c h w -> b t h w c')
-	vid_a = rearrange(vid_a, 'b c t h w -> b t h w c')
+	b,t,c,_,_ = vid_d.size()
+	vid_d_batch = resize_back(rearrange(vid_d, "b t c h w -> (b t) c h w"), w, h)
+	vid_a_batch = resize_back(rearrange(vid_a, "b c t h w -> (b t) c h w"), w, h)
+	
+	vid_d = rearrange(vid_d_batch, "(b t) c h w -> b t h w c", b=b) # B T H W C
+	vid_a = rearrange(vid_a_batch, "(b t) c h w -> b t h w c", b=b) # B T H W C	
 	vid_all = torch.cat([vid_d, vid_a], dim=3)
 
 	vid_a_np = (vid_denorm(vid_a[0]).numpy() * 255).astype('uint8')
@@ -126,14 +139,14 @@ def vid_edit(gen, chunk_size, device):
 	@torch.no_grad()
 	def edit_img(video, *selected_s):
 
-		vid_target_tensor, fps = vid_preprocessing(video, 512)
+		vid_target_tensor, fps, w, h = vid_preprocessing(video, 512)
 		video_target_tensor = vid_target_tensor.to(device)
 		image_tensor = video_target_tensor[:,0,:,:,:]
 
 		edited_image_tensor = gen.edit_img(image_tensor, labels_v, selected_s)
 
 		# de-norm
-		edited_image = img_postprocessing(edited_image_tensor)
+		edited_image = img_postprocessing(edited_image_tensor, w, h)
 
 		return edited_image
 
@@ -141,15 +154,15 @@ def vid_edit(gen, chunk_size, device):
 	@torch.no_grad()
 	def edit_vid(video, *selected_s):
 
-		video_target_tensor, fps = vid_preprocessing(video, 512)
+		video_target_tensor, fps, w, h = vid_preprocessing(video, 512)
 		video_target_tensor = video_target_tensor.to(device)
 
 		edited_video_tensor = gen.edit_vid_batch(video_target_tensor, labels_v, selected_s, chunk_size)
 		edited_image_tensor = edited_video_tensor[:,:,0,:,:]       
 
 		# de-norm
-		animated_video, animated_all_video = vid_all_save(video_target_tensor, edited_video_tensor, fps)
-		edited_image = img_postprocessing(edited_image_tensor)
+		animated_video, animated_all_video = vid_all_save(video_target_tensor, edited_video_tensor, w, h, fps)
+		edited_image = img_postprocessing(edited_image_tensor, w, h)
 
 		return edited_image, animated_video, animated_all_video
 
@@ -187,16 +200,13 @@ def vid_edit(gen, chunk_size, device):
 
 				with gr.Row():
 					with gr.Accordion(open=True, label="Edited First Frame"):
-						#image_output.render()
 						image_output = gr.Image(label="Image", elem_id="output_img", type='numpy', interactive=False, width=512)
 
 					with gr.Accordion(open=True, label="Edited Video"):
-						#video_output.render()
 						video_output = gr.Video(label="Video", elem_id="output_vid", width=512)
 
 				with gr.Row():
 					with gr.Accordion(open=True, label="Original & Edited Videos"):
-						#video_all_output.render()
 						video_all_output = gr.Video(label="Videos", elem_id="output_vid_all")
 
 			with gr.Column(scale=1):