Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,12 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+
+COPY src/ ./src/
+COPY models/ ./models/
+
+EXPOSE 8000 8501
+
+CMD ["sh", "-c", "uvicorn src.app:app --host 0.0.0.0 --port 8000 & streamlit run src/dashboard.py --server.port 8501 --server.address 0.0.0.0"]

README.md

@@ -1,3 +0,0 @@
----
-license: apache-2.0
----

requirements.txt

@@ -0,0 +1,14 @@
+# requirements.txt
+torch==2.0.1
+torchvision==0.15.2
+opencv-python==4.8.1.78
+fastapi==0.104.1
+uvicorn==0.24.0
+streamlit==1.28.1
+onnxruntime==1.16.1
+numpy==1.24.3
+pillow==10.0.1
+transformers==4.35.0  # Optional, but useful for some models
+scikit-learn==1.3.0
+matplotlib==3.7.2
+requests==2.31.0

src/app.py

@@ -0,0 +1,41 @@
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import JSONResponse
+import onnxruntime as ort
+import numpy as np
+from src.preprocess import preprocess_image
+from src.gradcam import GradCAM  # Note: GradCAM uses PyTorch, so for ONNX we approximate or skip full CAM
+import cv2
+from PIL import Image
+import io
+
+app = FastAPI(title="AutoVision API")
+
+# Load ONNX model
+ort_session = ort.InferenceSession('../models/resnet18_anomaly.onnx')
+classes = ['crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale', 'scratches']
+
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents))
+    image_np = np.array(image)
+    
+    input_data = preprocess_image(image_np)
+    
+    ort_inputs = {ort_session.get_inputs()[0].name: input_data.astype(np.float32)}
+    ort_outs = ort_session.run(None, ort_inputs)
+    pred = np.argmax(ort_outs[0])
+    confidence = np.max(ort_outs[0])
+    
+    # For Grad-CAM, we'd need to load PyTorch model separately for explainability
+    # Here, return pred and conf; overlay handled in UI
+    
+    return JSONResponse({
+        "prediction": classes[pred],
+        "confidence": float(confidence),
+        "class_id": int(pred)
+    })
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

src/dashboard.py

@@ -0,0 +1,56 @@
+import streamlit as st
+import cv2
+import numpy as np
+from src.preprocess import preprocess_image, camera_stream, overlay_heatmap
+from src.gradcam import GradCAM
+import torch
+from PIL import Image
+import requests
+
+st.title("AutoVision: Real-Time Defect Detection")
+
+# Load models
+@st.cache_resource
+def load_models():
+    gradcam = GradCAM('../models/resnet18_anomaly.pth')
+    return gradcam
+
+gradcam = load_models()
+classes = ['normal', 'crazing', 'inclusion', 'pits', 'pitted_surface', 'rolled-in_scale', 'scratches']
+
+# Real-time camera feed
+frame_placeholder = st.empty()
+prediction_placeholder = st.empty()
+
+# Optional: Use API for prediction (if backend running)
+use_api = st.checkbox("Use FastAPI Backend for Inference")
+
+for frame in camera_stream():
+    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    
+    # Preprocess for inference
+    input_data = preprocess_image(rgb_frame)
+    input_tensor = torch.from_numpy(input_data).float()
+    
+    if use_api:
+        # Upload to API (simplified; in practice, use multipart)
+        # For demo, use local PyTorch
+        pass
+    
+    # Local inference with PyTorch (for Grad-CAM compatibility)
+    with torch.no_grad():
+        output = gradcam.model(input_tensor)
+        pred = output.argmax().item()
+        confidence = torch.softmax(output, dim=1).max().item()
+    
+    # Generate Grad-CAM
+    heatmap = gradcam.generate(input_tensor, pred)
+    
+    # Overlay
+    overlaid = overlay_heatmap(rgb_frame, heatmap)
+    
+    frame_placeholder.image(overlaid, channels="RGB")
+    
+    prediction_placeholder.markdown(f"**Prediction:** {classes[pred]} ({confidence:.2%})")
+
+st.info("Press Ctrl+C to stop camera.")

