Add additional check

app.py

@@ -12,8 +12,10 @@ from scipy.special import softmax
 logger = logging.getLogger("basebody")
 CLIP_MODEL_NAME = "ViT-B/16"
 
+TEXT_PROMPTS_OLD_FILE_NAME = "text_prompts.json"
 TEXT_PROMPTS_FILE_NAME = "text_prompts2.json"
 HAIR_TEXT_PROMPTS_FILE_NAME = "text_prompts_hair.json"
+LOGISTIC_REGRESSION_OLD_MODEL_FILE_NAME = "logistic_regression_l1.pkl"
 LOGISTIC_REGRESSION_MODEL_FILE_NAME = "logistic_regression_l1_oct_26.pkl"
 HAIR_RF_CLASSIFIER_MODEL_FILE_NAME = "hairclassifier_rf.pkl"
 
@@ -36,6 +38,16 @@ with open(
 ) as f:
     hair_text_prompts = json.load(f)
 
+
+with open(
+    os.path.join(
+        os.path.dirname(__file__), TEXT_PROMPTS_OLD_FILE_NAME
+    ),
+    "r",
+) as f:
+    text_prompts_old = json.load(f)
+
+
 with open(
     os.path.join(
         os.path.dirname(__file__), LOGISTIC_REGRESSION_MODEL_FILE_NAME
@@ -44,6 +56,16 @@ with open(
 ) as f:
     lr_model = pickle.load(f)
 
+
+with open(
+    os.path.join(
+        os.path.dirname(__file__), LOGISTIC_REGRESSION_OLD_MODEL_FILE_NAME
+    ),
+    "rb",
+) as f:
+    lr_old_model = pickle.load(f)
+
+
 with open(
     os.path.join(
         os.path.dirname(__file__), HAIR_RF_CLASSIFIER_MODEL_FILE_NAME
@@ -55,27 +77,21 @@ with open(
 
 logger.info("Logistic regression model loaded, coefficients: ")
 
+def get_text_features(text_prompts):
+    all_text_features = []
+    with torch.no_grad():
+        for k, prompts in text_prompts.items():
+            assert len(prompts) == 2
+            inputs = clip.tokenize(prompts)
+            outputs = clip_model.encode_text(inputs)
+            all_text_features.append(outputs)
+        all_text_features = torch.cat(all_text_features, dim=0)
+        all_text_features = all_text_features.cpu()
+    return all_text_features
 
-all_text_features = []
-with torch.no_grad():
-    for k, prompts in text_prompts.items():
-        assert len(prompts) == 2
-        inputs = clip.tokenize(prompts)
-        outputs = clip_model.encode_text(inputs)
-        all_text_features.append(outputs)
-    all_text_features = torch.cat(all_text_features, dim=0)
-    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()
+all_text_features = get_text_features(text_prompts)
+hair_text_features = get_text_features(hair_text_prompts)
+old_text_features = get_text_features(text_prompts_old)
 
 
 def get_cosine_similarities(image_features, text_features, text_prompts):
@@ -101,6 +117,9 @@ def predict_fn(input_img):
         hair_cosine_simlarities = get_cosine_similarities(
             image_features, hair_text_features, hair_text_prompts
         )
+        old_cosine_simlarities = get_cosine_similarities(
+            image_features, old_text_features, text_prompts_old
+        )
         # logger.info(f"cosine_simlarities shape: {cosine_simlarities.shape}")
         logger.info(f"cosine_simlarities: {base_body_cosine_simlarities}")
     probabilities = lr_model.predict_proba(
@@ -109,9 +128,13 @@ def predict_fn(input_img):
     hair_probabilities = hair_rf_model.predict_proba(
         hair_cosine_simlarities.reshape(1, -1)
     )
+    old_lr_probabilities = lr_old_model.predict_proba(
+        old_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))
+    old_result_probabilty = float(old_lr_probabilities[0][1].round(3))
     # get decision string
     if result_probabilty > 0.77:
         if hair_result_probabilty < 0.5:
@@ -120,7 +143,11 @@ def predict_fn(input_img):
             decision = "AUTO REJECT"
         else:
             decision = "AUTO ACCEPT"
-    elif result_probabilty < 0.4:
+    elif result_probabilty < 0.2:
+        decision = "AUTO REJECT"
+    elif old_result_probabilty < 0.06:
+        logger.info("old_result_probabilty < 0.06")
+        result_probabilty = old_result_probabilty
         decision = "AUTO REJECT"
     else:
         decision = "MODERATION"