| 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 | |