src/gradcam.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn.functional as F
+import cv2
+import numpy as np
+from src.model import get_model
+
+class GradCAM:
+    def __init__(self, model_path, target_layer='layer4'):
+        self.model = get_model(pretrained=False)
+        self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
+        self.model.eval()
+        self.target_layer = target_layer
+        self.gradients = None
+        self.activations = None
+        self.hooks = []
+        self._register_hooks()
+    
+    def _register_hooks(self):
+        def backward_hook(module, grad_input, grad_output):
+            self.gradients = grad_output[0]
+        
+        def forward_hook(module, input, output):
+            self.activations = output
+        
+        for name, module in self.model.named_modules():
+            if target_layer in name:
+                self.hooks.append(module.register_forward_hook(forward_hook))
+                self.hooks.append(module.register_backward_hook(backward_hook))
+    
+    def generate(self, input_tensor, class_idx=None):
+        self.model.zero_grad()
+        output = self.model(input_tensor)
+        if class_idx is None:
+            class_idx = output.argmax().item()
+        score = output[0, class_idx]
+        score.backward()
+        
+        gradients = self.gradients[0]
+        activations = self.activations[0]
+        weights = torch.mean(gradients, dim=(1, 2), keepdim=True)
+        
+        cam = torch.sum(weights * activations, dim=1, keepdim=True)
+        cam = F.relu(cam)
+        cam = cam.squeeze().detach().cpu().numpy()
+        cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam) + 1e-8)
+        
+        # Remove hooks
+        for hook in self.hooks:
+            hook.remove()
+        
+        return cam
+    
+    def __del__(self):
+        for hook in self.hooks:
+            hook.remove()
+
+# Usage example:
+# gradcam = GradCAM('../models/resnet18_anomaly.pth')
+# input_tensor = torch.randn(1, 3, 224, 224)  # From preprocess
+# heatmap = gradcam.generate(input_tensor)

src/model.py

@@ -0,0 +1,19 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+from torchvision import transforms
+
+NUM_CLASSES = 7  # 6 defects + normal
+
+def get_model(pretrained=True):
+    model = models.resnet18(pretrained=pretrained)
+    num_ftrs = model.fc.in_features
+    model.fc = nn.Linear(num_ftrs, NUM_CLASSES)
+    return model
+
+def get_transforms():
+    return transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])

src/preprocess.py

@@ -0,0 +1,35 @@
+import cv2
+import numpy as np
+from PIL import Image
+from src.model import get_transforms
+
+def preprocess_image(image_path_or_np):
+    """Preprocess single image for inference."""
+    if isinstance(image_path_or_np, str):
+        image = cv2.imread(image_path_or_np)
+    else:
+        image = image_path_or_np
+    
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    image = Image.fromarray(image)
+    transform = get_transforms()
+    return transform(image).unsqueeze(0).numpy()  # To numpy for ONNX
+
+def camera_stream():
+    """Generator for real-time camera feed."""
+    cap = cv2.VideoCapture(0)  # Use default camera
+    while True:
+        ret, frame = cap.read()
+        if ret:
+            yield frame
+        else:
+            break
+    cap.release()
+
+def overlay_heatmap(frame, heatmap, alpha=0.4):
+    """Overlay Grad-CAM heatmap on frame."""
+    heatmap = cv2.resize(heatmap, (frame.shape[1], frame.shape[0]))
+    heatmap = np.uint8(255 * heatmap)
+    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+    superimposed = cv2.addWeighted(frame, 1 - alpha, heatmap, alpha, 0)
+    return superimposed

src/train.py

@@ -0,0 +1,74 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from src.model import get_model, get_transforms
+import numpy as np
+from sklearn.metrics import accuracy_score
+
+# Assume data/ has train/ and val/ folders with subfolders for classes: normal, crazing, inclusion, etc.
+# https://www.kaggle.com/datasets/kaustubhdikshit/neu-surface-defect-database
+DATA_DIR = '../data/neu_surface_defect_database'
+BATCH_SIZE = 32
+EPOCHS = 10
+LEARNING_RATE = 0.001
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def main():
+    transform = get_transforms()
+    
+    train_dataset = datasets.ImageFolder(os.path.join(DATA_DIR, 'train'), transform=transform)
+    val_dataset = datasets.ImageFolder(os.path.join(DATA_DIR, 'val'), transform=transform)
+    
+    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
+    
+    model = get_model(pretrained=True).to(DEVICE)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
+    
+    for epoch in range(EPOCHS):
+        model.train()
+        running_loss = 0.0
+        for inputs, labels in train_loader:
+            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            running_loss += loss.item()
+        
+        print(f'Epoch {epoch+1}/{EPOCHS}, Loss: {running_loss / len(train_loader):.4f}')
+        
+        # Validation
+        model.eval()
+        preds, trues = [], []
+        with torch.no_grad():
+            for inputs, labels in val_loader:
+                inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
+                outputs = model(inputs)
+                _, predicted = torch.max(outputs, 1)
+                preds.extend(predicted.cpu().numpy())
+                trues.extend(labels.cpu().numpy())
+        
+        acc = accuracy_score(trues, preds)
+        print(f'Validation Accuracy: {acc:.4f}')
+    
+    # Save PyTorch model
+    torch.save(model.state_dict(), '../models/resnet18_anomaly.pth')
+    
+    # Export to ONNX
+    model.eval()
+    dummy_input = torch.randn(1, 3, 224, 224).to(DEVICE)
+    torch.onnx.export(model, dummy_input, '../models/resnet18_anomaly.onnx',
+                      export_params=True, opset_version=11,
+                      do_constant_folding=True,
+                      input_names=['input'], output_names=['output'])
+    
+    print('Model trained and exported to ONNX!')
+
+if __name__ == '__main__':
+    main()