add detection for colored persons

app.py

@@ -8,6 +8,7 @@ from contextlib import nullcontext
 import requests
 import functools
 import random
+import timm
 
 from ldm.models.diffusion.ddim import DDIMSampler
 from ldm.models.diffusion.plms import PLMSSampler
@@ -28,10 +29,13 @@ model = load_model_from_config(config, ckpt, device=device, verbose=False)
 model = model.to(device).half()
 
 clip_model, preprocess = clip.load("ViT-L/14", device=device)
+
 gender_learn = load_learner('gender_model.pkl')
 gender_labels = gender_learn.dls.vocab
 beard_learn = load_learner('facial_hair_model.pkl')
 beard_labels = beard_learn.dls.vocab
+ethnic_learn = load_learner('ethnic_model.pkl')
+ethnic_labels = ethnic_learn.dls.vocab
 
 n_inputs = 5
 
@@ -128,7 +132,11 @@ def is_female(img):
     
 def has_beard(img):
     pred,pred_idx,probs = beard_learn.predict(img)
-    return float(probs[0]) > float(probs[1])
+    return float(probs[1]) > float(probs[0])
+    
+def ethnicity(img):
+    pred,pred_idx,probs = ethnic_learn.predict(img)
+    return pred
 
 
 import gradio
@@ -136,15 +144,24 @@ import gradio
 def boutsify(person):
     portrait_path = "bouts_m1.jpg"
     female_detected = is_female(person)
+    ethnicity_prediction = ethnicity(person)
     
-    if female_detected:
-        print("This is a female portrait")
-        portrait_path = "bouts_f1.jpg"
-    else:
-        print("This is a male portrait")
-        if has_beard(person):
-            print("The person has a beard")
-            portrait_path = "bouts_mb1.jpg"
+    if ethnicity_prediction == "Black":
+        print("Colored person")
+        if female_detected:
+            print("This is a female portrait")
+            portrait_path = "bouts_fc1.jpg"
+        else:
+            portrait_path = "bouts_mc1.jpg"            
+    else:    
+        if female_detected:
+            print("This is a female portrait")
+            portrait_path = "bouts_f1.jpg"
+        else:
+            print("This is a male portrait, checking facial hair")
+            if has_beard(person):
+                print("The person has a beard")
+                portrait_path = "bouts_mb1.jpg"
     
     person_image = Image.fromarray(person)
     
@@ -160,6 +177,7 @@ def boutsify(person):
         3.0, 1, random.randrange(0, 10000), 50,
     ]
     
+    #return person
     return run_image_mixer(inputs)
 
 gradio_interface = gradio.Interface(

requirements.txt

@@ -20,6 +20,7 @@ fire==0.4.0
 diffusers==0.3.0
 datasets[vision]==2.4.0
 fastai==2.7.11
+timm==0.6.13
 -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
 -e git+https://github.com/openai/CLIP.git@main#egg=clip
 -e git+https://github.com/justinpinkney/nomi.git@e9ded23b7e2269cc64d39683e1bf3c0319f552ab#egg=nomi