Update src/inference.py

src/inference.py

@@ -17,7 +17,6 @@ from src.utils import (
     tensor_lab2rgb
 )
 import numpy as np
-from tqdm import tqdm
 
 class SwinTExCo:
     def __init__(self, weights_path, swin_backbone='swinv2-cr-t-224', device=None):
@@ -62,13 +61,13 @@ class SwinTExCo:
                                                                 size=(H,W), 
                                                                 mode="bilinear", 
                                                                 align_corners=False)
-        large_IA_l = torch.cat((large_IA_l, large_current_ab_predict.cpu()), dim=1)
+        large_IA_l = torch.cat((large_IA_l, large_current_ab_predict), dim=1)
         large_current_rgb_predict = tensor_lab2rgb(large_IA_l)
-        return large_current_rgb_predict
+        return large_current_rgb_predict.cpu()
     
     def __proccess_sample(self, curr_frame, I_last_lab_predict,  I_reference_lab, features_B):
         large_IA_lab = ToTensor()(RGB2Lab()(curr_frame)).unsqueeze(0)
-        large_IA_l = large_IA_lab[:, 0:1, :, :]
+        large_IA_l = large_IA_lab[:, 0:1, :, :].to(self.device)
         
         IA_lab = self.processor(curr_frame)
         IA_lab = IA_lab.unsqueeze(0).to(self.device)
@@ -113,9 +112,7 @@ class SwinTExCo:
             I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1)).to(self.device)
             features_B = self.embed_net(I_reference_rgb)
         
-        #PBAR = tqdm(total=int(video.get(cv2.CAP_PROP_FRAME_COUNT)), desc="Colorizing video", unit="frame")
         while video.isOpened():
-            #PBAR.update(1)
             ret, curr_frame = video.read()
             
             if not ret:
@@ -130,7 +127,6 @@ class SwinTExCo:
             
             yield IA_predict_rgb
 
-        #PBAR.close()
         video.release()
 
     def predict_image(self, image, ref_image):