c

testVideo.py

@@ -107,6 +107,10 @@ def img2Comix(args,net,device,image):
         image = to_tensor(image).unsqueeze(0) * 2 - 1
         out = net(image.to(device), args.upsample_align).cpu()
         out = out.squeeze(0).clip(-1, 1) * 0.5 + 0.5
+
+        out = (out).numpy()
+        # switch channels -> (c, h, w) -> (h, w, c)
+        out = np.rollaxis(out, 0, 3)
         #out = to_pil_image(out)
     return out
 
@@ -183,10 +187,10 @@ def video2comix(args):
                 #style_frame.save(os.path.join(args.output_dir, imgname + str(i) + "_c.jpg"))
                 save(style_frame, args.output_dir, imgname + str(i) + "_c")
 
-            #img = style_frame.astype('uint8')
+            img = style_frame.astype('uint8')
             #print('---------style_frame type=' , type(style_frame))
             print('----------one frame time:',stylization_time,', style_frame type:',type(style_frame))
-            img = np.asarray(style_frame)
+            #img = np.asarray(style_frame)
             flow_video.write(img)
             # flow_video.write(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
             # cv2.waitKey(1)