Upload tuberculosis ViT model with complete configuration

README.md

@@ -0,0 +1,115 @@
+---
+language: en
+license: apache-2.0
+tags:
+- medical
+- chest-xray
+- tuberculosis
+- vision-transformer
+- pytorch
+- image-classification
+metrics:
+- accuracy
+- precision
+- recall
+- f1
+pipeline_tag: image-classification
+---
+
+# Tuberculosis Detection ViT Model
+
+This repository contains a Vision Transformer (ViT) model for classifying chest X-ray images as **Normal** or **Tuberculosis**.
+
+## Model Details
+
+- **Model Type**: Vision Transformer (ViT)
+- **Framework**: PyTorch 2.0.1
+- **Architecture**: 
+  - Patch size: 16
+  - Number of patches: 196
+  - Projection dimension: 768
+  - Number of attention heads: 12
+  - Number of encoder blocks: 12
+- **Input Size**: 224x224 RGB images
+- **Classes**: Normal, Tuberculosis
+
+## Dataset
+- **Source**: TB Chest Radiography Database
+- **Preprocessing**: Grayscale conversion, CLAHE enhancement, Gaussian blur, RGB conversion, standardization
+- **Training**: Early stopping with validation loss monitoring
+
+## Usage
+
+### Quick Start
+```python
+from huggingface_hub import hf_hub_download
+import torch
+import cv2
+import numpy as np
+
+# Download the handler
+handler_path = hf_hub_download(repo_id="sukhmani1303/tuberculosis-vit-model", filename="handler.py")
+
+# Import and use
+exec(open(handler_path).read())
+classifier = TBClassifier()
+
+# Load and predict
+image = cv2.imread("path/to/chest_xray.jpg")
+result = classifier.predict(image)
+print(f"Prediction: {result['prediction']}, Confidence: {result['confidence']:.4f}")
+```
+
+### Loading the PyTorch Model
+```python
+from huggingface_hub import hf_hub_download
+import torch
+import json
+
+# Download model files
+config_path = hf_hub_download(repo_id="sukhmani1303/tuberculosis-vit-model", filename="config.json")
+model_path = hf_hub_download(repo_id="sukhmani1303/tuberculosis-vit-model", filename="pytorch_model.bin")
+
+# Load configuration
+with open(config_path, 'r') as f:
+    config = json.load(f)
+
+# Initialize your ViT model class with config
+# model = ViT(**config)  # You'll need to have your ViT class available
+# model.load_state_dict(torch.load(model_path, map_location='cpu'))
+# model.eval()
+```
+
+## Performance
+- **Training**: Early stopping with validation monitoring
+- **Evaluation**: Tested on held-out validation set
+- **Metrics**: Accuracy, Precision, Recall, F1-score available upon request
+
+## Preprocessing Pipeline
+1. Convert RGB to grayscale
+2. Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)
+3. Apply Gaussian blur (5x5 kernel)
+4. Resize to 224x224
+5. Convert back to RGB format
+6. Normalize with z-score standardization
+
+## Medical Disclaimer
+⚠️ **IMPORTANT**: This model is for educational and research purposes only. It should not be used for actual medical diagnosis. Always consult qualified healthcare professionals for medical advice and diagnosis.
+
+## Citation
+If you use this model in your research, please cite:
+```
+@misc{tuberculosis-vit-model,
+  author = {Sukhmani},
+  title = {Tuberculosis Detection using Vision Transformer},
+  year = {2024},
+  publisher = {Hugging Face},
+  url = {https://huggingface.co/sukhmani1303/tuberculosis-vit-model}
+}
+```
+
+## License
+Apache 2.0
+
+## Contact
+For questions or issues, please open an issue in the repository.

config.json

@@ -0,0 +1,16 @@
+{
+  "model_type": "vit",
+  "patch_size": 16,
+  "num_patches": 196,
+  "projection_dim": 768,
+  "num_heads": 12,
+  "num_encoder": 12,
+  "input_size": 224,
+  "num_classes": 1,
+  "class_names": [
+    "Normal",
+    "Tuberculosis"
+  ],
+  "framework": "pytorch",
+  "torch_dtype": "float32"
+}

handler.py

