Upload 9 files

data_processing/__init__.py

@@ -0,0 +1,2 @@
+# __init__.py
+# print("Data processing")

data_processing/split_train_test.py

@@ -0,0 +1,168 @@
+import pandas as pd
+import numpy as np
+import os
+import torch
+from sklearn.model_selection import train_test_split
+import logging
+
+# cross_dataset
+def process_cross_dataset(train_data_name, test_data_name, metadata_path, feature_path, network_name):
+    metadata_name1 = f"{train_data_name.replace('_all', '').upper()}_metadata.csv"
+    metadata_name2 = f"{test_data_name.replace('_all', '').upper()}_metadata.csv"
+    # load CSV data
+    train_df = pd.read_csv(f'{metadata_path}/{metadata_name1}')
+    test_df = pd.read_csv(f'{metadata_path}/{metadata_name2}')
+
+    # split videonames into train and test sets
+    train_vids = train_df.iloc[:, 0]
+    test_vids = test_df.iloc[:, 0]
+
+    # scores (1-100) map to 1-5
+    train_scores = train_df['mos'].tolist()
+    test_scores = test_df['mos'].tolist()
+    if train_data_name == 'konvid_1k_all' or train_data_name == 'youtube_ugc_all':
+        train_mos_list = ((np.array(train_scores) - 1) * (99/4) + 1.0).tolist()
+    else:
+        train_mos_list = train_scores
+    if test_data_name == 'konvid_1k_all' or test_data_name == 'youtube_ugc_all':
+        test_mos_list = ((np.array(test_scores) - 1) * (99/4) + 1.0).tolist()
+    else:
+        test_mos_list = test_scores
+
+    # reorder columns
+    sorted_train_df = pd.DataFrame({'vid': train_df['vid'],  'framerate': train_df['framerate'], 'MOS': train_mos_list, 'MOS_raw': train_df['mos']})
+    sorted_test_df = pd.DataFrame({'vid': test_df['vid'], 'framerate': test_df['framerate'], 'MOS': test_mos_list, 'MOS_raw': test_df['mos']})
+
+    # use indices from the train and test DataFrames to split features
+    train_features = torch.load(f"{feature_path}/{network_name}_{train_data_name.replace('_all', '')}_features.pt")
+    test_features = torch.load(f"{feature_path}/{network_name}_{test_data_name.replace('_all', '')}_features.pt")
+
+    # save the files
+    sorted_train_df.to_csv(f'{metadata_path}mos_files/{train_data_name}_MOS_train.csv', index=False)
+    sorted_test_df.to_csv(f'{metadata_path}mos_files/{test_data_name}_MOS_test.csv', index=False)
+    os.makedirs(os.path.join(feature_path, "split_train_test"), exist_ok=True)
+    torch.save(train_features, f'{feature_path}/split_train_test/{network_name}_{train_data_name}_cross_train_features.pt')
+    torch.save(test_features, f'{feature_path}/split_train_test/{network_name}_{test_data_name}_cross_test_features.pt')
+
+    return train_features, test_features, test_vids
+
+#NR: original
+def process_lsvq(train_data_name, test_data_name, metadata_path, feature_path, network_name):
+    train_df = pd.read_csv(f'{metadata_path}/{train_data_name.upper()}_metadata.csv')
+    test_df = pd.read_csv(f'{metadata_path}/{test_data_name.upper()}_metadata.csv')
+
+    # grayscale videos, do not consider them for fair comparison
+    # grey_df_train = pd.read_csv(f'{metadata_path}/greyscale_report/{train_data_name.upper()}_greyscale_metadata.csv')
+    # grey_df_test = pd.read_csv(f'{metadata_path}/greyscale_report/{test_data_name.upper()}_greyscale_metadata.csv')
+    # grey_indices_train = grey_df_train.iloc[:, 0].tolist()
+    # grey_indices_test = grey_df_test.iloc[:, 0].tolist()
+    # train_df = train_df.drop(index=grey_indices_train).reset_index(drop=True)
+    # test_df = test_df.drop(index=grey_indices_test).reset_index(drop=True)
+    test_vids = test_df['vid']
+
+    # mos scores
+    train_scores = train_df['mos'].tolist()
+    test_scores = test_df['mos'].tolist()
+    train_mos_list = train_scores
+    test_mos_list = test_scores
+
+    # reorder columns
+    sorted_train_df = pd.DataFrame({'vid': train_df['vid'],  'framerate': train_df['framerate'], 'MOS': train_mos_list, 'MOS_raw': train_df['mos']})
+    sorted_test_df = pd.DataFrame({'vid': test_df['vid'], 'framerate': test_df['framerate'], 'MOS': test_mos_list, 'MOS_raw': test_df['mos']})
+
+    # use indices from the train and test DataFrames to split features
+    train_features = torch.load(f'{feature_path}/{network_name}_{train_data_name}_features.pt')
+    print(f"loaded {train_data_name}: dimensions are {train_features.shape}")
+    test_features = torch.load(f'{feature_path}/{network_name}_{test_data_name}_features.pt')
+
+    # grayscale videos
+    # train_mask = torch.ones(train_features.size(0), dtype=torch.bool, device=train_features.device)
+    # test_mask = torch.ones(test_features.size(0), dtype=torch.bool, device=test_features.device)
+    # train_mask[grey_indices_train] = False
+    # test_mask[grey_indices_test] = False
+    # train_features = train_features[train_mask]
+    # test_features = test_features[test_mask]
+    print(len(train_features))
+    print(len(test_features))
+
+    # save the files
+    sorted_train_df.to_csv(f'{metadata_path}mos_files/{train_data_name}_MOS_train.csv', index=False)
+    sorted_test_df.to_csv(f'{metadata_path}mos_files/{train_data_name}_MOS_test.csv', index=False)
+    os.makedirs(os.path.join(feature_path, "split_train_test"), exist_ok=True)
+    torch.save(train_features, f'{feature_path}/split_train_test/{network_name}_{train_data_name}_train_features.pt')
+    torch.save(test_features, f'{feature_path}/split_train_test/{network_name}_{test_data_name}_test_features.pt')
+
+    return train_features, test_features, test_vids
+
+def process_other(data_name, test_size, random_state, metadata_path, feature_path, network_name):
+    metadata_name = f'{data_name.upper()}_metadata.csv'
+    # load CSV data
+    df = pd.read_csv(f'{metadata_path}/{metadata_name}')
+
+    # if data_name == 'youtube_ugc':
+    #     # grayscale videos, do not consider them for fair comparison
+    #     grey_df = pd.read_csv(f'{metadata_path}/greyscale_report/{data_name.upper()}_greyscale_metadata.csv')
+    #     grey_indices = grey_df.iloc[:, 0].tolist()
+    #     df = df.drop(index=grey_indices).reset_index(drop=True)
+
+    # get unique vids
+    unique_vids = df['vid'].unique()
+
+    # split videonames into train and test sets
+    train_vids, test_vids = train_test_split(unique_vids, test_size=test_size, random_state=random_state)
+
+    # split all_dfs into train and test based on vids
+    train_df = df[df['vid'].isin(train_vids)]
+    test_df = df[df['vid'].isin(test_vids)]
+
+    # mos scores
+    train_scores = train_df['mos'].tolist()
+    test_scores = test_df['mos'].tolist()
+    train_mos_list = train_scores
+    test_mos_list = test_scores
+
+    # reorder columns
+    sorted_train_df = pd.DataFrame({'vid': train_df['vid'],  'framerate': train_df['framerate'], 'MOS': train_mos_list, 'MOS_raw': train_df['mos']})
+    sorted_test_df = pd.DataFrame({'vid': test_df['vid'], 'framerate': test_df['framerate'], 'MOS': test_mos_list, 'MOS_raw': test_df['mos']})
+
+    # use indices from the train and test DataFrames to split features
+    features = torch.load(f'{feature_path}/{network_name}_{data_name}_features.pt')
+    # if data_name == 'youtube_ugc':
+    #     # features = np.delete(features, grey_indices, axis=0)
+    #     mask = torch.ones(features.size(0), dtype=torch.bool, device=features.device)
+    #     mask[grey_indices] = False
+    #     features = features[mask]
+    train_features = features[train_df.index]
+    test_features = features[test_df.index]
+
+    # save the files
+    sorted_train_df.to_csv(f'{metadata_path}mos_files/{data_name}_MOS_train.csv', index=False)
+    sorted_test_df.to_csv(f'{metadata_path}mos_files/{data_name}_MOS_test.csv', index=False)
+    os.makedirs(os.path.join(feature_path, "split_train_test"), exist_ok=True)
+    torch.save(train_features, f'{feature_path}/split_train_test/{network_name}_{data_name}_train_features.pt')
+    torch.save(test_features, f'{feature_path}/split_train_test/{network_name}_{data_name}_test_features.pt')
+
+    return train_features, test_features, test_vids
+
+
+if __name__ == '__main__':
+    network_name = "slowfast"
+    data_name = "test"
+    metadata_path = '../../metadata/'
+    feature_path = '../../features/konvid_1k_test/slowfast/'
+
+    # train test split
+    test_size = 0.2
+    random_state = None
+
+    if data_name == 'lsvq_train':
+        test_data_name = 'lsvq_test'
+        process_lsvq(data_name, test_data_name, metadata_path, feature_path, network_name)
+
+    elif data_name == 'cross_dataset':
+        train_data_name = 'youtube_ugc_all'
+        test_data_name = 'cvd_2014_all'
+        _, _, test_vids = process_cross_dataset(train_data_name, test_data_name, metadata_path, feature_path, network_name)
+
+    else:
+        process_other(data_name, test_size, random_state, metadata_path, feature_path, network_name)

