Add optical flow generation, tensorboard and debugs of train.py

calib_generate_flow.py

@@ -0,0 +1,127 @@
+import sys
+sys.path.append('core')
+
+import os
+import numpy as np
+import torch
+import cv2
+from tqdm import tqdm
+from PIL import Image
+import matplotlib.pyplot as plt
+import argparse
+from pathlib import Path
+
+from utils import flow_viz
+from utils.raft import RAFT
+from utils.utils import InputPadder
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '5'
+DEVICE = 'cuda' if torch.cuda.is_available else 'cpu'
+
+images_dir = "/data_4T/EV_RAFT/calib_frame/4/"
+output_dir = "/data_4T/EV_RAFT/calib_flow/4/"
+
+outfig_dir = '/home/ljw/projects/vehicle-speed-estimation/result/vis_calib/'
+model_path = '/home/ljw/projects/vehicle-speed-estimation/model/raft_models/raft-things.pth'
+videos_path = '/home/ljw/projects/vehicle-speed-estimation/speedchallenge/data/'
+
+
+
+def generate_frames(videos_path, images_dir):
+    #vidcap = cv2.VideoCapture(videos_path+'train.mp4')
+    vidcap = cv2.VideoCapture(videos_path+'4.hevc')  
+
+    frames_rate=vidcap.get(5)
+    frame_num = int(vidcap.get(7))
+    print("frame rate", frames_rate)
+
+    for i in tqdm(range(frame_num)):
+        success,image = vidcap.read()
+        if image.size == 0: 
+            pass
+        else:
+            cv2.imwrite(images_dir+"frame%d.png" % i, image)
+
+
+def vis(npy_dir, output_dir):
+    npy_dir = Path(npy_dir)
+    output_dir = Path(output_dir)
+
+    npy_files = list(npy_dir.glob('*.npy'))
+
+    for i, npy_file in enumerate(npy_files):
+        f = str(npy_file)
+        of = np.load(f)
+        of = torch.from_numpy(of)
+        of = of[0].permute(1,2,0).numpy()
+        of = flow_viz.flow_to_image(of)
+        img = Image.fromarray(of)
+        output_f = output_dir / npy_file.stem
+        output_f = str(output_f) + '.jpg'
+        img.save(output_f)
+
+        if i % 20 == 0: print(f'{i}/{len(npy_files)}')
+
+def load_image(imfile):
+    img = np.array(Image.open(imfile)).astype(np.uint8)
+    img = torch.from_numpy(img).permute(2, 0, 1).float()
+    return img[None].to(DEVICE)
+
+def run(args):
+    model = torch.nn.DataParallel(RAFT(args))
+    model.load_state_dict(torch.load(model_path))
+
+    model = model.module
+    model.to(DEVICE)
+    model.eval()
+
+    output_dir = Path(args.output_dir)
+    images_dir = Path(args.images_dir)
+    images = list(images_dir.glob('frame*.png'))
+
+    with torch.no_grad():
+        # Important bug: the img_path need to be sorted!
+        images.sort(key=lambda x: int(x.name[5:-4]))
+        print(images[1010])
+
+        for i in tqdm(range(len(images)-1)):
+        # run first 200
+        #for i in tqdm(range(200)):
+            im_f1 = str(images[i])
+            im_f2 = str(images[i+1])
+
+            image1 = load_image(im_f1)
+            image2 = load_image(im_f2)
+
+            padder = InputPadder(image1.shape)
+            image1, image2 = padder.pad(image1, image2)
+
+            flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
+
+            # 2.4 MB each npy
+            of_f_name = output_dir / f'{i}.npy' 
+            np.save(of_f_name, flow_up.cpu())
+
+    if args.out_vis:
+        vis(output_dir, outfig_dir)
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--images_dir',
+                        default=os.path.abspath(images_dir),
+                        help="directory with your images")
+    parser.add_argument('--output_dir',
+                        default=os.path.abspath(output_dir),
+                        help="optical flow images will be stored here as .npy files")
+    parser.add_argument('--out_vis',
+                        action='store_true',
+                        help="output optical flow images")
+    args = parser.parse_args()
+    # transfer videos to frames
+    # generate_frames(videos_path,images_dir)
+
+    run(args)
+
+

calib_train.py

