first commit

rtdetr_pytorch/tools/README.md

@@ -0,0 +1,24 @@
+
+
+Train/test script examples
+- `CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master-port=8989 tools/train.py -c path/to/config &> train.log 2>&1 &`
+- `-r path/to/checkpoint`
+- `--amp`
+- `--test-only` 
+
+
+Tuning script examples
+- `torchrun --master_port=8844 --nproc_per_node=4 tools/train.py -c configs/rtdetr/rtdetr_r18vd_6x_coco.yml -t https://github.com/lyuwenyu/storage/releases/download/v0.1/rtdetr_r18vd_5x_coco_objects365_from_paddle.pth` 
+
+
+Export script examples
+- `python tools/export_onnx.py -c path/to/config -r path/to/checkpoint --check`
+
+
+GPU do not release memory
+- `ps aux | grep "tools/train.py" | awk '{print $2}' | xargs kill -9`
+
+
+Save all logs
+- Appending `&> train.log 2>&1 &` or `&> train.log 2>&1`
+

rtdetr_pytorch/tools/export_onnx.py

@@ -0,0 +1,147 @@
+"""by lyuwenyu
+"""
+
+import os 
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
+
+import argparse
+import numpy as np 
+
+from src.core import YAMLConfig
+
+import torch
+import torch.nn as nn 
+
+
+def main(args, ):
+    """main
+    """
+    cfg = YAMLConfig(args.config, resume=args.resume)
+
+    if args.resume:
+        checkpoint = torch.load(args.resume, map_location='cpu') 
+        if 'ema' in checkpoint:
+            state = checkpoint['ema']['module']
+        else:
+            state = checkpoint['model']
+    else:
+        raise AttributeError('only support resume to load model.state_dict by now.')
+
+    # NOTE load train mode state -> convert to deploy mode
+    cfg.model.load_state_dict(state)
+
+    class Model(nn.Module):
+        def __init__(self, ) -> None:
+            super().__init__()
+            self.model = cfg.model.deploy()
+            self.postprocessor = cfg.postprocessor.deploy()
+            print(self.postprocessor.deploy_mode)
+            
+        def forward(self, images, orig_target_sizes):
+            outputs = self.model(images)
+            return self.postprocessor(outputs, orig_target_sizes)
+    
+
+    model = Model()
+
+    dynamic_axes = {
+        'images': {0: 'N', },
+        'orig_target_sizes': {0: 'N'}
+    }
+
+    data = torch.rand(1, 3, 640, 640)
+    size = torch.tensor([[640, 640]])
+
+    torch.onnx.export(
+        model, 
+        (data, size), 
+        args.file_name,
+        input_names=['images', 'orig_target_sizes'],
+        output_names=['labels', 'boxes', 'scores'],
+        dynamic_axes=dynamic_axes,
+        opset_version=16, 
+        verbose=False
+    )
+
+
+    if args.check:
+        import onnx
+        onnx_model = onnx.load(args.file_name)
+        onnx.checker.check_model(onnx_model)
+        print('Check export onnx model done...')
+
+
+    if args.simplify:
+        import onnxsim
+        dynamic = True 
+        input_shapes = {'images': data.shape, 'orig_target_sizes': size.shape} if dynamic else None
+        onnx_model_simplify, check = onnxsim.simplify(args.file_name, input_shapes=input_shapes, dynamic_input_shape=dynamic)
+        onnx.save(onnx_model_simplify, args.file_name)
+        print(f'Simplify onnx model {check}...')
+
+
+    # import onnxruntime as ort 
+    # from PIL import Image, ImageDraw, ImageFont
+    # from torchvision.transforms import ToTensor
+    # from src.data.coco.coco_dataset import mscoco_category2name, mscoco_category2label, mscoco_label2category
+
+    # # print(onnx.helper.printable_graph(mm.graph))
+
+    # # Load the original image without resizing
+    # original_im = Image.open('./hongkong.jpg').convert('RGB')
+    # original_size = original_im.size
+
+    # # Resize the image for model input
+    # im = original_im.resize((640, 640))
+    # im_data = ToTensor()(im)[None]
+    # print(im_data.shape)
+
+    # sess = ort.InferenceSession(args.file_name)
+    # output = sess.run(
+    #     # output_names=['labels', 'boxes', 'scores'],
+    #     output_names=None,
+    #     input_feed={'images': im_data.data.numpy(), "orig_target_sizes": size.data.numpy()}
+    # )
+
+    # # print(type(output))
+    # # print([out.shape for out in output])
+
+    # labels, boxes, scores = output
+
+    # draw = ImageDraw.Draw(original_im)  # Draw on the original image
+    # thrh = 0.6
+
+    # for i in range(im_data.shape[0]):
+
+    #     scr = scores[i]
+    #     lab = labels[i][scr > thrh]
+    #     box = boxes[i][scr > thrh]
+
+    #     print(i, sum(scr > thrh))
+
+    #     for b, l in zip(box, lab):
+    #         # Scale the bounding boxes back to the original image size
+    #         b = [coord * original_size[j % 2] / 640 for j, coord in enumerate(b)]
+    #         # Get the category name from the label
+    #         category_name = mscoco_category2name[mscoco_label2category[l]]
+    #         draw.rectangle(list(b), outline='red', width=2)
+    #         font = ImageFont.truetype("Arial.ttf", 15)
+    #         draw.text((b[0], b[1]), text=category_name, fill='yellow', font=font)
+
+    # # Save the original image with bounding boxes
+    # original_im.save('test.jpg')
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', '-c', type=str, )
+    parser.add_argument('--resume', '-r', type=str, )
+    parser.add_argument('--file-name', '-f', type=str, default='model.onnx')
+    parser.add_argument('--check',  action='store_true', default=False,)
+    parser.add_argument('--simplify',  action='store_true', default=False,)
+
+    args = parser.parse_args()
+
+    main(args)

