Create classifier.py

classifier.py

@@ -0,0 +1,79 @@
+import torch
+import torch.nn.functional as F
+from datetime import datetime
+import os
+
+def few_shot_fault_classification(
+    model,
+    test_images,
+    test_image_filenames,
+    nominal_images,
+    nominal_descriptions,
+    defective_images,
+    defective_descriptions,
+    num_few_shot_nominal_imgs: int,
+    device="cpu",
+    file_path: str = '.',
+    file_name: str = 'image_classification_results.csv',
+    print_one_liner: bool = False
+):
+    if not isinstance(test_images, list):
+        test_images = [test_images]
+    if not isinstance(test_image_filenames, list):
+        test_image_filenames = [test_image_filenames]
+    if not isinstance(nominal_images, list):
+        nominal_images = [nominal_images]
+    if not isinstance(defective_images, list):
+        defective_images = [defective_images]
+
+    csv_file = os.path.join(file_path, file_name)
+    results = []
+
+    with torch.no_grad():
+        nominal_features = torch.stack([model.encode_image(img.to(device)) for img in nominal_images])
+        nominal_features /= nominal_features.norm(dim=-1, keepdim=True)
+
+        defective_features = torch.stack([model.encode_image(img.to(device)) for img in defective_images])
+        defective_features /= defective_features.norm(dim=-1, keepdim=True)
+
+        for idx, test_img in enumerate(test_images):
+            test_features = model.encode_image(test_img.to(device))
+            test_features /= test_features.norm(dim=-1, keepdim=True)
+
+            max_nominal_similarity = max_defective_similarity = -float('inf')
+            max_nominal_idx = max_defective_idx = -1
+
+            for i in range(nominal_features.shape[0]):
+                similarity = (test_features @ nominal_features[i].T).item()
+                if similarity > max_nominal_similarity:
+                    max_nominal_similarity = similarity
+                    max_nominal_idx = i
+
+            for j in range(defective_features.shape[0]):
+                similarity = (test_features @ defective_features[j].T).item()
+                if similarity > max_defective_similarity:
+                    max_defective_similarity = similarity
+                    max_defective_idx = j
+
+            similarities = torch.tensor([max_nominal_similarity, max_defective_similarity])
+            probabilities = F.softmax(similarities, dim=0).tolist()
+
+            classification = "Defective" if probabilities[1] > probabilities[0] else "Nominal"
+
+            result = {
+                "datetime_of_operation": datetime.now().isoformat(),
+                "image_path": test_image_filenames[idx],
+                "classification_result": classification,
+                "non_defect_prob": round(probabilities[0], 3),
+                "defect_prob": round(probabilities[1], 3),
+                "nominal_description": nominal_descriptions[max_nominal_idx],
+                "defective_description": defective_descriptions[max_defective_idx],
+            }
+
+            results.append(result)
+
+            if print_one_liner:
+                print(f"{test_image_filenames[idx]} → {classification} "
+                      f"(Nominal: {probabilities[0]:.3f}, Defective: {probabilities[1]:.3f})")
+
+    return results