change model

config.py

···
 
     model_save_path: str = "./models/nsfw_vit_classifier"
 
-    model_name: str = "google/vit-base-patch16-224"
+    model_name: str = "Falconsai/nsfw_image_detection"
     num_classes: int = 2
 
-    batch_size: int = 32
+    batch_size: int = 16
     num_epochs: int = 10
-    learning_rate: float = 2e-5
+    learning_rate: float = 5e-6
     weight_decay: float = 0.01
 
     train_ratio: float = 0.7

main.py

···
 import logging
-import shutil
 import random
+from io import BytesIO
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
···
             self.transform = transforms.Compose(
                 [
                     transforms.RandomHorizontalFlip(p=0.5),
-                    transforms.RandomRotation(20),
+                    transforms.RandomRotation(15),
                     transforms.ColorJitter(
-                        brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1
+                        brightness=0.2, contrast=0.2, saturation=0.2
                     ),
-                    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
-                    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
-                    transforms.RandomGrayscale(p=0.1),
+                    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
                 ]
             )
         else:
···
     nsfw_path = Path(nsfw_dir)
     if nsfw_path.exists():
         for f in nsfw_path.iterdir():
-            if f.suffix.lower() in valid_extensions or f.name.startswith("did:"):
+            if f.suffix.lower() in valid_extensions:
                 nsfw_paths.append(str(f))
     else:
         logger.warning(f"NSFW directory not found: {nsfw_dir}")
···
     sfw_humans_path = Path(sfw_human_dir)
     if sfw_humans_path.exists():
         for f in sfw_humans_path.iterdir():
-            if f.suffix.lower() in valid_extensions or f.name.startswith("did:"):
+            if f.suffix.lower() in valid_extensions:
                 sfw_paths.append(str(f))
     else:
         logger.warning(f"SFW humans directory not found: {sfw_human_dir}")
···
     sfw_anime_path = Path(sfw_anime_dir)
     if sfw_anime_path.exists():
         for f in sfw_anime_path.iterdir():
-            if f.suffix.lower() in valid_extensions or f.name.startswith("did:"):
+            if f.suffix.lower() in valid_extensions:
                 sfw_paths.append(str(f))
     else:
         logger.warning(f"SFW anime directory not found: {sfw_anime_dir}")
···
         target_ratio = 1.2
 
         if needs_balancing:
-            logger.info(f"🔄 Auto-balancing to target ratio: {target_ratio:.2f}")
-
             if current_ratio > target_ratio:
                 target_nsfw = int(sfw_count * target_ratio)
                 target_sfw = sfw_count
···
             sfw_count = len(sfw_paths)
             new_ratio = nsfw_count / sfw_count
 
-            logger.info(
-                f"✅ Dataset after balancing: {nsfw_count} NSFW, {sfw_count} SFW"
-            )
+            logger.info(f"Dataset after balancing: {nsfw_count} NSFW, {sfw_count} SFW")
             logger.info(f"New NSFW/SFW ratio: {new_ratio:.2f}")
 
-        else:
-            logger.info("✅ Dataset is already balanced, no adjustment needed")
-
-    elif nsfw_count == 0:
-        logger.error("❌ No NSFW images found! Check your data directory.")
-    elif sfw_count == 0:
-        logger.error("❌ No SFW images found! Check your data directory.")
-
-    # Combine and create labels
     image_paths: List[str] = []
     labels: List[int] = []
 
-    # Add NSFW (label = 1)
     for path in nsfw_paths:
         image_paths.append(path)
         labels.append(1)
 
-    # Add SFW (label = 0)
     for path in sfw_paths:
         image_paths.append(path)
         labels.append(0)
···
         image_paths.append(path)
         labels.append(0)
 
-    # Shuffle to mix NSFW and SFW
     combined = list(zip(image_paths, labels))
     random.shuffle(combined)
     image_paths, labels = zip(*combined)
     image_paths = list(image_paths)
     labels = list(labels)
 
-    logger.info(f"📊 Final dataset: {len(image_paths)} total images")
+    logger.info(f"Final dataset: {len(image_paths)} total images")
 
     return image_paths, labels
 
···
         ignore_mismatched_sizes=True,
     )
 
-    for param in model.vit.embeddings.parameters():
-        param.requires_grad = False
-    for param in model.vit.encoder.layer[:6].parameters():
-        param.requires_grad = False
-
     model = model.to(config.device)
 
     return model
···
         optimizer.zero_grad()
 
         outputs = model(pixel_values=images)