rtdetr_pytorch/tools/infer.py

@@ -0,0 +1,203 @@
+import torch
+import torch.nn as nn 
+import torchvision.transforms as T
+from torch.cuda.amp import autocast
+import numpy as np 
+from PIL import Image, ImageDraw, ImageFont
+import os 
+import sys 
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
+import argparse
+import src.misc.dist as dist 
+from src.core import YAMLConfig 
+from src.solver import TASKS
+import numpy as np
+
+def postprocess(labels, boxes, scores, iou_threshold=0.55):
+    def calculate_iou(box1, box2):
+        x1, y1, x2, y2 = box1
+        x3, y3, x4, y4 = box2
+        xi1 = max(x1, x3)
+        yi1 = max(y1, y3)
+        xi2 = min(x2, x4)
+        yi2 = min(y2, y4)
+        inter_width = max(0, xi2 - xi1)
+        inter_height = max(0, yi2 - yi1)
+        inter_area = inter_width * inter_height
+        box1_area = (x2 - x1) * (y2 - y1)
+        box2_area = (x4 - x3) * (y4 - y3)
+        union_area = box1_area + box2_area - inter_area
+        iou = inter_area / union_area if union_area != 0 else 0
+        return iou
+    merged_labels = []
+    merged_boxes = []
+    merged_scores = []
+    used_indices = set()
+    for i in range(len(boxes)):
+        if i in used_indices:
+            continue
+        current_box = boxes[i]
+        current_label = labels[i]
+        current_score = scores[i]
+        boxes_to_merge = [current_box]
+        scores_to_merge = [current_score]
+        used_indices.add(i)
+        for j in range(i + 1, len(boxes)):
+            if j in used_indices:
+                continue
+            if labels[j] != current_label:
+                continue  
+            other_box = boxes[j]
+            iou = calculate_iou(current_box, other_box)
+            if iou >= iou_threshold:
+                boxes_to_merge.append(other_box.tolist())  
+                scores_to_merge.append(scores[j])
+                used_indices.add(j)
+        xs = np.concatenate([[box[0], box[2]] for box in boxes_to_merge])
+        ys = np.concatenate([[box[1], box[3]] for box in boxes_to_merge])
+        merged_box = [np.min(xs), np.min(ys), np.max(xs), np.max(ys)]
+        merged_score = max(scores_to_merge)
+        merged_boxes.append(merged_box)
+        merged_labels.append(current_label)
+        merged_scores.append(merged_score)
+    return [np.array(merged_labels)], [np.array(merged_boxes)], [np.array(merged_scores)]
+def slice_image(image, slice_height, slice_width, overlap_ratio):
+    img_width, img_height = image.size
+    
+    slices = []
+    coordinates = []
+    step_x = int(slice_width * (1 - overlap_ratio))
+    step_y = int(slice_height * (1 - overlap_ratio))
+    
+    for y in range(0, img_height, step_y):
+        for x in range(0, img_width, step_x):
+            box = (x, y, min(x + slice_width, img_width), min(y + slice_height, img_height))
+            slice_img = image.crop(box)
+            slices.append(slice_img)
+            coordinates.append((x, y))
+    return slices, coordinates
+def merge_predictions(predictions, slice_coordinates, orig_image_size, slice_width, slice_height, threshold=0.30):
+    merged_labels = []
+    merged_boxes = []
+    merged_scores = []
+    orig_height, orig_width = orig_image_size
+    for i, (label, boxes, scores) in enumerate(predictions):
+        x_shift, y_shift = slice_coordinates[i]
+        scores = np.array(scores).reshape(-1)
+        valid_indices = scores > threshold
+        valid_labels = np.array(label).reshape(-1)[valid_indices]
+        valid_boxes = np.array(boxes).reshape(-1, 4)[valid_indices]
+        valid_scores = scores[valid_indices]
+        for j, box in enumerate(valid_boxes):
+            box[0] = np.clip(box[0] + x_shift, 0, orig_width)  
+            box[1] = np.clip(box[1] + y_shift, 0, orig_height)
+            box[2] = np.clip(box[2] + x_shift, 0, orig_width)  
+            box[3] = np.clip(box[3] + y_shift, 0, orig_height) 
+            valid_boxes[j] = box
+        merged_labels.extend(valid_labels)
+        merged_boxes.extend(valid_boxes)
+        merged_scores.extend(valid_scores)
+    return np.array(merged_labels), np.array(merged_boxes), np.array(merged_scores)
+def draw(images, labels, boxes, scores, thrh = 0.6, path = ""):
+    for i, im in enumerate(images):
+        draw = ImageDraw.Draw(im)
+        scr = scores[i]
+        lab = labels[i][scr > thrh]
+        box = boxes[i][scr > thrh]
+        scrs = scores[i][scr > thrh]
+        for j,b in enumerate(box):
+            draw.rectangle(list(b), outline='red',)
+            draw.text((b[0], b[1]), text=f"label: {lab[j].item()} {round(scrs[j].item(),2)}", font=ImageFont.load_default(), fill='blue')
+        if path == "":
+            im.save(f'results_{i}.jpg')
+        else:
+            im.save(path)
+            
+def main(args, ):
+    """main
+    """
+    cfg = YAMLConfig(args.config, resume=args.resume)
+    if args.resume:
+        checkpoint = torch.load(args.resume, map_location='cpu') 
+        if 'ema' in checkpoint:
+            state = checkpoint['ema']['module']
+        else:
+            state = checkpoint['model']
+    else:
+        raise AttributeError('Only support resume to load model.state_dict by now.')
+    # NOTE load train mode state -> convert to deploy mode
+    cfg.model.load_state_dict(state)
+    class Model(nn.Module):
+        def __init__(self, ) -> None:
+            super().__init__()
+            self.model = cfg.model.deploy()
+            self.postprocessor = cfg.postprocessor.deploy()
+            
+        def forward(self, images, orig_target_sizes):
+            outputs = self.model(images)
+            outputs = self.postprocessor(outputs, orig_target_sizes)
+            return outputs
+    
+    model = Model().to(args.device)
+    im_pil = Image.open(args.im_file).convert('RGB')
+    w, h = im_pil.size
+    orig_size = torch.tensor([w, h])[None].to(args.device)
+    
+    transforms = T.Compose([
+        T.Resize((640, 640)),  
+        T.ToTensor(),
+    ])
+    im_data = transforms(im_pil)[None].to(args.device)
+    if args.sliced:
+        num_boxes = args.numberofboxes
+        
+        aspect_ratio = w / h
+        num_cols = int(np.sqrt(num_boxes * aspect_ratio)) 
+        num_rows = int(num_boxes / num_cols)
+        slice_height = h // num_rows
+        slice_width = w // num_cols
+        overlap_ratio = 0.2
+        slices, coordinates = slice_image(im_pil, slice_height, slice_width, overlap_ratio)
+        predictions = []
+        for i, slice_img in enumerate(slices):
+            slice_tensor = transforms(slice_img)[None].to(args.device)
+            with autocast():  # Use AMP for each slice
+                output = model(slice_tensor, torch.tensor([[slice_img.size[0], slice_img.size[1]]]).to(args.device))
+            torch.cuda.empty_cache() 
+            labels, boxes, scores = output
+            
+            labels = labels.cpu().detach().numpy()
+            boxes = boxes.cpu().detach().numpy()
+            scores = scores.cpu().detach().numpy()
+            predictions.append((labels, boxes, scores))
+        
+        merged_labels, merged_boxes, merged_scores = merge_predictions(predictions, coordinates, (h, w), slice_width, slice_height)
+        labels, boxes, scores = postprocess(merged_labels, merged_boxes, merged_scores)
+    else:
+        output = model(im_data, orig_size)
+        labels, boxes, scores = output
+        
+    draw([im_pil], labels, boxes, scores, 0.6)
+  
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-c', '--config', type=str, )
+    parser.add_argument('-r', '--resume', type=str, )
+    parser.add_argument('-f', '--im-file', type=str, )
+    parser.add_argument('-s', '--sliced', type=bool, default=False)
+    parser.add_argument('-d', '--device', type=str, default='cpu')
+    parser.add_argument('-nc', '--numberofboxes', type=int, default=25)
+    args = parser.parse_args()
+    main(args)
+
+
+
+
+
+
+
+
+
+
+

