Rahaf1/dithar-api
0
1import os 2from fastapi import FastAPI, UploadFile, File3from fastapi.middleware.cors import CORSMiddleware4from PIL import Image, Image as _PILImage5import io, torch, numpy as np, colorsys6import open_clip7from skimage import color as skcolor8 9app = FastAPI()10app.add_middleware(11 CORSMiddleware,12 allow_origins=["*"],13 allow_methods=["*"],14 allow_headers=["*"],15)16 17 18device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")19model_name = 'ViT-B-32'20pretrained = 'laion2b_s34b_b79k'21 22 23 24CACHE_DIR = os.getenv("HF_HOME", "/tmp/hf")25os.makedirs(CACHE_DIR, exist_ok=True)26 27 28model, _, preprocess = open_clip.create_model_and_transforms(29 model_name,30 pretrained=pretrained,31 device=device,32 cache_dir=CACHE_DIR, 33)34model.eval()35 36 37CATEGORY_PROMPTS = [38 ("thobe (saudi thobe)", "thobe"),39 ("saudi thobe", "thobe"),40 ("gulf thobe", "thobe"),41 ("dishdasha", "thobe"),42 ("kandura", "thobe"),43 ("jalabiya (men's)", "thobe"),44 45 ("abaya (women's cloak)", "abaya"),46 ("black abaya", "abaya"),47 ("saudi abaya", "abaya"),48 ("abaya cloak", "abaya"),49 50 ("t-shirt", "tshirt"),51 ("short-sleeve t-shirt", "tshirt"),52 ("long-sleeve t-shirt", "tshirt"),53 ("shirt", "shirt"),54 ("button-up shirt", "shirt"),55 ("blouse", "blouse"),56 ("sweater", "sweater"),57 ("pullover knit", "sweater"),58 ("coat", "coat"),59 ("overcoat", "coat"),60 ("dress", "dress"),61 ("scarf", "scarf"),62 ("pants", "pants"),63 ("trousers", "pants"),64 ("jeans", "pants"),65 ("skirt", "skirt"),66 ("shorts", "shorts"),67 ("tank top", "tanktop"),68 ("camisole", "tanktop"),69 ("undershirt", "tanktop"),70 ("sneakers", "sneakers"),71 ("running shoes", "sneakers"),72 ("dress shoes", "dress_shoes"),73 ("oxford shoes", "dress_shoes"),74 ("sandals", "sandals"),75 ("high heels", "heels"),76 ("boots", "boots"),77 ("necklace", "necklace"),78 ("bracelet", "bracelet"),79 ("earrings", "earrings"),80 ("ring", "ring"),81 ("watch", "watch"),82 ("glasses", "glasses"),83 ("sunglasses", "glasses"),84 ("handbag", "bag"),85 ("bag", "bag"),86 ("belt", "belt"),87 ("hat", "hat"),88 ("cap", "hat"),89]90 91CATEGORY_AR = {92 "shirt" : "قميص",93 "blouse" : "بلوزة",94 "sweater" : "كنزة",95 "coat" : "معطف",96 "dress" : "فستان",97 "tshirt" : "تيشيرت",98 "scarf" : "وشاح",99 100 "pants" : "بنطال",101 "skirt" : "تنورة",102 "shorts" : "شورت",103 "tanktop" : "شيال",104 105 "sneakers" : "حذاء رياضي",106 "dress_shoes" : "حذاء رسمي",107 "sandals" : "صندل",108 "heels" : "كعب",109 "boots" : "بوت",110 111 "necklace" : "سلسال",112 "bracelet" : "اسورة",113 "earrings" : "حلق",114 "ring" : "خاتم",115 "watch" : "ساعة",116 "glasses" : "نظارة",117 "bag" : "حقيبة",118 "belt" : "حزام",119 "hat" : "قبعة",120 121 "thobe" : "ثوب",122 "abaya" : "عباية",123}124 125CATEGORY_LABELS_EN = [p[0] for p in CATEGORY_PROMPTS]126CATEGORY_KEYS = [p[1] for p in CATEGORY_PROMPTS]127 128PATTERN_PROMPTS = [129 ("solid/plain", "plain"),130 ("striped", "striped"),131 ("floral", "floral"),132 ("plaid", "plaid"),133 ("checkered", "checkered"),134 ("polka dot", "polka"),135 ("animal print", "animal"),136 ("geometric", "geometric"),137 ("camouflage", "camo"),138 ("lace", "lace"),139 ("crochet", "crochet"),140]141PATTERN_LABELS_EN = [p[0] for p in PATTERN_PROMPTS]142PATTERN_KEYS = [p[1] for p in PATTERN_PROMPTS]143 144PATTERN_AR = {145 "plain" : "سادة",146 "striped" : "مخطط",147 "floral" : "مورد",148 "plaid" : "كاروهات",149 "checkered" : "مربعات",150 "polka" : "منقط",151 "geometric" : "اشكال هندسية",152 "camo" : "تمويه",153 "lace" : "دانتيل",154 "crochet" : "كروشيه",155}156 157def center_crop(img: Image.Image, frac: float) -> Image.Image:158 w, h = img.size159 cw, ch = int(w * frac), int(h * frac)160 x0, y0 = (w - cw) // 2, (h - ch) // 2161 return img.crop((x0, y0, x0 + cw, y0 + ch))162 163@torch.no_grad()164def clip_scores(img: Image.Image, labels_en, templates):165 scores = torch.zeros(len(labels_en), device=device)166 image = preprocess(img).unsqueeze(0).to(device)167 for t in templates:168 texts = [t.format(l) for l in labels_en]169 text = open_clip.tokenize(texts).to(device)170 img_f = model.encode_image(image)171 txt_f = model.encode_text(text)172 img_f /= img_f.norm(dim=-1, keepdim=True)173 txt_f /= txt_f.norm(dim=-1, keepdim=True)174 logits = (100.0 * img_f @ txt_f.T).squeeze(0).softmax(-1)175 scores += logits176 return scores / len(templates)177 178 179 180def rgb_to_hex(rgb):181 r, g, b = rgb182 return f"#{r:02X}{g:02X}{b:02X}"183 184def _lab_from_rgb_pixel(rgb):185 arr = np.array([[[rgb[0] / 255.0, rgb[1] / 255.0, rgb[2] / 255.0]]], dtype=float)186 return skcolor.rgb2lab(arr)[0, 0, :]187 188def _hsv_from_rgb(rgb):189 r, g, b = rgb190 return colorsys.rgb_to_hsv(r / 255.0, g / 255.0, b / 255.0)191 192COLOR_AR = {193 "White": "أبيض", "Black": "أسود", "Gray": "رمادي", "Brown": "بني", "Beige": "بيج",194 "Red": "أحمر", "Pink": "وردي", "Purple": "بنفسجي", "Orange": "برتقالي", "Yellow": "أصفر",195 "Green": "أخضر", "Cyan": "سماوي", "Blue": "أزرق", "Gold": "ذهبي", "Silver": "فضي"196}197 198def _base_color_name(rgb):199 L, a, b = _lab_from_rgb_pixel(rgb)200 chroma = (a * a + b * b) ** 0.5201 h, s, v = _hsv_from_rgb(rgb)202 h_deg = h * 360.0203 204 if L > 85 and chroma < 20:205 return "White"206 207 if (L < 22 and chroma < 22) or (v < 0.22 and s < 0.35 and chroma < 20):208 return "Black"209 210 211 if (68 <= h_deg < 160) and ((L < 60 and s >= 0.07) or (0.10 <= s < 0.28)):212 return "Green"213 214 if (48 <= h_deg <= 68) and (s > 0.18) and (L > 55):215 return "Yellow"216 217 if (30 <= h_deg <= 75) and (55 <= L <= 95) and (v > 0.50) and (s <= 0.38) and (chroma <= 22):218 return "Beige"219 220 if chroma < 14:221 if (60 <= L <= 95) and (s < 0.35) and (v > 0.55) and (b > 0.5) and ((b - a) > 2):222 return "Beige"223 return "Gray"224 225 if (15 <= h_deg < 55) and (a > 5) and (b > 10):226 if (L >= 72 or v >= 0.78) and (s >= 0.52) and (chroma >= 28):227 return "Orange"228 return "Brown"229 230 if (h_deg >= 345 or h_deg < 15): return "Red"231 if 15 <= h_deg < 45: return "Orange"232 if 45 <= h_deg < 68: return "Yellow"233 if 68 <= h_deg < 160: return "Green"234 if 160 <= h_deg < 200: return "Cyan"235 if 200 <= h_deg < 255: return "Blue"236 if 255 <= h_deg < 290: return "Purple"237 if 290 <= h_deg < 345: return "Pink"238 239 if L > 82: return "White"240 return "Gray"241 242 243 244def _map_to_allowed(name):245 """246 حصر المخرجات في قائمتك فقط.247 أي اسم خارج المجموعة يُقرّب للأقرب المنطقي.248 """249 if name in COLOR_AR:250 return name251 if name in ("Cream", "Ivory", "Taupe", "Khaki"): return "Beige"252 if name in ("Navy",): return "Blue"253 if name in ("Teal", "Turquoise"): return "Cyan"254 if name in ("Maroon", "Burgundy", "Wine"): return "Red"255 return "Gray"256 257def main_color(image: Image.Image, is_accessory: bool):258 W, H = image.size259 focus = image.crop((int(W * 0.15), int(H * 0.15), int(W * 0.85), int(H * 0.85)))260 261 focus.thumbnail((200, 200), _PILImage.Resampling.LANCZOS)262 focus_np = np.array(focus)263 264 rgb01 = focus_np.reshape(-1, 3).astype(np.float32) / 255.0265 hsv = np.array([colorsys.rgb_to_hsv(r, g, b) for r, g, b in rgb01])266 V = hsv[:, 2]; S = hsv[:, 1]; H = hsv[:, 0]267 268 is_dark_item = (V < 0.35).mean() > 0.40269 270 mask = ((V > 0.35) & (S > 0.08) & ~((S < 0.10) & (V < 0.60))) | (V < 0.25)271 if mask.sum() < 500:272 mask = ((S > 0.04) & (V > 0.20)) | (V < 0.28)273 274 rgb_masked = rgb01[mask]275 if rgb_masked.shape[0] == 0:276 rgb_masked = rgb01277 278 lab = skcolor.rgb2lab(rgb_masked.reshape(-1, 1, 3)).reshape(-1, 3)279 Xw = np.stack([lab[:, 0] * 1.0, lab[:, 1] * 0.6, lab[:, 2] * 0.6], axis=1)280 281 def kmeans_np(data, k=5, iters=12, seed=42):282 rng = np.random.default_rng(seed)283 cent = data[rng.choice(len(data), size=k, replace=False)]284 for _ in range(iters):285 d = np.linalg.norm(data[:, None, :] - cent[None, :, :], axis=2)286 lbl = d.argmin(axis=1)287 newc = np.array([data[lbl == i].mean(axis=0) if np.any(lbl == i) else cent[i] for i in range(k)])288 if np.allclose(newc, cent): break289 cent = newc290 return lbl, cent291 292 labels, centers_lab = kmeans_np(Xw, k=5)293 counts = np.bincount(labels)294 props = counts / labels.size295 296 centers_rgb_raw = skcolor.lab2rgb(centers_lab / np.array([1.0, 0.6, 0.6]))297 centers_rgb_255 = (centers_rgb_raw * 255).clip(0, 255).astype(np.uint8)298 299 centers_lab_unweighted = centers_lab / np.array([1.0, 0.6, 0.6])300 a_vals = centers_lab_unweighted[:, 1]301 b_vals = centers_lab_unweighted[:, 2]302 chroma_vals = np.sqrt(a_vals * a_vals + b_vals * b_vals)303 L_vals = centers_lab[:, 0]304 305 def name_of_idx(i: int):306 rgb_i = tuple(map(int, centers_rgb_255[i]))307 return _base_color_name(rgb_i)308 309 if is_dark_item:310 candidates = np.where(props > 0.10)[0]311 if len(candidates) == 0:312 darkest_idx = int(np.argmin(L_vals))313 else:314 darkest_idx = int(candidates[np.argmin(L_vals[candidates])])315 316 hs_all = [ _hsv_from_rgb(tuple(map(int, centers_rgb_255[i]))) for i in range(len(centers_lab)) ]317 hdeg_all = np.array([h*360.0 for (h,s,v) in hs_all])318 s_all = np.array([s for (h,s,v) in hs_all])319 320 blue_cands = np.where((hdeg_all >= 200) & (hdeg_all < 255) & (s_all >= 0.08) & (props >= 0.06))[0]321 if len(blue_cands):322 blue_dark_idx = int(blue_cands[np.argmin(L_vals[blue_cands])])323 L_darkest = float(L_vals[darkest_idx])324 L_blue = float(L_vals[blue_dark_idx])325 best_idx = blue_dark_idx if (L_darkest - L_blue) <= 8.0 else darkest_idx326 else:327 best_idx = darkest_idx328 else:329 best_idx = int(np.argmax(counts))330 331 best_color_rgb = tuple(map(int, centers_rgb_255[best_idx]))332 main_rgb = best_color_rgb333 334 prelim = _base_color_name(main_rgb)335 if prelim == "Black":336 H_mask = H[mask]; S_mask = S[mask]; V_mask = V[mask]337 338 blue_band = (H_mask >= (200.0/360.0)) & (H_mask < (255.0/360.0))339 blueish_ratio = float(((blue_band) & (S_mask > 0.07)).sum()) / float(H_mask.size + 1e-6)340 if is_dark_item and blueish_ratio >= 0.30:341 prelim = "Blue"342 343 if prelim == "Black":344 red_band = (((H_mask >= (345.0/360.0)) | (H_mask < (20.0/360.0))) & (S_mask > 0.10))345 redish_ratio = float(red_band.sum()) / float(H_mask.size + 1e-6)346 s_med = float(np.median(S_mask)); v_med = float(np.median(V_mask))347 if is_dark_item and (redish_ratio >= 0.22 or (redish_ratio >= 0.16 and s_med >= 0.14 and v_med >= 0.12)):348 prelim = "Red"349 350 if (prelim == "Cyan") and (not is_dark_item) and (props[best_idx] < 0.20):351 major_idx = int(np.argmax(counts))352 main_rgb = tuple(map(int, centers_rgb_255[major_idx]))353 prelim = _base_color_name(main_rgb)354 355 if is_accessory:356 h, s, v = _hsv_from_rgb(main_rgb)357 h_deg = h * 360.0358 highlight_ratio = float(((hsv[:, 2] > 0.92) & mask).sum()) / float(mask.sum() + 1e-6)359 if (15 <= h_deg < 70) and (v > 0.45) and (s > 0.15) and (highlight_ratio > 0.01):360 prelim = "Gold"361 elif prelim in ("Gray", "White") and v > 0.55 and s < 0.20 and (highlight_ratio > 0.02):362 prelim = "Silver"363 364 final_name = _map_to_allowed(prelim)365 return COLOR_AR[final_name], rgb_to_hex(main_rgb)366 367@app.post("/classify")368async def classify(file: UploadFile = File(...)):369 raw = await file.read()370 img = Image.open(io.BytesIO(raw)).convert("RGB")371 372 cat_templates = [373 "a product photo of {}",374 "a clothing item: {}",375 "a {}",376 ]377 s_full = clip_scores(img, CATEGORY_LABELS_EN, cat_templates)378 s_cent = clip_scores(center_crop(img, 0.85), CATEGORY_LABELS_EN, cat_templates)379 scores_cat = (s_full + s_cent) / 2.0380 best_idx = int(torch.argmax(scores_cat).item())381 cat_key = CATEGORY_KEYS[best_idx]382 cat_ar = CATEGORY_AR.get(cat_key, "غير محدد")383 cat_conf = float(scores_cat[best_idx].item())384 385 is_accessory = cat_key in {"necklace","bracelet","earrings","ring","watch","glasses","bag","belt","hat"}386 387 pat_templates = ["a fabric pattern that is {}", "a clothing pattern: {}"]388 p_full = clip_scores(img, PATTERN_LABELS_EN, pat_templates)389 p_cent = clip_scores(center_crop(img, 0.70), PATTERN_LABELS_EN, pat_templates)390 p_scores = (p_full + p_cent) / 2.0391 p_idx = int(torch.argmax(p_scores).item())392 pattern_key = PATTERN_KEYS[p_idx]393 pattern_ar = PATTERN_AR.get(pattern_key, "سادة")394 pattern_conf = float(p_scores[p_idx].item())395 396 color_ar, color_hex = main_color(img, is_accessory=is_accessory)397 398 return {399 "category": cat_ar,400 "color": color_ar,401 "pattern": pattern_ar,402 "color_hex": color_hex,403 "scores": {404 "category": cat_conf,405 "pattern": pattern_conf406 },407 "device": device,408 "model": f"{model_name}/{pretrained}"409 }410 411 412 413 414 415@app.get("/")416def health():417 return {"ok": True, "model": f"{model_name}/{pretrained}"}418 419if __name__ == "__main__":420 import uvicorn421 port = int(os.environ.get("PORT", 7860)) 422 uvicorn.run("app.main:app", host="0.0.0.0", port=port)423 424