Update app.py

app.py

@@ -1,20 +1,35 @@
-import gradio as gr
-import tensorflow as tf
-from tensorflow.keras.preprocessing import image
-import matplotlib.pyplot as plt
+import os
+env_var = os.environ.get('env')
+import torch
+import time 
 import numpy as np
-model = tf.keras.models.load_model('dogcat_model_bak.h5')
-def image_classifier(img):
-    img1 = image.load_img(str(img), target_size=(64, 64))
-    img1 = image.img_to_array(img1)
-    img1 = img1/255
-    img1 = np.expand_dims(img1, axis=0)
-    res = model.predict(img, batch_size=None,steps=1)
-    if(res[:,:]>0.5):
-      value ='Dog :%1.2f'%(prediction[0,0])
-    else:
-      value ='Cat :%1.2f'%(1.0-prediction[0,0])
-    return value
+import gradio as gr
+from PIL import Image
+import torchvision 
+from torchvision import transforms
+
+device = 'cpu'
+model = torch.load('model.pkl').to(device).eval()
+transform = transforms.Resize(size=500)
+labels = ['Cat', 'Dog']
+
+def predict(image):
+    start = time.time()
+    with torch.no_grad():
+        image = Image.fromarray(np.uint8(image)).convert('RGB')
+        image = transform(image)
+        image = np.array(image)
+        image = torch.from_numpy(image).permute(2,0,1).float()
+        image = image.unsqueeze(0)
+        prediction = model(image.to(device))
+        pred_idx = np.argmax(prediction.to(device))
+        pred_label = "Cat" if pred_idx == 0 else "Dog"
+        label = [l for l in labels if l!=pred_label]
+        confidences = {pred_label: float(prediction[0][pred_idx])/100, label[len(label)-1]: 1-(float(prediction[0][pred_idx]))/100 }
+        infer = time.time()-start
+    return confidences, infer
 
-demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
-demo.launch()
+gr.Interface(fn=predict,
+             inputs=gr.inputs.Image(shape=(512, 512)), 
+             outputs=[gr.outputs.Label(num_top_classes=3), gr.outputs.Textbox('infer',label='Inference Time')],
+             examples='1.jpg 2.jpg'.split(' ')).launch()