rimjhimittal/final_final

1from transformers import pipeline2import numpy as np3import cv24import insightface5from insightface.app import FaceAnalysis6from PIL import Image, ImageDraw7 8 9# Initialize face detection10app = FaceAnalysis(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])11app.prepare(ctx_id=0, det_size=(640, 640))12 13# Initialize segmentation pipeline14segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")15 16 17def remove_face(img, mask):18    # Convert image to numpy array19    img_arr = np.asarray(img)20    21    # Run face detection22    faces = app.get(img_arr)23    24    # Get the first face25    faces = faces[0]['bbox']26 27    # Width and height of face28    w = faces[2] - faces[0]29    h = faces[3] - faces[1]30 31    # Make face locations bigger32    faces[0] = faces[0] - (w*0.5) # x left33    faces[2] = faces[2] + (w*0.5) # x right34    faces[1] = faces[1] - (h*0.5) # y top35    faces[3] = faces[3] + (h*0.2) # y bottom36 37    # Convert to [(x_left, y_top), (x_right, y_bottom)]38    face_locations = [(faces[0], faces[1]), (faces[2], faces[3])]39 40    # Draw black rect onto mask41    img1 = ImageDraw.Draw(mask)42    img1.rectangle(face_locations, fill=0)43 44    return mask45 46 47 48def segment_body(original_img, face=True):49    # Make a copy50    img = original_img.copy()51    52    # Segment image53    segments = segmenter(img)54 55    # Create list of masks56    segment_include = ["Hat", "Hair", "Sunglasses", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Face", "Left-leg", "Right-leg", "Left-arm", "Right-arm", "Bag","Scarf"]57    mask_list = []58    for s in segments:59        if(s['label'] in segment_include):60            mask_list.append(s['mask'])61 62 63    # Paste all masks on top of eachother 64    final_mask = np.array(mask_list[0])65    for mask in mask_list:66        current_mask = np.array(mask)67        final_mask = final_mask + current_mask68            69    # Convert final mask from np array to PIL image70    final_mask = Image.fromarray(final_mask)71 72    # Remove face73    if(face==False):74        final_mask = remove_face(img.convert('RGB'), final_mask)75 76    # Apply mask to original image77    img.putalpha(final_mask)78 79    return img, final_mask80 81 82 83def segment_torso(original_img):84    # Make a copy85    img = original_img.copy()86    87    # Segment image88    segments = segmenter(img)89 90    # Create list of masks91    segment_include = ["Upper-clothes", "Dress", "Belt", "Face", "Left-arm", "Right-arm"]92    mask_list = []93    for s in segments:94        if(s['label'] in segment_include):95            mask_list.append(s['mask'])96 97 98    # Paste all masks on top of eachother 99    final_mask = np.array(mask_list[0])100    for mask in mask_list:101        current_mask = np.array(mask)102        final_mask = final_mask + current_mask103            104    # Convert final mask from np array to PIL image105    final_mask = Image.fromarray(final_mask)106 107    # Remove face108    final_mask = remove_face(img.convert('RGB'), final_mask)109 110    # Apply mask to original image111    img.putalpha(final_mask)112 113    return img, final_mask114