Update app.py

app.py

@@ -1,27 +1,63 @@
 from transformers import pipeline
 import gradio
 import base64
-from PIL import Image
+from PIL import Image, ImageDraw
 from io import BytesIO
 from sentence_transformers import SentenceTransformer, util
 
-
 backgroundPipe = pipeline("image-segmentation", model="facebook/maskformer-swin-large-coco")
 PersonPipe = pipeline("image-segmentation", model="mattmdjaga/segformer_b2_clothes")
 sentenceModal = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+personDetailsPipe = pipeline("image-segmentation", model="yolo12138/segformer-b2-human-parse-24")
+faceModal = pipeline("image-segmentation", model="jonathandinu/face-parsing")
+faceDetectionModal = pipeline("object-detection", model="aditmohan96/detr-finetuned-face")
+PersonDetectionpipe = pipeline("object-detection", model="hustvl/yolos-tiny")
+
+def getPersonDetail(image):
+    data = PersonDetectionpipe(image)
+    persn = []
+    for per in data:
+        if per["label"].lower() == "person":
+            persn.append(per["box"])
+    n = 1
+    ret = {}
+    for cord in persn:
+        crop_box = (cord['xmin'], cord['ymin'], cord['xmax'], cord['ymax'])
+        cropped_image = image.crop(crop_box)
+        personData = personDetailsPipe(cropped_image)
+        for dt in personData:
+            if len(persn) > 1:
+                ret[(f'Person {n} {dt["label"]}').lower()] = cbiwm(image, dt["mask"], cord)
+            else:
+                ret[dt["label"].lower()] = cbiwm(image, dt["mask"], cord)
+        n = n + 1
+    return ret
+
+def cbiwm(image, mask, coordinates):
+    black_image = Image.new("RGBA", image.size, (0, 0, 0, 255))
+    black_image.paste(mask, (coordinates['xmin'], coordinates['ymin']), mask)
+    return black_image
+
+def processFaceDetails(image):
+    ret = getPersonDetail(image)
+    data = faceDetectionModal(image)
+    cordinates = data[1]["box"]
+    crop_box = (data[1]["box"]['xmin'], data[1]["box"]['ymin'], data[1]["box"]['xmax'], data[1]["box"]['ymax'])
+    cropped_image = image.crop(crop_box)
+    facedata = faceModal(cropped_image)
+    for imask in facedata:
+        ret[imask["label"].replace(".png", "").lower()] = cbiwm(image, imask["mask"], cordinates)
+    return ret
 
 def getImageDetails(image) -> dict:
+    ret = processFaceDetails(image)
     person = PersonPipe(image)
     bg = backgroundPipe(image)
-    ret = {}
-    labs = []
     for imask in bg:
-        ret[imask["label"]] = imask["mask"] # Apply base64 image converter here if needed
-        labs.append(imask["label"])
+        ret[imask["label"].lower()] = imask["mask"] # Apply base64 image converter here if needed
     for mask in person:
-        ret[mask["label"]] = mask["mask"] # Apply base64 image converter here if needed
-        labs.append(mask["label"])
-    return ret, labs
+        ret[mask["label"].lower()] = mask["mask"] # Apply base64 image converter here if needed
+    return ret
 
 def processSentence(sentence: str, semilist: list):
     query_embedding = sentenceModal.encode(sentence)
@@ -51,9 +87,12 @@ def process_image(image):
     return processed_image
 
 def processAndGetMask(image: str, text: str):
-    datas, labs = getImageDetails(image)
+    datas = getImageDetails(image)
+    labs = list(datas.keys())
+    print(labs)
     selector = processSentence(text, labs)
     imageout = datas[selector]
+    print(f"Selected : {selector}")
     return process_image(imageout)
 
 gr = gradio.Interface(
@@ -61,4 +100,4 @@ gr = gradio.Interface(
     [gradio.Image(type="pil"), gradio.Text()],
     gradio.Image(type="pil")
 )
-gr.launch()
+gr.launch(debug=True)