Background Removal

U-Net Like(VGG Encoder) Pretrained Model For Human Body Detection (pytorch, Semantic Segmentation)

Quick Links

• Dependencies
• Implementation
• Usage
• Examples

Dependencies

• Install Dependencies $ pip install -r requirements.txt
• Download model weights Here, place it in model_weights/

Implementation

VGG Encoder Implementation

  def VGGEncoder():
    weights = VGG16_Weights.DEFAULT
    base_model = vgg16(weights=weights)
    base_model.training = False
        
    encoder_seq =  nn.ModuleList()
    moduls = nn.Sequential()
    for layer in list(base_model.features.children()):
        if isinstance(layer, nn.modules.pooling.MaxPool2d):
            encoder_seq.append(moduls)
            moduls = nn.Sequential()
        else:
            moduls.append(layer)
    return encoder_seq

Usage:

  detector = BodyDetector("model_weights/bgrm-bh.pth")

Load model.

  fname = RandomSample("background folder", '*')
  bg = LoadImage(fname)

Read background image.

  fname = RandomSample("image folder")
  img = LoadImage(fname)
  img_resize = cv2.resize(img, (224, 224), interpolation = cv2.INTER_AREA)

Read image in bgr mode and resize it to 224*224.

  mask = detector.DetectBody(img_resize)

Detect object(human body) area.

  res = ReplaceBG(img, mask, bg)

Replace current background with loaded background image.

Examples

🛡️ License

Project is distributed under MIT License