@@ -0,0 +1,148 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.sampler import SubsetRandomSampler
+import torch.optim as optim
+from efficientnet_pytorch import EfficientNet
+from pathlib import Path
+import numpy as np
+import multiprocessing
+from torch.utils.tensorboard import SummaryWriter
+
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '6'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+v = 0     # model version
+in_c = 2  # number of input channels
+num_c = 2 # number of classes to predict
+
+
+# directory with the optical flow images
+of_dir = '/data_4T/EV_RAFT/calib_flow/'
+# labels as txt file
+labels_f = '/data_4T/EV_RAFT/speedchallenge/calib_challenge/labeled/'
+model_f = '/home/ljw/projects/vehicle-speed-estimation/model/calib_efnet_b0.pth'
+
+
+class OFDataset(Dataset):
+    '''File 0~3 for training, file 4 for testing'''
+
+    def __init__(self, of_dir, label_file):
+        self.of_dirlist, self.label_dirlist = [], []
+        for i in range(4):
+            self.of_dirlist.append(os.path.join(of_dir, str(i)))
+            self.label_dirlist.append(os.path.join(label_file,f'{i}.txt'))
+        self.len = sum([len(list(Path(self.of_dirlist[i]).glob('*.npy'))) 
+                        for i in range(len(self.of_dirlist))])
+
+        #print(self.len, self.of_dirlist, self.label_dirlist)
+        #self.label_file = open(label_f).readlines()
+    def __len__(self): return self.len
+
+    def __getitem__(self, idx):
+        num = 0
+        while idx >= 0:
+            max_id = len(list(Path(self.of_dirlist[num]).glob('*.npy')))-1
+            if idx <= max_id:
+                of_array = np.load(Path(self.of_dirlist[num])/f'{idx}.npy')
+                #print(Path(self.of_dirlist[num])/f'{idx}.npy')
+                of_tensor = torch.squeeze(torch.Tensor(of_array))
+                label_file = open(self.label_dirlist[num]).readlines()
+                #print(self.label_dirlist[num], idx)
+                label = [float(i) for i in label_file[idx].split()[:2]]
+                label = 100 * torch.tensor(label)
+                return [of_tensor, label]
+            else:
+                idx = idx - max_id - 1
+            num += 1
+
+ds = OFDataset(of_dir, labels_f)
+print(torch.isnan(torch.mean(ds[784+2400][1])))         #torch.Size([2, 880, 1168])
+
+
+#of = torch.randn(1,2,640,480)  # input shape (1,2,640,480)
+
+model = EfficientNet.from_pretrained(f'efficientnet-b{v}', in_channels=in_c, num_classes=num_c)
+#state = torch.load(MODEL_F)
+#model.load_state_dict(state)
+model.to(device)
+
+#of = of.to(device)
+#model(of).item()
+
+
+# 80% of data for training
+# 20% of data for validation
+train_split = .9
+
+ds_size = len(ds)
+indices = list(range(ds_size))
+split = int(np.floor(train_split * ds_size))
+train_idx, val_idx = indices[:split], indices[split:]
+#sample = ds[3]
+
+
+train_set, val_set = torch.utils.data.random_split(ds,[split,ds_size-split])
+train_dl = DataLoader(train_set, batch_size=4, shuffle=True, num_workers=12)
+val_dl = DataLoader(val_set, batch_size=4, shuffle=False, num_workers=12)
+
+print(len(train_dl), len(val_dl))
+
+
+
+epochs = 100
+#log_train_steps = 100
+
+writer = SummaryWriter('/home/ljw/projects/vehicle-speed-estimation/result/calib')
+
+
+criterion = nn.MSELoss()
+opt = optim.Adam(model.parameters(),lr=0.0001)
+
+for epoch in range(epochs):
+    model.train()
+    running_loss = 0.0
+    for i, sample in enumerate(train_dl):
+        if torch.isnan(torch.mean(sample[1])):
+            #print(i)
+            continue
+
+        of_tensor = sample[0].cuda()
+        label = sample[1].cuda()
+
+        opt.zero_grad()
+        pred = torch.squeeze(model(of_tensor))
+        loss = criterion(pred, label)
+        loss.backward()
+        opt.step()
+        if (i+1) % 100 == 0:
+            print('Epoch: [{}/{}], Step: [{}/{}], Loss: {}'
+                .format(epoch+1, epochs, i+1, len(train_dl), loss.item()))
+            writer.add_scalar('Training loss', loss.item(), global_step=i+epoch*len(train_dl))
+
+    # validation
+    model.eval()
+    val_losses = []
+    with torch.no_grad():
+        for j, val_sample in enumerate(val_dl):
+            if torch.isnan(torch.mean(val_sample[1])):
+                #print(i)
+                continue
+            of_tensor = val_sample[0].cuda()
+            label = val_sample[1].float().cuda()
+            pred = torch.squeeze(model(of_tensor))
+            loss = criterion(pred, label)
+            val_losses.append(loss)
+        print('Validation: Epoch: [{}/{}], mean Loss: {}, last loss:{}'
+                .format(epoch+1, epochs, sum(val_losses)/len(val_losses), loss.item()))
+        #print(f'{epoch}: {sum(val_losses)/len(val_losses)}')
+        writer.add_scalar('Validation loss', sum(val_losses)/len(val_losses), global_step=epoch)
+
+# test
+
+
+
+torch.save(model, model_f)