Spaces:
Sleeping
Sleeping
File size: 2,771 Bytes
20d7fde | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import os
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
import torch
torch.set_num_threads(1)
import torchvision.transforms as T
from PIL import Image, ImageOps
from fastai.vision.all import load_learner
import pickle
import plum._resolver
# --- THE MODEL LOADING PATCH ---
class PatchedResolver(plum._resolver.Resolver):
def __setstate__(self, state):
if isinstance(state, dict):
for k, v in state.items():
try: setattr(self, k, v)
except AttributeError: pass
elif isinstance(state, tuple):
for item in state:
if isinstance(item, dict):
for k, v in item.items():
try: setattr(self, k, v)
except AttributeError: pass
class SafeUnpickler(pickle.Unpickler):
def find_class(self, module, name):
if name == 'Resolver' and 'plum' in module:
return PatchedResolver
return super().find_class(module, name)
class SafePickle:
Unpickler = SafeUnpickler
# -------------------------------
# 1. Load the model using our safe unpickler
print("Loading model...")
learn = load_learner('handwriting_classifier_best (1).pkl', pickle_module=SafePickle)
print("Model loaded successfully!")
# 2. Extract the raw PyTorch model and the classes (vocab)
pytorch_model = learn.model
pytorch_model.eval() # Set model to evaluation mode
vocab = list(learn.dls.vocab) # Get the list of your 163 classes
# 3. Standard PyTorch Image Preprocessing (Bypasses FastAI transforms entirely)
# This perfectly matches the Resize and standard ImageNet normalization used by ResNet
preprocess = T.Compose([
T.Resize((40, 80)),
T.ToTensor(),
T.Normalize(
mean=[0.485, 0.456, 0.406], # Standard ImageNet stats
std=[0.229, 0.224, 0.225]
)
])
# 4. Open the image using Pillow and preprocess it
img = Image.open('IMG_1521 (1) (1) (1).jpg')
img = ImageOps.exif_transpose(img).convert('L').convert('RGB') # Fix iPhone rotation bug by applying EXIF orientation
input_tensor = preprocess(img).unsqueeze(0) # Add batch dimension [1, 3, 40, 80]
# 5. Raw PyTorch Forward Pass
print("Running raw PyTorch inference...")
with torch.no_grad():
outputs = pytorch_model(input_tensor)
# Apply softmax to convert raw outputs to probabilities
probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
# 6. Extract the top 3 predictions
top_prob, top_catid = torch.topk(probabilities, 3)
print("\nPrediction Results:")
for i in range(top_prob.size(0)):
class_name = vocab[top_catid[i].item()]
confidence = top_prob[i].item() * 100
print(f"Class {class_name}: {confidence:.2f}%") |