CoolFace
Modelpublic

atomdev-ibktommy/garment-stacking-classifier

sourceHugging Faceupdated 2mo agoView on Hugging Face
0likes
predict.py65 linesDownload Raw Back to root
1import os2import cv23import json4import joblib5import numpy as np6from skimage.feature import hog, local_binary_pattern7 8class GarmentClassifierPipeline:9    def __init__(self, artifact_dir="."):10        self.artifact_dir = artifact_dir11        12        with open(os.path.join(artifact_dir, "class_mapping.json"), "r") as f:13            self.class_mapping = {int(k): v for k, v in json.load(f).items()}14            15        self.scaler_base = joblib.load(os.path.join(artifact_dir, "scaler_base.joblib"))16        self.svm_base = joblib.load(os.path.join(artifact_dir, "svm_baseline.joblib"))17        self.rf_base = joblib.load(os.path.join(artifact_dir, "rf_baseline.joblib"))18        self.meta_base = joblib.load(os.path.join(artifact_dir, "meta_xgb_baseline.joblib"))19 20    def extract_features(self, image_bgr):21        img_resized = cv2.resize(image_bgr, (224, 224))22        23        # 1. HOG Feature Vector24        gray = cv2.cvtColor(img_resized, cv2.COLOR_BGR2GRAY)25        hog_feat = hog(26            gray, orientations=9, pixels_per_cell=(16, 16),27            cells_per_block=(2, 2), block_norm='L2-Hys', visualize=False28        )29        30        # 2. LBP Texture Histogram31        lbp = local_binary_pattern(gray, P=24, R=3, method='uniform')32        lbp_hist, _ = np.histogram(lbp.ravel(), bins=26, range=(0, 26), density=True)33        34        # 3. HSV Color Histogram35        hsv = cv2.cvtColor(img_resized, cv2.COLOR_BGR2HSV)36        hsv_hist = cv2.calcHist([hsv], [0, 1, 2], None, [8, 8, 8], [0, 180, 0, 256, 0, 256])37        hsv_hist = cv2.normalize(hsv_hist, hsv_hist).flatten()38        39        return np.hstack([hog_feat, lbp_hist, hsv_hist]).reshape(1, -1)40 41    def predict(self, image_bgr):42        raw_feat = self.extract_features(image_bgr)43        norm_feat = self.scaler_base.transform(raw_feat)44        45        p_svm = self.svm_base.predict_proba(norm_feat)46        p_rf  = self.rf_base.predict_proba(norm_feat)47        meta_feat = np.hstack([p_svm, p_rf])48        49        pred_label = int(self.meta_base.predict(meta_feat)[0])50        confidence_vector = self.meta_base.predict_proba(meta_feat)[0]51        52        predicted_class = self.class_mapping[pred_label]53        class_confidence = float(confidence_vector[pred_label])54        55        all_class_scores = {56            self.class_mapping[i]: round(float(prob), 4) 57            for i, prob in enumerate(confidence_vector)58        }59        60        return {61            "predicted_category": predicted_class,62            "confidence_score": round(class_confidence, 4),63            "all_class_probabilities": all_class_scores64        }65