@@ -0,0 +1,159 @@
+
+import torch
+import cv2
+import numpy as np
+import os
+
+class TBClassifier:
+    """
+    Tuberculosis classifier using Vision Transformer
+    """
+    def __init__(self, model_path="model.pt", config_path="config.json"):
+        """
+        Initialize the classifier
+        
+        Args:
+            model_path: Path to the TorchScript model file
+            config_path: Path to the configuration file
+        """
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Load configuration if available
+        self.config = None
+        if os.path.exists(config_path):
+            import json
+            with open(config_path, 'r') as f:
+                self.config = json.load(f)
+        
+        # Load model
+        try:
+            self.model = torch.jit.load(model_path, map_location=self.device)
+            self.model.eval()
+            print(f"Model loaded successfully on {self.device}")
+        except Exception as e:
+            raise RuntimeError(f"Failed to load model: {str(e)}")
+        
+        # Model configuration
+        self.class_names = self.config.get('class_names', ['Normal', 'Tuberculosis']) if self.config else ['Normal', 'Tuberculosis']
+        self.img_size = self.config.get('input_size', 224) if self.config else 224
+        
+        print(f"Classifier initialized with classes: {self.class_names}")
+
+    def preprocess(self, image):
+        """
+        Preprocess input image for model inference
+        
+        Args:
+            image: Input image as numpy array (BGR or RGB)
+            
+        Returns:
+            Preprocessed tensor ready for model inference
+        """
+        try:
+            # Handle different input formats
+            if isinstance(image, str):
+                image = cv2.imread(image)
+            
+            if image is None:
+                raise ValueError("Invalid image input")
+            
+            # Convert to grayscale if needed
+            if len(image.shape) == 3 and image.shape[2] == 3:
+                # Assume BGR format from cv2
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+            
+            # Apply CLAHE for contrast enhancement
+            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
+            image = clahe.apply(image)
+            
+            # Apply Gaussian blur for noise reduction
+            image = cv2.GaussianBlur(image, (5, 5), 0)
+            
+            # Resize to model input size
+            image = cv2.resize(image, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR)
+            
+            # Convert back to RGB format
+            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+            
+            # Convert to tensor format (C, H, W)
+            image = np.moveaxis(image, -1, 0).astype(np.float32)
+            
+            # Normalize
+            image = (image - image.mean()) / (image.std() + 1e-8)
+            
+            # Add batch dimension and move to device
+            tensor = torch.tensor(image).unsqueeze(0).to(self.device)
+            
+            return tensor
+            
+        except Exception as e:
+            raise RuntimeError(f"Preprocessing failed: {str(e)}")
+
+    def predict(self, image, return_probs=False):
+        """
+        Predict tuberculosis from chest X-ray image
+        
+        Args:
+            image: Input image (file path, numpy array, or PIL image)
+            return_probs: Whether to return raw probabilities
+            
+        Returns:
+            Dictionary with prediction results
+        """
+        try:
+            # Preprocess image
+            processed_image = self.preprocess(image)
+            
+            # Model inference
+            with torch.no_grad():
+                output = self.model(processed_image)
+                
+                # Handle different output formats
+                if len(output.shape) > 1:
+                    output = output.squeeze(-1)
+                
+                prob = torch.sigmoid(output).item()  # Ensure probability is in [0,1]
+                
+                # Determine class
+                class_id = 1 if prob > 0.5 else 0
+                confidence = prob if class_id == 1 else 1 - prob
+                prediction = self.class_names[class_id]
+                
+                result = {
+                    "prediction": prediction,
+                    "confidence": float(confidence),
+                    "class_id": class_id
+                }
+                
+                if return_probs:
+                    result["raw_probability"] = float(prob)
+                    result["probabilities"] = {
+                        self.class_names[0]: float(1 - prob),
+                        self.class_names[1]: float(prob)
+                    }
+                
+                return result
+                
+        except Exception as e:
+            return {
+                "error": str(e),
+                "prediction": None,
+                "confidence": None
+            }
+
+    def batch_predict(self, images, return_probs=False):
+        """
+        Predict on multiple images
+        
+        Args:
+            images: List of images
+            return_probs: Whether to return raw probabilities
+            
+        Returns:
+            List of prediction results
+        """
+        results = []
+        for img in images:
+            result = self.predict(img, return_probs=return_probs)
+            results.append(result)
+        return results

metadata.json

@@ -0,0 +1,17 @@
+{
+  "model_name": "Tuberculosis ViT Classifier",
+  "version": "1.0.0",
+  "author": "sukhmani1303",
+  "description": "Vision Transformer for tuberculosis detection in chest X-rays",
+  "tags": [
+    "medical",
+    "chest-xray",
+    "tuberculosis",
+    "vision-transformer",
+    "pytorch"
+  ],
+  "license": "apache-2.0",
+  "framework": "pytorch",
+  "model_size_mb": 654.1593074798584,
+  "created_date": "2025-06-08"
+}

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ceea603a3fd53322fcefff6a15cf0d326f55e80235f33bb98d12a62f41cba00
+size 343282923

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa2914422445a137de9749565af1a8d5e29327de2abb5411b6c44ac214f36578
+size 342643362

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.0.0,<3.0.0
+torchvision>=0.15.0
+opencv-python-headless>=4.8.0
+numpy>=1.24.0
+huggingface-hub>=0.16.0