Commit 43b5d3f8 authored by E144069X's avatar E144069X

Added preprocessing for marseille

parent 0275e2a9
......@@ -164,7 +164,7 @@ class SeqTrDataset(torch.utils.data.Dataset):
class PreProcess():
def __init__(self,maskTimeOnImage,mask,origImgSize,resizeImage,resizeTorchFunc,sobel=False,augmentData=False,gridShuffle=False,transfFunc=None):
def __init__(self,maskTimeOnImage,mask,origImgSize,resizeImage,resizeTorchFunc,sobel=False,augmentData=False,gridShuffle=False,transfFunc=None,dataset="big"):
self.maskTimeOnImage = maskTimeOnImage
self.origImgSize = origImgSize
......@@ -172,19 +172,34 @@ class PreProcess():
self.resizeTorchFunc = resizeTorchFunc
self.transfFunc = transfFunc
self.applyTransf = augmentData or gridShuffle
if dataset == "marseille":
self.contrastFunc = transforms.Compose([transforms.ToPILImage(),\
transforms.ColorJitter(brightness=(0.35,0.35), contrast=(2,2), saturation=0, hue=0),\
transforms.ToTensor()])
#self.contrastFunc = None
else:
self.contrastFunc = None
self.toTensorFunc = torchvision.transforms.ToTensor()
self.normalizeFunc = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
self.mask = mask
self.sobel = sobel
self.dataset = dataset
def maskTimeOnImageFunc(self,x):
if self.maskTimeOnImage:
if self.maskTimeOnImage and self.dataset != "marseille":
x[-50:] = 0
return x
def removeTopFunc(self,x):
#Removing the top part where the name of the video is written
x = x[:,x.shape[1]-self.origImgSize:,:]
def removeBorderFunc(self,x):
if self.dataset != "marseille":
#Removing the top part where the name of the video is written
x = x[:,x.shape[1]-self.origImgSize:,:]
else:
#x = x[:,x.shape[1]-self.origImgSize:,:]
x = x[:,:,(x.shape[2]-x.shape[1])//2:-(x.shape[2]-x.shape[1])//2]
return x
def resizeFunc(self,x):
......@@ -203,6 +218,11 @@ class PreProcess():
mag[-60:] = 0
return mag[np.newaxis]
def contrast(self,x):
if self.contrastFunc:
x = self.contrastFunc(x)
return x
class TestLoader():
'''
The dataset to sample sequence of frames from videos. As the video contains a great number of frame,
......@@ -247,7 +267,7 @@ class TestLoader():
self.transf = None
self.preproc = PreProcess(self.maskTimeOnImage,self.mask,self.origImgSize,self.resizeImage,self.reSizeTorchFunc,\
sobel=sobel,augmentData=False,gridShuffle=gridShuffle,transfFunc=self.transf)
sobel=sobel,augmentData=False,gridShuffle=gridShuffle,transfFunc=self.transf,dataset=self.dataset)
def __iter__(self):
self.videoInd = 0
......@@ -303,7 +323,7 @@ class TestLoader():
def loadFrames_and_process(frameInds,gt,timeElapsed,vidName,video,preproc):
#Building the frame sequence, remove the top of the video (if required)
frameSeq = np.concatenate(list(map(preproc.removeTopFunc,map(lambda x:video[x][np.newaxis],np.array(frameInds)))),axis=0)
frameSeq = np.concatenate(list(map(preproc.removeBorderFunc,map(lambda x:video[x][np.newaxis],np.array(frameInds)))),axis=0)
#Resize the images (if required) and mask the time (if required)
frameSeq = np.concatenate(list(map(preproc.resizeFunc,map(preproc.maskTimeOnImageFunc,frameSeq))),axis=0)
......@@ -322,8 +342,12 @@ def loadFrames_and_process(frameInds,gt,timeElapsed,vidName,video,preproc):
# T x H x W
frameSeq = frameSeq.unsqueeze(1)
# T x 1 x H x W
frameSeq = torch.cat(list(map(lambda x:preproc.contrast(x).unsqueeze(0),frameSeq)),dim=0)
frameSeq = frameSeq.expand(frameSeq.size(0),3,frameSeq.size(2),frameSeq.size(3))
# T x 3 x H x W
frameSeq = torch.cat(list(map(lambda x:preproc.normalizeFunc(x).unsqueeze(0),frameSeq.float())),dim=0)
timeElapsed = torch.tensor(timeElapsed).float().unsqueeze(0)/maxTime if not timeElapsed is None else None
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment