archaiveproject/CCR_OCR

1import os2import json3from PIL import Image4import torch5import torch.nn as nn6from torchvision import models, transforms7from huggingface_hub import snapshot_download8import gradio as gr9 10# -------- Model Definition --------11class ChineseClassifier(nn.Module):12    def __init__(self, embed_dim, num_classes, pretrainedEncoder=True, unfreezeEncoder=True):13        super().__init__()14        resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) if pretrainedEncoder else models.resnet50()15        self.resnet = nn.Sequential(*list(resnet.children())[:-1])16        for param in self.resnet.parameters():17            param.requires_grad = unfreezeEncoder18        self.fc = nn.Linear(resnet.fc.in_features, embed_dim)19        self.batch_norm = nn.BatchNorm1d(embed_dim)20        self.dropout = nn.Dropout(0.3)21        self.classifier = nn.Linear(embed_dim, num_classes)22 23    def forward(self, x, return_embedding=False):24        x = self.resnet(x)25        x = torch.flatten(x, 1)26        x = self.fc(x)27        x = self.batch_norm(x)28        x = self.dropout(x)29        if return_embedding:30            return x31        x = self.classifier(x)32        return x33 34# -------- Utility Functions --------35def get_sorted_classes(labels_dict):36    """Extract sorted unique classes from labels dictionary"""37    return sorted(set(labels_dict.values()))38 39def load_labels_json(labels_json_path):40    """Load and normalize labels JSON"""41    with open(labels_json_path, "r", encoding="utf-8") as f:42        labels_dict = json.load(f)43    # Normalize paths and remove directory prefixes44    return {os.path.basename(k).replace("\\", "/"): v for k, v in labels_dict.items()}45 46def prepare_transforms():47    return transforms.Compose([48        transforms.Resize((224, 224)),49        transforms.ToTensor(),50        transforms.Normalize(mean=[0.485, 0.456, 0.406],51                             std=[0.229, 0.224, 0.225]),52    ])53 54def load_model(model_path, embed_dim, num_classes, device, pretrained=True, unfreeze=True):55    model = ChineseClassifier(embed_dim, num_classes, pretrainedEncoder=pretrained, unfreezeEncoder=unfreeze).to(device)56    checkpoint = torch.load(model_path, map_location=device)57    if "model_state_dict" in checkpoint:58        try:59            model.load_state_dict(checkpoint["model_state_dict"])60        except RuntimeError as e:61            print("Warning:", e)62            print("Loading partial weights, skipping classifier layer...")63            filtered_state_dict = {k: v for k, v in checkpoint["model_state_dict"].items() if not k.startswith("classifier.")}64            model.load_state_dict(filtered_state_dict, strict=False)65    else:66        model.load_state_dict(checkpoint)67    model.eval()68    return model69 70# -------- Setup --------71DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")72EMBED_DIM = 51273LABELS_JSON_PATH = "labels.json"74 75# 1. Load labels and extract sorted classes76labels_dict = load_labels_json(LABELS_JSON_PATH)77classes = get_sorted_classes(labels_dict)78idx_to_class = {idx: c for idx, c in enumerate(classes)}79num_classes = len(classes)80 81# Verify class count matches training82print(f"Loaded {num_classes} classes")83print(f"First 5 classes: {classes[:5]}")84 85# 2. Download model86REPO_ID = "JJJHHHH/CCR_EthicalSplit_Finetune"87print("Downloading model from repo...")88repo_dir = snapshot_download(repo_id=REPO_ID)89model_path = os.path.join(repo_dir, "CCR_EthicalSplit_Finetune.pth")90print(f"Model path: {model_path}")91 92# 3. Load model93model = load_model(model_path, EMBED_DIM, num_classes, DEVICE)94transform = prepare_transforms()95 96# -------- Prediction Function --------97def predict(pil_img):98    """Predict character from PIL image"""99    img_t = transform(pil_img).unsqueeze(0).to(DEVICE)100    with torch.no_grad():101        output = model(img_t)102        pred_idx = output.argmax(dim=1).item()103        pred_label = idx_to_class[pred_idx]104    return pred_label105 106# -------- Gradio Interface --------107gr.Interface(108    fn=predict,109    inputs=gr.Image(type="pil", label="Upload Handwritten Chinese Character"),110    outputs=gr.Text(label="Predicted Character"),111    title="Chinese Character Recognition",112    description="Recognizes handwritten Chinese characters with 80% accuracy",113).launch()