romal , devngari classifier added

main.py

@@ -71,7 +71,7 @@ from pydantic import BaseModel
 import shutil
 
 # Import from optimized utils
-from utils import dev_number, roman_number, dev_letter, roman_letter
+from utils import dev_number, roman_number, dev_letter, roman_letter, predict_ne
 
 app = FastAPI(
     title="OCR API",
@@ -150,7 +150,8 @@ async def extract_text(
         "dev_number": dev_number,
         "roman_number": roman_number,
         "dev_letter": dev_letter,
-        "roman_letter": roman_letter
+        "roman_letter": roman_letter,
+        
     }
     
     if model_type not in ocr_functions:
@@ -179,6 +180,21 @@ async def extract_roman_letter(image: UploadFile = File(...)):
     """Extract Roman letters from an image"""
     return await process_ocr_request(image, roman_letter)
 
+@app.post("/predict_ne")
+async def classify_ne(image: UploadFile = File(...)):
+    """Predict Named Entities from an image"""
+    # Placeholder for Named Entity Recognition logic
+    image_path  = await save_upload_file_tmp(image)
+    prediction = predict_ne(
+        image_path=image_path,
+        # model="models/nepali_english_classifier.pth",  # Update with actual model path
+        device="cpu"  # or "cpu"
+    )
+
+
+
+    # Implement the logic as per your requirements
+    return JSONResponse(content={"predicted": prediction})
 # Health check endpoint
 @app.get("/health")
 async def health_check():

models/nepali_english_classifier.pth

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

utils.py

@@ -164,9 +164,11 @@ from PIL import Image
 import numpy as np
 import torchvision.transforms as transforms
 from doctr.io import DocumentFile
+from torchvision import models
 from doctr.models import recognition_predictor
 import os
 from functools import lru_cache
+import pickle
 
 # Character sets
 CHARACTER_NUM = "0123456789-"
@@ -176,11 +178,28 @@ CHARACTER_LETTER = ''' "()-./0123456789:?ABCDEFGHIKLMNOPQRSTUWYabcdefghijklmnopr
 MODEL_PATHS = {
     'dev_digits': "models/devnagri_digits_20k_v2.pth",
     'roman_digits': "models/roman_digits_20k_v5.pth",
-    'dev_letter': "models/small_devnagari_letter.pth"
+    'dev_letter': "models/small_devnagari_letter.pth",
+    'classify_ne': "models/nepali_english_classifier.pth"
 }
 
+
 # Use GPU if available
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+class ResNetClassifier(nn.Module):
+    def __init__(self, num_classes=2):
+        super(ResNetClassifier, self).__init__()
+        self.base_model = models.resnet50(weights='IMAGENET1K_V2')  # Pre-trained ResNet-50
+        for param in self.base_model.parameters():
+            param.requires_grad = False  # Freeze base model
+        num_ftrs = self.base_model.fc.in_features
+        self.base_model.fc = nn.Sequential(
+            nn.Linear(num_ftrs, 128),
+            nn.ReLU(),
+            nn.Linear(128, num_classes)
+        )
+    
+    def forward(self, x):
+        return self.base_model(x)
 
 # Define the CRNN model
 class CRNN(nn.Module):
@@ -326,6 +345,8 @@ class OCRModelManager:
         result = self.roman_letter_model(img)
         # print(result)
         return result[0][0]
+    
+    
 
 
 # Initialize the model manager as a singleton
@@ -346,4 +367,25 @@ def dev_letter(image_path):
 
 def roman_letter(image_path):
     """Recognize Roman letters in an image"""
-    return ocr_manager.predict_roman_letter(image_path)
+    return ocr_manager.predict_roman_letter(image_path)
+
+
+def predict_ne(image_path, device="cpu"):
+        # load label encoder
+        
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        model = ResNetClassifier(num_classes=2).to(device)
+        # model.eval()
+        transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+        image = Image.open(image_path).convert('RGB')
+        image_tensor = transform(image).unsqueeze(0).to(device)
+        model.load_state_dict(torch.load('models/nepali_english_classifier.pth', map_location=device))
+        model.eval()
+        with torch.no_grad():
+            output = model(image_tensor)
+            _, predicted = torch.max(output, 1)
+        return predicted