Create app.py

app.py

@@ -0,0 +1,337 @@
+import streamlit as st
+from PIL import Image
+import torch.nn as nn
+import numpy as np
+import torch
+import cv2
+import pandas as pd
+import os
+
+from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
+from transformers import AutoImageProcessor, DetrForObjectDetection
+# segmentation
+processor_seg = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
+model_seg = AutoModelForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
+#object detection
+processor_obj = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50")
+model_obj = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+
+
+def center_image(image_path,width=700):
+    st.markdown(
+        f'<style>img {{ display: block; margin-left: auto; margin-right: auto; }} </style>',
+        unsafe_allow_html=True
+    )
+    st.image(image_path,width = width)
+
+
+### INTRO ###
+st.header('👚 오늘 뭐입지?! 👕')
+st.markdown('💬 : 🚨 **설마 너 지금.. 그렇게 입고 나가게?** 🚨')
+st.markdown(' **패션센스가 2% 부족한 당신을 위해 준비했습니다!** 사진 이미지만 입력하면, 요즘 트렌디한 스타일과 여러분의 TPO를 고려하여 코디를 추천해드립니다. 무신사와 온더룩의 패셔니스타들의 코디를 지금 바로 참고해보세요! ')
+center_image('./intro_img/fashionista.jpg')
+
+st.markdown('--------------------------------------------------------------------------------------')
+st.subheader('PROCESS')
+center_image('./intro_img/process.png')
+st.markdown('--------------------------------------------------------------------------------------')
+
+
+## INPUT ###
+st.subheader(' ✅ 의류 이미지 업로드 ')
+input_image = st.file_uploader(" **의류 이미지를 업로드하세요. (배경이 깔끔한 사진이라면 더 좋습니다!)** ", type=['png', 'jpg', 'jpeg'])
+if not input_image :
+        con = st.container()
+        st.stop()
+center_image(input_image,400)
+st.markdown('--------------------------------------------------------------------------------------')
+
+st.subheader(' ✅ 업로드한 의류 이미지 카테고리 선택 ')
+input_cat = st.radio(
+    "**귀하가 업로드한 의류 이미지의 카테고리를 골라주세요.**",
+    ['top👕', 'bottom👖', 'shoes👞', 'hat🧢', 'sunglasses🕶️', 'scarf🧣', 'bag👜'],
+    index=None,
+    horizontal = True)
+
+
+if not input_cat :
+        con = st.container()
+        st.stop()
+input_cat = input_cat[:-1]
+st.write('You selected:', input_cat)
+st.markdown('--------------------------------------------------------------------------------------')
+
+st.subheader(' ✅ 추천받고 싶은 의류 카테고리 선택 ')
+output_cat = st.radio(
+    '**추천받고 싶은 의류 카테고리를 선택해주세요.**',
+    ['top👕', 'bottom👖', 'shoes👞', 'hat🧢', 'sunglasses🕶️', 'scarf🧣', 'bag👜'],
+    index=None,
+    horizontal = True)
+
+if not output_cat :
+        con = st.container()
+        st.write('🚫 주의: 업로드한 의류 카테고리와 다른 카테고리를 선택해주세요.')
+        st.stop()
+output_cat = output_cat[:-1]
+st.write('You selected:', output_cat)
+st.write(' ')
+st.markdown('--------------------------------------------------------------------------------------')
+
+
+st.subheader(' ✅ 상황 카테고리 선택 ')
+situation = st.radio(
+    "**상황 카테고리를 선택해주세요.**",
+    ['여행🌊', '카페☕️', '전시회🖼️', '캠퍼스🏫 & 출근💼', '급추위🤧', '운동💪'],
+    captions = ['(바다,여행)','(카페, 데일리)','(데이트, 결혼식)','','',''],
+    index=None,
+    horizontal = True)
+
+# 선택된 상황 카테고리를 영어로 변환해서 변수 저장
+situation_mapping = {
+    '여행🌊': 'travel',
+    '카페☕️': 'cafe',
+    '전시회🖼️': 'exhibit',
+    '캠퍼스🏫 & 출근💼': 'campus_work',
+    '급추위🤧': 'cold',
+    '운동💪': 'exercise'}
+
+if not situation:
+        con = st.container()
+        st.stop()
+situation= situation_mapping[situation]
+st.write('You selected:', situation)
+
+## 변수 명
+# input_img
+# input_cat : 입은 옷 카테고리
+# output_cat  : 추천 받을 카테고리
+# situation : 상황
+
+st.markdown('--------------------------------------------------------------------------------------')
+
+
+### 입력받은 이미지 segmentation & detection & vector변환 ###
+image = Image.open(input_image)
+
+# object detection & cropping 함수
+def cropping(images,st = 1,
+  fi = 0.0,
+  step = -0.05):
+  image_1 = Image.fromarray(images)
+  inputs = processor_obj(images=image_1, return_tensors="pt")
+  outputs = model_obj(**inputs)
+  for tre in np.arange(st,fi,step):
+    try:
+        # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
+        target_sizes = torch.tensor([image_1.size[::-1]])
+        results = processor_obj.post_process_object_detection(outputs, threshold=tre, target_sizes=target_sizes)[0]
+        
+        img = None
+        for idx, (score, label, box) in enumerate(zip(results["scores"], results["labels"], results["boxes"])):
+            box = [round(i, 2) for i in box.tolist()]
+            xmin, ymin, xmax, ymax = box
+            img = image_1.crop((xmin, ymin, xmax, ymax))
+
+        poss = np.array(img).sum().sum()
+        return img
+        break
+    except:
+        continue
+  return images
+
+# vector 변환 함수
+default_path = './'
+def image_to_vector(image,resize_size=(256,256)):  # 이미지 size 변환 resize(256,256)
+    #image = Image.fromarray(image)
+    #image = image.resize(resize_size)
+    image = Image.fromarray(np.copy(image))
+    image = image.resize(resize_size)
+    image_array = np.array(image, dtype=np.float32)
+    image_vector = image_array.flatten()
+    return image_vector
+
+# 전체 통합 함수
+def final_image(image):    
+    if len(np.array(image).shape) == 2:
+        image = Image.fromarray(image).convert('RGB')
+    # segmentation
+    inputs = processor_seg(images=image, return_tensors="pt")
+    outputs = model_seg(**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]
+    segments = torch.unique(pred_seg)
+    default_path = './'
+    
+    for i in segments:
+        if int(i) == 0:
+            continue
+        if int(i) == 1:
+            cloth = 'hat'
+            cloths = 'hat'
+            mask = pred_seg == i
+            image = np.array(image)
+            mask_np = (mask * 255).numpy().astype(np.uint8)
+            result = cv2.bitwise_and(image.astype(np.uint8), image.astype(np.uint8), mask=mask_np)
+            img = cropping(result)
+            img_vector = image_to_vector(img)
+        elif int(i) == 3:
+            cloth= 'sunglasses'
+            cloths= 'sunglasses'
+            mask = pred_seg == i
+            image = np.array(image)
+            mask_np = (mask * 255).numpy().astype(np.uint8)
+            result = cv2.bitwise_and(image.astype(np.uint8), image.astype(np.uint8), mask=mask_np)
+            img = cropping(result)
+            img_vector = image_to_vector(img)
+        elif int(i) == 4:
+            cloth = 'top'
+            cloths = 'top'
+            mask = pred_seg == i
+            image = np.array(image)
+            mask_np = (mask * 255).numpy().astype(np.uint8)
+            result = cv2.bitwise_and(image.astype(np.uint8), image.astype(np.uint8), mask=mask_np)
+            img = cropping(result)
+            img_vector = image_to_vector(img)
+        elif int(i) in [5,6,7]:
+            cloth= ['pants','skirt','dress']
+            cloths= 'bottom'
+            mask  = (pred_seg == torch.tensor(5)) | (pred_seg == torch.tensor(6)) | (pred_seg == torch.tensor(7))
+            image = np.array(image)
+            mask_np = (mask * 255).numpy().astype(np.uint8)
+            result = cv2.bitwise_and(image.astype(np.uint8), image.astype(np.uint8), mask=mask_np)
+            img = cropping(result)
+            img_vector = image_to_vector(img)
+        elif int(i) == 8:
+            cloth = 'belt'
+            cloths = 'belt'
+            mask = pred_seg == torch.tensor(8)
+            image = np.array(image)
+            mask_np = (mask * 255).numpy().astype(np.uint8)
+            result = cv2.bitwise_and(image.astype(np.uint8), image.astype(np.uint8), mask=mask_np)
+            img = cropping(result)
+            img_vector = image_to_vector(img)
+        elif (int(i) == 9):
+            cloth = 'shoes'
+            cloths = 'shoes'
+            mask = (pred_seg == torch.tensor(9)) | (pred_seg == torch.tensor(10))
+            image = np.array(image)
+            mask_np = (mask * 255).numpy().astype(np.uint8)
+            result = cv2.bitwise_and(image.astype(np.uint8), image.astype(np.uint8), mask=mask_np)
+            img = cropping(result)
+            img_vector = image_to_vector(img)
+        elif int(i) == 16:
+            cloth = 'bag'
+            cloths = 'bag'
+            mask = pred_seg == torch.tensor(16)
+            image = np.array(image)
+            mask_np = (mask * 255).numpy().astype(np.uint8)
+            result = cv2.bitwise_and(image.astype(np.uint8), image.astype(np.uint8), mask=mask_np)
+            img = cropping(result)
+            img_vector = image_to_vector(img)
+        elif int(i) == 17:
+            cloth = 'scarf'
+            cloths = 'scarf'
+            mask = pred_seg == torch.tensor(17)
+            image = np.array(image)
+            mask_np = (mask * 255).numpy().astype(np.uint8)
+            result = cv2.bitwise_and(image.astype(np.uint8), image.astype(np.uint8), mask=mask_np)
+            img = cropping(result)
+            img_vector = image_to_vector(img)
+        return img_vector
+    
+# 입력받은 이미지 전처리 완료
+input_img = final_image(image)
+
+### 유사도 분석 ###
+# 하나는 이미���, 다른 하나는 경로로 받는 경우
+def cosine_similarity(vec1, vec2_path):
+    vec2 = np.loadtxt(vec2_path)
+    dot_product = np.dot(vec1, vec2)
+    norm_vec1 = np.linalg.norm(vec1)
+    norm_vec2 = np.linalg.norm(vec2)
+    similarity = dot_product / (norm_vec1 * norm_vec2)   
+    return similarity
+
+# 둘 다 경로로 받는 경우
+def cosine_similarity_2(vec1_path, vec2_path):
+    vec1 = np.loadtxt(vec1_path)
+    vec2 = np.loadtxt(vec2_path)
+    dot_product = np.dot(vec1, vec2)
+    norm_vec1 = np.linalg.norm(vec1)
+    norm_vec2 = np.linalg.norm(vec2)
+    similarity = dot_product / (norm_vec1 * norm_vec2)
+    return similarity
+
+with st.spinner('Wait for it...'):
+    # 입력받은 이미지 & 동일 카테고리 폴더에 저장된 스타일 이미지
+    sim_list = []
+    file_path = './style/' + situation + '/' + input_cat + '/'   # ex) './cafe/top/'
+    cloths = os.listdir('./style/' + situation + '/' + input_cat + '/')
+    for cloth in cloths:
+        sim_list.append(cosine_similarity(input_img, file_path + cloth))
+    max_idx = np.argmax(sim_list)
+
+    # target_image 정의
+    target_image = './style/' + situation + '/' + output_cat + '/' + cloths[max_idx]
+    # 유사도 분석 완료된 스타일seg 이미지와 product_seg 유사도분석
+    sim_list = []
+    file_path = './product/' + output_cat + '/'
+    cloths = os.listdir('./product/' + output_cat + '/')
+    for cloth in cloths:
+        sim_list.append(cosine_similarity_2(target_image, file_path + cloth))
+    max_idx = np.argmax(sim_list)
+    output_name  = cloths[max_idx]
+    ## 예시 출력값: 'bottom_1883.txt'
+
+    # name 로드
+    acc_name = pd.read_csv('acc_name.csv')
+    bottom_name =pd.read_csv('bottom_name.csv')
+    outer_name =pd.read_csv('outer_name.csv')
+    shoes_name =pd.read_csv('shoes_name.csv')
+    top_name =pd.read_csv('top_name.csv')
+
+    #상품 데이터 로드
+    outer = pd.read_csv('outer.csv')
+    top = pd.read_csv('top.csv')
+    bottom = pd.read_csv('bottom.csv')
+    shoes = pd.read_csv('shoes.csv')
+    acc = pd.read_csv('acc.csv')
+
+    if output_cat == 'bottom':
+        df = bottom.copy()
+        df_name = bottom_name.copy()
+    elif output_cat == 'top':
+        df = top.copy()
+        df_name = top_name.copy()
+    elif output_cat == 'shoes':
+        df = shoes.copy()
+        df_name = shoes_name.copy()
+    elif (output_cat == 'hat') or (output_cat == 'sunglasses') or (output_cat == 'scarf') or (output_cat == 'bag') or (output_cat == 'belt'):
+        df = acc.copy()
+        df_name = acc_name.copy()
+
+    output_name = output_name.split('.')[0]
+    file_name = df_name[df_name['index']==output_name].iloc[0,1] #3049906_16754112975667_500.jpg
+    final = df[df['id'] == file_name]
+
+    name = final['name'].values[0].split('\n')[-1] # 상품명
+    price = final['price'].values[0] # 상품가격
+
+image_path = './product/img/'
+
+st.subheader('OUTPUT')
+
+img = Image.open(image_path+output_cat+'/'+file_name)
+
+col1, col2, col3 = st.columns(3)
+with col1:
+    st.image(img,width=400)
+with col3:
+    st.caption('상풍명 : ' + name)
+    st.caption('가격 : ' + price)