# helper functions to get segmented mask import gradio as gr import os from PIL import Image import urllib import shutil url = "http://static.okkular.io/scripted.model" output_file = "./scripted.model" with urllib.request.urlopen(url) as response, open(output_file, 'wb') as out_file: shutil.copyfileobj(response, out_file) def get_stl(input_sku): preds=shop_the_look(f'./data/dress_{input_sku}.jpg') ret_bag = preds['./segs/bag.jpg'][1] ret_shoes = preds['./segs/shoe.jpg'][1] return Image.open(f'./data/dress_{input_sku}.jpg'), Image.open(ret_bag), Image.open(ret_shoes) sku = gr.Dropdown( ["1", "2", "3", '4', '5'], label="Dress Sku", ), demo = gr.Interface(get_stl, gr.Dropdown( ["1", "2", "3", '4', '5'], label="Dress Sku"), ["image", "image", "image"]) demo.launch(root_path=f"/{os.getenv('TOKEN')}") from PIL import Image, ImageChops import numpy as np from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation from PIL import Image import requests import matplotlib.pyplot as plt import torch import torch.nn as nn import os import nmslib from fastai.vision.all import * def get_segment(image, num,ret=False): extractor = AutoFeatureExtractor.from_pretrained("mattmdjaga/segformer_b2_clothes") model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes") inputs = extractor(images=image, return_tensors="pt") outputs = model(**inputs) logits = outputs.logits.cpu() upsampled_logits = nn.functional.interpolate( logits, size=image.size[::-1], mode="bilinear", align_corners=False, ) #pred_seg = upsampled_logits.argmax(dim=1)[0] pred_seg = upsampled_logits.argmax(dim=1)[0] np_im = np.array(image) pred_seg[pred_seg != num] = 0 mask = pred_seg.detach().cpu().numpy() # masked region np_im[mask.squeeze()==0] = 0 # white bg np_im[np.where((np_im==[0,0,0]).all(axis=2))] = [255,255,255] # trim extra whitespace im = Image.fromarray(np.uint8(np_im)).convert('RGB') im = trim(im) if ret==False: plt.imshow(im) plt.show() elif ret==True: print('here and returning', im) return im def trim(im): bg = Image.new(im.mode, im.size, im.getpixel((0,0))) diff = ImageChops.difference(im, bg) diff = ImageChops.add(diff, diff, 2.0, -100) bbox = diff.getbbox() if bbox: return im.crop(bbox) def get_pred_seg(image_url): extractor = AutoFeatureExtractor.from_pretrained("mattmdjaga/segformer_b2_clothes") model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes") image = Image.open(image_url) inputs = extractor(images=image, return_tensors="pt") outputs = model(**inputs) logits = outputs.logits.cpu() upsampled_logits = nn.functional.interpolate( logits, size=image.size[::-1], mode="bilinear", align_corners=False, ) pred_seg = upsampled_logits.argmax(dim=1)[0] #plt.imshow(pred_seg) return upsampled_logits,pred_seg #### get predictions and neighbours def get_predictions(feed): pairs = [[x["sku"], x["category"]] for x in feed] skus, labels = zip(*pairs) labels = np.array(labels) skus = np.array(skus) categories = list(set(labels)) def get_image_fpath(x): return x[0] data = DataBlock( blocks=(ImageBlock, CategoryBlock), get_x = get_image_fpath, get_y = ItemGetter(1), item_tfms=[Resize(256)], batch_tfms=[Normalize.from_stats(*imagenet_stats)], splitter=IndexSplitter([]) ) dls = data.dataloaders( pairs, device=default_device(), shuffle_fn=lambda x:x, drop_last=False ) #model = torch.jit.load("../inference_script/scripted.model").cpu() #model = torch.jit.load("scripted.model").cuda() with open('./scripted.model', 'rb') as f: buffer = io.BytesIO(f.read()) # Load all tensors to the original device model = torch.jit.load(buffer, map_location=torch.device('cpu')) preds_list = [] with torch.no_grad(): for x,y in progress_bar(iter(dls.train), total=len(dls.train)): pred = model(x) preds_list.append(pred) preds = torch.cat(preds_list) preds = to_np(preds) predictions_json = {} for cat in categories: filtered_preds = preds[labels == cat] filtered_skus = skus[labels==cat] neighbours,dists = get_neighbours(filtered_preds) #neighbours = neighbours[:,1:] for i, sku in enumerate(filtered_skus): predictions_json[sku] = [filtered_skus[j] for j in neighbours[i]] return predictions_json INDEX_TIME_PARAMS = {'M': 100, 'indexThreadQty': 8, 'efConstruction': 2000, 'post': 0} QUERY_TIME_PARAMS = {"efSearch": 2000} N_NEIGHBOURS = 4 def get_neighbours(embeddings): index = nmslib.init(method='hnsw', space='l2') index.addDataPointBatch(embeddings) index.createIndex(INDEX_TIME_PARAMS) index.setQueryTimeParams(QUERY_TIME_PARAMS) res = index.knnQueryBatch( embeddings, k=min(N_NEIGHBOURS+1, embeddings.shape[0]), num_threads=8) proc_res = [l[None] for (l, d) in res] neighbours = np.concatenate(proc_res).astype(np.int32) dists = np.array([d for (_, d) in res]).astype(np.float32) return neighbours , dists def shop_the_look(prod): #upsampled_logits, all_segs = get_pred_seg(prod) bag_segment=get_segment(Image.open(prod), 16, ret=True) bag_segment.save('./segs/bag.jpg') shoe_l = get_segment(Image.open(prod), 9, True)#left shoe shoe_r = get_segment(Image.open(prod), 10, True) #right shoe shoe_segment = concat_h(shoe_l, shoe_r) shoe_segment.save('./segs/shoe.jpg') feed= [] main_prods=os.listdir('./data') for sku in main_prods: if 'checkpoint' not in sku: cat = sku.split('_')[0] x={'sku':f'./data/{sku}', 'category':cat} feed.append(x) feed.extend([{'sku':'./segs/shoe.jpg', 'category':'shoes'}, {'sku':'./segs/bag.jpg', 'category':'bag'}]) preds=get_predictions(feed) return preds def concat_h(image1,image2): #resize, first image image1 = image1.resize((426, 240)) image1_size = image1.size image2_size = image2.size new_image = Image.new('RGB',(2*image1_size[0], image1_size[1]), (250,250,250)) new_image.paste(image1,(0,0)) new_image.paste(image2,(image1_size[0],0)) return new_image