-        logits: torch.Tensor = outputs.logits  # Shape: (batch_size, num_classes)
+        logits: torch.Tensor = outputs.logits
 
         loss: torch.Tensor = criterion(logits, labels)
 
···
     val_loader: DataLoader[torch.Tensor],
     config: Config,
 ):
-    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
+    criterion = nn.CrossEntropyLoss()
     optimizer = optim.AdamW(
         model.parameters(),
         lr=config.learning_rate,
···
         history["val_acc"].append(val_acc)
 
         print(f"\nEpoch {epoch + 1} Summary:")
-        print(f"    Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}")
-        print(f"    Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}")
+        print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}")
+        print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}")
 
         if val_acc > best_val_acc:
             best_val_acc = val_acc
···
             "nsfw_prob": probs[1].item(),
         }
 
-    def predict_batch(
-        self, folder_path: str, extensions: Optional[List[str]] = None
-    ) -> List[Dict[str, Union[str, float]]]:
-        if extensions is None:
-            extensions = [".jpg", ".jpeg", ".png", ".bmp", ".webp"]
-
-        folder: Path = Path(folder_path)
-        results: List[Dict[str, Union[str, float]]] = []
-
-        # Get all image files
-        image_files: List[Path] = []
-        for ext in extensions:
-            image_files.extend(folder.glob(f"*{ext}"))
-            image_files.extend(folder.glob(f"*{ext.upper()}"))
-
-        print(f"Found {len(image_files)} images")
-
-        for img_path in image_files:
-            try:
-                result: Dict[str, Union[str, float]] = self.predict_single(
-                    str(img_path)
-                )
-                result["filename"] = img_path.name
-                results.append(result)
-
-                print(
-                    f"  {img_path.name}: {result['prediction']} "
-                    f"({result['confidence']:.2%})"
-                )
-            except Exception as e:
-                print(f"  Error processing {img_path.name}: {e}")
-
-        return results
-
-    def predict_with_threshold(
-        self, image_path: str, nsfw_threshold: float = 0.5
-    ) -> Dict[str, Union[str, float]]:
-        result: Dict[str, Union[str, float]] = self.predict_single(image_path)
+    def predict_single_bytes(self, image_bytes: bytes) -> Dict[str, Union[str, float]]:
+        image = Image.open(BytesIO(image_bytes)).convert("RGB")
+        inputs: Dict[str, torch.Tensor] = self.processor(
+            images=image, return_tensors="pt"
+        )
+        pixel_values: torch.Tensor = inputs["pixel_values"].to(self.device)
 
-        if result["nsfw_prob"] >= nsfw_threshold:
-            result["prediction"] = "NSFW"
-        else:
-            result["prediction"] = "SFW"
+        with torch.no_grad():
+            outputs = self.model(pixel_values=pixel_values)
+            logits: torch.Tensor = outputs.logits
+            probs: torch.Tensor = torch.softmax(logits, dim=1)[0]
+            predicted_class: int = torch.argmax(probs).item()
 
-        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")
+        return {
+            "prediction": self.class_names[predicted_class],
+            "confidence": probs[predicted_class].item(),
+            "sfw_prob": probs[0].item(),
+            "nsfw_prob": probs[1].item(),
+        }
 
 
 @click.group()
···
 
     model, history = train_model(model, train_loader, val_loader, config)
 
-    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)
 
     logger.info("Saving model")
     save_model(model, processor, config)
 
-    print("\n" + "=" * 60)
-    print("✓ Training pipeline complete!")
-    print("=" * 60)
-
-    print("\nExample: Testing inference on a sample image...")
-    if len(image_paths) > 0:
-        sample_image: str = image_paths[0]
-        prediction, confidence = predict_image(
-            sample_image, model, processor, config.device, config
-        )
-        print(f"Sample image: {sample_image}")
-        print(f"Prediction: {prediction} (confidence: {confidence:.2%})")
-
 
 @cli.command()
 @click.argument("image", type=str, required=True)
···
         CONFIG.device,
     )
 
-    resp = requests.get(url, stream=True)
+    resp = requests.get(url)
     resp.raise_for_status()
 
+    # we can use predict_bytes_single, but we wont so that we can do something with this image like
+    # move it to a directory for retraining
     with open("./testimg.jpeg", "wb") as f:
-        resp.raw.decode_content = True
-        shutil.copyfileobj(resp.raw, f)
+        f.write(resp.content)
 
     print(f"{classifier.predict_single('./testimg.jpeg')}")
 
 
 if __name__ == "__main__":
     cli()
+