From eec74aa7f22c93901be19e2e2e021331c656d3dc Mon Sep 17 00:00:00 2001 From: aspaul20 Date: Tue, 25 Feb 2025 16:26:12 +0500 Subject: [PATCH] added multihead and multilabel support --- .../resnet50tsm_memory_tcn_gtea_s1.yaml | 144 +++++++ .../resnet50tsm_memory_tcn_split1.yaml | 144 +++++++ .../resnet50tsm_memory_tcn_split132sw.yaml | 144 +++++++ config/gtea/TSM_memory_tcn/test.yaml | 158 ++++++++ .../raw_frame_stream_segmentation_dataset.py | 378 +++++++++++++++++- loader/sampler/frame_sampler.py | 201 +++++++++- ...temporal_action_segmentation_base_class.py | 33 +- .../temporal_action_segmentation_metric.py | 11 +- ...mporal_action_segmentation_metric_utils.py | 3 +- .../segmentation/stream_segmentation2d.py | 189 ++++++++- ...stream_segmentation2d_with_backboneloss.py | 90 ----- .../stream_segmentation2d_with_neck.py | 88 ---- ...stream_segmentation3d_with_backboneloss.py | 100 ----- model/backbones/__init__.py | 3 +- model/backbones/image/resnet.py | 5 +- model/backbones/video/resnet_tsm.py | 5 +- model/heads/segmentation/lstm_head.py | 2 +- model/losses/segmentation_loss.py | 210 +++++++++- model/necks/cross_attn.py | 84 ++++ .../stream_score_post_processing.py | 95 ++++- optimizer/tsm_adam_optimizer.py | 7 + tasks/runner.py | 165 +++++--- tasks/test.py | 1 - tasks/train.py | 7 +- tools/avi2mp4.py | 23 ++ tools/find_unused_files.py | 62 +++ tools/plot.py | 66 +++ tools/plot_multi.py | 75 ++++ tools/rotate.py | 47 +++ tools/split_labels.py | 49 +++ tools/split_mapping.py | 49 +++ tools/verify_splits.py | 15 + utils/recorder.py | 26 ++ 33 files changed, 2303 insertions(+), 376 deletions(-) create mode 100644 config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_gtea_s1.yaml create mode 100644 config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_split1.yaml create mode 100644 config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_split132sw.yaml create mode 100644 config/gtea/TSM_memory_tcn/test.yaml delete mode 100644 model/architectures/segmentation/stream_segmentation2d_with_backboneloss.py delete mode 100644 model/architectures/segmentation/stream_segmentation2d_with_neck.py delete mode 100644 model/architectures/segmentation/stream_segmentation3d_with_backboneloss.py create mode 100644 model/necks/cross_attn.py create mode 100644 tools/avi2mp4.py create mode 100644 tools/find_unused_files.py create mode 100644 tools/plot.py create mode 100644 tools/plot_multi.py create mode 100644 tools/rotate.py create mode 100644 tools/split_labels.py create mode 100644 tools/split_mapping.py create mode 100644 tools/verify_splits.py diff --git a/config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_gtea_s1.yaml b/config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_gtea_s1.yaml new file mode 100644 index 0000000..46a4f8d --- /dev/null +++ b/config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_gtea_s1.yaml @@ -0,0 +1,144 @@ +MODEL: + architecture: "StreamSegmentation2D" + backbone: + name: "ResNetTSM" + pretrained: "data/tsm_r50_dense_256p_1x1x8_100e_kinetics400_rgb_20200727-e1e0c785.pth" + clip_seg_num: 32 + shift_div: 8 + out_indices: (3, ) + neck: + name: "AvgPoolNeck" + num_classes: 11 + in_channels: 2048 + clip_seg_num: 32 + drop_ratio: 0.5 + need_pool: True + head: + name: "MemoryTCNHead" + num_stages: 1 + num_layers: 4 + num_f_maps: 64 + dim: 2048 + num_classes: 11 + sample_rate: 4 + loss: + name: "SegmentationLoss" + num_classes: 11 + sample_rate: 4 + smooth_weight: 0.15 + ignore_index: -100 + +POSTPRECESSING: + name: "StreamScorePostProcessing" + num_classes: 11 + clip_seg_num: 32 + sliding_window: 128 + sample_rate: 4 + ignore_index: -100 + +COLLATE: + name: "StreamBatchCompose" + to_tensor_keys: ["imgs", "labels", "masks", "precise_sliding_num"] + +DATASET: #DATASET field + temporal_clip_batch_size: 3 + video_batch_size: 2 + num_workers: 0 + train: + name: "RawFrameStreamSegmentationDataset" + data_prefix: "./" #Mandatory, train data root path + file_path: "data/gtea/splits/train.split1.bundle" #Mandatory, train data index file path + videos_path: "data/gtea/Videos" + gt_path: "data/gtea/groundTruth" + actions_map_file_path: "data/gtea/mapping.txt" + dataset_type: "gtea" + train_mode: True + sliding_window: 128 + clip_seg_num: 32 + sample_rate: 4 + test: + name: "RawFrameStreamSegmentationDataset" + data_prefix: "./" #Mandatory, train data root path + file_path: "data/gtea/splits/test.split1.bundle" #Mandatory, train data index file path + videos_path: "./data/gtea/Videos" + gt_path: "./data/gtea/groundTruth" + actions_map_file_path: "data/gtea/mapping.txt" + dataset_type: "gtea" + train_mode: False + sliding_window: 128 + clip_seg_num: 32 + sample_rate: 4 + +PIPELINE: #PIPELINE field + train: + name: "BasePipline" + decode: + name: "VideoDecoder" + backend: "decord" + sample: + name: "VideoStreamSampler" + is_train: True + sample_rate: 4 + clip_seg_num: 32 + sliding_window: 128 + sample_mode: "uniform" + transform: #Mandotary, image transform operator. + name: "VideoStreamTransform" + transform_list: + - Resize: + size: [256, 320] + - RandomCrop: + size: 224 + - RandomHorizontalFlip: + - ToTensor: + - Normalize: + mean: [0.551, 0.424, 0.179] + std: [0.133, 0.141, 0.124] + + test: + name: "BasePipline" + decode: + name: "VideoDecoder" + backend: "decord" + sample: + name: "VideoStreamSampler" + is_train: "" + sample_rate: 4 + clip_seg_num: 32 + sliding_window: 128 + sample_mode: "uniform" + transform: + name: "VideoStreamTransform" + transform_list: + - Resize: + size: [256, 320] + - CenterCrop: + size: 224 + - ToTensor: + - Normalize: + mean: [0.551, 0.424, 0.179] + std: [0.133, 0.141, 0.124] + +OPTIMIZER: + name: "TSMAdamOptimizer" + learning_rate: 0.0005 + weight_decay: 1e-4 + betas: (0.9, 0.999) + +LRSCHEDULER: + name: "MultiStepLR" + step_size: [50] + gamma: 0.1 + +METRIC: + name: "TASegmentationMetric" + overlap: [.1, .25, .5] + actions_map_file_path: "./data/gtea/mapping.txt" + file_output: True + score_output: True + + +model_name: "ResnetTSM_Memory_TCN_gtea_latest" +log_interval: 4 #Optional, the interal of logger, default:10 +epochs: 50 #Mandatory, total epoch +save_interval: 50 diff --git a/config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_split1.yaml b/config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_split1.yaml new file mode 100644 index 0000000..fe6dcd3 --- /dev/null +++ b/config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_split1.yaml @@ -0,0 +1,144 @@ +MODEL: + architecture: "StreamSegmentation2D" + backbone: + name: "ResNetTSM" + pretrained: "clsLess_tsm_r50_256p_1x1x8_kinetics400_rgb_.pth" + clip_seg_num: 32 + shift_div: 8 + out_indices: (3, ) + neck: + name: "AvgPoolNeck" + num_classes: 5 + in_channels: 2048 + clip_seg_num: 32 + drop_ratio: 0.5 + need_pool: True + head: + name: "MemoryTCNHead" + num_stages: 1 + num_layers: 4 + num_f_maps: 64 + dim: 2048 + num_classes: 5 + sample_rate: 4 + loss: + name: "SegmentationLoss" + num_classes: 5 + sample_rate: 4 + smooth_weight: 0.15 + ignore_index: -100 + +POSTPRECESSING: + name: "StreamScorePostProcessing" + num_classes: 5 + clip_seg_num: 32 + sliding_window: 128 + sample_rate: 4 + ignore_index: -100 + +COLLATE: + name: "StreamBatchCompose" + to_tensor_keys: ["imgs", "labels", "masks", "precise_sliding_num"] + +DATASET: #DATASET field + temporal_clip_batch_size: 3 + video_batch_size: 2 + num_workers: 0 + train: + name: "RawFrameStreamSegmentationDataset" + data_prefix: "./" #Mandatory, train data root path + file_path: "./data/thal/splits/train.split1.bundle" #Mandatory, train data index file path + videos_path: "./data/thal/Videos" + gt_path: "./data/thal/groundTruth" + actions_map_file_path: "./data/thal/mapping.txt" + dataset_type: "gtea" + train_mode: True + sliding_window: 128 + clip_seg_num: 32 + sample_rate: 4 + test: + name: "RawFrameStreamSegmentationDataset" + data_prefix: "./" #Mandatory, train data root path + file_path: "./data/thal/splits/test.split1.bundle" #Mandatory, train data index file path + videos_path: "./data/thal/Videos" + gt_path: "./data/thal/groundTruth" + actions_map_file_path: "./data/thal/mapping.txt" + dataset_type: "gtea" + train_mode: False + sliding_window: 128 + clip_seg_num: 32 + sample_rate: 4 + +PIPELINE: #PIPELINE field + train: + name: "BasePipline" + decode: + name: "VideoDecoder" + backend: "decord" + sample: + name: "VideoStreamSampler" + is_train: True + sample_rate: 4 + clip_seg_num: 32 + sliding_window: 128 + sample_mode: "uniform" + transform: #Mandotary, image transform operator. + name: "VideoStreamTransform" + transform_list: + - Resize: + size: [256, 320] + - RandomCrop: + size: 224 + - RandomHorizontalFlip: + - ToTensor: + - Normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + + test: + name: "BasePipline" + decode: + name: "VideoDecoder" + backend: "decord" + sample: + name: "VideoStreamSampler" + is_train: False + sample_rate: 4 + clip_seg_num: 32 + sliding_window: 128 + sample_mode: "uniform" + transform: + name: "VideoStreamTransform" + transform_list: + - Resize: + size: [256, 320] + - CenterCrop: + size: 224 + - ToTensor: + - Normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +OPTIMIZER: + name: "TSMAdamOptimizer" + learning_rate: 0.0005 + weight_decay: 1e-4 + betas: (0.9, 0.999) + +LRSCHEDULER: + name: "MultiStepLR" + step_size: [50] + gamma: 0.1 + +METRIC: + name: "TASegmentationMetric" + overlap: [.1, .25, .5] + actions_map_file_path: "./data/thal/mapping.txt" + file_output: True + score_output: True + + +model_name: "ResnetMemTCN_Thal_Areas_45_4_32SW" +log_interval: 1 #Optional, the interal of logger, default:10 +epochs: 150 #Mandatory, total epoch +save_interval: 50 diff --git a/config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_split132sw.yaml b/config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_split132sw.yaml new file mode 100644 index 0000000..c74e3d1 --- /dev/null +++ b/config/gtea/TSM_memory_tcn/resnet50tsm_memory_tcn_split132sw.yaml @@ -0,0 +1,144 @@ +MODEL: + architecture: "StreamSegmentation2D" + backbone: + name: "ResNetTSM" + pretrained: "clsLess_tsm_r50_256p_1x1x8_kinetics400_rgb_.pth" + clip_seg_num: 8 + shift_div: 8 + out_indices: (3, ) + neck: + name: "AvgPoolNeck" + num_classes: 30 + in_channels: 2048 + clip_seg_num: 8 + drop_ratio: 0.5 + need_pool: True + head: + name: "MemoryTCNHead" + num_stages: 1 + num_layers: 4 + num_f_maps: 64 + dim: 2048 + num_classes: 30 + sample_rate: 4 + loss: + name: "SegmentationLoss" + num_classes: 30 + sample_rate: 4 + smooth_weight: 0.15 + ignore_index: -100 + +POSTPRECESSING: + name: "StreamScorePostProcessing" + num_classes: 30 + clip_seg_num: 8 + sliding_window: 32 + sample_rate: 4 + ignore_index: -100 + +COLLATE: + name: "StreamBatchCompose" + to_tensor_keys: ["imgs", "labels", "masks", "precise_sliding_num"] + +DATASET: #DATASET field + temporal_clip_batch_size: 3 + video_batch_size: 2 + num_workers: 0 + train: + name: "RawFrameStreamSegmentationDataset" + data_prefix: "./" #Mandatory, train data root path + file_path: "./data/thal/splits/train.split1.bundle" #Mandatory, train data index file path + videos_path: "./data/thal/Videos" + gt_path: "./data/thal/groundTruth" + actions_map_file_path: "./data/thal/mapping.txt" + dataset_type: "gtea" + train_mode: True + sliding_window: 32 + clip_seg_num: 8 + sample_rate: 4 + test: + name: "RawFrameStreamSegmentationDataset" + data_prefix: "./" #Mandatory, train data root path + file_path: "./data/thal/splits/test.split1.bundle" #Mandatory, train data index file path + videos_path: "./data/thal/Videos" + gt_path: "./data/thal/groundTruth" + actions_map_file_path: "./data/thal/mapping.txt" + dataset_type: "gtea" + train_mode: False + sliding_window: 32 + clip_seg_num: 8 + sample_rate: 4 + +PIPELINE: #PIPELINE field + train: + name: "BasePipline" + decode: + name: "VideoDecoder" + backend: "decord" + sample: + name: "VideoStreamSampler" + is_train: True + sample_rate: 4 + clip_seg_num: 8 + sliding_window: 32 + sample_mode: "uniform" + transform: #Mandotary, image transform operator. + name: "VideoStreamTransform" + transform_list: + - Resize: + size: [256, 320] + - RandomCrop: + size: 224 + - RandomHorizontalFlip: + - ToTensor: + - Normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + + test: + name: "BasePipline" + decode: + name: "VideoDecoder" + backend: "decord" + sample: + name: "VideoStreamSampler" + is_train: False + sample_rate: 4 + clip_seg_num: 8 + sliding_window: 32 + sample_mode: "uniform" + transform: + name: "VideoStreamTransform" + transform_list: + - Resize: + size: [256, 320] + - CenterCrop: + size: 224 + - ToTensor: + - Normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +OPTIMIZER: + name: "TSMAdamOptimizer" + learning_rate: 0.0005 + weight_decay: 1e-4 + betas: (0.9, 0.999) + +LRSCHEDULER: + name: "MultiStepLR" + step_size: [50] + gamma: 0.1 + +METRIC: + name: "TASegmentationMetric" + overlap: [.1, .25, .5] + actions_map_file_path: "./data/thal/mapping.txt" + file_output: True + score_output: True + + +model_name: "Thal_Branchwise_HighSampletest" +log_interval: 1 #Optional, the interal of logger, default:10 +epochs: 150 #Mandatory, total epoch +save_interval: 50 diff --git a/config/gtea/TSM_memory_tcn/test.yaml b/config/gtea/TSM_memory_tcn/test.yaml new file mode 100644 index 0000000..e7eca2e --- /dev/null +++ b/config/gtea/TSM_memory_tcn/test.yaml @@ -0,0 +1,158 @@ +MODEL: + architecture: "StreamSegmentation2DMultiLabel" + backbone: + name: "ResNetTSM" + pretrained: "data/clsLess_tsm_r50_256p_1x1x8_kinetics400_rgb_.pth" + clip_seg_num: 32 + shift_div: 8 + out_indices: (3, ) + neck: + name: "AvgPoolNeck" + num_classes: 21 + in_channels: 2048 + clip_seg_num: 32 + drop_ratio: 0.5 + need_pool: True + head: + action_head: + name: "MemoryTCNHead" + num_stages: 1 + num_layers: 4 + num_f_maps: 64 + dim: 2048 + num_classes: 5 + sample_rate: 1 + branch_head: + name: "MemoryTCNHead" + num_stages: 1 + num_layers: 4 + num_f_maps: 64 + dim: 2048 + num_classes: 16 + sample_rate: 1 + sample_rate: 1 + loss: + name: "SegmentationLossMultiLabel" + num_classes_action: 5 + num_classes_branches: 16 + sample_rate: 1 + smooth_weight: 0.15 + ignore_index: -100 + +POSTPRECESSING: + name: "StreamScorePostProcessingMultiLabel" + num_action_classes: 5 + num_branch_classes: 16 + clip_seg_num: 32 + sliding_window: 32 + sample_rate: 1 + ignore_index: -100 + +COLLATE: + name: "StreamBatchCompose" + to_tensor_keys: ["imgs", "labels", "masks", "precise_sliding_num"] + +DATASET: #DATASET field + temporal_clip_batch_size: 1 + video_batch_size: 1 + num_workers: 0 + train: + name: "RawFrameStreamSegmentationDatasetMultiLabel" + data_prefix: "./" #Mandatory, train data root path + file_path: "./data/thal/splits/train.split2.bundle" #Mandatory, train data index file path + videos_path: "./data/thal/Videos" + gt_path: "./data/thal/groundTruth_split" + actions_map_file_path: "./data/thal/mapping_tasks.txt" + branches_map_file_path: "./data/thal/mapping_branches.txt" + dataset_type: "gtea" + train_mode: True + sliding_window: 32 + clip_seg_num: 32 + sample_rate: 1 + test: + name: "RawFrameStreamSegmentationDatasetMultiLabel" + data_prefix: "./" #Mandatory, train data root path + file_path: "./data/thal/splits/test.split2.bundle" #Mandatory, train data index file path + videos_path: "./data/thal/Videos" + gt_path: "./data/thal/groundTruth_split" + actions_map_file_path: "./data/thal/mapping_tasks.txt" + branches_map_file_path: "./data/thal/mapping_branches.txt" + dataset_type: "gtea" + train_mode: False + sliding_window: 32 + clip_seg_num: 32 + sample_rate: 1 + +PIPELINE: #PIPELINE field + train: + name: "BasePipline" + decode: + name: "VideoDecoder" + backend: "decord" + sample: + name: "VideoStreamSamplerMultiLabel" + is_train: True + sample_rate: 1 + clip_seg_num: 32 + sliding_window: 32 + sample_mode: "uniform" + transform: #Mandotary, image transform operator. + name: "VideoStreamTransform" + transform_list: + - Resize: + size: [224, 224] # 256 320 + - RandomCrop: + size: 224 + - RandomHorizontalFlip: + - ToTensor: + - Normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + + test: + name: "BasePipline" + decode: + name: "VideoDecoder" + backend: "decord" + sample: + name: "VideoStreamSamplerMultiLabel" + is_train: False + sample_rate: 1 + clip_seg_num: 32 + sliding_window: 32 + sample_mode: "uniform" + transform: + name: "VideoStreamTransform" + transform_list: + - Resize: + size: [224, 224] + - CenterCrop: + size: 224 + - ToTensor: + - Normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +OPTIMIZER: + name: "TSMAdamOptimizer" + learning_rate: 0.0005 + weight_decay: 1e-4 + betas: (0.9, 0.999) + +LRSCHEDULER: + name: "MultiStepLR" + step_size: [50] + gamma: 0.1 + +METRIC: + name: "TASegmentationMetric" + overlap: [.1, .25, .5] + actions_map_file_path: "./data/thal/mapping_split.txt" + file_output: True + score_output: True + + +model_name: "ThalMultiHead_Colored_32Frames" +log_interval: 1 #Optional, the interal of logger, default:10 +epochs: 200 #Mandatory, total epoch +save_interval: 50 diff --git a/loader/dataset/raw_frame_stream_segmentation_dataset.py b/loader/dataset/raw_frame_stream_segmentation_dataset.py index f8d88ac..1573624 100644 --- a/loader/dataset/raw_frame_stream_segmentation_dataset.py +++ b/loader/dataset/raw_frame_stream_segmentation_dataset.py @@ -166,12 +166,14 @@ def load_file(self, sample_videos_list): if self.dataset_type in ['gtea', '50salads', 'thumos14', 'egtea']: video_name = video_segment.split('.')[0] label_path = os.path.join(self.gt_path, video_name + '.txt') - + import pdb; pdb.set_trace() video_path = os.path.join(self.videos_path, video_name + '.mp4') if not osp.isfile(video_path): + video_path = os.path.join(self.videos_path, video_name + '.avi') if not osp.isfile(video_path): - raise NotImplementedError + import pdb; pdb.set_trace() + raise IndexError(f'Could not find file: {video_path}') elif self.dataset_type in ['breakfast']: video_segment_name, video_segment_path = video_segment video_name = video_segment_name.split('.')[0] @@ -186,7 +188,11 @@ def load_file(self, sample_videos_list): content = file_ptr.read().split('\n')[:-1] classes = np.zeros(len(content), dtype='int64') for i in range(len(content)): - classes[i] = self.actions_dict[content[i]] + + ### for GTEA ### + # classes[i] = self.actions_dict[content[i]] + ### for THAL ### + classes[i] = self.actions_dict[content[i].split(':')[-1].strip()] # caculate sliding num if max_len < len(content): @@ -313,6 +319,372 @@ def __iter__(self): sample_iterator = self._genrate_sampler(woker_id, num_workers) return sample_iterator +class VideoSamplerDataset(data.Dataset): + def __init__(self, + file_path): + super().__init__() + + self.file_path = file_path + self.info = self.load_file() + + def parse_file_paths(self, input_path): + file_ptr = open(input_path, 'r') + info = file_ptr.read().split('\n')[:-1] + file_ptr.close() + return info + + def load_file(self): + """Load index file to get video information.""" + video_segment_lists = self.parse_file_paths(self.file_path) + return video_segment_lists + + def __len__(self): + """get the size of the dataset.""" + return len(self.info) + + def __getitem__(self, idx): + """ Get the sample for either training or testing given index""" + return idx + + + + +@DATASET.register() +class RawFrameStreamSegmentationDatasetMultiLabel(data.IterableDataset): + """Video dataset for action recognition + The dataset loads raw videos and apply specified transforms on them. + The index file is a file with multiple lines, and each line indicates + a sample video with the filepath and label, which are split with a whitesapce. + Example of a inde file: + file tree: + ─── gtea + ├── Videos + │ ├── S1_Cheese_C1.mp4 + │ ├── S1_Coffee_C1.mp4 + │ ├── S1_CofHoney_C1.mp4 + │ └── ... + ├── groundTruth + │ ├── S1_Cheese_C1.txt + │ ├── S1_Coffee_C1.txt + │ ├── S1_CofHoney_C1.txt + │ └── ... + ├── splits + │ ├── test.split1.bundle + │ ├── test.split2.bundle + │ ├── test.split3.bundle + │ └── ... + └── mapping.txt + Args: + file_path(str): Path to the index file. + pipeline(XXX): A sequence of data transforms. + **kwargs: Keyword arguments for ```BaseDataset```. + """ + def __init__(self, + file_path, + videos_path, + gt_path, + pipeline, + actions_map_file_path, + branches_map_file_path, + temporal_clip_batch_size, + video_batch_size, + sliding_window=60, + clip_seg_num=15, + sample_rate=4, + suffix='', + dataset_type='gtea', + data_prefix=None, + train_mode=True, + drap_last=False, + local_rank=-1, + nprocs=1, + ): + super().__init__() + self.suffix = suffix + self.videos_path = videos_path + self.gt_path = gt_path + self.actions_map_file_path = actions_map_file_path + self.branches_map_file_path = branches_map_file_path + self.dataset_type = dataset_type + self.sliding_window = sliding_window + self.clip_seg_num = clip_seg_num + self.sample_rate = sample_rate + + self.file_path = file_path + self.data_prefix = osp.realpath(data_prefix) if \ + data_prefix is not None and osp.isdir(data_prefix) else data_prefix + self.pipeline = pipeline + + # distribute + self.local_rank = local_rank + self.nprocs = nprocs + self.drap_last = drap_last + # if self.nprocs > 1: + # self.drap_last = True + self.video_batch_size = video_batch_size + + # actions dict generate + file_ptr = open(self.actions_map_file_path, 'r') + actions = file_ptr.read().split('\n')[:-1] + file_ptr.close() + self.actions_dict = dict() + for a in actions: + self.actions_dict[a.split()[1]] = int(a.split()[0]) + br_file_ptr = open(self.branches_map_file_path, 'r') + branches = br_file_ptr.read().split('\n')[:-1] + br_file_ptr.close() + self.branches_dict = dict() + for a in branches: + self.branches_dict[a.split()[1]] = int(a.split()[0]) + + # construct sampler + self.video_sampler_dataloader = torch.utils.data.DataLoader( + VideoSamplerDataset(file_path=file_path), + batch_size=video_batch_size, + num_workers=0, + shuffle=train_mode) + if train_mode == False: + self._viodeo_sample_shuffle() + + # iterable + self.temporal_clip_batch_size = temporal_clip_batch_size + + def _viodeo_sample_shuffle(self): + # sampler video order + self.info_list = self._stream_order_sample(self.video_sampler_dataloader) + + def _stream_order_sample(self, video_sampler_dataloader): + sample_videos_list = [] + self.step_num = len(video_sampler_dataloader) + for step, sample_videos in enumerate(video_sampler_dataloader): + if self.drap_last is True and len(list(sample_videos)) < self.video_batch_size: + break + sample_videos_list.append([step, list(sample_videos)]) + + info_list = self.load_file(sample_videos_list).copy() + return info_list + + def parse_file_paths(self, input_path): + if self.dataset_type in ['gtea', '50salads', 'thumos14', 'egtea']: + file_ptr = open(input_path, 'r') + info = file_ptr.read().split('\n')[:-1] + file_ptr.close() + elif self.dataset_type in ['breakfast']: + file_ptr = open(input_path, 'r') + info = file_ptr.read().split('\n')[:-1] + file_ptr.close() + refine_info = [] + for info_name in info: + video_ptr = info_name.split('.')[0].split('_') + file_name = '' + for j in range(2): + if video_ptr[j] == 'stereo01': + video_ptr[j] = 'stereo' + file_name = file_name + video_ptr[j] + '/' + file_name = file_name + video_ptr[2] + '_' + video_ptr[3] + if 'stereo' in file_name: + file_name = file_name + '_ch0' + refine_info.append([info_name, file_name]) + info = refine_info + return info + + def load_file(self, sample_videos_list): + """Load index file to get video information.""" + video_segment_lists = self.parse_file_paths(self.file_path) + info_list = [[] for i in range(self.nprocs)] + # sample step + for step, sample_idx_list in sample_videos_list: + # sample step clip + video_sample_segment_lists = [[] for i in range(self.nprocs)] + for sample_idx_list_idx in range(len(sample_idx_list)): + nproces_idx = sample_idx_list_idx % self.nprocs + sample_idx = sample_idx_list[sample_idx_list_idx] + video_sample_segment_lists[nproces_idx].append(video_segment_lists[sample_idx]) + + max_len = 0 + info_proc = [[] for i in range(self.nprocs)] + for proces_idx in range(self.nprocs): + # convert sample + info = [] + for video_segment in video_sample_segment_lists[proces_idx]: + if self.dataset_type in ['gtea', '50salads', 'thumos14', 'egtea']: + video_name = video_segment.split('.')[0] + label_path = os.path.join(self.gt_path, video_name + '.txt') + branch_label_path = os.path.join(self.gt_path, video_name + '_branch.txt') + + video_path = os.path.join(self.videos_path, video_name + '.mp4') + if not osp.isfile(video_path): + + video_path = os.path.join(self.videos_path, video_name + '.avi') + if not osp.isfile(video_path): + import pdb; pdb.set_trace() + raise IndexError(f'Could not find file: {video_path}') + elif self.dataset_type in ['breakfast']: + video_segment_name, video_segment_path = video_segment + video_name = video_segment_name.split('.')[0] + label_path = os.path.join(self.gt_path, video_name + '.txt') + + video_path = os.path.join(self.videos_path, video_segment_path + '.mp4') + if not osp.isfile(video_path): + video_path = os.path.join(self.videos_path, video_segment_path + '.avi') + if not osp.isfile(video_path): + raise NotImplementedError + file_ptr = open(label_path, 'r') + content = file_ptr.read().split('\n')[:-1] + classes = np.zeros(len(content), dtype='int64') + for i in range(len(content)): + + ### for GTEA ### + # classes[i] = self.actions_dict[content[i]] + ### for THAL ### + try: + classes[i] = self.actions_dict[content[i].split(':')[-1].strip()] + except: + import pdb; pdb.set_trace() + file_ptr_branch = open(branch_label_path, 'r') + branch_content = file_ptr_branch.read().split('\n')[:-1] + branch_classes = np.zeros(len(branch_content), dtype='int64') + for i in range(len(branch_content)): + + ### for GTEA ### + # classes[i] = self.actions_dict[content[i]] + ### for THAL ### + branch_classes[i] = self.branches_dict[branch_content[i].split(':')[-1].strip()] + + # caculate sliding num + if max_len < len(content): + max_len = len(content) + precise_sliding_num = len(content) // self.sliding_window + if len(content) % self.sliding_window != 0: + precise_sliding_num = precise_sliding_num + 1 + + info.append( + dict(filename=video_path, + raw_labels=classes, + raw_branch_labels=branch_classes, + video_name=video_name, + precise_sliding_num=precise_sliding_num)) + + info_proc[proces_idx] = info + + # construct sliding num + sliding_num = max_len // self.sliding_window + if max_len % self.sliding_window != 0: + sliding_num = sliding_num + 1 + + # nprocs sync + for proces_idx in range(self.nprocs): + info_list[proces_idx].append([step, sliding_num, info_proc[proces_idx]]) + return info_list + + def _get_one_videos_clip(self, idx, info): + imgs_list = [] + labels_list = [] + masks_list = [] + vid_list = [] + branch_labels_list = [] + precise_sliding_num_list = [] + + for single_info in info: + sample_segment = single_info.copy() + sample_segment['sample_sliding_idx'] = idx + sample_segment = self.pipeline(sample_segment) + # imgs: tensor labels: ndarray mask: ndarray vid_list : str list + imgs_list.append(copy.deepcopy(sample_segment['imgs'].unsqueeze(0))) + branch_labels_list.append(np.expand_dims(sample_segment['branch_labels'], axis=0).copy()) + labels_list.append(np.expand_dims(sample_segment['labels'], axis=0).copy()) + masks_list.append(np.expand_dims(sample_segment['mask'], axis=0).copy()) + vid_list.append(copy.deepcopy(sample_segment['video_name'])) + precise_sliding_num_list.append(np.expand_dims(sample_segment['precise_sliding_num'], axis=0).copy()) + + imgs = copy.deepcopy(torch.concat(imgs_list, dim=0)) + try: + labels = copy.deepcopy(np.concatenate(labels_list, axis=0).astype(np.int64)) + branch_labels = copy.deepcopy(np.concatenate(branch_labels_list, axis=0).astype(np.int64)) + except: + import pdb; pdb.set_trace() + masks = copy.deepcopy(np.concatenate(masks_list, axis=0).astype(np.float32)) + precise_sliding_num = copy.deepcopy(np.concatenate(precise_sliding_num_list, axis=0).astype(np.float32)) + # compose result + data_dict = {} + data_dict['imgs'] = imgs + data_dict['labels'] = labels + data_dict['branch_labels'] = torch.tensor(branch_labels) + data_dict['masks'] = masks + data_dict['precise_sliding_num'] = precise_sliding_num + data_dict['vid_list'] = vid_list + return data_dict + + def _get_end_videos_clip(self): + # compose result + data_dict = {} + data_dict['imgs'] = 0 + data_dict['labels'] = 0 + data_dict['branch_labels'] = 0 + + data_dict['masks'] = 0 + data_dict['vid_list'] = [] + data_dict['sliding_num'] = 0 + data_dict['precise_sliding_num'] = 0 + data_dict['step'] = self.step_num + data_dict['current_sliding_cnt'] = -1 + return data_dict + + def _genrate_sampler(self, woker_id, num_workers): + if self.local_rank < 0: + # single gpu train + return self._step_sliding_sampler(woker_id=woker_id, num_workers=num_workers, info_list=self.info_list[0]) + else: + # multi gpu train + sample_info_list = self.info_list[self.local_rank] + return self._step_sliding_sampler(woker_id=woker_id, num_workers=num_workers, info_list=sample_info_list) + + def _step_sliding_sampler(self, woker_id, num_workers, info_list): + # dispatch function + current_sliding_cnt = woker_id * self.temporal_clip_batch_size + mini_sliding_cnt = 0 + next_step_flag = False + for step, sliding_num, info in info_list: + while current_sliding_cnt < sliding_num and len(info) > 0: + while mini_sliding_cnt < self.temporal_clip_batch_size: + if current_sliding_cnt < sliding_num: + data_dict = self._get_one_videos_clip(current_sliding_cnt, info) + data_dict['sliding_num'] = sliding_num + data_dict['step'] = step + data_dict['current_sliding_cnt'] = current_sliding_cnt + yield data_dict + current_sliding_cnt = current_sliding_cnt + 1 + mini_sliding_cnt = mini_sliding_cnt + 1 + else: + next_step_flag = True + break + if current_sliding_cnt <= sliding_num and next_step_flag == False: + current_sliding_cnt = current_sliding_cnt + (num_workers - 1) * self.temporal_clip_batch_size + + if mini_sliding_cnt >= self.temporal_clip_batch_size: + mini_sliding_cnt = 0 + + # modify num_worker + current_sliding_cnt = current_sliding_cnt - sliding_num + next_step_flag = False + yield self._get_end_videos_clip() + + def __len__(self): + """get the size of the dataset.""" + return self.step_num + + def __iter__(self): + """ Get the sample for either training or testing given index""" + worker_info = torch.utils.data.get_worker_info() + # single worker + if worker_info is None: + sample_iterator = self._genrate_sampler(0, 1) + else: # multiple workers + woker_id = worker_info.id + num_workers = int(worker_info.num_workers) + sample_iterator = self._genrate_sampler(woker_id, num_workers) + return sample_iterator + class VideoSamplerDataset(data.Dataset): def __init__(self, file_path): diff --git a/loader/sampler/frame_sampler.py b/loader/sampler/frame_sampler.py index f55d395..e974e7e 100644 --- a/loader/sampler/frame_sampler.py +++ b/loader/sampler/frame_sampler.py @@ -9,6 +9,8 @@ import numpy as np import random from PIL import Image +import cv2 +import albumentations as A from ..builder import SAMPLER class VideoFrameSample(): @@ -59,7 +61,8 @@ def __init__(self, sliding_window=60, ignore_index=-100, channel_mode="RGB", - sample_mode='random' + sample_mode='random', + aug=[] ): self.sample_rate = sample_rate self.is_train = is_train @@ -68,7 +71,12 @@ def __init__(self, self.ignore_index = ignore_index self.channel_mode = channel_mode self.sample = VideoFrameSample(mode = sample_mode) - + self.transform = A.ReplayCompose([ + A.RandomBrightnessContrast(p=0.6), + A.HueSaturationValue(hue_shift_limit=(-10, 10), sat_shift_limit=(-10, 10), val_shift_limit=(-10, 10), p=1), + A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.01, rotate_limit=4, p=1), + A.ImageCompression(quality_lower=75, quality_upper=95, p=1), + ]) def _all_valid_frames(self, start_frame, end_frame, video_len, container, labels): imgs = [] vid_end_frame = end_frame @@ -79,6 +87,17 @@ def _all_valid_frames(self, start_frame, end_frame, video_len, container, labels frames_select = container.get_batch(frames_idx) # dearray_to_img np_frames = frames_select.asnumpy() + + ### augmentations ### + first_frame = np_frames[0] + replayed = self.transform(image=first_frame) + replay_data = replayed['replay'] + augmented_video = np.empty_like(np_frames) + for i in range(np_frames.shape[0]): + augmented_frame = A.ReplayCompose.replay(replay_data, image=np_frames[i])['image'] + augmented_video[i] = augmented_frame + np_frames = augmented_video + import pdb; pdb.set_trace() for i in range(np_frames.shape[0]): imgbuf = np_frames[i].copy() imgs.append(Image.fromarray(imgbuf, mode=self.channel_mode)) @@ -164,6 +183,184 @@ def __call__(self, results): results['mask'] = mask.copy() return results + + +@SAMPLER.register() +class VideoStreamSamplerMultiLabel(): + """ + Sample frames id. + Returns: + frames_idx: the index of sampled #frames. + """ + + def __init__(self, + is_train=False, + sample_rate=4, + clip_seg_num=15, + sliding_window=60, + ignore_index=-100, + channel_mode="RGB", + sample_mode='random', + aug=[] + ): + self.sample_rate = sample_rate + self.is_train = is_train + self.clip_seg_num = clip_seg_num + self.sliding_window = sliding_window + self.ignore_index = ignore_index + self.channel_mode = channel_mode + self.sample = VideoFrameSample(mode = sample_mode) + self.transform = A.ReplayCompose([ + A.RandomBrightnessContrast(p=0.6), + A.HueSaturationValue(hue_shift_limit=(-30, 30), sat_shift_limit=(-30, 30), val_shift_limit=(-30, 30), p=0.6), + A.ShiftScaleRotate(shift_limit=0.25, scale_limit=0.4, rotate_limit=15, p=0.6), + A.ImageCompression(quality_lower=75, quality_upper=95, p=0.6), + ]) + self.visualize = False + + def save_video(self, np_frames: np.ndarray, output_path: str, fps: int = 30) -> None: + """ + Saves a numpy array of shape (num_frames, height, width, channels) as a video file. + Args: + np_frames (np.ndarray): Array containing video frames. + output_path (str): Path to save the output video file. + fps (int, optional): Frames per second. Defaults to 30. + """ + # Extract frame dimensions + num_frames, height, width, channels = np_frames.shape + assert channels == 3, "Frames should have 3 channels (RGB)." + + # Define the video codec and create a VideoWriter object + fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Codec for .mp4 + out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) + + # Write each frame to the video + for frame in np_frames: + # Convert RGB (numpy default) to BGR (OpenCV default) + frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) + out.write(frame_bgr) + + # Release the video writer + out.release() + print(f'Video saved to {output_path}') + + def _all_valid_frames(self, start_frame, end_frame, video_len, container, labels, branch_labels): + imgs = [] + vid_end_frame = end_frame + if end_frame > video_len: + vid_end_frame = video_len + frames_idx = self.sample(start_frame, vid_end_frame, self.sample_rate) + labels = self._labels_sample(labels, start_frame=start_frame, end_frame=end_frame, samples_idx=frames_idx).copy() + branch_labels= self._labels_sample(branch_labels, start_frame=start_frame, end_frame=end_frame, samples_idx=frames_idx).copy() + + frames_select = container.get_batch(frames_idx) + # dearray_to_img + np_frames = frames_select.asnumpy() + + ### augmentations ### + first_frame = np_frames[0] + replayed = self.transform(image=first_frame) + replay_data = replayed['replay'] + augmented_video = np.empty_like(np_frames) + for i in range(np_frames.shape[0]): + augmented_frame = A.ReplayCompose.replay(replay_data, image=np_frames[i])['image'] + augmented_video[i] = augmented_frame + np_frames = augmented_video + + if self.visualize: + self.save_video(np_frames=np_frames, output_path='test.mp4', fps=1) + + for i in range(np_frames.shape[0]): + imgbuf = np_frames[i].copy() + imgs.append(Image.fromarray(imgbuf, mode=self.channel_mode)) + + if len(imgs) < self.clip_seg_num: + np_frames = np_frames[-1].copy() + pad_len = self.clip_seg_num - len(imgs) + for i in range(pad_len): + imgs.append(Image.fromarray(np_frames, mode=self.channel_mode)) + + mask = np.ones((labels.shape[0])) + + return imgs, labels, mask, branch_labels + + def _some_valid_frames(self, start_frame, end_frame, video_len, frames_len, container, labels, branch_labels): + imgs = [] + small_frames_video_len = min(frames_len, video_len) + frames_idx = self.sample(start_frame, video_len, self.sample_rate) + label_frames_idx = self.sample(start_frame, small_frames_video_len, self.sample_rate) + labels = self._labels_sample(labels, start_frame=start_frame, end_frame=small_frames_video_len, samples_idx=label_frames_idx).copy() + branch_labels = self._labels_sample(branch_labels, start_frame=start_frame, end_frame=small_frames_video_len, samples_idx=label_frames_idx).copy() + + frames_select = container.get_batch(frames_idx) + # dearray_to_img + np_frames = frames_select.asnumpy() + if np_frames.shape[0] > 0: + for i in range(np_frames.shape[0]): + imgbuf = np_frames[i].copy() + imgs.append(Image.fromarray(imgbuf, mode=self.channel_mode)) + np_frames = np.zeros_like(np_frames[0]) + else: + np_frames = np.zeros((240, 320, 3)) + pad_len = self.clip_seg_num - len(imgs) + for i in range(pad_len): + imgs.append(Image.fromarray(np_frames, mode=self.channel_mode)) + vaild_mask = np.ones((labels.shape[0])) + mask_pad_len = self.clip_seg_num * self.sample_rate - labels.shape[0] + void_mask = np.zeros((mask_pad_len)) + mask = np.concatenate([vaild_mask, void_mask], axis=0) + labels = np.concatenate([labels, np.full((mask_pad_len), self.ignore_index)]) + branch_labels = np.concatenate([branch_labels, np.full((mask_pad_len), self.ignore_index)]) + + return imgs, labels, mask, branch_labels + + def _labels_sample(self, labels, start_frame=0, end_frame=0, samples_idx=[]): + if self.is_train: + sample_labels = labels[samples_idx] + sample_labels = np.repeat(sample_labels, repeats=self.sample_rate, axis=-1) + else: + sample_labels = labels[start_frame:end_frame] + return sample_labels + + def __call__(self, results): + """ + Args: + frames_len: length of frames. + return: + sampling id. + """ + frames_len = int(results['frames_len']) + video_len = int(results['video_len']) + results['frames_len'] = frames_len + container = results['frames'] + labels = results['raw_labels'] + branch_labels = results['raw_branch_labels'] + small_frames_video_len = min(frames_len, video_len) + + # generate sample index + start_frame = results['sample_sliding_idx'] * self.sliding_window + end_frame = start_frame + self.clip_seg_num * self.sample_rate + if start_frame < small_frames_video_len and end_frame < small_frames_video_len: + imgs, labels, mask, branch_labels = self._all_valid_frames(start_frame, end_frame, video_len, container, labels, branch_labels) + elif start_frame < small_frames_video_len and end_frame >= small_frames_video_len: + imgs, labels, mask, branch_labels = self._some_valid_frames(start_frame, end_frame, video_len, frames_len, container, labels, branch_labels) + else: + imgs = [] + np_frames = np.zeros((240, 320, 3)) + pad_len = self.clip_seg_num + for i in range(pad_len): + imgs.append(Image.fromarray(np_frames, mode=self.channel_mode)) + mask = np.zeros((self.clip_seg_num * self.sample_rate)) + labels = np.full((self.clip_seg_num * self.sample_rate), self.ignore_index) + branch_labels = np.full((self.clip_seg_num * self.sample_rate), self.ignore_index) + + results['imgs'] = imgs[:self.clip_seg_num].copy() + results['labels'] = labels.copy() + results['mask'] = mask.copy() + results['branch_labels'] = branch_labels.copy() + + return results + @SAMPLER.register() class RGBFlowVideoStreamSampler(): """ diff --git a/metric/temporal_action_segmentation/temporal_action_segmentation_base_class.py b/metric/temporal_action_segmentation/temporal_action_segmentation_base_class.py index 6cb64b2..462ce39 100644 --- a/metric/temporal_action_segmentation/temporal_action_segmentation_base_class.py +++ b/metric/temporal_action_segmentation/temporal_action_segmentation_base_class.py @@ -61,7 +61,6 @@ def __init__(self, self.actions_dict = dict() for a in actions: self.actions_dict[a.split()[1]] = int(a.split()[0]) - # cls score self.overlap = overlap self.overlap_len = len(overlap) @@ -178,31 +177,39 @@ def _write_seg_file(self, input_data, write_name, write_path): f.writelines(recog_content) f.close() - def _transform_model_result(self, vid, outputs_np, gt_np, outputs_arr): + def _transform_model_result(self, vid, outputs_np, gt_np, outputs_arr, action_dict=None): recognition = [] + + if action_dict is not None: + act_dict = action_dict + else: + act_dict = self.actions_dict + for i in range(outputs_np.shape[0]): recognition = np.concatenate((recognition, [ - list(self.actions_dict.keys())[list( - self.actions_dict.values()).index(outputs_np[i])] + list(act_dict.keys())[list( + act_dict.values()).index(outputs_np[i])] ])) recog_content = list(recognition) gt_content = [] for i in range(gt_np.shape[0]): gt_content = np.concatenate((gt_content, [ - list(self.actions_dict.keys())[list( - self.actions_dict.values()).index(gt_np[i])] + list(act_dict.keys())[list( + act_dict.values()).index(gt_np[i])] ])) gt_content = list(gt_content) if self.file_output is True and self.train_mode is False: - self._write_seg_file(gt_content, vid + '-gt', self.output_dir) - self._write_seg_file(recog_content, vid + '-pred', self.output_dir) - - pred_detection = get_labels_scores_start_end_time( - outputs_arr, recog_content, self.actions_dict) - gt_detection = get_labels_scores_start_end_time( - np.ones(outputs_arr.shape), gt_content, self.actions_dict) + self._write_seg_file(gt_content, vid + f'-{list(act_dict.keys())[0]}-gt', self.output_dir) + self._write_seg_file(recog_content, vid + f'-{list(act_dict.keys())[0]}-pred', self.output_dir) + try: + pred_detection = get_labels_scores_start_end_time( + outputs_arr, recog_content, act_dict) + gt_detection = get_labels_scores_start_end_time( + np.ones(outputs_arr.shape), gt_content, act_dict) + except: + import pdb; pdb.set_trace() return [recog_content, gt_content, pred_detection, gt_detection] def update(self, outputs): diff --git a/metric/temporal_action_segmentation/temporal_action_segmentation_metric.py b/metric/temporal_action_segmentation/temporal_action_segmentation_metric.py index c0bc579..082bf74 100644 --- a/metric/temporal_action_segmentation/temporal_action_segmentation_metric.py +++ b/metric/temporal_action_segmentation/temporal_action_segmentation_metric.py @@ -33,9 +33,16 @@ def __init__(self, max_proposal, tiou_thresholds, file_output, score_output, output_dir, score_output_dir) - def update(self, vid, ground_truth_batch, outputs): + def update(self, vid, ground_truth_batch, outputs, action_dict_path=None): """update metrics during each iter """ + if action_dict_path: + action_dict = open(action_dict_path, 'r') + acts = action_dict.readlines() + action_dict = {act.strip().split(' ')[1]: int(act.strip().split(' ')[0]) for act in acts} + else: + action_dict = None + # list [N, T] predicted_batch = outputs['predict'] # list [N, C, T] @@ -61,7 +68,7 @@ def update(self, vid, ground_truth_batch, outputs): score_output_path = os.path.join(self.score_output_dir, vid[bs] + ".npy") np.save(score_output_path, output_np) - result = self._transform_model_result(vid[bs], outputs_np, gt_np, outputs_arr) + result = self._transform_model_result(vid[bs], outputs_np, gt_np, outputs_arr, action_dict) recog_content, gt_content, pred_detection, gt_detection = result single_f1, acc = self._update_score([vid[bs]], recog_content, gt_content, pred_detection, gt_detection) diff --git a/metric/temporal_action_segmentation/temporal_action_segmentation_metric_utils.py b/metric/temporal_action_segmentation/temporal_action_segmentation_metric_utils.py index 7a5d55c..dab9076 100644 --- a/metric/temporal_action_segmentation/temporal_action_segmentation_metric_utils.py +++ b/metric/temporal_action_segmentation/temporal_action_segmentation_metric_utils.py @@ -10,7 +10,7 @@ import pandas as pd from joblib import Parallel, delayed -ignore_bg_class = ["background", "None"] +ignore_bg_class = ["bg", "None"] def get_labels_scores_start_end_time(input_np, frame_wise_labels, @@ -377,7 +377,6 @@ def wrapper_compute_average_precision(prediction, ground_truth, tiou_thresholds, label_name, cidx), tiou_thresholds=tiou_thresholds, ) for label_name, cidx in activity_dict.items()) - for i, cidx in enumerate(activity_dict.values()): ap[:, cidx] = results[i] diff --git a/model/architectures/segmentation/stream_segmentation2d.py b/model/architectures/segmentation/stream_segmentation2d.py index 55a6ffb..c590fdb 100644 --- a/model/architectures/segmentation/stream_segmentation2d.py +++ b/model/architectures/segmentation/stream_segmentation2d.py @@ -17,6 +17,8 @@ from ...builder import build_head from ...builder import ARCHITECTURE +from model.backbones.image.resnet import ResNet + @ARCHITECTURE.register() class StreamSegmentation2D(nn.Module): @@ -27,16 +29,48 @@ def __init__(self, loss=None): super().__init__() self.backbone = build_backbone(backbone) + # backbone['pretrained'] = 'data/cleaned.pth' + # self.det_backbone = build_backbone(backbone) + + self.det_weights = torch.load('data/ew_cleaned.pth') + self.det_backbone = ResNet(depth=50, pretrained='data/new_cleaned.pth') + # self.det_backbone = ResNetModel.from_pretrained("microsoft/resnet-50") + # self.det_weights = {key.replace('backbone.', ''):val for key, val in torch.load('data/last_epoch_model.pth').items() if key.startswith('backbone.')} + # self.det_weights = torch.load('efficientEgoNet_bb_512x512.pth', map_location='cpu') + # self.det_backbone = torchvision.models.efficientnet_v2_s(weights=EfficientNet_V2_S_Weights) + # self.det_backbone = torch.nn.Sequential(*(list(self.det_backbone.children())[:-2])) + + self.neck = build_neck(neck) self.head = build_head(head) + # self.backbone.init_weights() + # self.det_backbone.init_weights() + + self.feat_combine = torch.nn.Sequential( + conv_bn_relu(in_channels=4096, out_channels=2048, kernel_size=3, padding=1, dilation=1), + torch.nn.Conv2d(in_channels=2048, out_channels=2048, kernel_size=1), + ) + + self.feat_refine = torch.nn.Sequential( + conv_bn_relu(in_channels=2048, out_channels=2048, kernel_size=3, padding=1, dilation=1) + ) + self.init_weights() - self.sample_rate = head.sample_rate def init_weights(self): if self.backbone is not None: self.backbone.init_weights(child_model=False, revise_keys=[(r'backbone.', r'')]) + + if self.det_backbone is not None: + # self.det_backbone.init_weights(child_model=False, revise_keys=[(r'backbone.', r''), (r'conv.net', 'conv')]) + self.det_backbone.load_state_dict(self.det_weights, strict=True) + print('det backbone loaded') + for param in self.det_backbone.parameters(): + param.requires_grad=False + print('det backbone freezed') + if self.neck is not None: self.neck.init_weights() if self.head is not None: @@ -49,30 +83,43 @@ def _clear_memory_buffer(self): self.neck._clear_memory_buffer() if self.head is not None: self.head._clear_memory_buffer() + + def feature_fusion(self, temporal_fet, spatial_feat): + + x_1 = temporal_fet + x_2 = spatial_feat + x_concat = torch.cat([x_1, x_2], dim=1) + x_3 = self.feat_combine(x_concat) + feature = self.feat_refine(x_3) + + return feature def forward(self, input_data): masks = input_data['masks'] imgs = input_data['imgs'] - # masks.shape=[N,T] masks = masks.unsqueeze(1) - + import pdb; pdb.set_trace() # x.shape=[N,T,C,H,W], for most commonly case imgs = torch.reshape(imgs, [-1] + list(imgs.shape[2:])) # x [N * T, C, H, W] - if self.backbone is not None: # masks.shape [N * T, 1, 1, 1] backbone_masks = torch.reshape(masks[:, :, ::self.sample_rate], [-1]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) feature = self.backbone(imgs, backbone_masks) + feature_det =self.det_backbone(imgs, backbone_masks)#self.det_backbone(imgs).last_hidden_state # + + # feature_det = self.det_backbone(imgs) + else: feature = imgs - + + feature = self.feature_fusion(feature, feature_det) + # feature = 0.5*feature+0.5*feature_det # feature [N * T , F_dim, 7, 7] # step 3 extract memory feature if self.neck is not None: - seg_feature = self.neck( - feature, masks[:, :, ::self.sample_rate]) + seg_feature = self.neck(feature, masks[:, :, ::self.sample_rate]) else: seg_feature = feature @@ -86,4 +133,130 @@ def forward(self, input_data): head_score = seg_feature # seg_score [stage_num, N, C, T] # cls_score [N, C, T] - return head_score \ No newline at end of file + return head_score + + +@ARCHITECTURE.register() +class StreamSegmentation2DMultiLabel(nn.Module): + def __init__(self, + backbone=None, + neck=None, + head=None, + loss=None): + super().__init__() + self.backbone = build_backbone(backbone) + # backbone['pretrained'] = 'data/cleaned.pth' + # self.det_backbone = build_backbone(backbone) + + self.det_weights = torch.load('data/new_cleaned.pth') + self.det_backbone = ResNet(depth=50, pretrained='data/new_cleaned.pth') + # self.det_backbone = ResNetModel.from_pretrained("microsoft/resnet-50") + # self.det_weights = {key.replace('backbone.', ''):val for key, val in torch.load('data/last_epoch_model.pth').items() if key.startswith('backbone.')} + # self.det_weights = torch.load('efficientEgoNet_bb_512x512.pth', map_location='cpu') + # self.det_backbone = torchvision.models.efficientnet_v2_s(weights=EfficientNet_V2_S_Weights) + # self.det_backbone = torch.nn.Sequential(*(list(self.det_backbone.children())[:-2])) + + self.neck = build_neck(neck) + self.action_head = build_head(head['action_head']) + self.branch_head = build_head(head['branch_head']) + + # self.backbone.init_weights() + # self.det_backbone.init_weights() + + self.feat_combine = torch.nn.Sequential( + conv_bn_relu(in_channels=4096, out_channels=2048, kernel_size=3, padding=1, dilation=1), + torch.nn.Conv2d(in_channels=2048, out_channels=2048, kernel_size=1), + ) + + self.feat_refine = torch.nn.Sequential( + conv_bn_relu(in_channels=2048, out_channels=2048, kernel_size=3, padding=1, dilation=1) + ) + + self.init_weights() + self.sample_rate = head.sample_rate + + def init_weights(self): + if self.backbone is not None: + self.backbone.init_weights(child_model=False, revise_keys=[(r'backbone.', r'')]) + + if self.det_backbone is not None: + # self.det_backbone.init_weights(child_model=False, revise_keys=[(r'backbone.', r''), (r'conv.net', 'conv')]) + self.det_backbone.load_state_dict(self.det_weights, strict=True) + print('det backbone loaded') + for param in self.det_backbone.parameters(): + param.requires_grad=False + print('det backbone freezed') + + if self.neck is not None: + self.neck.init_weights() + if self.action_head is not None: + self.action_head.init_weights() + if self.branch_head is not None: + self.branch_head.init_weights() + + def _clear_memory_buffer(self): + if self.backbone is not None: + self.backbone._clear_memory_buffer() + if self.neck is not None: + self.neck._clear_memory_buffer() + if self.action_head is not None: + self.action_head._clear_memory_buffer() + if self.branch_head is not None: + self.branch_head._clear_memory_buffer() + def feature_fusion(self, temporal_fet, spatial_feat): + + x_1 = temporal_fet + x_2 = spatial_feat + x_concat = torch.cat([x_1, x_2], dim=1) + x_3 = self.feat_combine(x_concat) + feature = self.feat_refine(x_3) + + return feature + + def forward(self, input_data): + masks = input_data['masks'] + imgs = input_data['imgs'] + # masks.shape=[N,T] + masks = masks.unsqueeze(1) + # x.shape=[N,T,C,H,W], for most commonly case + imgs = torch.reshape(imgs, [-1] + list(imgs.shape[2:])) + # x [N * T, C, H, W] + if self.backbone is not None: + # masks.shape [N * T, 1, 1, 1] + backbone_masks = torch.reshape(masks[:, :, ::self.sample_rate], [-1]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + feature = self.backbone(imgs, backbone_masks) + feature_det =self.det_backbone(imgs, backbone_masks)#self.det_backbone(imgs).last_hidden_state # + + # feature_det = self.det_backbone(imgs) + + else: + feature = imgs + + feature = self.feature_fusion(feature, feature_det) + # feature = 0.5*feature+0.5*feature_det + # feature [N * T , F_dim, 7, 7] + # step 3 extract memory feature + if self.neck is not None: + seg_feature = self.neck(feature, masks[:, :, ::self.sample_rate]) + + else: + seg_feature = feature + + # step 5 segmentation + # seg_feature [N, H_dim, T] + # cls_feature [N, F_dim, T] + + if self.action_head is not None: + action_head_score = self.action_head(seg_feature, masks) + else: + action_head_score = seg_feature + + if self.branch_head is not None: + branch_head_score = self.branch_head(seg_feature, masks) + else: + branch_head_score = seg_feature + + # seg_score [stage_num, N, C, T] + # cls_score [N, C, T] + + return {'branch_score':branch_head_score, 'action_score':action_head_score} \ No newline at end of file diff --git a/model/architectures/segmentation/stream_segmentation2d_with_backboneloss.py b/model/architectures/segmentation/stream_segmentation2d_with_backboneloss.py deleted file mode 100644 index d0322f0..0000000 --- a/model/architectures/segmentation/stream_segmentation2d_with_backboneloss.py +++ /dev/null @@ -1,90 +0,0 @@ -''' -Author: Thyssen Wen -Date: 2022-03-25 10:29:10 -LastEditors : Thyssen Wen -LastEditTime : 2022-07-06 20:38:09 -Description: etesvs model framework -FilePath : /ETESVS/model/architectures/segmentation/stream_segmentation2d_with_backboneloss.py -''' -import torch -import torch.nn as nn -from mmcv.runner import load_checkpoint - -from utils.logger import get_logger - -from ...builder import build_backbone -from ...builder import build_neck -from ...builder import build_head - -from ...builder import ARCHITECTURE - -@ARCHITECTURE.register() -class StreamSegmentation2DWithBackbone(nn.Module): - def __init__(self, - backbone=None, - neck=None, - head=None, - loss=None): - super().__init__() - self.backbone = build_backbone(backbone) - self.neck = build_neck(neck) - self.head = build_head(head) - - self.init_weights() - - self.sample_rate = head.sample_rate - - def init_weights(self): - if self.backbone is not None: - self.backbone.init_weights(child_model=False, revise_keys=[(r'backbone.', r'')]) - if self.neck is not None: - self.neck.init_weights() - if self.head is not None: - self.head.init_weights() - - def _clear_memory_buffer(self): - if self.backbone is not None: - self.backbone._clear_memory_buffer() - if self.neck is not None: - self.neck._clear_memory_buffer() - if self.head is not None: - self.head._clear_memory_buffer() - - def forward(self, input_data): - masks = input_data['masks'] - imgs = input_data['imgs'] - - # masks.shape=[N,T] - masks = masks.unsqueeze(1) - - # x.shape=[N,T,C,H,W], for most commonly case - imgs = torch.reshape(imgs, [-1] + list(imgs.shape[2:])) - # x [N * T, C, H, W] - - if self.backbone is not None: - # masks.shape [N * T, 1, 1, 1] - backbone_masks = torch.reshape(masks[:, :, ::self.sample_rate], [-1]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) - feature = self.backbone(imgs, backbone_masks) - else: - feature = imgs - - # feature [N * T , F_dim, 7, 7] - # step 3 extract memory feature - if self.neck is not None: - seg_feature, backbone_score = self.neck( - feature, masks[:, :, ::self.sample_rate]) - - else: - seg_feature = feature - backbone_score = None - - # step 5 segmentation - # seg_feature [N, H_dim, T] - # cls_feature [N, F_dim, T] - if self.head is not None: - head_score = self.head(seg_feature, masks) - else: - head_score = seg_feature - # seg_score [stage_num, N, C, T] - # cls_score [N, C, T] - return backbone_score, head_score \ No newline at end of file diff --git a/model/architectures/segmentation/stream_segmentation2d_with_neck.py b/model/architectures/segmentation/stream_segmentation2d_with_neck.py deleted file mode 100644 index f487a13..0000000 --- a/model/architectures/segmentation/stream_segmentation2d_with_neck.py +++ /dev/null @@ -1,88 +0,0 @@ -''' -Author: Thyssen Wen -Date: 2022-03-25 10:29:10 -LastEditors : Thyssen Wen -LastEditTime : 2022-06-15 21:12:13 -Description: etesvs model framework -FilePath : /ETESVS/model/architectures/segmentation/stream_segmentation2d_with_neck.py -''' -import torch -import torch.nn as nn - -from ...builder import build_backbone -from ...builder import build_neck -from ...builder import build_head - -from ...builder import ARCHITECTURE - -@ARCHITECTURE.register() -class StreamSegmentation2DWithNeck(nn.Module): - def __init__(self, - backbone=None, - neck=None, - head=None, - loss=None): - super().__init__() - self.backbone = build_backbone(backbone) - self.neck = build_neck(neck) - self.head = build_head(head) - - self.init_weights() - - self.sample_rate = head.sample_rate - - def init_weights(self): - if self.backbone is not None: - self.backbone.init_weights(child_model=False, revise_keys=[(r'backbone.', r'')]) - if self.neck is not None: - self.neck.init_weights() - if self.head is not None: - self.head.init_weights() - - def _clear_memory_buffer(self): - if self.backbone is not None: - self.backbone._clear_memory_buffer() - if self.neck is not None: - self.neck._clear_memory_buffer() - if self.head is not None: - self.head._clear_memory_buffer() - - def forward(self, input_data): - masks = input_data['masks'] - imgs = input_data['imgs'] - - # masks.shape=[N,T] - masks = masks.unsqueeze(1) - - # x.shape=[N,T,C,H,W], for most commonly case - imgs = torch.reshape(imgs, [-1] + list(imgs.shape[2:])) - # x [N * T, C, H, W] - - if self.backbone is not None: - # masks.shape [N * T, 1, 1, 1] - backbone_masks = torch.reshape(masks[:, :, ::self.sample_rate], [-1]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) - feature = self.backbone(imgs, backbone_masks) - else: - feature = imgs - - # feature [N * T , F_dim, 7, 7] - # step 3 extract memory feature - if self.neck is not None: - seg_feature, backbone_score, neck_score = self.neck( - feature, masks[:, :, ::self.sample_rate]) - - else: - seg_feature = feature - backbone_score = None - neck_score = None - - # step 5 segmentation - # seg_feature [N, H_dim, T] - # cls_feature [N, F_dim, T] - if self.head is not None: - head_score = self.head(seg_feature, masks) - else: - head_score = None - # seg_score [stage_num, N, C, T] - # cls_score [N, C, T] - return backbone_score, neck_score, head_score \ No newline at end of file diff --git a/model/architectures/segmentation/stream_segmentation3d_with_backboneloss.py b/model/architectures/segmentation/stream_segmentation3d_with_backboneloss.py deleted file mode 100644 index 7cb10b6..0000000 --- a/model/architectures/segmentation/stream_segmentation3d_with_backboneloss.py +++ /dev/null @@ -1,100 +0,0 @@ -''' -Author: Thyssen Wen -Date: 2022-03-25 10:29:10 -LastEditors : Thyssen Wen -LastEditTime : 2022-07-06 20:37:59 -Description: etesvs model framework -FilePath : /ETESVS/model/architectures/segmentation/stream_segmentation3d_with_backboneloss.py -''' -import torch -import torch.nn as nn -from mmcv.runner import load_checkpoint - -from utils.logger import get_logger - -from ...builder import build_backbone -from ...builder import build_neck -from ...builder import build_head - -from ...builder import ARCHITECTURE - -@ARCHITECTURE.register() -class StreamSegmentation3DWithBackbone(nn.Module): - def __init__(self, - backbone=None, - neck=None, - head=None, - loss=None): - super().__init__() - self.backbone = build_backbone(backbone) - self.neck = build_neck(neck) - self.head = build_head(head) - - self.init_weights() - - self.sample_rate = head.sample_rate - - def init_weights(self): - if self.backbone is not None: - self.backbone.init_weights(child_model=False, revise_keys=[(r'backbone.', r'')]) - if self.neck is not None: - self.neck.init_weights() - if self.head is not None: - self.head.init_weights() - - def _clear_memory_buffer(self): - if self.backbone is not None: - self.backbone._clear_memory_buffer() - if self.neck is not None: - self.neck._clear_memory_buffer() - if self.head is not None: - self.head._clear_memory_buffer() - - def forward(self, input_data): - masks = input_data['masks'] - imgs = input_data['imgs'] - - # masks.shape=[N,T] - masks = masks.unsqueeze(1) - - # x.shape=[N,T,C,H,W], for most commonly case - imgs = imgs.transpose(1, 2).contiguous() - # x [N * T, C, H, W] - - if self.backbone is not None: - # masks.shape [N, 1, T, 1, 1] - backbone_masks = masks[:, :, ::self.sample_rate].unsqueeze(-1).unsqueeze(-1) - # [N, C, T, H, W] or [N T D] - feature = self.backbone(imgs, backbone_masks) - - if len(feature.shape) == 5: - out_channels = feature.shape[1] - # [N, T, C, H, W] - feature = feature.transpose(1, 2) - # [N * T, C, H, W] - feature = torch.reshape(feature, shape=[-1, out_channels] + list(feature.shape[-2:])) - elif len(feature.shape) == 3: - feature = torch.reshape(feature, shape=[-1, feature.shape[-1]]) - else: - feature = imgs - - # feature [N * T , F_dim, 7, 7] or [N * T, D] - # step 3 extract memory feature - if self.neck is not None: - seg_feature, backbone_score = self.neck( - feature, masks[:, :, ::self.sample_rate]) - - else: - seg_feature = feature - backbone_score = None - - # step 5 segmentation - # seg_feature [N, H_dim, T] - # cls_feature [N, F_dim, T] - if self.head is not None: - head_score = self.head(seg_feature, masks) - else: - head_score = seg_feature - # seg_score [stage_num, N, C, T] - # cls_score [N, C, T] - return backbone_score, head_score \ No newline at end of file diff --git a/model/backbones/__init__.py b/model/backbones/__init__.py index 3f77798..ba71362 100644 --- a/model/backbones/__init__.py +++ b/model/backbones/__init__.py @@ -7,7 +7,6 @@ FilePath : /ETESVS/model/backbones/__init__.py ''' from .image import ResNet, MobileNetV2, MobileViT, ViT -from .flow import FastFlowNet, RAFT, LiteFlowNetV3 from .video import (ResNet2Plus1d, ResNet3d, PredRNNV2, I3D, MobileNetV2TSM, MoViNet, TimeSformer, ResNetTSM) from .language import TransducerTextEncoder @@ -16,7 +15,7 @@ __all__ = [ 'ResNet', 'ResNetTSM', 'MobileNetV2', 'MobileNetV2TSM', 'MobileNetV2TMM' - 'ResNet3d', 'FastFlowNet', 'RAFT', 'I3D' + 'ResNet3d', 'I3D' 'MoViNet', 'MobileViT', 'ViT', 'TimeSformer', 'ResNet3d', 'ResNet2Plus1d', diff --git a/model/backbones/image/resnet.py b/model/backbones/image/resnet.py index e5647f7..41f8ce2 100644 --- a/model/backbones/image/resnet.py +++ b/model/backbones/image/resnet.py @@ -18,11 +18,13 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch.nn as nn +import torch from mmcv.cnn import ConvModule, constant_init, kaiming_init from mmcv.runner import _load_checkpoint, load_checkpoint from mmcv.utils import _BatchNorm from torch.utils import checkpoint as cp from utils.logger import get_logger +import torchvision from ...builder import BACKBONES @@ -348,7 +350,7 @@ class ResNet(nn.Module): } def __init__(self, - depth, + depth=50, pretrained=None, torchvision_pretrain=False, in_channels=3, @@ -367,6 +369,7 @@ def __init__(self, super().__init__() if depth not in self.arch_settings: raise KeyError(f'invalid depth {depth} for resnet') + self.depth = depth self.in_channels = in_channels self.pretrained = pretrained diff --git a/model/backbones/video/resnet_tsm.py b/model/backbones/video/resnet_tsm.py index ec300c0..4ce58d4 100644 --- a/model/backbones/video/resnet_tsm.py +++ b/model/backbones/video/resnet_tsm.py @@ -157,7 +157,7 @@ class ResNetTSM(ResNet): """ def __init__(self, - depth, + depth=50, modality="RGB", clip_seg_num=8, is_shift=True, @@ -317,8 +317,7 @@ def init_weights(self, child_model=False, revise_keys=[(r'^module\.', '')]): self._load_torchvision_checkpoint(logger) else: # ours - load_checkpoint( - self, self.pretrained, strict=False, logger=logger, revise_keys=revise_keys) + load_checkpoint(self, self.pretrained, strict=False, logger=logger, revise_keys=revise_keys) elif self.pretrained is None: for m in self.modules(): if isinstance(m, nn.Conv2d): diff --git a/model/heads/segmentation/lstm_head.py b/model/heads/segmentation/lstm_head.py index 9d74a82..ca023c1 100644 --- a/model/heads/segmentation/lstm_head.py +++ b/model/heads/segmentation/lstm_head.py @@ -36,13 +36,13 @@ def __init__(self, self.sample_rate = sample_rate self.num_directions = 2 if self.bidirectional else 1 - self.lstm = nn.LSTM(input_size=in_channels, hidden_size=hidden_channels, num_layers=num_layers, batch_first=self.batch_first, dropout=self.dropout, bidirectional=self.bidirectional) + self.fc_cls = nn.Linear(self.hidden_channels, self.num_classes) self.memory_hidden_list = None diff --git a/model/losses/segmentation_loss.py b/model/losses/segmentation_loss.py index 767e32c..bd29f2f 100644 --- a/model/losses/segmentation_loss.py +++ b/model/losses/segmentation_loss.py @@ -10,9 +10,50 @@ import numpy as np import torch.nn as nn import torch.nn.functional as F +import torchvision +import json from ..builder import LOSSES +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Variable + +class FocalLoss(nn.Module): + def __init__(self, gamma=0, alpha=None, size_average=True): + super(FocalLoss, self).__init__() + self.gamma = gamma + self.alpha = alpha + if isinstance(alpha,(float,int)): self.alpha = torch.Tensor([alpha,1-alpha]) + if isinstance(alpha,list): self.alpha = torch.Tensor(alpha) + self.size_average = size_average + + def forward(self, input, target): + if input.dim()>2: + input = input.view(input.size(0),input.size(1),-1) # N,C,H,W => N,C,H*W + input = input.transpose(1,2) # N,C,H*W => N,H*W,C + input = input.contiguous().view(-1,input.size(2)) # N,H*W,C => N*H*W,C + target = target.view(-1,1) + + logpt = F.log_softmax(input) + logpt = logpt.gather(1,target) + logpt = logpt.view(-1) + pt = Variable(logpt.data.exp()) + + if self.alpha is not None: + if self.alpha.type()!=input.data.type(): + self.alpha = self.alpha.type_as(input.data) + at = self.alpha.gather(0,target.data.view(-1)) + logpt = logpt * Variable(at) + + loss = -1 * (1-pt)**self.gamma * logpt + return loss + import pdb; pdb.set_trace() + if self.size_average: return loss.mean() + else: return loss.sum() + + @LOSSES.register() class SegmentationLoss(nn.Module): def __init__(self, @@ -34,25 +75,178 @@ def __init__(self, self.ce = nn.CrossEntropyLoss(weight=class_weight, ignore_index=self.ignore_index, reduction='none') else: self.ce = nn.CrossEntropyLoss(ignore_index=self.ignore_index, reduction='none') - self.mse = nn.MSELoss(reduction='none') - + self.mse = nn.MSELoss(reduction='none')# + # self.focal_loss = FocalLoss(gamma=2)#torchvision.ops.focal_loss.sigmoid_focal_loss# + self.iter = {'inp': None, 'lbl':None, 'loss':None} + def focal_loss(self, predictions, targets, ignore_idx=None): + """ + Compute Focal Loss. + + Args: + - predictions: Tensor of shape (N, C), probabilities (after softmax). + - targets: Tensor of shape (N, C), one-hot encoded. + - alpha: Weighting factor for class imbalance. + - gamma: Focusing parameter for hard-to-classify examples. + + Returns: + - Loss: Scalar loss value. + """ + # Avoid zero probabilities (numerical stability) + alpha=1 + gamma=2 + predictions = F.softmax(predictions.clamp(min=1e-7, max=1-1e-7),dim=1) + + if ignore_idx is not None: + targets = targets.clone() # Avoid modifying the original targets + targets[ignore_idx, :] = 0 # Set the target class at ignore_idx to 0 for all samples + + # Focal Loss formula + focal_loss = -alpha * (1 - predictions) ** gamma * targets * predictions.log() + return focal_loss.mean(dim=1) + def forward(self, model_output, input_data): # score shape [stage_num N C T] # masks shape [N T] + import pdb; pdb.set_trace() head_score = model_output masks, labels, precise_sliding_num = input_data["masks"], input_data["labels"], input_data['precise_sliding_num'] - _, b, _, t = head_score.shape loss = 0. for p in head_score: + + + #### focal loss ##### + # inp = p.transpose(2, 1).contiguous().view(-1, self.num_classes) + # lbl = labels.view(-1) + # lbl = lbl.clamp(0, self.num_classes - 1) + # assert lbl.min() >= 0 and lbl.max() < self.num_classes, "Labels out of range" + # assert lbl.dtype == torch.int64, "Labels must be int64" + # lbl_one_hot = F.one_hot(lbl, num_classes=5) + # # self.iter['inp'] = inp + # # self.iter['lbl'] = lbl + # seg_cls_loss = self.focal_loss(inp, lbl_one_hot, self.ignore_index) + # if torch.isnan(seg_cls_loss).any() or torch.isinf(seg_cls_loss).any(): + # print("NaN or Inf detected") + + # json.dump(self.iter, open('recent.json', 'w')) + + # seg_cls_loss = torch.mean(self.focal_loss(inp.float(), lbl_one_hot.float()), dim=1) + + #### ce loss #### seg_cls_loss = self.ce(p.transpose(2, 1).contiguous().view(-1, self.num_classes), labels.view(-1)) loss += torch.sum(torch.sum(torch.reshape(seg_cls_loss, shape=[b, t]), dim=-1) / (precise_sliding_num + self.elps)) / (torch.sum(labels != -100) + self.elps) - loss += self.smooth_weight * torch.mean( - (torch.mean(torch.reshape(torch.clamp( - self.mse(F.log_softmax(p[:, :, 1:], dim=1), F.log_softmax(p.detach()[:, :, :-1], dim=1)) - , min=0, max=16) * masks[:, 1:].unsqueeze(1), [b, -1]), dim=-1) / (precise_sliding_num + self.elps))) - + loss += self.smooth_weight * torch.mean((torch.mean(torch.reshape(torch.clamp( self.mse(F.log_softmax(p[:, :, 1:], dim=1), F.log_softmax(p.detach()[:, :, :-1], dim=1)) , min=0, max=16) * masks[:, 1:].unsqueeze(1), [b, -1]), dim=-1) / (precise_sliding_num + self.elps))) loss_dict={} loss_dict["loss"] = loss * self.loss_weight + return loss_dict + + +@LOSSES.register() +class SegmentationLossMultiLabel(nn.Module): + def __init__(self, + num_classes_action, + num_classes_branches, + loss_weight=1.0, + sample_rate=1, + smooth_weight=0.5, + ignore_index=-100, + class_weight=None): + super().__init__() + self.smooth_weight = smooth_weight + self.ignore_index = ignore_index + self.num_classes_action = num_classes_action + self.num_classes_branches = num_classes_branches + self.sample_rate = sample_rate + self.elps = 1e-10 + self.loss_weight = loss_weight + if class_weight is not None: + class_weight = torch.tensor(class_weight) + self.ce = nn.CrossEntropyLoss(weight=class_weight, ignore_index=self.ignore_index, reduction='none') + else: + self.ce = nn.CrossEntropyLoss(ignore_index=self.ignore_index, reduction='none') + self.mse = nn.MSELoss(reduction='none')# + # self.focal_loss = FocalLoss(gamma=2)#torchvision.ops.focal_loss.sigmoid_focal_loss# + self.iter = {'inp': None, 'lbl':None, 'loss':None} + def focal_loss(self, predictions, targets, num_classes, ignore_idx=None): + """ + Compute Focal Loss. + + Args: + - predictions: Tensor of shape (N, C), probabilities (after softmax). + - targets: Tensor of shape (N, C), one-hot encoded. + - alpha: Weighting factor for class imbalance. + - gamma: Focusing parameter for hard-to-classify examples. + + Returns: + - Loss: Scalar loss value. + """ + # Avoid zero probabilities (numerical stability) + alpha=1 + gamma=2 + predictions = F.softmax(predictions.clamp(min=1e-7, max=1-1e-7),dim=1) + + if ignore_idx is not None: + targets = targets.clone() # Avoid modifying the original targets + targets[ignore_idx, :] = 0 # Set the target class at ignore_idx to 0 for all samples + + # Focal Loss formula + focal_loss = -alpha * (1 - predictions) ** gamma * targets * predictions.log() + return focal_loss.mean(dim=1) + + def compute_loss(self, model_output, labels, input_data, num_classes): + """ + Compute the loss for a given model output and input data. + + Parameters: + model_output (torch.Tensor): Output from the model. + input_data (dict): Dictionary containing "masks", "labels", and "precise_sliding_num". + + Returns: + torch.Tensor: The computed loss value. + """ + head_score = model_output + labels = labels + masks, precise_sliding_num = input_data["masks"], input_data["precise_sliding_num"] + _, b, _, t = head_score.shape + + loss = 0.0 + + for p in head_score: + # Cross-entropy loss + seg_cls_loss = self.ce( + p.transpose(2, 1).contiguous().view(-1, num_classes), + labels.view(-1) + ) + loss += torch.sum( + torch.sum( + torch.reshape(seg_cls_loss, shape=[b, t]), dim=-1 + ) / (precise_sliding_num + self.elps) + ) / (torch.sum(labels != -100) + self.elps) + + # Smoothness loss + smooth_loss = self.mse( + F.log_softmax(p[:, :, 1:], dim=1), + F.log_softmax(p.detach()[:, :, :-1], dim=1) + ) + smooth_loss = torch.clamp(smooth_loss, min=0, max=16) * masks[:, 1:].unsqueeze(1) + smooth_loss = torch.mean( + torch.mean( + torch.reshape(smooth_loss, [b, -1]), dim=-1 + ) / (precise_sliding_num + self.elps) + ) + loss += self.smooth_weight * smooth_loss + + return loss + + def forward(self, model_output, input_data): + # score shape [stage_num N C T] + # masks shape [N T] + action_loss = self.compute_loss(model_output['action_score'], input_data['labels'], input_data, self.num_classes_action) + branch_loss = self.compute_loss(model_output['branch_score'], input_data['branch_labels'], input_data, self.num_classes_branches) + + loss_dict={} + loss_dict["action_loss"] = action_loss * self.loss_weight + loss_dict["branch_loss"] = branch_loss * self.loss_weight + loss_dict["loss"] = action_loss + branch_loss return loss_dict \ No newline at end of file diff --git a/model/necks/cross_attn.py b/model/necks/cross_attn.py new file mode 100644 index 0000000..c6bbbf7 --- /dev/null +++ b/model/necks/cross_attn.py @@ -0,0 +1,84 @@ +import torch +import torch.nn as nn + +class CrossAttention(nn.Module): + def __init__(self, dim, num_heads): + super(CrossAttention, self).__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + assert dim % num_heads == 0, "Dimension must be divisible by the number of heads." + self.scale = self.head_dim ** -0.5 + + self.query_proj = nn.Linear(dim, dim) + self.key_proj = nn.Linear(dim, dim) + self.value_proj = nn.Linear(dim, dim) + self.output_proj = nn.Linear(dim, dim) + + def forward(self, x1, x2): + B, C, H, W = x1.shape + HW = H * W + + # Flatten spatial dimensions and project queries, keys, and values + Q = self.query_proj(x1.view(B, C, HW).permute(0, 2, 1)) # (B, HW, C) + K = self.key_proj(x2.view(B, C, HW).permute(0, 2, 1)) # (B, HW, C) + V = self.value_proj(x2.view(B, C, HW).permute(0, 2, 1)) # (B, HW, C) + + Q = Q.view(B, HW, self.num_heads, self.head_dim).permute(0, 2, 1, 3) # (B, num_heads, HW, head_dim) + K = K.view(B, HW, self.num_heads, self.head_dim).permute(0, 2, 1, 3) # (B, num_heads, HW, head_dim) + V = V.view(B, HW, self.num_heads, self.head_dim).permute(0, 2, 1, 3) # (B, num_heads, HW, head_dim) + + # Compute attention scores: (B, num_heads, HW, head_dim) x (B, num_heads, head_dim, HW) + attention_scores = torch.matmul(Q, K.transpose(-1, -2)) * self.scale # (B, num_heads, HW, HW) + attention_probs = torch.softmax(attention_scores, dim=-1) # (B, num_heads, HW, HW) + + # Compute weighted sum of values + attention_output = torch.matmul(attention_probs, V) # (B, num_heads, HW, head_dim) + + # Concatenate heads and project back + attention_output = attention_output.permute(0, 2, 1, 3).reshape(B, HW, -1) # (B, HW, C) + output = self.output_proj(attention_output).permute(0, 2, 1).view(B, C, H, W) # (B, C, H, W) + + return output + +# class CrossAttention(nn.Module): +# def __init__(self, dim, num_heads): +# super(CrossAttention, self).__init__() +# self.num_heads = num_heads +# self.head_dim = dim // num_heads +# self.scale = self.head_dim ** -0.5 + +# # Projection layers for queries, keys, and values +# self.query_proj = nn.Linear(dim, dim) +# self.key_proj = nn.Linear(dim, dim) +# self.value_proj = nn.Linear(dim, dim) +# self.output_proj = nn.Linear(dim, dim) + +# def forward(self, x1, x2): +# # Flatten spatial dimensions +# B, C, H, W = x1.shape +# A_flat = x1.view(B, C, H * W).permute(0, 2, 1) # (B, HW, C) +# B_flat = x2.view(B, C, H * W).permute(0, 2, 1) # (B, HW, C) + +# # Project queries, keys, and values +# Q_A = self.query_proj(A_flat) # (B, HW, C) +# K_B = self.key_proj(B_flat) # (B, HW, C) +# V_B = self.value_proj(B_flat) # (B, HW, C) + +# # Compute attention scores +# attention_scores = torch.bmm(Q_A, K_B.transpose(1, 2)) * self.scale # (B, HW, HW) +# attention_probs = torch.softmax(attention_scores, dim=-1) # (B, HW, HW) + +# # Compute weighted sum of values +# attention_output = torch.bmm(attention_probs, V_B) # (B, HW, C) + +# # Project output back and reshape to original dimensions +# output = self.output_proj(attention_output).permute(0, 2, 1) # (B, C, HW) +# return output.view(B, C, H, W) + +# # Example usage +# batch, channels, height, width = 1, 2048, 7, 7 +# A = torch.randn(batch, channels, height, width) # Feature map A +# B = torch.randn(batch, channels, height, width) # Feature map B + +# cross_attention = CrossAttention(dim=channels, num_heads=8) +# output_A_to_B = cross_attention(A, B) # Cross-attention output A -> B diff --git a/model/post_precessings/stream_score_post_processing.py b/model/post_precessings/stream_score_post_processing.py index 28b7ba3..7b1a708 100644 --- a/model/post_precessings/stream_score_post_processing.py +++ b/model/post_precessings/stream_score_post_processing.py @@ -66,4 +66,97 @@ def output(self): pred_score_list.append(self.pred_scores[bs, :, :ignore_start].copy()) ground_truth_list.append(self.video_gt[bs, :ignore_start].copy()) - return pred_score_list, pred_cls_list, ground_truth_list \ No newline at end of file + return pred_score_list, pred_cls_list, ground_truth_list + +@POSTPRECESSING.register() +class StreamScorePostProcessingMultiLabel(): + def __init__(self, + num_action_classes, + num_branch_classes, + clip_seg_num, + sliding_window, + sample_rate, + ignore_index=-100): + self.clip_seg_num = clip_seg_num + self.sliding_window = sliding_window + self.sample_rate = sample_rate + self.num_action_classes = num_action_classes + self.num_branch_classes = num_branch_classes + self.ignore_index = ignore_index + self.init_flag = False + self.epls = 1e-10 + + def init_scores(self, sliding_num, batch_size): + max_temporal_len = sliding_num * self.sliding_window + self.sample_rate * self.clip_seg_num + sample_videos_max_len = max_temporal_len + \ + ((self.clip_seg_num * self.sample_rate) - max_temporal_len % (self.clip_seg_num * self.sample_rate)) + if sample_videos_max_len % self.sliding_window != 0: + sample_videos_max_len = sample_videos_max_len + \ + (self.sliding_window - (sample_videos_max_len % self.sliding_window)) + self.sample_videos_max_len = sample_videos_max_len + + # Initialize scores for action and branch outputs + self.pred_action_scores = np.zeros((batch_size, self.num_action_classes, sample_videos_max_len)) + self.pred_branch_scores = np.zeros((batch_size, self.num_branch_classes, sample_videos_max_len)) + + self.video_action_gt = np.full((batch_size, sample_videos_max_len), self.ignore_index) + self.video_branch_gt = np.full((batch_size, sample_videos_max_len), self.ignore_index) + + self.init_flag = True + + def update(self, action_scores, branch_scores, action_gt, branch_gt, idx): + # action_scores, branch_scores [stage_num N C T] + # action_gt, branch_gt [N C T] + with torch.no_grad(): + start_frame = idx * self.sliding_window + if start_frame < 0: + start_frame = 0 + end_frame = start_frame + (self.clip_seg_num * self.sample_rate) + + # Update scores and ground truths for action + self.pred_action_scores[:, :, start_frame:end_frame] = action_scores[-1, :].detach().cpu().numpy().copy() + self.video_action_gt[:, start_frame:end_frame] = action_gt.detach().cpu().numpy().copy() + + # Update scores and ground truths for branch + self.pred_branch_scores[:, :, start_frame:end_frame] = branch_scores[-1, :].detach().cpu().numpy().copy() + self.video_branch_gt[:, start_frame:end_frame] = branch_gt.detach().cpu().numpy().copy() + + # Calculate accuracy for action and branch + pred_action = np.argmax(action_scores[-1, :].detach().cpu().numpy(), axis=-2) + acc_action = np.mean((np.sum(pred_action == action_gt.detach().cpu().numpy(), axis=1) / + (np.sum(action_gt.detach().cpu().numpy() != self.ignore_index, axis=1) + self.epls))) + + pred_branch = np.argmax(branch_scores[-1, :].detach().cpu().numpy(), axis=-2) + acc_branch = np.mean((np.sum(pred_branch == branch_gt.detach().cpu().numpy(), axis=1) / + (np.sum(branch_gt.detach().cpu().numpy() != self.ignore_index, axis=1) + self.epls))) + + return acc_action, acc_branch + + def output(self): + action_pred_score_list = [] + action_pred_cls_list = [] + action_ground_truth_list = [] + + branch_pred_score_list = [] + branch_pred_cls_list = [] + branch_ground_truth_list = [] + + for bs in range(self.pred_action_scores.shape[0]): + # Action outputs + action_index = np.where(self.video_action_gt[bs, :] == self.ignore_index) + action_ignore_start = min(list(action_index[0]) + [self.sample_videos_max_len]) + action_predicted = np.argmax(self.pred_action_scores[bs, :, :action_ignore_start], axis=0) + action_pred_cls_list.append(action_predicted.squeeze().copy()) + action_pred_score_list.append(self.pred_action_scores[bs, :, :action_ignore_start].copy()) + action_ground_truth_list.append(self.video_action_gt[bs, :action_ignore_start].copy()) + + # Branch outputs + branch_index = np.where(self.video_branch_gt[bs, :] == self.ignore_index) + branch_ignore_start = min(list(branch_index[0]) + [self.sample_videos_max_len]) + branch_predicted = np.argmax(self.pred_branch_scores[bs, :, :branch_ignore_start], axis=0) + branch_pred_cls_list.append(branch_predicted.squeeze().copy()) + branch_pred_score_list.append(self.pred_branch_scores[bs, :, :branch_ignore_start].copy()) + branch_ground_truth_list.append(self.video_branch_gt[bs, :branch_ignore_start].copy()) + + return (action_pred_score_list, action_pred_cls_list, action_ground_truth_list), \ + (branch_pred_score_list, branch_pred_cls_list, branch_ground_truth_list) diff --git a/optimizer/tsm_adam_optimizer.py b/optimizer/tsm_adam_optimizer.py index a8a4d8b..868e81d 100644 --- a/optimizer/tsm_adam_optimizer.py +++ b/optimizer/tsm_adam_optimizer.py @@ -77,8 +77,15 @@ def get_optim_policies(self, model): for param in list(m.parameters()): if param.requires_grad: bn.append(param) + elif isinstance(m, torch.nn.LSTM): + for param in list(m.parameters()): + if param.requires_grad: + bn.append(param) + elif len(m._modules) == 0: if len(list(m.parameters())) > 0: + import pdb; pdb.set_trace() + raise ValueError("New atomic module type: {}. Need to give it a learning policy".format(type(m))) # pop the cls_head fc layer params diff --git a/tasks/runner.py b/tasks/runner.py index 766327e..c4ecd8e 100644 --- a/tasks/runner.py +++ b/tasks/runner.py @@ -13,7 +13,7 @@ from utils.logger import get_logger import numpy as np - +import cv2 try: from apex import amp except: @@ -56,7 +56,7 @@ def __init__(self, # self.writer = get_logger(name="SVTAS", tensorboard=True) # self.step = 1 # self.cnt = 1 - + self.segs = 0 assert runner_mode in ['train', 'validation', 'test'], "Not support this runner mode: " + runner_mode self.runner_mode = runner_mode @@ -83,7 +83,9 @@ def _log_loss_dict(self): def epoch_init(self): # batch videos sampler self._init_loss_dict() - self.seg_acc = 0. + self.action_seg_acc = 0. + self.branch_seg_acc = 0. + self.post_processing.init_flag = False self.current_step = 0 self.current_step_vid_list = None @@ -120,39 +122,78 @@ def batch_end_step(self, sliding_num, vid_list, step, epoch): self.model._clear_memory_buffer() # get pred result - pred_score_list, pred_cls_list, ground_truth_list = self.post_processing.output() - outputs = dict(predict=pred_cls_list, - output_np=pred_score_list) - vid = self.current_step_vid_list - - if self.runner_mode in ['validation', 'test']: - # distribution - if self.nprocs > 1: - collect_dict = dict( - predict=pred_cls_list, - output_np=pred_score_list, - ground_truth=ground_truth_list, - vid=self.current_step_vid_list - ) - gather_objects = [collect_dict for _ in range(self.nprocs)] # any picklable object - output_list = [None for _ in range(self.nprocs)] - dist.all_gather_object(output_list, gather_objects[dist.get_rank()]) - # collect - pred_cls_list_i = [] - pred_score_list_i = [] - ground_truth_list_i = [] - vid_i = [] - for output_dict in output_list: - pred_cls_list_i = pred_cls_list_i + output_dict["predict"] - pred_score_list_i = pred_score_list_i + output_dict["output_np"] - ground_truth_list_i = ground_truth_list_i + output_dict["ground_truth"] - vid_i = vid_i + output_dict["vid"] - outputs = dict(predict=pred_cls_list_i, - output_np=pred_score_list_i) - ground_truth_list = ground_truth_list_i - vid = vid_i - - f1, acc = self.Metric.update(vid, ground_truth_list, outputs) + # pred_score_list, pred_cls_list, ground_truth_list, = self.post_processing.output() + + def process_post_output(postprocessed, runner_mode, nprocs, current_step_vid_list): + """ + Process post-processing outputs and optionally handle distributed gathering. + + Args: + post_processing: An object with a method `output()` that returns + pred_score_list, pred_cls_list, ground_truth_list. + runner_mode: Mode of the runner, either 'validation', 'test', or others. + nprocs: Number of processes in distributed training. + current_step_vid_list: Current step video ID list. + + Returns: + tuple: (outputs, ground_truth_list, vid) + """ + # Get the initial output from post-processing + pred_score_list, pred_cls_list, ground_truth_list = postprocessed + outputs = dict( + predict=pred_cls_list, + output_np=pred_score_list + ) + vid = current_step_vid_list + + # Handle distributed gathering if in validation or test mode + # if runner_mode in ['validation', 'test']: + # if nprocs > 1: + # # Prepare collect_dict for distributed gathering + # collect_dict = dict( + # predict=pred_cls_list, + # output_np=pred_score_list, + # ground_truth=ground_truth_list, + # vid=current_step_vid_list + # ) + # gather_objects = [collect_dict for _ in range(nprocs)] # Prepare objects to gather + # output_list = [None for _ in range(nprocs)] + + # # Perform distributed all_gather_object + # dist.all_gather_object(output_list, gather_objects[dist.get_rank()]) + + # # Collect gathered data from all processes + # pred_cls_list_i = [] + # pred_score_list_i = [] + # ground_truth_list_i = [] + # vid_i = [] + + # for output_dict in output_list: + # pred_cls_list_i += output_dict["predict"] + # pred_score_list_i += output_dict["output_np"] + # ground_truth_list_i += output_dict["ground_truth"] + # vid_i += output_dict["vid"] + + # # Update outputs and ground_truth_list with gathered data + # outputs = dict( + # predict=pred_cls_list_i, + # output_np=pred_score_list_i + # ) + # ground_truth_list = ground_truth_list_i + # vid = vid_i + + return outputs, ground_truth_list, vid + + out = self.post_processing.output() + outputs, ground_truth_list, vid = process_post_output(out[0], self.runner_mode, self.nprocs, self.current_step_vid_list) + if not vid: + import pdb; pdb.set_trace() + f1_action, acc_action = self.Metric.update(vid, ground_truth_list, outputs, action_dict_path='data/thal/mapping_tasks.txt') + outputs, ground_truth_list, vid = process_post_output(out[1], self.runner_mode, self.nprocs, self.current_step_vid_list) + if not vid: + import pdb; pdb.set_trace() + f1_branch, acc_branch = self.Metric.update(vid, ground_truth_list, outputs, action_dict_path='data/thal/mapping_branches.txt') + self.current_step_vid_list = vid_list if len(self.current_step_vid_list) > 0: @@ -169,12 +210,16 @@ def batch_end_step(self, sliding_num, vid_list, step, epoch): self._log_loss_dict() self.record_dict['batch_time'].update(time.time() - self.b_tic) - self.record_dict['F1@0.5'].update(f1, self.video_batch_size) - self.record_dict['Acc'].update(acc, self.video_batch_size) - self.record_dict['Seg_Acc'].update(self.seg_acc, self.video_batch_size) + self.record_dict['F1Action@0.5'].update(f1_action, self.video_batch_size) + self.record_dict['F1Branch@0.5'].update(f1_branch, self.video_batch_size) + self.record_dict['ActionAcc'].update(acc_action, self.video_batch_size) + self.record_dict['BranchAcc'].update(acc_branch, self.video_batch_size) + self.record_dict['ActionSeg_Acc'].update(self.action_seg_acc, self.video_batch_size) + self.record_dict['BranchSeg_Acc'].update(self.branch_seg_acc, self.video_batch_size) self._init_loss_dict() - self.seg_acc = 0. + self.action_seg_acc = 0. + self.branch_seg_acc = 0. if self.current_step % self.cfg.get("log_interval", 10) == 0: ips = "ips: {:.5f} instance/sec.".format( @@ -189,17 +234,27 @@ def batch_end_step(self, sliding_num, vid_list, step, epoch): self.b_tic = time.time() self.current_step = step - + + def postprocess_frame(self, frame_list, wh = (1280, 720)): + mean = torch.tensor([0.485, 0.456, 0.406]) + std = torch.tensor([0.229, 0.224, 0.225]) + + original_frames = [] + for frame in frame_list: + unnormalized_frame = frame * std[:, None, None] + mean[:, None, None] + uint8_frame = (unnormalized_frame.clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8) + original_frames.append(cv2.resize(uint8_frame, wh)) + + return original_frames def _model_forward(self, data_dict): # move data input_data = {} for key, value in data_dict.items(): if torch.is_tensor(value): input_data[key] = value.cuda() - + outputs = self.model(input_data) loss_dict = self.criterion(outputs, input_data) - loss = loss_dict["loss"] if self.runner_mode in ['train']: if self.use_amp is True: @@ -207,8 +262,7 @@ def _model_forward(self, data_dict): scaled_loss.backward() else: loss.backward() - - if not torch.is_tensor(outputs): + if not torch.is_tensor(outputs) and not isinstance(outputs, dict): outputs = outputs[-1] return outputs, loss_dict @@ -218,7 +272,20 @@ def run_one_clip(self, data_dict): sliding_num = data_dict['sliding_num'] idx = data_dict['current_sliding_cnt'] labels = data_dict['labels'] + branch_labels = data_dict['branch_labels'] # train segment + # self.segs += 1 + # mean = torch.tensor([0.485, 0.456, 0.406]) + # std = torch.tensor([0.229, 0.224, 0.225]) + # fourcc = cv2.VideoWriter_fourcc(*'XVID') + # out = cv2.VideoWriter(f'input_segments/{self.segs}.mp4', fourcc, 30, (224, 224)) + # for frame in data_dict['imgs'][0]: + # unnorm = frame + # norm = unnorm * std[:, None, None] + mean[:, None, None] + # norm = (norm.numpy()*255).astype(np.uint8).transpose(1,2,0)[:,:,::-1] + # out.write(norm) + + # out.release() if self.nprocs > 1 and idx < sliding_num - 1 and self.use_amp is False: with self.model.no_sync(): # multi-gpus @@ -232,12 +299,14 @@ def run_one_clip(self, data_dict): # input=score, # scale_factor=[1, 4], # mode="nearest") - with torch.no_grad(): if self.post_processing.init_flag is not True: self.post_processing.init_scores(sliding_num, len(vid_list)) self.current_step_vid_list = vid_list - self.seg_acc += self.post_processing.update(score, labels, idx) / sliding_num + accs = [acc / sliding_num for acc in list(self.post_processing.update(score['action_score'],score['branch_score'], labels, branch_labels, idx))] + self.action_seg_acc += accs[0] + self.branch_seg_acc += accs[1] + # logger loss self._update_loss_dict(loss_dict) @@ -247,11 +316,13 @@ def run_one_iter(self, data, r_tic=None, epoch=None): self.record_dict['reader_time'].update(time.time() - r_tic) for sliding_seg in data: + step = sliding_seg['step'] vid_list = sliding_seg['vid_list'] sliding_num = sliding_seg['sliding_num'] idx = sliding_seg['current_sliding_cnt'] # wheather next step + if self.current_step != step or (len(vid_list) <= 0 and step == 1): self.batch_end_step(sliding_num=sliding_num, vid_list=vid_list, step=step, epoch=epoch) diff --git a/tasks/test.py b/tasks/test.py index a8cb91d..2c9ca40 100644 --- a/tasks/test.py +++ b/tasks/test.py @@ -107,7 +107,6 @@ def test(cfg, checkpoint = torch.load(weights, map_location=map_location) state_dicts = checkpoint['model_state_dict'] - if nprocs > 1: model.module.load_state_dict(state_dicts) else: diff --git a/tasks/train.py b/tasks/train.py index ea792b5..a79d97f 100644 --- a/tasks/train.py +++ b/tasks/train.py @@ -103,9 +103,9 @@ def train(cfg, # Resume resume_epoch = cfg.get("resume_epoch", 0) - if resume_epoch: - path_checkpoint = osp.join(output_dir, - model_name + f"_epoch_{resume_epoch:05d}" + ".pkl") + if weights: + path_checkpoint = weights#osp.join(output_dir, + # model_name + f"_epoch_{resume_epoch:05d}" + ".pkl") if local_rank < 0: checkpoint = torch.load(path_checkpoint) @@ -247,7 +247,6 @@ def evaluate(best): if local_rank <= 0: # metric output Metric_dict = runner.Metric.accumulate() - if Metric_dict["Acc"] > best: best = Metric_dict["Acc"] best_flag = True diff --git a/tools/avi2mp4.py b/tools/avi2mp4.py new file mode 100644 index 0000000..dd3aa9e --- /dev/null +++ b/tools/avi2mp4.py @@ -0,0 +1,23 @@ +import os +import subprocess + +def convert_avi_to_mp4(directory): + """Convert all .avi files in the specified directory to .mp4 using ffmpeg.""" + for filename in os.listdir(directory): + if filename.lower().endswith(".avi"): + input_path = os.path.join(directory, filename) + output_path = os.path.join(directory, os.path.splitext(filename)[0] + ".mp4") + + # Run ffmpeg command to copy streams without re-encoding + command = [ + "ffmpeg", "-i", input_path, "-c:v", "copy", "-c:a", "copy", output_path + ] + + print(f"Converting: {input_path} -> {output_path}") + subprocess.run(command, check=True) + + print(f"Conversion completed: {output_path}") + +if __name__ == "__main__": + directory = 'data/thal/Videos'# Use current working directory + convert_avi_to_mp4(directory) diff --git a/tools/find_unused_files.py b/tools/find_unused_files.py new file mode 100644 index 0000000..cf1a260 --- /dev/null +++ b/tools/find_unused_files.py @@ -0,0 +1,62 @@ +import os + +# Define paths for the folder and the two text files +folder_path = "/home/multi-gpu/amur/copilot/copilot-svtas/data/thal/groundTruth" +old_text_file_1 = "/home/multi-gpu/amur/copilot/copilot-svtas/data/thal/splits/testLongOnly_test.bundle" ### test file +old_text_file_2 = "/home/multi-gpu/amur/copilot/copilot-svtas/data/thal/splits/testLongOnly_train.bundle" ### train file + +# New text file names (to avoid overwriting the originals) +new_text_file_1 = "/home/multi-gpu/amur/copilot/copilot-svtas/data/thal/splits/ext_testLongOnly_test.bundle" +new_text_file_2 = "/home/multi-gpu/amur/copilot/copilot-svtas/data/thal/splits/ext_testLongOnly_train.bundle" + +# Get the set of filenames in the folder +folder_files = set(os.listdir(folder_path)) + +# Read the contents of the old text files into sets and lists +with open(old_text_file_1, 'r') as f1: + old_entries_1 = [line.strip() for line in f1 if line.strip()] +with open(old_text_file_2, 'r') as f2: + old_entries_2 = [line.strip() for line in f2 if line.strip()] + +# Create a set of already listed files from both text files +already_listed = set(old_entries_1) | set(old_entries_2) + +# Identify files that are in the folder but not already listed +files_not_listed = folder_files - already_listed + +# Prepare new entries based on the file's number of lines +new_entries_1 = [] # for files with > 1800 lines +new_entries_2 = [] # for files with <= 1800 lines + +for filename in files_not_listed: + file_path = os.path.join(folder_path, filename) + + # Only process if it's a file + if os.path.isfile(file_path): + try: + with open(file_path, 'r') as file: + num_lines = sum(1 for _ in file) + except Exception as e: + print(f"Error reading {file_path}: {e}") + continue + + if num_lines > 1800: + new_entries_1.append(filename) + else: + new_entries_2.append(filename) + +# Combine old and new entries for each new text file +combined_entries_1 = old_entries_1 + new_entries_1 +combined_entries_2 = old_entries_2 + new_entries_2 + +# Write the combined entries to new text files +with open(new_text_file_1, 'w') as nt1: + for name in combined_entries_1: + nt1.write(name + "\n") + +with open(new_text_file_2, 'w') as nt2: + for name in combined_entries_2: + nt2.write(name + "\n") + +print(f"Total entries in {new_text_file_1}: {len(combined_entries_1)}") +print(f"Total entries in {new_text_file_2}: {len(combined_entries_2)}") \ No newline at end of file diff --git a/tools/plot.py b/tools/plot.py new file mode 100644 index 0000000..0d4d09b --- /dev/null +++ b/tools/plot.py @@ -0,0 +1,66 @@ +import cv2 + +def plot_labels_on_video(input_video_path, labels_txt_path1, labels_txt_path2, output_video_path, font_scale=1, font_thickness=2): + # Open the video + cap = cv2.VideoCapture(input_video_path) + if not cap.isOpened(): + print("Error: Unable to open the video file.") + return + + # Read labels from the text files + with open(labels_txt_path1, 'r') as f1, open(labels_txt_path2, 'r') as f2: + labels1 = f1.read().splitlines()[:-1] + labels2 = f2.read().splitlines()[:-1] + + # Get video properties + frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = int(cap.get(cv2.CAP_PROP_FPS)) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + # Check label count matches frame count + if len(labels1) != total_frames - 1 or len(labels2) != total_frames - 1: + print(len(labels1), len(labels2), total_frames-1) + print(f"Error: The number of labels in the files does not match the number of frames in the video.") + return + + # Define codec and create VideoWriter object + fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Use 'XVID' for AVI + out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height)) + + # Define text properties + font = cv2.FONT_HERSHEY_SIMPLEX + text_position1 = (50, 50) # Position for the first set of labels (x, y) + text_position2 = (50, 100) # Position for the second set of labels (x, y) + + frame_index = 0 + while True: + ret, frame = cap.read() + if not ret: + break + + # Overlay the labels on the frame + if frame_index == total_frames - 1: + break + label1 = 'Pred: ' + labels1[frame_index] + label2 = 'GT: ' + labels2[frame_index] + + cv2.putText(frame, label1, text_position1, font, font_scale, (0, 255, 0), font_thickness, lineType=cv2.LINE_AA) + cv2.putText(frame, label2, text_position2, font, font_scale, (0, 0, 255), font_thickness, lineType=cv2.LINE_AA) + + # Write the frame to the output video + out.write(frame) + + frame_index += 1 + + # Release resources + cap.release() + out.release() + print(f"Plotted video saved at {output_video_path}") + +# Example usage +input_video = 'data/thal/Videos/IMG_4579_counter.mp4' +labels_txt1 = 'output/results/pred_gt_list/IMG_4579_counter-gt.txt' +labels_txt2 = 'output/results/pred_gt_list/IMG_4579_counter-pred.txt' +output_video = 'vis/IMG_4579_counter.mp4' +plot_labels_on_video(input_video, labels_txt1, labels_txt2, output_video) diff --git a/tools/plot_multi.py b/tools/plot_multi.py new file mode 100644 index 0000000..182f2d7 --- /dev/null +++ b/tools/plot_multi.py @@ -0,0 +1,75 @@ +import cv2 + +def plot_labels_on_video(input_video_path, gt_txt_path1, gt_txt_path2, pred_txt_path1, pred_txt_path2, output_video_path, font_scale=1, font_thickness=2): + # Open the video + cap = cv2.VideoCapture(input_video_path) + if not cap.isOpened(): + print("Error: Unable to open the video file.") + return + + # Read and concatenate labels from the text files + with open(gt_txt_path1, 'r') as gt1, open(gt_txt_path2, 'r') as gt2: + gt_labels1 = gt1.read().splitlines()[:-1] + gt_labels2 = gt2.read().splitlines()[:-1] + gt_labels = [f"{g1} {g2}" for g1, g2 in zip(gt_labels1, gt_labels2)] + + with open(pred_txt_path1, 'r') as pred1, open(pred_txt_path2, 'r') as pred2: + pred_labels1 = pred1.read().splitlines()[:-1] + pred_labels2 = pred2.read().splitlines()[:-1] + pred_labels = [f"{p1} {p2}" for p1, p2 in zip(pred_labels1, pred_labels2)] + + # Get video properties + frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = int(cap.get(cv2.CAP_PROP_FPS)) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + # Check label count matches frame count + if len(gt_labels) != total_frames - 1 or len(pred_labels) != total_frames - 1: + print(len(gt_labels), len(pred_labels), total_frames - 1) + print(f"Error: The number of labels in the files does not match the number of frames in the video.") + return + + # Define codec and create VideoWriter object + fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Use 'XVID' for AVI + out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height)) + + # Define text properties + font = cv2.FONT_HERSHEY_SIMPLEX + text_position_gt = (50, 50) # Position for ground truth labels (x, y) + text_position_pred = (50, 100) # Position for prediction labels (x, y) + + frame_index = 0 + while True: + ret, frame = cap.read() + if not ret: + break + + # Overlay the labels on the frame + if frame_index == total_frames - 1: + break + gt_label = 'GT: ' + gt_labels[frame_index] + pred_label = 'Pred: ' + pred_labels[frame_index] + + cv2.putText(frame, gt_label, text_position_gt, font, font_scale, (0, 255, 0), font_thickness, lineType=cv2.LINE_AA) + cv2.putText(frame, pred_label, text_position_pred, font, font_scale, (0, 0, 255), font_thickness, lineType=cv2.LINE_AA) + + # Write the frame to the output video + out.write(frame) + + frame_index += 1 + + # Release resources + cap.release() + out.release() + print(f"Plotted video saved at {output_video_path}") + +# Example usage +input_video = 'data/thal/Videos/IMG_4742_counter.mp4' +gt_txt1 = 'output/results/pred_gt_list/IMG_4742_counter-b1-gt.txt' +gt_txt2 = 'output/results/pred_gt_list/IMG_4742_counter-tag-gt.txt' +pred_txt1 = 'output/results/pred_gt_list/IMG_4742_counter-b1-pred.txt' +pred_txt2 = 'output/results/pred_gt_list/IMG_4742_counter-tag-pred.txt' +output_video = 'vis/IMG_4742_counter.mp4' + +plot_labels_on_video(input_video, gt_txt1, gt_txt2, pred_txt1, pred_txt2, output_video) diff --git a/tools/rotate.py b/tools/rotate.py new file mode 100644 index 0000000..7439e2e --- /dev/null +++ b/tools/rotate.py @@ -0,0 +1,47 @@ +import os +import subprocess + +def convert_and_rotate_videos(input_dir, output_dir): + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + for filename in os.listdir(input_dir): + if filename.endswith(".MOV"): + input_path = os.path.join(input_dir, filename) + base_name = os.path.splitext(filename)[0] + output_path = os.path.join(output_dir, base_name + ".mp4") + + if "counter" in filename.lower(): + rotate_filter = "transpose=1" # Rotate 90 degrees clockwise + elif "double" in filename.lower(): + rotate_filter = "transpose=2,transpose=2" # Rotate 180 degrees + else: + rotate_filter = None # No rotation + + # Build the ffmpeg command + command = [ + "ffmpeg", + "-i", input_path, + ] + + if rotate_filter: + command += ["-vf", rotate_filter] + + command += [ + "-c:v", "libx264", # Use H.264 codec for output + "-preset", "fast", # Use a fast preset + "-crf", "23", # Set quality + output_path + ] + + # Run the command + try: + subprocess.run(command, check=True) + print(f"Converted and rotated: {filename} -> {output_path}") + except subprocess.CalledProcessError as e: + print(f"Error processing {filename}: {e}") + +# Example usage +input_directory = "data/Colored Thal Dataset" +output_directory = "data/Colored Thal Dataset/rot" +convert_and_rotate_videos(input_directory, output_directory) diff --git a/tools/split_labels.py b/tools/split_labels.py new file mode 100644 index 0000000..1e6dd52 --- /dev/null +++ b/tools/split_labels.py @@ -0,0 +1,49 @@ +import os +import re + +def process_label_files(input_folder, output_folder): + # Ensure the output folder exists + os.makedirs(output_folder, exist_ok=True) + + # Iterate over all text files in the input folder + for filename in os.listdir(input_folder): + if filename.endswith(".txt"): + input_path = os.path.join(input_folder, filename) + + # Output file paths + base_name = os.path.splitext(filename)[0] + output_b_number_path = os.path.join(output_folder, f"{base_name}_branch.txt") + output_word_path = os.path.join(output_folder, f"{base_name}.txt") + + with open(input_path, "r") as infile, \ + open(output_b_number_path, "w") as b_number_file, \ + open(output_word_path, "w") as word_file: + + for line in infile: + # Split the line into index and label + parts = line.strip().split(": ", 1) + if len(parts) == 2: + label = parts[1] + + if label == "background": + # Write "background" to both files + b_number_file.write("background\n") + word_file.write("background\n") + else: + # Match the pattern for "bNUMBERWORD" + match = re.match(r'(b\d+)([a-zA-Z]+)', label) + if match: + b_number = match.group(1) + word = match.group(2) + # Write to respective output files + b_number_file.write(f"{b_number}\n") + word_file.write(f"{word}\n") + + print(f"Processing complete. Output files saved in {output_folder}") + +# Specify the input and output folder paths +input_folder = "data/thal_new/groundTruth" +output_folder = "data/thal_new/groundTruth_split" + +# Process the files +process_label_files(input_folder, output_folder) diff --git a/tools/split_mapping.py b/tools/split_mapping.py new file mode 100644 index 0000000..5241296 --- /dev/null +++ b/tools/split_mapping.py @@ -0,0 +1,49 @@ +import os +import re + +def process_mapping_file(input_file, output_folder): + # Ensure the output folder exists + os.makedirs(output_folder, exist_ok=True) + + # Output file path + base_name = os.path.splitext(os.path.basename(input_file))[0] + output_file_path = os.path.join(output_folder, f"{base_name}_unique_labels.txt") + + # Sets to store unique labels + unique_b_numbers = set() + unique_words = set() + + with open(input_file, "r") as infile, open(output_file_path, "w") as outfile: + for line in infile: + # Split the line into index and label + parts = line.strip().split(" ", 1) + if len(parts) == 2: + label = parts[1] + + if label == "background": + # Ensure "background" is only written once + if "background" not in unique_b_numbers: + outfile.write("background\n") + unique_b_numbers.add("background") + else: + # Match the pattern for "bNUMBERWORD" + match = re.match(r'(b\d+)([a-zA-Z]+)', label) + if match: + b_number = match.group(1) + word = match.group(2) + # Write unique bNUMBER and word + if b_number not in unique_b_numbers: + outfile.write(f"{b_number}\n") + unique_b_numbers.add(b_number) + if word not in unique_words: + outfile.write(f"{word}\n") + unique_words.add(word) + + print(f"Processing complete. Output file saved at {output_file_path}") + +# Specify the input file and output folder paths +input_file = "data/thal/mapping.txt" +output_folder = "data/thal/mapping_split.txt" + +# Process the mapping file +process_mapping_file(input_file, output_folder) diff --git a/tools/verify_splits.py b/tools/verify_splits.py new file mode 100644 index 0000000..28186b9 --- /dev/null +++ b/tools/verify_splits.py @@ -0,0 +1,15 @@ +# Read the filenames from both files into sets +with open('data/THAL_CLSWISE/thal/splits/test.split1.bundle', 'r') as f1, open('data/THAL_CLSWISE/thal/splits/train.split1.bundle', 'r') as f2: + set1 = set(f1.read().splitlines()) + set2 = set(f2.read().splitlines()) + +# Find the common files +print(set1) +print('\n') +print(set2) +common_files = set1.intersection(set2) + +if common_files: + print("Common files:", common_files) +else: + print("No common files.") diff --git a/utils/recorder.py b/utils/recorder.py index cc77cd2..055746d 100644 --- a/utils/recorder.py +++ b/utils/recorder.py @@ -76,6 +76,32 @@ def build_recod(architecture_type, mode): 'Acc': AverageMeter("Acc", '.5f'), 'Seg_Acc': AverageMeter("Seg_Acc", '.5f') } + elif architecture_type in ['StreamSegmentation2DMultiLabel']: + if mode == "train": + return {'batch_time': AverageMeter('batch_cost', '.5f'), + 'reader_time': AverageMeter('reader_time', '.5f'), + 'action_loss': AverageMeter('action_loss', '7.5f'), + 'branch_loss': AverageMeter('branch_loss', '7.5f'), + 'lr': AverageMeter('lr', 'f', need_avg=False), + 'F1Action@0.5': AverageMeter("F1Action@0.5", '.5f'), + 'F1Branch@0.5': AverageMeter("F1Branch@0.5", '.5f'), + 'ActionAcc': AverageMeter("Acc", '.5f'), + 'ActionSeg_Acc': AverageMeter("ActionSeg_Acc", '.5f'), + 'BranchAcc': AverageMeter("Acc", '.5f'), + 'BranchSeg_Acc': AverageMeter("BranchSeg_Acc", '.5f') + } + elif mode == "validation": + return {'batch_time': AverageMeter('batch_cost', '.5f'), + 'reader_time': AverageMeter('reader_time', '.5f'), + 'action_loss': AverageMeter('action_loss', '7.5f'), + 'branch_loss': AverageMeter('branch_loss', '7.5f'), + 'F1Action@0.5': AverageMeter("F1Action@0.5", '.5f'), + 'F1Branch@0.5': AverageMeter("F1Branch@0.5", '.5f'), + 'ActionAcc': AverageMeter("Acc", '.5f'), + 'ActionSeg_Acc': AverageMeter("ActionSeg_Acc", '.5f'), + 'BranchAcc': AverageMeter("Acc", '.5f'), + 'BranchSeg_Acc': AverageMeter("BranchSeg_Acc", '.5f') + } elif architecture_type in ["SegmentationCLIP"]: if mode == "train": return {'batch_time': AverageMeter('batch_cost', '.5f'),