Update app.py

app.py

@@ -1,36 +1,28 @@
-import torch
 import gradio as gr
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-REPO_ID = "https://huggingface.co/spaces/Kr1n3/Fashion-Items-Classification"
-FILENAME = "best.pt"
-
-yolov5_weights = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+import torch
+from torch import nn
+from torch.nn import functional as F
+import torchvision
+from torchvision import transforms
 
-model = torch.hub.load('ultralytics/yolov5', 'custom', path=yolov5_weights, force_reload=True)  # local repo
+model= torch.jit.load('best.pt')
+data_transform1=transforms.Compose([
+    transforms.Resize((640,640)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
+])
 
-def object_detection(im, size=640):
-    results = model(im)  # inference
-    #results.print()  # print results to screen
-    #results.show()  # display results
-    #results.save()  # save as results1.jpg, results2.jpg... etc.
-    results.render()  # updates results.imgs with boxes and labels
-    return Image.fromarray(results.imgs[0])
+title = " Fashion Items Classification"
 
-title = "Fashion Items Classification"
-description = """Esse modelo é uma pequena demonstração baseada em uma análise de cerca de 60 imagens somente. Para resultados mais confiáveis e genéricos, são necessários mais exemplos (imagens).
-"""
+examples=[['https://github.com/Kr1n3/MPC_2022/blob/main/dataset/pants_33.jpg?raw=true'],['https://github.com/Kr1n3/MPC_2022/blob/main/dataset/pants_30.jpeg?raw=true'],['https://github.com/Kr1n3/MPC_2022/blob/main/dataset/bag_01.jpg?raw=true']]
 
-image = gr.inputs.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Imagem", optional=False)
-outputs = gr.outputs.Image(type="pil", label="Output Image")
+classes=['Bags','Dress','Pants','Shoes','Skirt']
+def predict(img):
+  imag=data_transform1(img)
+  inp =imag.unsqueeze(0)
+  outputs=model(inp)
+  pred=F.softmax(outputs[0], dim=0).cpu().data.numpy()
+  confidences = {classes[i]:(float(pred[i])) for i in range(5)}    
+  return confidences
 
-gr.Interface(
-    fn=object_classification,
-    inputs=image,
-    outputs=outputs,
-    title=title,
-    description=description,
-    examples=[["https://github.com/Kr1n3/MPC_2022/blob/main/dataset/bag_01.jpg"], ["https://github.com/Kr1n3/MPC_2022/blob/main/dataset/bag_18.JPG?raw=true"], 
-              ["https://github.com/Kr1n3/MPC_2022/blob/main/dataset/pants_30.jpeg?raw=true"], ["https://github.com/Kr1n3/MPC_2022/blob/main/dataset/pants_33.jpg?raw=true"]],
-).launch()
+gr.Interface(predict,gr.inputs.Image(type='pil'),title=title,examples=examples,outputs='label').launch(debug=True)