Merge pull request #22 from freds-dev/feat/four-channels

Feat/four channels

Author: jakobdanel

code/start_scene_cross_validation.py

@@ -17,7 +17,7 @@ def start_scene_cross_validation(dataset_name,video_event_name,video_rgb_name,co
         # Run the Bash script with arguments
         if id is None:
             id = -1
-        id = subprocess.check_output(['bash', "create_scene_split.sh", dataset_name, scene,video_event_name,video_rgb_name,config_name,str(id)]).strip().split("\n")[-1]
+        id = subprocess.check_output(['bash', "create_scene_split.sh", dataset_name, scene,video_event_name,video_rgb_name,config_name,str(id)]).decode('utf-8').strip().split("\n")[-1]
 
 
 if __name__ == "__main__":
@@ -30,4 +30,4 @@ def start_scene_cross_validation(dataset_name,video_event_name,video_rgb_name,co
 
 
     args = parser.parse_args()
-    start_scene_cross_validation(args.dataset,args.video_event_name,args.video_rgb_name,args.config_nameargs.exception_scenes)
+    start_scene_cross_validation(args.dataset,args.video_event_name,args.video_rgb_name,args.config_name,args.exception_scenes)

code/utils/build_sbatch.py

@@ -81,4 +81,4 @@ def main():
     write_file(args.script_location,gpu_script_content)
 
 if __name__ == "__main__":
-    main()
+    main()

code/utils/build_split_script.py

@@ -45,10 +45,10 @@ def main():
     parser.add_argument("--video_event_name", type=str, required=True,help="Directory where the event videos are located")
     parser.add_argument("--video_rgb_name", type=str, required=True,help="Directory where the rgb videos are located")
     parser.add_argument("--config_name",type=str,required=True,help="Name of the configuration file")
-    parser.add_argument("--cpus", type=int, default=18, help="Number of CPUs")
-    parser.add_argument("--memory", type=int, default=48, help="Memory in GB")
-    parser.add_argument("--hours", type=int, default=12, help="Wallclock time in hours")
-    parser.add_argument("--partition", default="normal", help="Partition for the job")
+    parser.add_argument("--cpus", type=int, default=36, help="Number of CPUs")
+    parser.add_argument("--memory", type=int, default=25, help="Memory in GB")
+    parser.add_argument("--hours", type=int, default=8, help="Wallclock time in hours")
+    parser.add_argument("--partition", default="normal,long", help="Partition for the job")
 
     args = parser.parse_args()
 

code/yolov7_custom/train.py

@@ -38,9 +38,48 @@
 from utils.plots import plot_lr_scheduler, plot_images, plot_labels, plot_results, plot_evolution
 from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, is_parallel
 from utils.wandb_logging.wandb_utils import WandbLogger, check_wandb_resume
-from ..utils.paths import get_data_yaml, get_result_dir
 logger = logging.getLogger(__name__)
 
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+IS_LOCAL = os.getenv('IS_LOCAL') == "TRUE"
+LOCAL_PATH_DATA = os.getenv('LOCAL_PATH_DATA')
+PALMA_PATH_DATA = os.getenv('PALMA_PATH_DATA')
+
+def get_base_dir()-> str:
+    return LOCAL_PATH_DATA if IS_LOCAL else PALMA_PATH_DATA
+
+def get_homography() -> str:
+    return os.path.join(get_base_dir(), "homography_calib.yaml")
+
+def get_calib_rgb() -> str:
+    return os.path.join(get_base_dir(), "calib_dng.yaml")
+
+def get_calib_event() -> str:
+    return os.path.join(get_base_dir(), "calib_raw.yaml")
+
+def get_config_path(config_name:str) -> str:
+    return os.path.join(get_base_dir(),"configs",config_name)
+
+def get_annotations_path()->str:
+    return os.path.join(get_base_dir(), "annotations.ndjson")
+
+def get_dataset_dir(dataset_name:str)-> str:
+    return os.path.join(get_base_dir(),"datasets",dataset_name)
+
+def get_video_dir(video_dir_name: str) -> str:
+    return os.path.join(get_base_dir(),"videos",video_dir_name)
+
+def get_data_yaml(dataset_name: str) -> str:
+    return os.path.join(get_base_dir(),"datasets",dataset_name,"data.yaml")
+    
+def get_result_dir(dataset_name: str) -> str:
+    dir = os.path.join(get_base_dir(),"results",dataset_name)
+    os.makedirs(dir,exist_ok=True)
+    return dir
 
 def train(hyp, opt, device, tb_writer=None):
     logger.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
@@ -88,18 +127,18 @@ def train(hyp, opt, device, tb_writer=None):
 
     # Model
     pretrained = weights.endswith('.pt')
-    if pretrained:
-        with torch_distributed_zero_first(rank):
-            attempt_download(weights)  # download if not found locally
-        ckpt = torch.load(weights, map_location=device)  # load checkpoint
-        model = Model(opt.cfg or ckpt['model'].yaml, ch=(4 if opt.four_channels else 3)*opt.multi_frame, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
-        exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else []  # exclude keys
-        state_dict = ckpt['model'].float().state_dict()  # to FP32
-        state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)  # intersect
-        model.load_state_dict(state_dict, strict=False)  # load
-        logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights))  # report
-    else:
-        model = Model(opt.cfg, ch=(4 if opt.four_channels else 3)*opt.multi_frame, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
+    #if pretrained:
+    #    with torch_distributed_zero_first(rank):
+    #        attempt_download(weights)  # download if not found locally
+    #    ckpt = torch.load(weights, map_location=device)  # load checkpoint
+    #    model = Model(opt.cfg or ckpt['model'].yaml, ch=(4 if opt.four_channels else 3)*opt.multi_frame, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
+    #    exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else []  # exclude keys
+    #    state_dict = ckpt['model'].float().state_dict()  # to FP32
+    #    state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)  # intersect
+    #    model.load_state_dict(state_dict, strict=False)  # load
+    #    logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights))  # report
+    #else:
+    model = Model(opt.cfg, ch=(4 if opt.four_channels else 3)*opt.multi_frame, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
     with torch_distributed_zero_first(rank):
         check_dataset(data_dict)  # check
     train_path = data_dict['train']
@@ -208,6 +247,7 @@ def train(hyp, opt, device, tb_writer=None):
 
     # Resume
     start_epoch, best_fitness = 0, 0.0
+    """
     if pretrained:
         # Optimizer
         if ckpt['optimizer'] is not None:
@@ -233,7 +273,7 @@ def train(hyp, opt, device, tb_writer=None):
             epochs += ckpt['epoch']  # finetune additional epochs
 
         del ckpt, state_dict
-
+"""
     # if opt.multi_frame > 1:
         # multi_train_path = stack_images(train_path, opt.multi_frame)
         # train_path = multi_train_path
@@ -477,8 +517,22 @@ def train(hyp, opt, device, tb_writer=None):
                     tb_writer.add_scalar(tag, x, epoch)  # tensorboard
                 if wandb_logger.wandb:
                     wandb_logger.log({tag: x})  # W&B
+           
+
+
+            def flatten_data(data):
+                flattened_data = []
+                for item in data:
+                    if isinstance(item, np.ndarray):
+                        flattened_data.extend(item.flatten())
+                    else:
+                        flattened_data.append(item)
+                return flattened_data
 
             # Update best mAP
+            print(results)
+            results = flatten_data(results)
+            print(results)
             fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
             if fi > best_fitness:
                 best_fitness = fi
@@ -571,13 +625,13 @@ def train(hyp, opt, device, tb_writer=None):
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', type=str, default='yolo7.pt', help='initial weights path')
-    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
-    parser.add_argument('--dataset', type=str, default='data/coco.yaml', help='data.yaml path')
-    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.p5.yaml', help='hyperparameters path')
+    parser.add_argument('--weights', type=str, default='/scratch/tmp/jdanel/data/best.pt', help='initial weights path')
+    parser.add_argument('--cfg', type=str, default='/home/j/jdanel/codespace/ML4IM/code/yolov7_custom/cfg/training/yolov7.yaml', help='model.yaml path')
+    parser.add_argument('--data', type=str, default='data/coco.yaml', help='data.yaml path')
+    parser.add_argument('--hyp', type=str, default='/home/j/jdanel/codespace/ML4IM/code/yolov7_custom/data/hyp.scratch.p5.yaml', help='hyperparameters path')
     parser.add_argument('--epochs', type=int, default=100)
-    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
-    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
+    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs') #TODO: Change default to 256
+    parser.add_argument('--img-size', nargs='+', type=int, default=[1280, 1280], help='[train, test] image sizes') # TODO: Change default to 1280,1280
     parser.add_argument('--rect', action='store_true', help='rectangular training')
     parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
     parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
@@ -613,14 +667,10 @@ def main():
     parser.add_argument('--multi-frame', type=int, default=1, choices=range(1,101), help='how many frames to load at once')
     parser.add_argument('--center-point', action='store_true', help='use center point metric instead of iou')
     opt = parser.parse_args()
-
+    print(opt.data)
     # Set DDP variables
     opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
     opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1
-    dataset = opt.dataset
-    opt.dataset = get_data_yaml(dataset)
-    opt.project = get_result_dir(dataset)
-
     set_logging(opt.global_rank)
     #if opt.global_rank in [-1, 0]:
     #    check_git_status()

code/yolov7_custom/utils/datasets.py

@@ -451,7 +451,7 @@ def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, r
                 x[:, 0] = 0
 
         n = len(shapes)  # number of images
-        bi = np.floor(np.arange(n) / batch_size).astype(np.int)  # batch index
+        bi = np.floor(np.arange(n) / batch_size).astype(int)  # batch index
         nb = bi[-1] + 1  # number of batches
         self.batch = bi  # batch index of image
         self.n = n
@@ -479,7 +479,7 @@ def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, r
                 elif mini > 1:
                     shapes[i] = [1, 1 / mini]
 
-            self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
+            self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(int) * stride
 
         # Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)
         self.imgs = [None] * n

code/yolov7_custom/utils/general.py

@@ -219,7 +219,7 @@ def labels_to_class_weights(labels, nc=80):
         return torch.Tensor()
 
     labels = np.concatenate(labels, 0)  # labels.shape = (866643, 5) for COCO
-    classes = labels[:, 0].astype(np.int)  # labels = [class xywh]
+    classes = labels[:, 0].astype(int)  # labels = [class xywh]
     weights = np.bincount(classes, minlength=nc)  # occurrences per class
 
     # Prepend gridpoint count (for uCE training)

code/yolov7_custom/utils/loss.py

@@ -642,7 +642,7 @@ def build_targets(self, p, targets, imgs):
         #indices, anch = self.find_4_positive(p, targets)
         #indices, anch = self.find_5_positive(p, targets)
         #indices, anch = self.find_9_positive(p, targets)
-
+        device = torch.device(targets.device)
         matching_bs = [[] for pp in p]
         matching_as = [[] for pp in p]
         matching_gjs = [[] for pp in p]
@@ -682,7 +682,7 @@ def build_targets(self, p, targets, imgs):
                 all_gj.append(gj)
                 all_gi.append(gi)
                 all_anch.append(anch[i][idx])
-                from_which_layer.append(torch.ones(size=(len(b),)) * i)
+                from_which_layer.append((torch.ones(size=(len(b),)) * i).to(device))
 
                 fg_pred = pi[b, a, gj, gi]                
                 p_obj.append(fg_pred[:, 4:5])
@@ -739,7 +739,7 @@ def build_targets(self, p, targets, imgs):
                 + 3.0 * pair_wise_iou_loss
             )
 
-            matching_matrix = torch.zeros_like(cost)
+            matching_matrix = torch.zeros_like(cost, device=device)
 
             for gt_idx in range(num_gt):
                 _, pos_idx = torch.topk(
@@ -753,7 +753,7 @@ def build_targets(self, p, targets, imgs):
                 _, cost_argmin = torch.min(cost[:, anchor_matching_gt > 1], dim=0)
                 matching_matrix[:, anchor_matching_gt > 1] *= 0.0
                 matching_matrix[cost_argmin, anchor_matching_gt > 1] = 1.0
-            fg_mask_inboxes = matching_matrix.sum(0) > 0.0
+            fg_mask_inboxes = (matching_matrix.sum(0) > 0.0).to(device)
             matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)
 
             from_which_layer = from_which_layer[fg_mask_inboxes]