Add hairclassifier to the model

app.py

@@ -13,7 +13,9 @@ logger = logging.getLogger("basebody")
 CLIP_MODEL_NAME = "ViT-B/16"
 
 TEXT_PROMPTS_FILE_NAME = "text_prompts.json"
+HAIR_TEXT_PROMPTS_FILE_NAME = "text_prompts_hair.json"
 LOGISTIC_REGRESSION_MODEL_FILE_NAME = "logistic_regression_l1_oct_2.pkl"
+HAIR_RF_CLASSIFIER_MODEL_FILE_NAME = "hairclassifier_rf.pkl"
 
 HF_TOKEN = os.getenv('HF_TOKEN')
 hf_writer = gr.HuggingFaceDatasetSaver(
@@ -28,6 +30,12 @@ with open(
     os.path.join(os.path.dirname(__file__), TEXT_PROMPTS_FILE_NAME), "r"
 ) as f:
     text_prompts = json.load(f)
+
+with open(
+    os.path.join(os.path.dirname(__file__), HAIR_TEXT_PROMPTS_FILE_NAME), "r"
+) as f:
+    hair_text_prompts = json.load(f)
+
 with open(
     os.path.join(
         os.path.dirname(__file__), LOGISTIC_REGRESSION_MODEL_FILE_NAME
@@ -36,6 +44,15 @@ with open(
 ) as f:
     lr_model = pickle.load(f)
 
+with open(
+    os.path.join(
+        os.path.dirname(__file__), HAIR_RF_CLASSIFIER_MODEL_FILE_NAME
+    ),
+    "rb",
+) as f:
+    hair_rf_model = pickle.load(f)
+
+
 logger.info("Logistic regression model loaded, coefficients: ")
 
 
@@ -50,6 +67,27 @@ with torch.no_grad():
     all_text_features = all_text_features.cpu()
 
 
+hair_text_features = []
+with torch.no_grad():
+    for k, prompts in hair_text_prompts.items():
+        assert len(prompts) == 2
+        inputs = clip.tokenize(prompts)
+        outputs = clip_model.encode_text(inputs)
+        hair_text_features.append(outputs)
+    hair_text_features = torch.cat(hair_text_features, dim=0)
+    hair_text_features = hair_text_features.cpu()
+
+
+def get_cosine_similarities(image_features, text_features):
+    cosine_simlarities = softmax(
+        (text_features @ image_features.cpu().T)
+        .squeeze()
+        .reshape(len(text_prompts), 2, -1),
+        axis=1,
+    )[:, 0, :]
+    return cosine_simlarities
+
+
 def predict_fn(input_img):
     input_img = Image.fromarray(input_img.astype("uint8"), "RGB")
     image = preprocess(
@@ -57,21 +95,28 @@ def predict_fn(input_img):
     ).unsqueeze(0)
     with torch.no_grad():
         image_features = clip_model.encode_image(image)
-        cosine_simlarities = softmax(
-            (all_text_features @ image_features.cpu().T)
-            .squeeze()
-            .reshape(len(text_prompts), 2, -1),
-            axis=1,
-        )[:, 0, :]
+        base_body_cosine_simlarities = get_cosine_similarities(
+            image_features, all_text_features
+        )
+        hair_cosine_simlarities = get_cosine_similarities(
+            image_features, hair_text_features
+        )
         # logger.info(f"cosine_simlarities shape: {cosine_simlarities.shape}")
-        logger.info(f"cosine_simlarities: {cosine_simlarities}")
+        logger.info(f"cosine_simlarities: {base_body_cosine_simlarities}")
     probabilities = lr_model.predict_proba(
-        cosine_simlarities.reshape(1, -1)
+        base_body_cosine_simlarities.reshape(1, -1)
+    )
+    hair_probabilities = hair_rf_model.predict_proba(
+        hair_cosine_simlarities.reshape(1, -1)
     )
     logger.info(f"probabilities: {probabilities}")
     result_probabilty = float(probabilities[0][1].round(3))
+    hair_result_probabilty = float(hair_probabilities[0][1].round(3))
     # get decision string
-    if result_probabilty > 0.95:
+    if result_probabilty > 0.77:
+        if hair_result_probabilty < 0.5:
+            result_probabilty = hair_result_probabilty
+            decision = "AUTO REJECT"
         decision = "AUTO ACCEPT"
     elif result_probabilty < 0.4:
         decision = "AUTO REJECT"

hairclassifier_rf.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ab2fb4a72a581f1943ee57e51031616a7a4d4c6e8be5fb7dfac7fc67cebd7c7
+size 83719733

text_prompts_hair.json

@@ -0,0 +1,74 @@
+{
+    "hair": [
+        "does character have hair on its head",
+        "does character not have hair on its head"
+    ],
+    "bald": [
+        "is character bald",
+        "is character not bald"
+    ],
+    "hat": [
+        "is character wearing a hat",
+        "is character not wearing a hat"
+    ],
+    "helmet": [
+        "is character wearing a helmet",
+        "is character not wearing a helmet"
+    ],
+    "headband": [
+        "is character wearing a headband",
+        "is character not wearing a headband"
+    ],
+    "tiara": [
+        "is character wearing a tiara",
+        "is character not wearing a tiara"
+    ],
+    "turban": [
+        "is character wearing a turban",
+        "is character not wearing a turban"
+    ],
+    "crown": [
+        "is character wearing a crown",
+        "is character not wearing a crown"
+    ],
+    "bandana": [
+        "is character wearing a bandana",
+        "is character not wearing a bandana"
+    ],
+    "hood": [
+        "is character wearing a hood",
+        "is character not wearing a hood"
+    ],
+    "wig": [
+        "is character wearing a wig",
+        "is character not wearing a wig"
+    ],
+    "headphones": [
+        "is character wearing headphones",
+        "is character not wearing headphones"
+    ],
+    "earmuffs": [
+        "is character wearing earmuffs",
+        "is character not wearing earmuffs"
+    ],
+    "veil": [
+        "is character wearing a veil",
+        "is character not wearing a veil"
+    ],
+    "feathers": [
+        "are there feathers on the character's head",
+        "there are no feathers on the character's head"
+    ],
+    "horns": [
+        "does character have horns on its head",
+        "does character not have horns on its head"
+    ],
+    "antenna": [
+        "does character have antenna on its head",
+        "does character not have antenna on its head"
+    ],
+    "head-decoration": [
+        "is there any decoration on the character's head",
+        "there is no decoration on the character's head"
+    ]
+}