Add complete inference pipeline script

inference_pipeline.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+"""
+KYC Document Cropping & Rotation - Complete Inference Pipeline
+
+Models:
+- Segmentation (corner detection): https://huggingface.co/Jwalit/kyc-document-corner-detector
+- Rotation classifier: https://huggingface.co/Jwalit/kyc-document-rotation-classifier
+
+Usage:
+    python inference_pipeline.py --image path/to/image.jpg --output out.jpg
+"""
+
+import argparse, warnings
+from pathlib import Path
+import numpy as np
+import cv2
+from PIL import Image
+import torch, torch.nn as nn
+from torchvision import transforms
+
+warnings.filterwarnings("ignore")
+
+# ── Model Classes ──
+class SegModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        from torchvision import models
+        self.enc = models.mobilenet_v3_small(weights=None).features
+        self.dec = nn.Sequential(
+            nn.Conv2d(576,256,3,padding=1), nn.ReLU(),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+            nn.Conv2d(256,128,3,padding=1), nn.ReLU(),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+            nn.Conv2d(128,64,3,padding=1), nn.ReLU(),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+            nn.Conv2d(64,32,3,padding=1), nn.ReLU(),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+            nn.Conv2d(32,16,3,padding=1), nn.ReLU(),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+            nn.Conv2d(16,1,3,padding=1),
+        )
+    def forward(self, x):
+        return self.dec(self.enc(x))
+
+class RotModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        from torchvision import models
+        self.m = models.mobilenet_v3_small(weights=None)
+        self.m.classifier[3] = nn.Linear(self.m.classifier[3].in_features, 4)
+    def forward(self, x):
+        return self.m(x)
+
+# ── Preprocessing ──
+seg_tf = transforms.Compose([
+    transforms.Resize((224,224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
+])
+rot_tf = transforms.Compose([
+    transforms.Resize((224,224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
+])
+
+# ── Helpers ──
+def mask_to_corners(mask):
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return None
+    cnt = max(contours, key=cv2.contourArea)
+    peri = cv2.arcLength(cnt, True)
+    for eps in [0.02, 0.05, 0.1]:
+        approx = cv2.approxPolyDP(cnt, eps*peri, True)
+        if len(approx) == 4:
+            pts = approx.reshape(4,2).astype(np.float32)
+            s, d = pts.sum(axis=1), np.diff(pts, axis=1).flatten()
+            rect = np.zeros((4,2), dtype=np.float32)
+            rect[0] = pts[np.argmin(s)]
+            rect[2] = pts[np.argmax(s)]
+            rect[1] = pts[np.argmin(d)]
+            rect[3] = pts[np.argmax(d)]
+            return rect
+    rect = cv2.minAreaRect(cnt)
+    return cv2.boxPoints(rect).astype(np.float32)
+
+def predict_corners(seg_model, img_path, device='cpu'):
+    img = Image.open(img_path).convert('RGB')
+    orig_w, orig_h = img.size
+    t = seg_tf(img).unsqueeze(0).to(device)
+    seg_model.eval()
+    with torch.no_grad():
+        pred = torch.sigmoid(seg_model(t))[0,0].cpu().numpy()
+    mask = (pred > 0.5).astype(np.uint8) * 255
+    mask = cv2.resize(mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
+    return mask_to_corners(mask)
+
+def predict_rotation(rot_model, img_path, device='cpu'):
+    img = Image.open(img_path).convert('RGB')
+    t = rot_tf(img).unsqueeze(0).to(device)
+    rot_model.eval()
+    with torch.no_grad():
+        pred = torch.argmax(rot_model(t), 1).item()
+    return [0, 90, 180, 270][pred]
+
+def crop_and_rotate(img_path, corners, angle):
+    img = cv2.imread(str(img_path))
+    if img is None:
+        return None
+    if corners is not None:
+        pts = corners.astype(np.float32)
+        w = int(max(np.linalg.norm(pts[1]-pts[0]), np.linalg.norm(pts[2]-pts[3])))
+        h = int(max(np.linalg.norm(pts[3]-pts[0]), np.linalg.norm(pts[2]-pts[1])))
+        dst = np.array([[0,0],[w-1,0],[w-1,h-1],[0,h-1]], dtype=np.float32)
+        M = cv2.getPerspectiveTransform(pts, dst)
+        cropped = cv2.warpPerspective(img, M, (w, h))
+    else:
+        cropped = img.copy()
+    if angle == 90:
+        cropped = cv2.rotate(cropped, cv2.ROTATE_90_COUNTERCLOCKWISE)
+    elif angle == 180:
+        cropped = cv2.rotate(cropped, cv2.ROTATE_180)
+    elif angle == 270:
+        cropped = cv2.rotate(cropped, cv2.ROTATE_90_CLOCKWISE)
+    return cropped
+
+def process_image(img_path, seg_weights, rot_weights, output_path=None, device='cpu'):
+    seg_model = SegModel()
+    seg_model.load_state_dict(torch.load(seg_weights, map_location=device))
+    seg_model.to(device)
+    rot_model = RotModel()
+    rot_model.load_state_dict(torch.load(rot_weights, map_location=device))
+    rot_model.to(device)
+    print(f"Processing: {img_path}")
+    corners = predict_corners(seg_model, img_path, device)
+    angle = predict_rotation(rot_model, img_path, device)
+    print(f"  Rotation: {angle}°, Corners: {corners is not None}")
+    result = crop_and_rotate(img_path, corners, angle)
+    if output_path:
+        cv2.imwrite(str(output_path), result)
+        print(f"  Saved: {output_path}")
+    return result
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image", required=True)
+    parser.add_argument("--seg-weights", default="pytorch_model.bin")
+    parser.add_argument("--rot-weights", required=True)
+    parser.add_argument("--output", default="output.jpg")
+    parser.add_argument("--device", default="cpu")
+    args = parser.parse_args()
+    process_image(args.image, args.seg_weights, args.rot_weights, args.output, args.device)
+
+if __name__ == "__main__":
+    main()