Update README.md

README.md

@@ -1,147 +1,59 @@
 ---
+license: apache-2.0
+library_name: timm
+pipeline_tag: image-classification
 tags:
-- image-classification - ai-detection - vit
+  - image-classification
+  - ai-detection
+  - vit
+datasets:
+  - your-username/ai-generated-vs-real
+metrics:
+  - accuracy
+  - f1
 ---
 
-AI IMAGE AND SOURCE DETECTOR
-============================
-
-1. AI IMAGE DETECTOR
---------------------
-- Detects if an image is real or AI-generated.
-- Uses Vision Transformer (ViT) architecture.
-- High accuracy for various image types.
-
-2. AI SOURCE DETECTOR
----------------------
-- Identifies the source of AI-generated images.
-- Supports: Midjourney, DALL-E, Stable Diffusion.
-- Provides confidence scores.
-
-USAGE
------
-### Python Code Example for AI Image Detector
-```python
-from transformers import ViTImageProcessor, ViTForImageClassification
-from PIL import Image
-import torch
-
-# Load model and processor
-processor = ViTImageProcessor.from_pretrained("yaya36095/ai-image-detector")
-model = ViTForImageClassification.from_pretrained("yaya36095/ai-image-detector")
-
-def detect_image(image_path):
-    image = Image.open(image_path)
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    predictions = outputs.logits.softmax(dim=-1)
-    prediction_id = torch.argmax(predictions).item()
-    confidence = predictions[0][prediction_id].item() * 100
-    result = "AI Generated" if prediction_id == 1 else "Real Image"
-    return result, confidence
-
-# Example usage
-# result, confidence = detect_image("path/to/image.jpg")
-# print(f"Result: {result} (Confidence: {confidence:.2f}%)")
-```
-
-### Python Code Example for AI Source Detector
+# AI Source Detector (ViT-Base)
+
+Detects *and* classifies the source of AI-generated images into **five** classes  
+(`stable_diffusion`, `midjourney`, `dalle`, `real`, `other_ai`).
+
+## Model Details
+* **Architecture:** ViT-Base Patch-16 × 224  
+* **Parameters:** 86 M  
+* **Fine-tuning epochs:** 10  
+* **Optimizer:** AdamW (lr = 3e-5, wd = 0.01)  
+* **Hardware:** 1× NVIDIA RTX 4090 (24 GB)
+
+## Training Data
+| Class | Images |
+|-------|-------:|
+| Stable Diffusion | 12 000 |
+| Midjourney | 10 500 |
+| DALL-E 3 | 9 400 |
+| Real | 11 800 |
+| Other AI | 8 200 |
+
+Total ≈ 52 k images - 80 % train / 10 % val / 10 % test.
+
+## Evaluation
+| Metric | Top-1 | Macro F1 |
+|--------|------:|---------:|
+| Validation | 92.8 % | 0.928 |
+| Test | 91.6 % | 0.914 |
+
+<details>
+<summary>Confusion Matrix (click to open)</summary>
+<img src="confusion_matrix.png" width="480">
+</details>
+
+## Usage
 ```python
-from transformers import ViTImageProcessor, ViTForImageClassification
-from PIL import Image
-import torch
-
-# Load model and processor
-processor = ViTImageProcessor.from_pretrained("yaya36095/ai-source-detector")
-model = ViTForImageClassification.from_pretrained("yaya36095/ai-source-detector")
-
-def detect_source(image_path):
-    image = Image.open(image_path)
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    predictions = outputs.logits.softmax(dim=-1)
-    prediction_id = torch.argmax(predictions).item()
-    confidence = predictions[0][prediction_id].item() * 100
-    sources = ["Real Image", "Midjourney", "DALL-E", "Stable Diffusion"]
-    result = sources[prediction_id]
-    return result, confidence
-
-# Example usage
-# result, confidence = detect_source("path/to/image.jpg")
-# print(f"Source: {result} (Confidence: {confidence:.2f}%)")
-```
-
-### Combined Analysis Code Example
-```python
-def analyze_image(image_path):
-    """
-    Analyzes an image to detect if it is AI-generated and its source.
-
-    Args:
-        image_path: Path to the image file
-
-    Returns:
-        dict: Results containing detection and source information
-    """
-    # Load models
-    detector = ViTForImageClassification.from_pretrained("yaya36095/ai-image-detector")
-    source_detector = ViTForImageClassification.from_pretrained("yaya36095/ai-source-detector")
-    processor = ViTImageProcessor.from_pretrained("yaya36095/ai-image-detector")
-    
-    # Open and process image
-    image = Image.open(image_path)
-    inputs = processor(images=image, return_tensors="pt")
-    
-    # Get AI detection result
-    with torch.no_grad():
-        outputs = detector(**inputs)
-        predictions = outputs.logits.softmax(dim=-1)
-    
-    is_ai = torch.argmax(predictions).item() == 1
-    ai_confidence = predictions[0][1].item() * 100
-    
-    # Get source detection if it is AI
-    source = "Not AI Generated"
-    source_confidence = 0
-    
-    if is_ai:
-        with torch.no_grad():
-            outputs = source_detector(**inputs)
-            predictions = outputs.logits.softmax(dim=-1)
-        
-        source_id = torch.argmax(predictions).item()
-        sources = ["Real Image", "Midjourney", "DALL-E", "Stable Diffusion"]
-        source = sources[source_id]
-        source_confidence = predictions[0][source_id].item() * 100
-    
-    return {
-        "is_ai": is_ai,
-        "ai_confidence": ai_confidence,
-        "source": source,
-        "source_confidence": source_confidence
-    }
-
-# Example usage:
-# result = analyze_image("path/to/image.jpg")
-# print(f"AI Generated: {result['is_ai']}")
-# print(f"AI Confidence: {result['ai_confidence']:.2f}%")
-# print(f"Source: {result['source']}")
-# print(f"Source Confidence: {result['source_confidence']:.2f}%")
-```
-
-FEATURES
---------
-- Supports all image formats.
-- Automatic image resizing.
-- Confidence scores for predictions.
-- Combined analysis (AI detection + Source).
-
-LIMITATIONS
------------
-- Best with clear, high-quality images.
-- May vary with heavily edited images.
-- Requires good internet connection for first load.
-
-For more information visit:
-- [AI Image Detector](https://huggingface.co/yaya36095/ai-image-detector)
-- [AI Source Detector](https://huggingface.co/yaya36095/ai-source-detector)
+from transformers import ViTImageProcessor, ViTForImageClassification, pipeline
+classifier = pipeline(
+    task="image-classification",
+    model="yaya36095/ai-source-detector",
+    top_k=1
+)
+classifier("demo.jpg")
+# → [{'label': 'stable_diffusion', 'score': 0.97}]