Update README.md

README.md

@@ -1,3 +1,83 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+metrics:
+- accuracy
+base_model:
+- google/efficientnet-b4
+pipeline_tag: image-classification
+library_name: timm
+tags:
+- art
+- pytorch
+- images
+- ai
+---
+
+# AI Image Detection
+
+## Dataset
+- **AI**: ≈100,000 Images
+- **Human**: ≈100,000 Images
+
+## Model
+- **Architecture**: EfficientNet-B4
+- **Framework**: PyTorch
+
+## Evaluation Metrics
+- **Training Accuracy**: 99.75%
+- **Validation Accuracy**: 98.59%
+- **Training Loss**: 0.0072
+- **Validation Loss**: 0.0553
+
+
+## Usage
+
+```
+pip install torch torchvision timm huggingface_hub pillow
+```
+
+### Example Code
+```python
+import torch
+from torchvision import transforms
+from PIL import Image
+from timm import create_model
+from huggingface_hub import hf_hub_download
+
+# Parameters
+IMG_SIZE = 380
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+LABEL_MAPPING = {1: "human", 0: "ai"}
+
+# Download model from HuggingFace Hub
+MODEL_PATH = hf_hub_download(repo_id="Dafilab/ai-vs-human-image-detection", filename="model_epoch_8_acc_0.9859.pth")
+
+# Preprocessing
+transform = transforms.Compose([
+    transforms.Resize(IMG_SIZE + 20),
+    transforms.CenterCrop(IMG_SIZE),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# Load model
+model = create_model('efficientnet_b4', pretrained=False, num_classes=2)
+model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
+model.to(DEVICE).eval()
+
+# Prediction function
+def predict_image(image_path):
+    img = Image.open(image_path).convert("RGB")
+    img = transform(img).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        logits = model(img)
+        probs = torch.nn.functional.softmax(logits, dim=1)
+        predicted_class = torch.argmax(probs, dim=1).item()
+        confidence = probs[0, predicted_class].item()
+    return LABEL_MAPPING[predicted_class], confidence
+
+# Example usage
+image_path = "path/to/image.jpg"
+label, confidence = predict_image(image_path)
+print(f"Label: {label}, Confidence: {confidence:.2f}")
+```