fix(bug): feature extractor added back in

app.py

@@ -61,10 +61,19 @@ def load_models():
 clf_1, clf_2, feature_extractor_3, model_3, feature_extractor_4, model_4, clf_5, clf_5b = load_models()
 
 @spaces.GPU(duration=10)
-def predict_with_model(img_pil, clf, class_names, confidence_threshold, model_name, model_id):
+def predict_with_model(img_pil, clf, class_names, confidence_threshold, model_name, model_id, feature_extractor=None):
     try:
-        prediction = clf(img_pil)
-        result = {pred['label']: pred['score'] for pred in prediction}
+        if feature_extractor:
+            inputs = feature_extractor(img_pil, return_tensors="pt").to(device)
+            with torch.no_grad():
+                outputs = clf(**inputs)
+            logits = outputs.logits
+            probabilities = softmax(logits.cpu().numpy()[0])
+            result = {class_names[i]: probabilities[i] for i in range(len(class_names))}
+        else:
+            prediction = clf(img_pil)
+            result = {pred['label']: pred['score'] for pred in prediction}
+
         result_output = [model_id, model_name, result.get(class_names[1], 0.0), result.get(class_names[0], 0.0)]
         logger.info(result_output)
         for class_name in class_names:
@@ -97,8 +106,8 @@ def predict_image(img, confidence_threshold):
 
     label_1, result_1output = predict_with_model(img_pil, clf_1, CLASS_NAMES["model_1"], confidence_threshold, "SwinV2-base", 1)
     label_2, result_2output = predict_with_model(img_pilvits, clf_2, CLASS_NAMES["model_2"], confidence_threshold, "ViT-base Classifier", 2)
-    label_3, result_3output = predict_with_model(img_pil, model_3, CLASS_NAMES["model_3"], confidence_threshold, "SDXL-Trained", 3)
-    label_4, result_4output = predict_with_model(img_pil, model_4, CLASS_NAMES["model_4"], confidence_threshold, "SDXL + FLUX", 4)
+    label_3, result_3output = predict_with_model(img_pil, model_3, CLASS_NAMES["model_3"], confidence_threshold, "SDXL-Trained", 3, feature_extractor_3)
+    label_4, result_4output = predict_with_model(img_pil, model_4, CLASS_NAMES["model_4"], confidence_threshold, "SDXL + FLUX", 4, feature_extractor_4)
     label_5, result_5output = predict_with_model(img_pilvits, clf_5, CLASS_NAMES["model_5"], confidence_threshold, "ViT-base Newcomer", 5)
     label_5b, result_5boutput = predict_with_model(img_pilvits, clf_5b, CLASS_NAMES["model_5b"], confidence_threshold, "ViT-base Newcomer", 6)
 
@@ -113,7 +122,6 @@ def predict_image(img, confidence_threshold):
 
     combined_outputs = [result_1output, result_2output, result_3output, result_4output, result_5output, result_5boutput]
     return img_pil, combined_outputs
-
 # Define a function to generate the HTML content
 def generate_results_html(results):
     def get_header_color(label):