Update app.py

app.py

@@ -1,32 +1,28 @@
-import torch as tf
-
 import gradio as gr
-
-inception_net = tf.keras.models.load_model('best.pt')  # load the model
-
-labels = ['bom', 'ruim']
-
-
-title = "Fashion Items Classification"
-description = """
-"""
-
-def classify_image(inp):
-    inp = inp.reshape((-1, 640, 640, 3))
-    inp = tf.keras.applications.mobilenet_v2.preprocess_input(inp)
-    prediction = inception_net.predict(inp).flatten()
-    return (labels[1] if float(prediction) >= 0 else labels[0])
-
-
-image = gr.inputs.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Imagem", optional=False)
-label = gr.outputs.Textbox(type="auto", label="Classificação")
-
-gr.Interface(
-    fn=classify_image,
-    inputs=image,
-    outputs=label,
-    title=title,
-    description=description,
-    examples=[["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"], 
-              ["https://github.com/Kr1n3/MPC_2022/blob/main/dataset/bag_14.JPG?raw=true"], ["https://github.com/Kr1n3/MPC_2022/blob/main/dataset/dress_45.JPG?raw=true"]],
-).launch()
+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', 'runs/train-cls/exp/weights/best.pt')
+data_transform1=transforms.Compose([
+    transforms.Resize((224,224)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
+])
+
+title = " Fashion Items Classification"
+
+examples=[['https://github.com/Kr1n3/MPC_2022/blob/main/dataset/bag_14.JPG?raw=true'],['https://github.com/Kr1n3/MPC_2022/blob/main/dataset/dress_45.JPG?raw=true'],['https://github.com/Kr1n3/MPC_2022/blob/main/dataset/pants_30.jpeg?raw=true']]
+
+classes=['Bag','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(predict,gr.inputs.Image(type='pil'),title=title,examples=examples,outputs='label').launch(debug=True)