save misclassifications

main.py

···
         return result
 
 
+def save_misclassified(
+    model: ViTForImageClassification,
+    test_loader: DataLoader[torch.Tensor],
+    config: Config,
+    output_dir: str = "misclassified",
+) -> None:
+    model.eval()
+
+    fp_dir = Path(output_dir) / "false_positives"
+    fn_dir = Path(output_dir) / "false_negatives"
+    fp_dir.mkdir(parents=True, exist_ok=True)
+    fn_dir.mkdir(parents=True, exist_ok=True)
+
+    dataset = test_loader.dataset
+
+    false_positives = []
+    false_negatives = []
+
+    with torch.no_grad():
+        sample_idx = 0
+        for images, labels in tqdm(test_loader, desc="Finding misclassified"):
+            images = images.to(config.device)
+
+            outputs = model(pixel_values=images)
+            logits = outputs.logits
+            probs = torch.softmax(logits, dim=1)
+            _, predicted = torch.max(logits, 1)
+
+            for i in range(len(labels)):
+                true_label = labels[i].item()
+                pred_label = predicted[i].item()
+                confidence = probs[i, pred_label].item()
+
+                if true_label != pred_label:
+                    img_path = dataset.image_paths[sample_idx]
+
+                    if true_label == 0 and pred_label == 1:
+                        false_positives.append(
+                            {
+                                "path": img_path,
+                                "confidence": confidence,
+                                "nsfw_prob": probs[i, 1].item(),
+                            }
+                        )
+                    elif true_label == 1 and pred_label == 0:
+                        false_negatives.append(
+                            {
+                                "path": img_path,
+                                "confidence": confidence,
+                                "sfw_prob": probs[i, 0].item(),
+                            }
+                        )
+
+                sample_idx += 1
+
+    import shutil
+
+    for i, fp in enumerate(false_positives):
+        src = Path(fp["path"])
+        dst = fp_dir / f"{fp['confidence']:.3f}_{src.name}"
+        shutil.copy(src, dst)
+
+    for i, fn in enumerate(false_negatives):
+        src = Path(fn["path"])
+        dst = fn_dir / f"{fn['confidence']:.3f}_{src.name}"
+        shutil.copy(src, dst)
+
+    print(f"\n{'=' * 60}")
+    print("Misclassification Analysis")
+    print(f"{'=' * 60}")
+    print(f"False Positives (SFW → NSFW): {len(false_positives)}")
+    print(f"False Negatives (NSFW → SFW): {len(false_negatives)}")
+    print(f"\nImages saved to: {output_dir}/")
+    print(f"- false_positives/  ({len(false_positives)} images)")
+    print(f"- false_negatives/  ({len(false_negatives)} images)")
+
+    report_path = Path(output_dir) / "report.txt"
+    with open(report_path, "w") as f:
+        f.write("FALSE POSITIVES (SFW classified as NSFW)\n")
+        f.write("=" * 50 + "\n")
+        for fp in sorted(false_positives, key=lambda x: x["confidence"], reverse=True):
+            f.write(f"{fp['confidence']:.3f} NSFW prob | {fp['path']}\n")
+
+        f.write("\n\nFALSE NEGATIVES (NSFW classified as SFW)\n")
+        f.write("=" * 50 + "\n")
+        for fn in sorted(false_negatives, key=lambda x: x["confidence"], reverse=True):
+            f.write(f"{fn['confidence']:.3f} SFW prob | {fn['path']}\n")
+
+
 @click.group()
 def cli():
     pass
···
     print("\nStep 6: Evaluating model...")
     logger.info("Evaluating model")
     evaluate_model(model, test_loader, config)
+
+    logger.info("Saving misclassified images")
+    save_misclassified(model, test_loader, config, output_dir="misclassified")
 
     logger.info("Generating plots")
     plot_training_history(history)