A Keras EfficientNet model for classifying real-world document images into structured categories. Includes a full validation pipeline covering image quality checks and AI/fake image detection.
1# Step 1 — Install dependencies
2# pip install huggingface_hub tensorflow opencv-python pillow
3
4# Step 2 — Copy and run this complete code
5
6from huggingface_hub import snapshot_download
7import tensorflow as tf
8import numpy as np
9import cv2
10import json
11from tensorflow.keras.applications.efficientnet import preprocess_input
12
13# Download model from Hugging Face (cached after first run)
14local_path = snapshot_download(repo_id="shailgsits/document-classifier")
15
16# Load model + class labels
17model = tf.saved_model.load(local_path)
18infer = model.signatures["serving_default"]
19
20with open(f"{local_path}/class_index.json") as f:
21 class_indices = json.load(f)
22LABELS = {int(v): k for k, v in class_indices.items()}
23
24DOCUMENT_TYPE_LABELS = {
25 "1_visiting_card": "Visiting Card",
26 "2_prescription": "Prescription",
27 "3_shop_banner": "Shop Banner",
28 "4_invalid_image": "Invalid",
29}
30
31def predict(image_path: str) -> dict:
32 img = cv2.imread(image_path)
33 if img is None:
34 return {"status": "ERROR", "message": "Could not read image"}
35
36 img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
37 resized = cv2.resize(img_rgb, (224, 224))
38 input_arr = np.expand_dims(resized.astype(np.float32), axis=0)
39 input_arr = preprocess_input(input_arr)
40
41 outputs = infer(tf.constant(input_arr))
42 preds = list(outputs.values())[0].numpy()[0]
43 class_id = int(np.argmax(preds))
44 confidence = float(np.max(preds))
45 label = LABELS.get(class_id, "unknown")
46 friendly = DOCUMENT_TYPE_LABELS.get(label, label)
47
48 return {
49 "status": "VALID" if confidence >= 0.75 else "LOW_CONFIDENCE",
50 "document_type": label,
51 "document_type_label": friendly,
52 "confidence": round(confidence * 100, 2),
53 "all_scores": {
54 DOCUMENT_TYPE_LABELS.get(LABELS[i], LABELS[i]): round(float(p) * 100, 2)
55 for i, p in enumerate(preds)
56 }
57 }
58
59# --- Run prediction ---
60result = predict("your_image.jpg")
61print(result)
62
63# Example output:
64# {
65# 'status': 'VALID',
66# 'document_type': '1_visiting_card',
67# 'document_type_label': 'Visiting Card',
68# 'confidence': 97.43,
69# 'all_scores': {'Visiting Card': 97.43, 'Prescription': 1.2, 'Shop Banner': 0.9, 'Invalid': 0.47}
70# }
1!pip install huggingface_hub tensorflow pillow opencv-python requests -q
2
3import tensorflow as tf, numpy as np, cv2, requests, json
4from PIL import Image
5from io import BytesIO
6from huggingface_hub import hf_hub_download
7from tensorflow.keras.applications.efficientnet import preprocess_input
8
9# Load model + class mapping
10model = tf.keras.models.load_model(
11 hf_hub_download("shailgsits/document-classifier", "document_classifier_final.keras")
12)
13with open(hf_hub_download("shailgsits/document-classifier", "class_index.json")) as f:
14 index_to_label = {v: k.split("_", 1)[1] for k, v in json.load(f).items()}
15
16# Predict from any image URL
17def predict_from_url(url: str):
18 img = np.array(Image.open(BytesIO(requests.get(url).content)).convert("RGB"))[:, :, ::-1]
19 h, w = img.shape[:2]
20 scale = min(224 / w, 224 / h)
21 nw, nh = int(w * scale), int(h * scale)
22 res = cv2.resize(img, (nw, nh))
23 canvas = np.ones((224, 224, 3), np.uint8) * 255
24 canvas[(224 - nh) // 2:(224 - nh) // 2 + nh, (224 - nw) // 2:(224 - nw) // 2 + nw] = res
25 input_arr = preprocess_input(np.expand_dims(canvas.astype(np.float32), 0))
26 pred = model.predict(input_arr)[0]
27 idx = int(np.argmax(pred))
28 return {"label": index_to_label[idx], "confidence": round(float(pred[idx]) * 100, 2)}
29
30# Test with a Google Drive image
31url = "https://drive.google.com/uc?export=download&id=YOUR_FILE_ID"
32print(predict_from_url(url))
33# {'label': 'visiting_card', 'confidence': 97.43}
1from google.colab import files
2uploaded = files.upload()
3image_path = list(uploaded.keys())[0]
4
5img = cv2.imread(image_path)
6h, w = img.shape[:2]
7scale = min(224 / w, 224 / h)
8nw, nh = int(w * scale), int(h * scale)
9res = cv2.resize(img, (nw, nh))
10canvas = np.ones((224, 224, 3), np.uint8) * 255
11canvas[(224 - nh) // 2:(224 - nh) // 2 + nh, (224 - nw) // 2:(224 - nw) // 2 + nw] = res
12input_arr = preprocess_input(np.expand_dims(canvas.astype(np.float32), 0))
13pred = model.predict(input_arr)[0]
14idx = int(np.argmax(pred))
15print({"label": index_to_label[idx], "confidence": round(float(pred[idx]) * 100, 2)})