extractor/__init__.py

@@ -0,0 +1,2 @@
+# __init__.py
+# print("Initializing extractor")

extractor/extract_rf_feats.py

@@ -0,0 +1,317 @@
+import os
+import cv2
+import pandas as pd
+import time
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+from torchvision import transforms
+from PIL import Image
+from torch.utils import data
+
+
+class VideoDataset_feature(data.Dataset):
+    def __init__(self, data_dir, filename_path, transform, resize, database, patch_size=16, target_size=224, top_n=196):
+        super(VideoDataset_feature, self).__init__()
+        if isinstance(filename_path, str):
+            dataInfo = pd.read_csv(filename_path)
+        elif isinstance(filename_path, pd.DataFrame):
+            dataInfo = filename_path
+        else:
+            raise ValueError("filename_path: CSV file or DataFrame")
+        self.video_names = dataInfo['vid'].tolist()
+        self.transform = transform
+        self.videos_dir = data_dir
+        self.resize = resize
+        self.database = database
+        self.patch_size = patch_size
+        self.target_size = target_size
+        self.top_n = top_n
+        self.length = len(self.video_names)
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, idx):
+        if self.database == 'konvid_1k' or self.database == 'test':
+            video_clip_min = 8
+            video_name = str(self.video_names[idx]) + '.mp4'
+        elif self.database == 'live_vqc':
+            video_clip_min = 10
+            video_name = str(self.video_names[idx]) + '.mp4'
+        elif self.database == 'cvd_2014':
+            video_clip_min = 12
+            video_name = str(self.video_names[idx]) + '.avi'
+        elif self.database == 'youtube_ugc':
+            video_clip_min = 20
+            video_name = str(self.video_names[idx]) + '.mkv'
+        elif self.database == 'youtube_ugc_h264':
+            video_clip_min = 20
+            video_name = str(self.video_names[idx]) + '.mp4'
+        elif self.database == 'lsvq_test_1080p' or self.database == 'lsvq_test' or self.database == 'lsvq_train':
+            video_clip_min = 8
+            video_name = str(self.video_names[idx]) + '.mp4'
+
+        filename = os.path.join(self.videos_dir, video_name)
+
+        video_capture = cv2.VideoCapture(filename)
+        video_capture.set(cv2.CAP_PROP_BUFFERSIZE, 3)
+        if not video_capture.isOpened():
+            raise RuntimeError(f"Failed to open video: {filename}")
+
+        video_channel = 3
+        video_length = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
+        video_frame_rate = int(round(video_capture.get(cv2.CAP_PROP_FPS)))
+        video_clip = int(video_length / video_frame_rate) if video_frame_rate != 0 else 10
+        video_length_clip = 32
+        # print(video_length)
+        # print(video_frame_rate)
+        # print(video_clip)
+
+        all_frame_tensor = torch.zeros((video_length, video_channel, self.resize, self.resize), dtype=torch.float32)
+        all_residual_frag_tensor = torch.zeros((video_length - 1, video_channel, self.resize, self.resize), dtype=torch.float32)
+        all_frame_frag_tensor = torch.zeros((video_length - 1, video_channel, self.resize, self.resize), dtype=torch.float32)
+
+        video_read_index = 0
+        prev_frame = None
+        for i in range(video_length):
+            has_frames, frame = video_capture.read()
+            if has_frames:
+                curr_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                # save_img(curr_frame, fig_path='original', img_title=f'original_{i}')
+
+                # frame features
+                curr_frame_tensor = self.transform(Image.fromarray(curr_frame))
+                all_frame_tensor[video_read_index] = curr_frame_tensor
+
+                # frame frag features
+                if prev_frame is not None:
+                    residual = cv2.absdiff(curr_frame, prev_frame)
+                    # save_img(residual, fig_path='residual', img_title=f'residual_{i}')
+
+                    diff = self.get_patch_diff(residual)
+                    # frame residual fragment
+                    imp_patches, positions = self.extract_important_patches(residual, diff)
+                    imp_patches_pil = Image.fromarray(imp_patches.astype('uint8'))
+                    # save_img(imp_patches_pil, fig_path='residual_frag', img_title=f'residual_frag_{i}')
+
+                    residual_frag_tensor = self.transform(imp_patches_pil)
+                    all_residual_frag_tensor[video_read_index] = residual_frag_tensor
+
+                    # current frame fragment
+                    ori_patches = self.get_original_frame_patches(curr_frame, positions)
+                    ori_patches_pil = Image.fromarray(ori_patches.astype('uint8'))
+                    # save_img(ori_patches_pil, fig_path='ori_frag', img_title=f'ori_frag_{i}')
+
+                    frame_frag_tensor = self.transform(ori_patches_pil)
+                    all_frame_frag_tensor[video_read_index] = frame_frag_tensor
+
+                    video_read_index += 1
+            prev_frame = curr_frame
+        video_capture.release()
+        # visualisation
+        visualise_image(curr_frame, 'Current Frame')
+        visualise_image(imp_patches_pil, 'Residual Fragment')
+        visualise_image(ori_patches_pil, 'Frame Fragment')
+
+        # Unfilled frames
+        self.fill_tensor(all_frame_tensor, video_read_index, video_length)
+        self.fill_tensor(all_residual_frag_tensor, video_read_index, video_length - 1)
+        self.fill_tensor(all_frame_frag_tensor, video_read_index, video_length - 1)
+
+        video_all = []
+        video_res_frag_all = []
+        video_frag_all = []
+        for i in range(video_clip):
+            clip_tensor = torch.zeros([video_length_clip, video_channel, self.resize, self.resize])
+            clip_res_frag_tensor = torch.zeros([video_length_clip, video_channel, self.resize, self.resize])
+            clip_frag_tensor = torch.zeros([video_length_clip, video_channel, self.resize, self.resize])
+
+            start_idx = i * video_frame_rate
+            end_idx = start_idx + video_length_clip
+            # frame features
+            if end_idx <= video_length:
+                clip_tensor = all_frame_tensor[start_idx:end_idx]
+            else:
+                clip_tensor[:(video_length - start_idx)] = all_frame_tensor[start_idx:]
+                clip_tensor[(video_length - start_idx):video_length_clip] = clip_tensor[video_length - start_idx - 1]
+
+            # frame frag features
+            if end_idx <= (video_length - 1):
+                clip_res_frag_tensor = all_residual_frag_tensor[start_idx:end_idx]
+                clip_frag_tensor = all_frame_frag_tensor[start_idx:end_idx]
+            else:
+                clip_res_frag_tensor[:(video_length - 1 - start_idx)] = all_residual_frag_tensor[start_idx:]
+                clip_frag_tensor[:(video_length - 1 - start_idx)] = all_frame_frag_tensor[start_idx:]
+                clip_res_frag_tensor[(video_length - 1 - start_idx):video_length_clip] = clip_res_frag_tensor[video_length - 1 - start_idx - 1]
+                clip_frag_tensor[(video_length - 1 - start_idx):video_length_clip] = clip_frag_tensor[video_length - 1 - start_idx - 1]
+
+            video_all.append(clip_tensor)
+            video_res_frag_all.append(clip_res_frag_tensor)
+            video_frag_all.append(clip_frag_tensor)
+
+        # Underfilling of clips
+        if video_clip < video_clip_min:
+            for i in range(video_clip, video_clip_min):
+                video_all.append(video_all[video_clip - 1])
+                video_res_frag_all.append(video_res_frag_all[video_clip - 1])
+                video_frag_all.append(video_frag_all[video_clip - 1])
+        return video_all, video_res_frag_all, video_frag_all, video_name
+
+    @staticmethod
+    # duplicat the final frames
+    def fill_tensor(tensor, read_index, length):
+        if read_index < length:
+            tensor[read_index:length] = tensor[read_index - 1]
+
+    def get_patch_diff(self, residual_frame):
+        h, w = residual_frame.shape[:2]
+        patch_size = self.patch_size
+        h_adj = (h // patch_size) * patch_size
+        w_adj = (w // patch_size) * patch_size
+        residual_frame_adj = residual_frame[:h_adj, :w_adj]
+        # calculate absolute patch difference
+        diff = np.zeros((h_adj // patch_size, w_adj // patch_size))
+        for i in range(0, h_adj, patch_size):
+            for j in range(0, w_adj, patch_size):
+                patch = residual_frame_adj[i:i+patch_size, j:j+patch_size]
+                # absolute sum
+                diff[i // patch_size, j // patch_size] = np.sum(np.abs(patch))
+        return diff
+
+    def extract_important_patches(self, residual_frame, diff):
+        patch_size = self.patch_size
+        target_size = self.target_size
+        top_n = self.top_n
+
+        # find top n patches indices
+        patch_idx = np.unravel_index(np.argsort(-diff.ravel()), diff.shape)
+        top_patches = list(zip(patch_idx[0][:top_n], patch_idx[1][:top_n]))
+        sorted_idx = sorted(top_patches, key=lambda x: (x[0], x[1]))
+
+        imp_patches_img = np.zeros((target_size, target_size, residual_frame.shape[2]), dtype=residual_frame.dtype)
+        patches_per_row = target_size // patch_size  # 14
+        # order the patch in the original location relation
+        positions = []
+        for idx, (y, x) in enumerate(sorted_idx):
+            patch = residual_frame[y * patch_size:(y + 1) * patch_size, x * patch_size:(x + 1) * patch_size]
+            # new patch location
+            row_idx = idx // patches_per_row
+            col_idx = idx % patches_per_row
+            start_y = row_idx * patch_size
+            start_x = col_idx * patch_size
+            imp_patches_img[start_y:start_y + patch_size, start_x:start_x + patch_size] = patch
+            positions.append((y, x))
+        return imp_patches_img, positions
+
+    def get_original_frame_patches(self, original_frame, positions):
+        patch_size = self.patch_size
+        target_size = self.target_size
+        imp_original_patches_img = np.zeros((target_size, target_size, original_frame.shape[2]), dtype=original_frame.dtype)
+        patches_per_row = target_size // patch_size
+
+        for idx, (y, x) in enumerate(positions):
+            start_y = y * patch_size
+            start_x = x * patch_size
+            end_y = start_y + patch_size
+            end_x = start_x + patch_size
+
+            patch = original_frame[start_y:end_y, start_x:end_x]
+            row_idx = idx // patches_per_row
+            col_idx = idx % patches_per_row
+            target_start_y = row_idx * patch_size
+            target_start_x = col_idx * patch_size
+
+            imp_original_patches_img[target_start_y:target_start_y + patch_size,
+                                     target_start_x:target_start_x + patch_size] = patch
+        return imp_original_patches_img
+
+def visualise_tensor(tensors, num_frames_to_visualise=5, img_title='Frag'):
+    np_feat = tensors.numpy()
+    fig, axes = plt.subplots(1, num_frames_to_visualise, figsize=(15, 5))
+    for i in range(num_frames_to_visualise):
+        # move channels to last dimension for visualisation: (height, width, channels)
+        frame = np_feat[i].transpose(1, 2, 0)
+        # normalize to [0, 1] for visualisation
+        frame = (frame - frame.min()) / (frame.max() - frame.min())
+        axes[i].imshow(frame)
+        axes[i].axis('off')
+        axes[i].set_title(f'{img_title} {i + 1}')
+
+    plt.tight_layout()
+    save_path =  f'../../figs/{img_title}.png'
+    plt.savefig(save_path, dpi=300)
+    plt.show()
+
+def visualise_image(frame, img_title='Residual Fragment', debug=False):
+    if debug:
+        plt.figure(figsize=(5, 5))
+        plt.imshow(frame)
+        plt.axis('off')
+        plt.title(img_title)
+        plt.show()
+
+def save_img(frame, fig_path, img_title):
+    from torchvision.transforms import ToPILImage
+    save_path = f'../../figs/{fig_path}/{img_title}.png'
+    if isinstance(frame, torch.Tensor):
+        if frame.dim() == 3 and frame.size(0) in [1, 3]:
+            frame = ToPILImage()(frame)
+        else:
+            raise ValueError("Unsupported tensor shape. Expected shape (C, H, W) with C=1 or C=3.")
+
+    if save_path:
+        if isinstance(frame, torch.Tensor):
+            frame = ToPILImage()(frame)
+        elif isinstance(frame, np.ndarray):
+            frame = Image.fromarray(frame)
+        frame.save(save_path)
+        print(f"Image saved to {save_path}")
+
+
+if __name__ == "__main__":
+    database = 'konvid_1k'
+    videos_dir = '../../ugc_original_videos/'
+    metadata_csv = '../../metadata/TEST_metadata.csv'
+    # videos_dir = '/home/xinyi/video_dataset/KoNViD_1k/KoNViD_1k_videos/'
+    # videos_dir = '/media/on23019/server/LSVQ/'
+    # metadata_csv = f'../../metadata/{database.upper()}_metadata.csv'
+
+    resize = 224 # 224, 384
+    start_time = time.time()
+    resize_transform = transforms.Compose([transforms.Resize([resize, resize]),
+                                           transforms.ToTensor(),
+                                           transforms.Normalize(mean=[0.45, 0.45, 0.45], std=[0.225, 0.225, 0.225])])
+
+    dataset = VideoDataset_feature(
+        data_dir=videos_dir,
+        filename_path=metadata_csv,
+        transform=resize_transform,
+        resize=resize,
+        database=database,
+        patch_size=16, # 8, 16, 32, 16, 32
+        target_size=224, # 224, 224, 224, 384, 384
+        top_n=14*14 # 28*28, 14*14, 7*7, 24*24, 12*12
+    )
+
+    # test
+    index = 0
+    video_segments, video_res_frag_all, video_frag_all, video_name = dataset[index]
+    print(f"Video Name: {video_name}")
+    print(f"Number of Video Segments: {len(video_segments)}")
+    print(f"Number of Video Residual Fragment Segments: {len(video_res_frag_all)}")
+    print(f"Number of Video Fragment Segments: {len(video_frag_all)}")
+    print(f"Shape of Each Segment: {video_segments[0].shape}")  # (video_length_clip, channels, height, width)
+    print(f"Shape of Each Residual Fragment Segments: {video_res_frag_all[0].shape}")
+    print(f"Shape of Each Fragment Segments: {video_frag_all[0].shape}")
+
+    # visualisation
+    first_segments = video_segments[0]
+    visualise_tensor(first_segments, num_frames_to_visualise=5, img_title='Frame')
+
+    first_segment_residuals = video_res_frag_all[0]
+    visualise_tensor(first_segment_residuals, num_frames_to_visualise=6, img_title='Residual Frag')
+
+    first_segment_fragments = video_frag_all[0]
+    visualise_tensor(first_segment_fragments, num_frames_to_visualise=5, img_title='Frame Frag')
+    print(f"Processed {time.time() - start_time:.2f} seconds")

extractor/extract_rf_subsampling.py

@@ -0,0 +1,301 @@
+import os
+import cv2
+import pandas as pd
+import time
+import torch
+import matplotlib.pyplot as plt
+import math
+import numpy as np
+from torchvision import transforms
+from PIL import Image
+from torch.utils import data
+
+
+class VideoDataset_feature(data.Dataset):
+    def __init__(self, data_dir, filename_path, transform, resize, database, patch_size=16, target_size=224, top_n=196):
+        super(VideoDataset_feature, self).__init__()
+        dataInfo = pd.read_csv(filename_path)
+        self.video_names = dataInfo['vid'].tolist()
+        self.transform = transform
+        self.videos_dir = data_dir
+        self.resize = resize
+        self.database = database
+        self.patch_size = patch_size
+        self.target_size = target_size
+        self.top_n = top_n
+        self.length = len(self.video_names)
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, idx):
+        if self.database == 'konvid_1k' or self.database == 'test':
+            video_clip_min = 4
+            video_name = str(self.video_names[idx]) + '.mp4'
+        elif self.database == 'live_vqc':
+            video_clip_min = 5
+            video_name = str(self.video_names[idx]) + '.mp4'
+        elif self.database == 'cvd_2014':
+            video_clip_min = 6
+            video_name = str(self.video_names[idx]) + '.avi'
+        elif self.database == 'youtube_ugc':
+            video_clip_min = 10
+            video_name = str(self.video_names[idx]) + '.mkv'
+        elif self.database == 'youtube_ugc_h264':
+            video_clip_min = 10
+            video_name = str(self.video_names[idx]) + '.mp4'
+        elif self.database == 'lsvq_test_1080p' or self.database == 'lsvq_test' or self.database == 'lsvq_train':
+            video_clip_min = 4
+            video_name = str(self.video_names[idx]) + '.mp4'
+
+        filename = os.path.join(self.videos_dir, video_name)
+
+        video_capture = cv2.VideoCapture(filename)
+        video_capture.set(cv2.CAP_PROP_BUFFERSIZE, 3)
+        if not video_capture.isOpened():
+            raise RuntimeError(f"Failed to open video: {filename}")
+
+        video_channel = 3
+        video_length = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
+        video_frame_rate = int(round(video_capture.get(cv2.CAP_PROP_FPS)))
+        video_clip = int(video_length / video_frame_rate) if video_frame_rate != 0 else 10
+        video_length_clip = 32
+        # print(video_length)
+        # print(video_frame_rate)
+        # print(video_clip)
+
+        # reduce FPS by half: every other frame
+        sub_video_length = math.ceil(video_length / 2)
+        all_frame_tensor = torch.zeros((sub_video_length, video_channel, self.resize, self.resize), dtype=torch.float32)
+        all_residual_frag_tensor = torch.zeros((sub_video_length - 1 , video_channel, self.resize, self.resize), dtype=torch.float32)
+        all_frame_frag_tensor = torch.zeros((sub_video_length - 1, video_channel, self.resize, self.resize), dtype=torch.float32)
+
+        video_read_index = 0
+        prev_frame = None
+        for i in range(0, video_length, 2):  # step by 2 to reduce FPS
+            video_capture.set(cv2.CAP_PROP_POS_FRAMES, i)  # set frame position
+            has_frames, frame = video_capture.read()
+            if has_frames:
+                curr_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+                # frame features
+                curr_frame_tensor = self.transform(Image.fromarray(curr_frame))
+                all_frame_tensor[video_read_index] = curr_frame_tensor
+
+                # frame frag features
+                if prev_frame is not None:
+                    residual = cv2.absdiff(curr_frame, prev_frame)
+                    diff = self.get_patch_diff(residual)
+                    # frame residual fragment
+                    imp_patches, positions = self.extract_important_patches(residual, diff)
+                    imp_patches_pil = Image.fromarray(imp_patches.astype('uint8'))
+                    residual_frag_tensor = self.transform(imp_patches_pil)
+                    all_residual_frag_tensor[video_read_index] = residual_frag_tensor
+
+                    # current frame fragment
+                    ori_patches = self.get_original_frame_patches(curr_frame, positions)
+                    ori_patches_pil = Image.fromarray(ori_patches.astype('uint8'))
+                    frame_frag_tensor = self.transform(ori_patches_pil)
+                    all_frame_frag_tensor[video_read_index] = frame_frag_tensor
+
+                    video_read_index += 1
+            prev_frame = curr_frame
+        video_capture.release()
+        # visualisation
+        visualise_image(curr_frame, 'Current Frame')
+        visualise_image(imp_patches_pil, 'Residual Fragment')
+        visualise_image(ori_patches_pil, 'Frame Fragment')
+
+        # Unfilled frames
+        self.fill_tensor(all_frame_tensor, video_read_index, sub_video_length)
+        self.fill_tensor(all_residual_frag_tensor, video_read_index, sub_video_length - 1)
+        self.fill_tensor(all_frame_frag_tensor, video_read_index, sub_video_length - 1)
+
+        # update video length
+        sub_video_length = len(all_frame_tensor)
+        sub_video_frame_rate = int(round(video_frame_rate/2))
+        sub_video_clip = int(video_clip/2)
+        # print(sub_video_length)
+        # print(sub_video_frame_rate)
+        # print(sub_video_clip)
+
+
+        video_all = []
+        video_res_frag_all = []
+        video_frag_all = []
+        for i in range(sub_video_clip):
+            clip_tensor = torch.zeros([video_length_clip, video_channel, self.resize, self.resize])
+            clip_res_frag_tensor = torch.zeros([video_length_clip, video_channel, self.resize, self.resize])
+            clip_frag_tensor = torch.zeros([video_length_clip, video_channel, self.resize, self.resize])
+
+            start_idx = i * sub_video_frame_rate * 2
+            end_idx = start_idx + video_length_clip
+            # frame features
+            if end_idx <= sub_video_length:
+                clip_tensor = all_frame_tensor[start_idx:end_idx]
+            else:
+                clip_tensor[:(sub_video_length - start_idx)] = all_frame_tensor[start_idx:]
+                clip_tensor[(sub_video_length - start_idx):video_length_clip] = clip_tensor[sub_video_length - start_idx - 1]
+
+            # frame frag features
+            if end_idx <= (sub_video_length - 1):
+                clip_res_frag_tensor = all_residual_frag_tensor[start_idx:end_idx]
+                clip_frag_tensor = all_frame_frag_tensor[start_idx:end_idx]
+            else:
+                clip_res_frag_tensor[:(sub_video_length - 1 - start_idx)] = all_residual_frag_tensor[start_idx:]
+                clip_frag_tensor[:(sub_video_length - 1 - start_idx)] = all_frame_frag_tensor[start_idx:]
+                clip_res_frag_tensor[(sub_video_length - 1 - start_idx):video_length_clip] = clip_res_frag_tensor[sub_video_length - 1 - start_idx - 1]
+                clip_frag_tensor[(sub_video_length - 1 - start_idx):video_length_clip] = clip_frag_tensor[sub_video_length - 1 - start_idx - 1]
+
+            video_all.append(clip_tensor)
+            video_res_frag_all.append(clip_res_frag_tensor)
+            video_frag_all.append(clip_frag_tensor)
+
+        # Underfilling of clips
+        if sub_video_clip < video_clip_min:
+            for i in range(video_clip, video_clip_min):
+                video_all.append(video_all[sub_video_clip - 1])
+                video_res_frag_all.append(video_res_frag_all[sub_video_clip - 1])
+                video_frag_all.append(video_frag_all[sub_video_clip - 1])
+        return video_all, video_res_frag_all, video_frag_all, video_name
+
+    @staticmethod
+    # duplicat the final frames
+    def fill_tensor(tensor, read_index, length):
+        if read_index < length:
+            tensor[read_index:length] = tensor[read_index - 1]
+
+    def get_patch_diff(self, residual_frame):
+        h, w = residual_frame.shape[:2]
+        patch_size = self.patch_size
+        h_adj = (h // patch_size) * patch_size
+        w_adj = (w // patch_size) * patch_size
+        residual_frame_adj = residual_frame[:h_adj, :w_adj]
+        # calculate absolute patch difference
+        diff = np.zeros((h_adj // patch_size, w_adj // patch_size))
+        for i in range(0, h_adj, patch_size):
+            for j in range(0, w_adj, patch_size):
+                patch = residual_frame_adj[i:i+patch_size, j:j+patch_size]
+                # absolute sum
+                diff[i // patch_size, j // patch_size] = np.sum(np.abs(patch))
+        return diff
+
+    def extract_important_patches(self, residual_frame, diff):
+        patch_size = self.patch_size
+        target_size = self.target_size
+        top_n = self.top_n
+
+        # find top n patches indices
+        patch_idx = np.unravel_index(np.argsort(-diff.ravel()), diff.shape)
+        top_patches = list(zip(patch_idx[0][:top_n], patch_idx[1][:top_n]))
+        sorted_idx = sorted(top_patches, key=lambda x: (x[0], x[1]))
+
+        imp_patches_img = np.zeros((target_size, target_size, residual_frame.shape[2]), dtype=residual_frame.dtype)
+        patches_per_row = target_size // patch_size  # 14
+        # order the patch in the original location relation
+        positions = []
+        for idx, (y, x) in enumerate(sorted_idx):
+            patch = residual_frame[y * patch_size:(y + 1) * patch_size, x * patch_size:(x + 1) * patch_size]
+            # new patch location
+            row_idx = idx // patches_per_row
+            col_idx = idx % patches_per_row
+            start_y = row_idx * patch_size
+            start_x = col_idx * patch_size
+            imp_patches_img[start_y:start_y + patch_size, start_x:start_x + patch_size] = patch
+            positions.append((y, x))
+        return imp_patches_img, positions
+
+    def get_original_frame_patches(self, original_frame, positions):
+        patch_size = self.patch_size
+        target_size = self.target_size
+        imp_original_patches_img = np.zeros((target_size, target_size, original_frame.shape[2]), dtype=original_frame.dtype)
+        patches_per_row = target_size // patch_size
+
+        for idx, (y, x) in enumerate(positions):
+            start_y = y * patch_size
+            start_x = x * patch_size
+            end_y = start_y + patch_size
+            end_x = start_x + patch_size
+
+            patch = original_frame[start_y:end_y, start_x:end_x]
+            row_idx = idx // patches_per_row
+            col_idx = idx % patches_per_row
+            target_start_y = row_idx * patch_size
+            target_start_x = col_idx * patch_size
+
+            imp_original_patches_img[target_start_y:target_start_y + patch_size,
+                                     target_start_x:target_start_x + patch_size] = patch
+        return imp_original_patches_img
+
+def visualise_tensor(tensors, num_frames_to_visualise=5, img_title='Frag'):
+    np_feat = tensors.numpy()
+    fig, axes = plt.subplots(1, num_frames_to_visualise, figsize=(15, 5))
+    for i in range(num_frames_to_visualise):
+        # move channels to last dimension for visualisation: (height, width, channels)
+        frame = np_feat[i].transpose(1, 2, 0)
+        # normalize to [0, 1] for visualisation
+        frame = (frame - frame.min()) / (frame.max() - frame.min())
+        axes[i].imshow(frame)
+        axes[i].axis('off')
+        axes[i].set_title(f'{img_title} {i + 1}')
+
+    plt.tight_layout()
+    save_path =  f'../../figs/{img_title}.png'
+    plt.savefig(save_path, dpi=300)
+    plt.show()
+
+def visualise_image(frame, img_title='Residual Fragment', debug=False):
+    if debug:
+        plt.figure(figsize=(5, 5))
+        plt.imshow(frame)
+        plt.axis('off')
+        plt.title(img_title)
+        plt.show()
+
+
+if __name__ == "__main__":
+    database = 'konvid_1k'
+    videos_dir = '../../ugc_original_videos/'
+    metadata_csv = '../../metadata/TEST_metadata.csv'
+    # videos_dir = '/home/xinyi/video_dataset/KoNViD_1k/KoNViD_1k_videos/'
+    # videos_dir = '/media/on23019/server/LSVQ/'
+    # metadata_csv = f'../../metadata/{database.upper()}_metadata.csv'
+
+    resize = 224 # 224, 384
+    start_time = time.time()
+    resize_transform = transforms.Compose([transforms.Resize([resize, resize]),
+                                           transforms.ToTensor(),
+                                           transforms.Normalize(mean=[0.45, 0.45, 0.45], std=[0.225, 0.225, 0.225])])
+
+    dataset = VideoDataset_feature(
+        data_dir=videos_dir,
+        filename_path=metadata_csv,
+        transform=resize_transform,
+        resize=resize,
+        database=database,
+        patch_size=16, # 8, 16, 32, 16, 32
+        target_size=224, # 224, 224, 224, 384, 384
+        top_n=14*14 # 28*28, 14*14, 7*7, 24*24, 12*12
+    )
+
+    # test
+    index = 0
+    video_segments, video_res_frag_all, video_frag_all, video_name = dataset[index]
+    print(f"Video Name: {video_name}")
+    print(f"Number of Video Segments: {len(video_segments)}")
+    print(f"Number of Video Residual Fragment Segments: {len(video_res_frag_all)}")
+    print(f"Number of Video Fragment Segments: {len(video_frag_all)}")
+    print(f"Shape of Each Segment: {video_segments[0].shape}")  # (video_length_clip, channels, height, width)
+    print(f"Shape of Each Residual Fragment Segments: {video_res_frag_all[0].shape}")
+    print(f"Shape of Each Fragment Segments: {video_frag_all[0].shape}")
+
+    # visualisation
+    first_segments = video_segments[0]
+    visualise_tensor(first_segments, num_frames_to_visualise=5, img_title='Frame')
+
+    first_segment_residuals = video_res_frag_all[0]
+    visualise_tensor(first_segment_residuals, num_frames_to_visualise=6, img_title='Residual Frag')
+
+    first_segment_fragments = video_frag_all[0]
+    visualise_tensor(first_segment_fragments, num_frames_to_visualise=5, img_title='Frame Frag')
+    print(f"Processed {time.time() - start_time:.2f} seconds")

extractor/extract_slowfast_clip.py

@@ -0,0 +1,110 @@
+import torch
+import torch.nn as nn
+from pytorchvideo.models.hub import slowfast_r50
+
+def pack_pathway_output(frames, device):
+    fast_pathway = frames
+    # temporal sampling from the fast pathway.
+    slow_pathway = torch.index_select(
+        frames,
+        2,
+        torch.linspace(0, frames.shape[2] - 1, frames.shape[2] // 4).long(),
+    )
+    return [slow_pathway.to(device), fast_pathway.to(device)]
+
+
+class SlowFast(torch.nn.Module):
+    def __init__(self):
+        super(SlowFast, self).__init__()
+        slowfast_pretrained_features = nn.Sequential(*list(slowfast_r50(pretrained=True).children())[0])
+
+        self.feature_extraction = torch.nn.Sequential()
+        self.slow_avg_pool = torch.nn.Sequential()
+        self.fast_avg_pool = torch.nn.Sequential()
+        self.adp_avg_pool = torch.nn.Sequential()
+
+        for x in range(0, 5):
+            self.feature_extraction.add_module(str(x), slowfast_pretrained_features[x])
+
+        self.slow_avg_pool.add_module('slow_avg_pool', slowfast_pretrained_features[5].pool[0])
+        self.fast_avg_pool.add_module('fast_avg_pool', slowfast_pretrained_features[5].pool[1])
+        self.adp_avg_pool.add_module('adp_avg_pool', slowfast_pretrained_features[6].output_pool)
+
+    def forward(self, x):
+        with torch.no_grad():
+            x = self.feature_extraction(x)
+            slow_feature = self.slow_avg_pool(x[0])
+            fast_feature = self.fast_avg_pool(x[1])
+            slow_feature = self.adp_avg_pool(slow_feature)
+            fast_feature = self.adp_avg_pool(fast_feature)
+        return slow_feature, fast_feature
+
+
+def extract_features_slowfast(video, model, device):
+    slow_features_list = []
+    fast_features_list = []
+
+    with torch.cuda.amp.autocast():
+        for idx, segment in enumerate(video):
+            segment = segment.permute(0, 2, 1, 3, 4)
+            inputs = pack_pathway_output(segment, device)
+            # print(f"Inputs shapes: slow={inputs[0].shape}, fast={inputs[1].shape}")
+
+            # extract features
+            slow_feature, fast_feature = model(inputs)
+            slow_features_list.append(slow_feature)
+            fast_features_list.append(fast_feature)
+
+    # concatenate and flatten features
+    slow_features = torch.cat(slow_features_list, dim=0).flatten()
+    fast_features = torch.cat(fast_features_list, dim=0).flatten()
+    return slow_features, fast_features
+
+
+def extract_features_slowfast_pool(video, model, device):
+    slow_features_list = []
+    fast_features_list = []
+
+    with torch.cuda.amp.autocast():
+        for idx, segment in enumerate(video):
+            segment = segment.permute(0, 2, 1, 3, 4)
+            inputs = pack_pathway_output(segment, device)
+            # print(f"Inputs shapes: slow={inputs[0].shape}, fast={inputs[1].shape}")
+
+            # extract features
+            slow_feature, fast_feature = model(inputs)
+            # global average pooling to reduce dimensions
+            slow_feature = slow_feature.mean(dim=[2, 3, 4])  # Pool over spatial and temporal dims
+            fast_feature = fast_feature.mean(dim=[2, 3, 4])
+            slow_features_list.append(slow_feature)
+            fast_features_list.append(fast_feature)
+
+    # concatenate pooled features
+    slow_features = torch.cat(slow_features_list, dim=0)
+    fast_features = torch.cat(fast_features_list, dim=0)
+    slowfast_features = torch.cat((slow_features, fast_features), dim=1)  # along feature dimension
+    return slow_features, fast_features, slowfast_features
+
+
+# slow_features, fast_features = extract_features_slowfast_pool(video, model, device)
+
+# extract_features_slowfast():
+# Segment shape: torch.Size([1, 3, 32, 224, 224])
+# Inputs shapes: slow=torch.Size([1, 3, 8, 224, 224]), fast=torch.Size([1, 3, 32, 224, 224])
+# Slow feature shape: torch.Size([1, 2048, 1, 1, 1])
+# Fast feature shape: torch.Size([1, 256, 1, 1, 1])
+# Slow features shape: torch.Size([16384])
+# Fast features shape: torch.Size([2048])
+# Combined features shape: torch.Size([18432])
+#
+# extract_features_slowfast_pool():
+# Segment shape: torch.Size([1, 3, 32, 224, 224])
+# Inputs shapes: slow=torch.Size([1, 3, 8, 224, 224]), fast=torch.Size([1, 3, 32, 224, 224])
+# Slow feature shape: torch.Size([1, 2048, 1, 1, 1])
+# Fast feature shape: torch.Size([1, 256, 1, 1, 1])
+# Pooled Slow feature shape: torch.Size([1, 2048])
+# Pooled Fast feature shape: torch.Size([1, 256])
+# Pooled Slow features shape: torch.Size([8, 2048])
+# Pooled Fast features shape: torch.Size([8, 256])
+# Combined features shape: torch.Size([8, 2304])
+# Averaged combined features shape: torch.Size([2304])

extractor/extract_swint_clip.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+from timm import create_model
+
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+
+class SwinT(nn.Module):
+    def __init__(self, model_name='swin_base_patch4_window7_224', global_pool='avg', pretrained=True):
+        super(SwinT, self).__init__()
+        self.swin_model = create_model(
+            model_name, pretrained=pretrained, global_pool=global_pool
+        )
+        self.swin_model.head = Identity()  # Remove classification head
+        self.global_pool = global_pool
+
+    def forward(self, x):
+        features = self.swin_model(x)  # Shape: (batch_size, 7, 7, 1024)
+        if self.global_pool == 'avg':
+            features = features.mean(dim=[1, 2])  # Global pool
+        return features
+
+def extract_features_swint_pool(video, model, device):
+    swint_feature_list = []
+
+    with torch.cuda.amp.autocast():
+        for segment in video:
+            # Flatten the segment into a batch of frames
+            frames = segment.squeeze(0).to(device)  # Shape: (32, 3, 224, 224)
+
+            swint_features = model(frames)  # Shape: (32, feature_dim)
+            swint_feature_list.append(swint_features)
+
+        # Concatenate features across segments
+        features = torch.cat(swint_feature_list, dim=0)  # Shape: (num_frames, feature_dim)
+    return features