rtdetr_pytorch/tools/train.py

@@ -0,0 +1,50 @@
+"""by lyuwenyu
+"""
+
+import os 
+import sys 
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
+import argparse
+
+import src.misc.dist as dist 
+from src.core import YAMLConfig 
+from src.solver import TASKS
+
+
+def main(args, ) -> None:
+    '''main
+    '''
+    dist.init_distributed()
+    if args.seed is not None:
+        dist.set_seed(args.seed)
+
+    assert not all([args.tuning, args.resume]), \
+        'Only support from_scrach or resume or tuning at one time'
+
+    cfg = YAMLConfig(
+        args.config,
+        resume=args.resume, 
+        use_amp=args.amp,
+        tuning=args.tuning
+    )
+
+    solver = TASKS[cfg.yaml_cfg['task']](cfg)
+    
+    if args.test_only:
+        solver.val()
+    else:
+        solver.fit()
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', '-c', type=str, )
+    parser.add_argument('--resume', '-r', type=str, )
+    parser.add_argument('--tuning', '-t', type=str, )
+    parser.add_argument('--test-only', action='store_true', default=False,)
+    parser.add_argument('--amp', action='store_true', default=False,)
+    parser.add_argument('--seed', type=int, help='seed',)
+    args = parser.parse_args()
+
+    main(args)